Pool queries should honor connection timeout (#1586) (#1588)

See #1232

If a connection timeout is defined in connect options, the pool uses it when users acquire a connection with pool.getConnection, but not when they execute a query with pool.query

Signed-off-by: Thomas Segismont <tsegismont@gmail.com>

Author: tsegismont

vertx-pg-client/src/test/java/io/vertx/pgclient/PgPoolTest.java

@@ -34,17 +34,14 @@
 import org.junit.Rule;
 import org.junit.Test;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collector;
 
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toList;
 
@@ -592,4 +589,17 @@ private void testConnectionClosedInProvider(TestContext ctx, boolean immediately
       }));
     }));
   }
+
+  @Test
+  public void testPooledQueryTimeout(TestContext ctx) {
+    Async async = ctx.async();
+    PoolOptions poolOptions = new PoolOptions().setMaxSize(1).setConnectionTimeout(1).setConnectionTimeoutUnit(SECONDS);
+    Pool pool = createPool(options, poolOptions);
+    pool.getConnection().onComplete(ctx.asyncAssertSuccess(conn -> {
+      pool.query("SELECT 1").execute().onComplete(ctx.asyncAssertFailure(t -> {
+        conn.close();
+        async.complete();
+      }));
+    }));
+  }
 }

vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/PoolImpl.java

@@ -25,7 +25,10 @@
 import io.vertx.core.impl.future.PromiseInternal;
 import io.vertx.core.spi.metrics.PoolMetrics;
 import io.vertx.core.spi.metrics.VertxMetrics;
-import io.vertx.sqlclient.*;
+import io.vertx.sqlclient.Pool;
+import io.vertx.sqlclient.PoolOptions;
+import io.vertx.sqlclient.SqlConnection;
+import io.vertx.sqlclient.TransactionRollbackException;
 import io.vertx.sqlclient.impl.command.CommandBase;
 import io.vertx.sqlclient.impl.pool.SqlConnectionPool;
 import io.vertx.sqlclient.spi.Driver;
@@ -179,7 +182,7 @@ public Future<SqlConnection> getConnection() {
 
   @Override
   public <R> Future<R> schedule(ContextInternal context, CommandBase<R> cmd) {
-    return pool.execute(context, cmd);
+    return pool.execute(context, cmd, connectionTimeout);
   }
 
   private void acquire(ContextInternal context, long timeout, Handler<AsyncResult<SqlConnectionPool.PooledConnection>> completionHandler) {

vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/pool/SqlConnectionPool.java

@@ -207,28 +207,53 @@ private Object endUse(Object metric) {
     return NO_METRICS;
   }
 
-  public <R> Future<R> execute(ContextInternal context, CommandBase<R> cmd) {
+  private static final Exception POOL_QUERY_TIMEOUT_EXCEPTION = new VertxException("Timeout waiting for connection", true);
+
+  public <R> Future<R> execute(ContextInternal context, CommandBase<R> cmd, long timeout) {
+    Promise<R> res = context.promise();
+    long timerId;
+    if (timeout > 0) {
+      timerId = context.setTimer(timeout, t -> res.fail(POOL_QUERY_TIMEOUT_EXCEPTION));
+    } else {
+      timerId = -1;
+    }
     Promise<Lease<PooledConnection>> p = context.promise();
     Object queueMetric = enqueueMetric();
     pool.acquire(context, 0, p);
-    return p.future().compose(lease -> {
-      Object useMetric = dequeueAndBeginUse(queueMetric);
+    p.future().compose(lease -> {
       PooledConnection pooled = lease.get();
-      Connection conn = pooled.conn;
+      Object useMetric;
       Future<R> future;
-      if (afterAcquire != null) {
-        future = afterAcquire.apply(conn)
-          .compose(v -> pooled.schedule(context, cmd))
-          .eventually(v -> beforeRecycle.apply(conn));
+      if (timerId != -1 && !vertx.cancelTimer(timerId)) {
+        // We want to make sure the connection is released properly below
+        // But we don't want to record begin/end pool metrics
+        dequeueAndReject(queueMetric);
+        useMetric = NO_METRICS;
+        future = Future.failedFuture(POOL_QUERY_TIMEOUT_EXCEPTION);
       } else {
-        future = pooled.schedule(context, cmd);
+        useMetric = dequeueAndBeginUse(queueMetric);
+        if (afterAcquire != null) {
+          Connection conn = pooled.conn;
+          future = afterAcquire.apply(conn)
+            .compose(v -> pooled.schedule(context, cmd))
+            .eventually(() -> beforeRecycle.apply(conn));
+        } else {
+          future = pooled.schedule(context, cmd);
+        }
       }
       return future.andThen(ar -> {
         endUse(useMetric);
         pooled.refresh();
         lease.recycle();
       });
+    }).onComplete(ar -> {
+      if (ar.succeeded()) {
+        res.complete(ar.result());
+      } else if (!POOL_QUERY_TIMEOUT_EXCEPTION.equals(ar.cause())) {
+        res.fail(ar.cause());
+      }
     });
+    return res.future();
   }
 
   public void acquire(ContextInternal context, long timeout, Handler<AsyncResult<PooledConnection>> handler) {