Bound Query.close() drain and fix the close() hang race

Query.close() used to drain an in-flight submit with an unbounded,
uninterruptible wait: while (!done) doneCondition.awaitUninterruptibly().
Because the lease model makes close() mandatory (try-with-resources),
ordinary code that bounded its own await(timeout) and gave up still hit
this drain and could block the caller for the full remaining query
duration when the server was slow to honor the cancel. Worse, when a
QuestDB.close() raced an in-flight lease close() and the client I/O
thread failed to join within shutdownJoinMs, QwpQueryClient.close()
skipped closePool() -- the synthetic-terminal source -- so done was
never set and the lease close() hung forever.

Make the drain bounded and interruptible, and fail safe on timeout:

- QueryImpl.close() now waits at most closeQueryTimeoutMillis via
  awaitNanos and aborts on interrupt (re-raising the flag). A caller is
  never pinned to the full query duration.
- On timeout or interrupt the worker is discarded, not returned: its
  connection may still carry late RESULT_* frames for the abandoned
  query, which would corrupt the next borrower's stream. The pool grows
  a fresh worker on the next borrow. QueryClientPool.discard() mirrors
  the ingest side's discardBroken (closed/stale-generation guarded).
- QueryWorker.shutdown() now interrupts the dispatch thread. takeEvent()
  (QwpSpscQueue.take) is interrupt-aware and executeOnce() turns the
  InterruptedException into a terminal -> signalDone, so a caller parked
  in close() is released even when the I/O thread is wedged and
  closePool() never runs. This closes the hang-forever race and makes
  QuestDB.close() more prompt.

Add the query_close_timeout_ms pool knob (default 5000ms, symmetric with
the ingest close_flush_timeout_millis) via
QuestDBBuilder.queryCloseTimeoutMillis(long), wired through QuestDBImpl
to the pool and guarded by PoolConfigHonoredTest's drift check. Update
the Query.close() Javadoc to state the bounded/interruptible/discard
contract.

QueryCloseDrainTest covers the new behavior deterministically with a
no-op connect hook: close() returns within the budget and discards a
worker that does not drain, an interrupt aborts the drain promptly, and
an already-drained worker is returned for reuse. The dispatch-thread
interrupt's effect on a genuinely stuck execute() is verified by
construction here; its end-to-end reproduction lives in the parent
repo's concurrency tests.

Author: bluestreak01

core/src/main/java/io/questdb/client/Query.java

@@ -64,14 +64,24 @@ public interface Query extends Closeable {
 
     /**
      * Releases the leased pooled query client back to the pool. The caller
-     * MUST call this (typically via try-with-resources). If a submit is still
-     * in flight, {@code close()} cancels it and waits for the terminal event
-     * before returning the client. A real disconnect only happens at
-     * {@link QuestDB#close()}. Idempotent.
+     * MUST call this (typically via try-with-resources). A real disconnect only
+     * happens at {@link QuestDB#close()}. Idempotent.
+     * <p>
+     * If a submit is still in flight (the caller never awaited, or its
+     * {@code await(timeout)} expired), {@code close()} cancels it and waits for
+     * the terminal event so the client is idle before it returns to the pool.
+     * That wait is bounded by {@code query_close_timeout_ms} (default 5000ms,
+     * see {@link QuestDBBuilder#queryCloseTimeoutMillis(long)}) and is
+     * interruptible -- interrupting the calling thread aborts it. If the query
+     * does not drain within the budget, the client is discarded rather than
+     * returned (its connection may carry late frames for the abandoned query),
+     * and the pool grows a fresh one on the next borrow. {@code close()}
+     * therefore never blocks the caller unbounded, even when the server is slow
+     * to honor the cancel.
      * <p>
      * Must NOT be called from a result handler: handlers run on the worker
-     * thread, so {@code close()} would block forever waiting for a terminal
-     * event that only that thread can deliver. Calling it there throws
+     * thread, so {@code close()} would block waiting for a terminal event that
+     * only that thread can deliver. Calling it there throws
      * {@link IllegalStateException}. Use {@link #cancel()} (non-blocking) to
      * stop a query from inside a handler.
      */

core/src/main/java/io/questdb/client/QuestDBBuilder.java

@@ -58,6 +58,7 @@ public final class QuestDBBuilder {
     static final long DEFAULT_MAX_LIFETIME_MILLIS = 30 * 60_000L;
     static final int DEFAULT_POOL_MAX = 4;
     static final int DEFAULT_POOL_MIN = 1;
+    static final long DEFAULT_QUERY_CLOSE_TIMEOUT_MILLIS = 5_000;
 
     // Every valid pool value is >= 0, so -1 unambiguously marks "not set
     // explicitly". The public pool setters are the only writers of these
@@ -73,6 +74,7 @@ public final class QuestDBBuilder {
     private String config;
     private long idleTimeoutMillis = UNSET;
     private long maxLifetimeMillis = UNSET;
+    private long queryCloseTimeoutMillis = UNSET;
     private int queryPoolMax = UNSET;
     private int queryPoolMin = UNSET;
     private int senderPoolMax = UNSET;
@@ -94,6 +96,21 @@ public QuestDBBuilder acquireTimeoutMillis(long millis) {
         return this;
     }
 
+    /**
+     * Maximum time {@link Query#close()} waits for an in-flight query to drain
+     * (after issuing a cancel) before discarding the leased query client and
+     * letting the pool grow a fresh one. Bounds the close of a handle whose
+     * {@code submit()} is still running -- e.g. when the caller's own
+     * {@code await(timeout)} expired and they gave up. Defaults to 5000ms.
+     */
+    public QuestDBBuilder queryCloseTimeoutMillis(long millis) {
+        if (millis < 0) {
+            throw new IllegalArgumentException("queryCloseTimeoutMillis must be >= 0");
+        }
+        this.queryCloseTimeoutMillis = millis;
+        return this;
+    }
+
     /**
      * Sets the async connection-event listener applied to every pooled ingest
      * {@link Sender}. The listener observes connect / disconnect / failover
@@ -174,6 +191,7 @@ public QuestDB build() {
         resolvePoolInt(queryPoolMin, "query_pool_min", view, lazyConnect ? 0 : DEFAULT_POOL_MIN, this::queryPoolMin);
         resolvePoolInt(queryPoolMax, "query_pool_max", view, DEFAULT_POOL_MAX, this::queryPoolMax);
         resolvePoolLong(acquireTimeoutMillis, "acquire_timeout_ms", view, DEFAULT_ACQUIRE_TIMEOUT_MILLIS, this::acquireTimeoutMillis);
+        resolvePoolLong(queryCloseTimeoutMillis, "query_close_timeout_ms", view, DEFAULT_QUERY_CLOSE_TIMEOUT_MILLIS, this::queryCloseTimeoutMillis);
         resolvePoolLong(idleTimeoutMillis, "idle_timeout_ms", view, DEFAULT_IDLE_TIMEOUT_MILLIS, this::idleTimeoutMillis);
         resolvePoolLong(maxLifetimeMillis, "max_lifetime_ms", view, DEFAULT_MAX_LIFETIME_MILLIS, this::maxLifetimeMillis);
         resolvePoolLong(housekeeperIntervalMillis, "housekeeper_interval_ms", view, DEFAULT_HOUSEKEEPER_INTERVAL_MILLIS, this::housekeeperIntervalMillis);
@@ -189,6 +207,7 @@ public QuestDB build() {
                 idleTimeoutMillis,
                 maxLifetimeMillis,
                 housekeeperIntervalMillis,
+                queryCloseTimeoutMillis,
                 errorHandler,
                 connectionListener
         );
@@ -372,6 +391,7 @@ public java.util.Map<String, Object> poolConfigSnapshotForTest() {
         m.put("query_pool_min", queryPoolMin);
         m.put("query_pool_max", queryPoolMax);
         m.put("acquire_timeout_ms", acquireTimeoutMillis);
+        m.put("query_close_timeout_ms", queryCloseTimeoutMillis);
         m.put("idle_timeout_ms", idleTimeoutMillis);
         m.put("max_lifetime_ms", maxLifetimeMillis);
         m.put("housekeeper_interval_ms", housekeeperIntervalMillis);

core/src/main/java/io/questdb/client/impl/ConfigSchema.java

@@ -109,6 +109,7 @@ public final class ConfigSchema {
         intRange("query_pool_min", Side.POOL, OPEN, OPEN_MAX, false, false);
         intRange("query_pool_max", Side.POOL, OPEN, OPEN_MAX, false, false);
         longRange("acquire_timeout_ms", Side.POOL, OPEN, OPEN_MAX, false, false);
+        longRange("query_close_timeout_ms", Side.POOL, OPEN, OPEN_MAX, false, false);
         longRange("idle_timeout_ms", Side.POOL, OPEN, OPEN_MAX, false, false);
         longRange("max_lifetime_ms", Side.POOL, OPEN, OPEN_MAX, false, false);
         longRange("housekeeper_interval_ms", Side.POOL, OPEN, OPEN_MAX, false, false);

core/src/main/java/io/questdb/client/impl/QueryClientPool.java

@@ -50,6 +50,12 @@
  */
 public final class QueryClientPool implements AutoCloseable {
 
+    // Default upper bound, in milliseconds, on how long Query.close() waits for
+    // an in-flight query to drain (after issuing a cancel) before discarding the
+    // worker. Mirrors the ingest side's close_flush_timeout_millis default so a
+    // close() can never block the caller unbounded. Tunable per pool via
+    // closeQueryTimeoutMillis(long).
+    static final long DEFAULT_CLOSE_QUERY_TIMEOUT_MILLIS = 5_000;
     private final long acquireTimeoutMillis;
     private final ArrayList<QueryWorker> all;
     private final ArrayDeque<QueryWorker> available;
@@ -76,6 +82,10 @@ public final class QueryClientPool implements AutoCloseable {
     private final AtomicInteger nextSlotIndex = new AtomicInteger();
     private final Condition workerReleased;
     private volatile boolean closed;
+    // Upper bound on the Query.close() drain wait; see
+    // DEFAULT_CLOSE_QUERY_TIMEOUT_MILLIS. Volatile because QuestDBImpl sets it
+    // once at build time on a different thread than the borrowers that read it.
+    private volatile long closeQueryTimeoutMillis = DEFAULT_CLOSE_QUERY_TIMEOUT_MILLIS;
     private int inFlightCreations;
 
     public QueryClientPool(
@@ -330,6 +340,62 @@ void cancelIfCurrent(QueryWorker w, long gen) {
         }
     }
 
+    long closeQueryTimeoutMillis() {
+        return closeQueryTimeoutMillis;
+    }
+
+    void closeQueryTimeoutMillis(long millis) {
+        this.closeQueryTimeoutMillis = millis;
+    }
+
+    /**
+     * Evicts a worker whose lease {@link QueryImpl#close(long)} could not drain
+     * the in-flight query within {@link #closeQueryTimeoutMillis} (the cancel
+     * was not honored in time, or the caller was interrupted). The worker's
+     * connection is left in an unknown protocol state -- a late {@code RESULT_*}
+     * frame for the abandoned query could corrupt the next borrower's stream --
+     * so it must NOT return to the pool. Removes it from {@code all} (freeing
+     * capacity for a fresh worker) and tears it down outside the lock via
+     * {@link QueryWorker#shutdown()}, which interrupts the dispatch thread so a
+     * stuck {@code execute()} returns promptly.
+     * <p>
+     * Bails when the pool is already closed: {@link #close()} owns the teardown
+     * of every worker via its snapshot loop, so mutating {@code all} here would
+     * race that iteration on a non-thread-safe {@code ArrayList}. Also bails on a
+     * stale generation -- the worker was already released/discarded and possibly
+     * re-borrowed, so discarding it would evict a worker a different borrower now
+     * owns. Mirrors {@link SenderPool#discardBroken} on the ingest side.
+     */
+    void discard(QueryWorker w, long gen) {
+        lock.lock();
+        try {
+            if (closed) {
+                return;
+            }
+            if (w.generation() != gen) {
+                return;
+            }
+            // Invalidate the lease so a duplicate close()/release with the same
+            // generation is dropped and the in-flight handle can no longer drive
+            // this worker.
+            w.bumpGeneration();
+            all.remove(w);
+            // Capacity freed -- a waiter in acquire() may now create a fresh
+            // worker in this slot's place.
+            workerReleased.signal();
+        } finally {
+            lock.unlock();
+        }
+        // Tear down outside the lock so a slow join doesn't keep the pool
+        // latched. shutdown() is best-effort and idempotent.
+        try {
+            w.shutdown();
+        } catch (Throwable ignored) {
+            // Best-effort: a teardown Error (e.g. an -ea AssertionError) must
+            // not propagate out of Query.close().
+        }
+    }
+
     void reapIdle() {
         if (closed) {
             return;

core/src/main/java/io/questdb/client/impl/QueryImpl.java

@@ -155,19 +155,24 @@ void close(long gen) {
         if (gen != worker.generation()) {
             return;
         }
-        // If a submit is still in flight (the caller did not await), cancel it
-        // and wait for the terminal event so the leased worker is idle before
-        // it returns to the pool -- otherwise the next borrower would inherit a
-        // running execute().
+        // If a submit is still in flight (the caller did not await, or its
+        // await timed out), cancel it and wait for the terminal event so the
+        // leased worker is idle before it returns to the pool -- otherwise the
+        // next borrower would inherit a running execute().
+        //
+        // The wait is bounded (closeQueryTimeoutMillis) and interruptible, so a
+        // caller that bounded its own await() is never pinned to the full
+        // remaining query duration here. If the query does NOT drain in time (a
+        // server slow to honor the cancel, or the caller interrupting), the
+        // worker is still running execute() on a connection whose protocol state
+        // is now uncertain -- a late RESULT_* for the abandoned query could
+        // corrupt the next borrower's stream -- so it is discarded rather than
+        // returned. The pool grows a fresh worker on the next borrow.
         if (!done) {
             worker.cancelInFlight(gen);
-            doneLock.lock();
-            try {
-                while (!done) {
-                    doneCondition.awaitUninterruptibly();
-                }
-            } finally {
-                doneLock.unlock();
+            if (!awaitDone(worker.closeQueryTimeoutMillis())) {
+                worker.discardFromPool(gen);
+                return;
             }
         }
         worker.releaseToPool(gen);
@@ -226,6 +231,34 @@ private void applyBinds(QwpBindValues binds) {
         }
     }
 
+    /**
+     * Waits up to {@code timeoutMillis} for the in-flight query's terminal
+     * event. Returns {@code true} once {@code done} is set, {@code false} on
+     * timeout or interrupt. Unlike an uninterruptible drain, an interrupt aborts
+     * the wait and re-raises the thread's interrupt flag, so {@code close()}
+     * stays responsive to a caller that wants to give up.
+     */
+    private boolean awaitDone(long timeoutMillis) {
+        long remaining = TimeUnit.MILLISECONDS.toNanos(timeoutMillis);
+        doneLock.lock();
+        try {
+            while (!done) {
+                if (remaining <= 0) {
+                    return false;
+                }
+                try {
+                    remaining = doneCondition.awaitNanos(remaining);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    return false;
+                }
+            }
+            return true;
+        } finally {
+            doneLock.unlock();
+        }
+    }
+
     private void checkLive(long gen) {
         if (gen != worker.generation()) {
             throw new IllegalStateException("query handle is not borrowed (closed or never leased)");

core/src/main/java/io/questdb/client/impl/QueryWorker.java

@@ -170,6 +170,21 @@ Query lease() {
         return new QueryLease(query, generation);
     }
 
+    long closeQueryTimeoutMillis() {
+        return pool.closeQueryTimeoutMillis();
+    }
+
+    /**
+     * Discards this worker from the pool instead of returning it. Called by
+     * {@link QueryImpl#close(long)} when the in-flight query could not be
+     * drained within the close budget, leaving the connection in an unknown
+     * protocol state. The captured lease {@code gen} lets the pool reject a
+     * stale discard whose worker has already been re-borrowed.
+     */
+    void discardFromPool(long gen) {
+        pool.discard(this, gen);
+    }
+
     /**
      * Returns this worker to the pool. Called by {@link QueryImpl#close(long)}
      * when the borrowed lease is released; the captured lease {@code gen} lets
@@ -188,10 +203,19 @@ void shutdown() {
             signalLock.unlock();
         }
         try {
-            // If a query is in flight on this worker, ask the client to abort so
-            // execute() returns promptly and the thread can exit before join
-            // times out. cancel() is documented as thread-safe and is a no-op
-            // when idle.
+            // If a query is in flight on this worker, force execute() to return
+            // promptly so the dispatch thread exits before the join below times
+            // out. Two nudges, strongest first:
+            //   1. Interrupt the dispatch thread. takeEvent() (QwpSpscQueue.take)
+            //      is interrupt-aware, and executeOnce() turns the resulting
+            //      InterruptedException into a terminal event -> signalDone. This
+            //      releases a caller parked in Query.close() even when the I/O
+            //      thread is wedged and client.close()'s synthetic terminal
+            //      (closePool()) never runs -- the race that would otherwise
+            //      strand the caller forever.
+            //   2. Ask the client to cancel on the wire so the server stops work.
+            //      Best-effort and a no-op when idle.
+            thread.interrupt();
             try {
                 client.cancel();
             } catch (Throwable ignored) {

core/src/main/java/io/questdb/client/impl/QuestDBImpl.java

@@ -62,12 +62,13 @@ public QuestDBImpl(
             long idleTimeoutMillis,
             long maxLifetimeMillis,
             long housekeeperIntervalMillis,
+            long queryCloseTimeoutMillis,
             SenderErrorHandler errorHandler,
             SenderConnectionListener connectionListener
     ) {
         this(ingestConfig, queryConfig, senderMin, senderMax, queryMin, queryMax,
                 acquireTimeoutMillis, idleTimeoutMillis, maxLifetimeMillis,
-                housekeeperIntervalMillis, null, null, errorHandler, connectionListener);
+                housekeeperIntervalMillis, queryCloseTimeoutMillis, null, null, errorHandler, connectionListener);
     }
 
     // Test-only constructor exposing the senderFactory and connectHook seams:
@@ -93,7 +94,8 @@ public QuestDBImpl(
     ) {
         this(ingestConfig, queryConfig, senderMin, senderMax, queryMin, queryMax,
                 acquireTimeoutMillis, idleTimeoutMillis, maxLifetimeMillis,
-                housekeeperIntervalMillis, senderFactory, connectHook, null, null);
+                housekeeperIntervalMillis, QueryClientPool.DEFAULT_CLOSE_QUERY_TIMEOUT_MILLIS,
+                senderFactory, connectHook, null, null);
     }
 
     // Full constructor adding the ingest-side errorHandler/connectionListener,
@@ -111,6 +113,7 @@ public QuestDBImpl(
             long idleTimeoutMillis,
             long maxLifetimeMillis,
             long housekeeperIntervalMillis,
+            long queryCloseTimeoutMillis,
             IntFunction<Sender> senderFactory,
             Consumer<QwpQueryClient> connectHook,
             SenderErrorHandler errorHandler,
@@ -131,6 +134,7 @@ public QuestDBImpl(
             builtQueryPool = new QueryClientPool(
                     queryConfig, queryMin, queryMax, acquireTimeoutMillis,
                     idleTimeoutMillis, maxLifetimeMillis, connectHook);
+            builtQueryPool.closeQueryTimeoutMillis(queryCloseTimeoutMillis);
             builtHousekeeper = new PoolHousekeeper(builtSenderPool, builtQueryPool, housekeeperIntervalMillis);
             builtHousekeeper.start();
         } catch (Throwable e) {

core/src/test/java/io/questdb/client/test/impl/PoolConfigHonoredTest.java

@@ -54,6 +54,7 @@ public void testEveryPoolKeyIsHonored() {
         expected.put("query_pool_min", 0);
         expected.put("query_pool_max", 5);
         expected.put("acquire_timeout_ms", 1234L);
+        expected.put("query_close_timeout_ms", 2468L);
         expected.put("idle_timeout_ms", 4321L);
         expected.put("max_lifetime_ms", 98765L);
         expected.put("housekeeper_interval_ms", 222L);