Pool queries should honor connection timeout (#1586)

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 <[email protected]>

Author: tsegismont

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

@@ -42,6 +42,7 @@
 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;
 
@@ -615,4 +616,17 @@ public void testConnectionClosedInHook(TestContext ctx) {
       }));
     }));
   }
+
+  @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/pool/PoolImpl.java

@@ -27,13 +27,13 @@
 import io.vertx.core.spi.metrics.VertxMetrics;
 import io.vertx.sqlclient.*;
 import io.vertx.sqlclient.impl.TransactionPropagationLocal;
-import io.vertx.sqlclient.spi.connection.Connection;
 import io.vertx.sqlclient.internal.SqlClientBase;
 import io.vertx.sqlclient.internal.SqlConnectionInternal;
+import io.vertx.sqlclient.spi.Driver;
+import io.vertx.sqlclient.spi.connection.Connection;
 import io.vertx.sqlclient.spi.connection.ConnectionContext;
-import io.vertx.sqlclient.spi.protocol.CommandBase;
 import io.vertx.sqlclient.spi.connection.ConnectionFactory;
-import io.vertx.sqlclient.spi.Driver;
+import io.vertx.sqlclient.spi.protocol.CommandBase;
 
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -179,7 +179,7 @@ public Future<SqlConnection> getConnection() {
 
   @Override
   public <R> void schedule(CommandBase<R> cmd, Completable<R> handler) {
-    pool.execute(cmd, handler);
+    pool.execute(cmd, handler, connectionTimeout);
   }
 
   private void acquire(ContextInternal context, long timeout, Completable<SqlConnectionPool.PooledConnection> completionHandler) {

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

@@ -205,34 +205,54 @@ private void endMetric(Object metric) {
     }
   }
 
+  private static final Exception POOL_QUERY_TIMEOUT_EXCEPTION = new VertxException("Timeout waiting for connection", true);
+
   // TODO : try optimize without promise
-  public <R> void execute(CommandBase<R> cmd, Completable<R> handler) {
+  public <R> void execute(CommandBase<R> cmd, Completable<R> handler, long timeout) {
     ContextInternal context = vertx.getOrCreateContext();
     Promise<Lease<PooledConnection>> p = context.promise();
+    long timerId;
+    if (timeout > 0) {
+      timerId = vertx.setTimer(timeout, t -> handler.fail(POOL_QUERY_TIMEOUT_EXCEPTION));
+    } else {
+      timerId = -1;
+    }
     Object metric = enqueueMetric();
     pool.acquire(context, 0, p);
     p.future().compose(lease -> {
       dequeueMetric(metric);
       PooledConnection pooled = lease.get();
-      pooled.timerMetric = beginMetric();
-      Connection conn = pooled.conn;
-
       Future<R> future;
-      if (afterAcquire != null) {
-        future = afterAcquire.apply(conn)
-          .compose(v -> Future.<R>future(d -> pooled.schedule(cmd, d)))
-          .eventually(() -> 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
+        pooled.timerMetric = NO_METRICS;
+        future = Future.failedFuture(POOL_QUERY_TIMEOUT_EXCEPTION);
       } else {
-        PromiseInternal<R> pp = context.promise();
-        pooled.schedule(cmd, pp);
-        future = pp;
+        pooled.timerMetric = beginMetric();
+        if (afterAcquire != null) {
+          Connection conn = pooled.conn;
+          future = afterAcquire.apply(conn)
+            .compose(v -> Future.<R>future(d -> pooled.schedule(cmd, d)))
+            .eventually(() -> beforeRecycle.apply(conn));
+        } else {
+          PromiseInternal<R> pp = context.promise();
+          pooled.schedule(cmd, pp);
+          future = pp;
+        }
       }
       return future.andThen(ar -> {
         endMetric(pooled.timerMetric);
         pooled.refresh();
         lease.recycle();
       });
-    }).onComplete(handler);
+    }).onComplete(ar -> {
+      if (ar.succeeded()) {
+        handler.succeed(ar.result());
+      } else if (!POOL_QUERY_TIMEOUT_EXCEPTION.equals(ar.cause())) {
+        handler.fail(ar.cause());
+      }
+    });
   }
 
   public void acquire(ContextInternal context, long timeout, Completable<PooledConnection> handler) {