Reject blocking lease calls from inside a result handler

Result handlers (onBatch/onEnd/onError) run inline on the worker's
dispatch thread: QueryWorker.runLoop -> QueryImpl.runOn ->
QwpQueryClient.execute consumes the I/O thread's events and invokes the
handler on the worker thread, then signalDone (which sets done) runs only
when that same thread loops back. So a handler that called the lease's
blocking close() or await() parked the worker thread waiting for a done
that only it could later produce -- a permanent, uninterruptible
self-deadlock that also leaked the worker and hung any other thread
awaiting the same handle. close()'s awaitUninterruptibly made it
unrecoverable, and close() is now mandatory via try-with-resources, so the
foot-gun was one stray call away.

Add a worker-thread reentrancy guard: QueryWorker.isCurrentThreadWorker()
compares Thread.currentThread() to the worker's dispatch thread, and
QueryImpl.close()/await()/await(timeout) throw IllegalStateException up
front when called on it. The exception unwinds at the user's call site
with a message pointing to cancel() (the non-blocking in-handler stop);
the worker is released normally by the app-thread close() afterwards, so
no deadlock and no leak.

Also fix the docs the bug exposed. Query.handler, Completion, and the
QueryWorker class javadoc all claimed the handler runs on "the I/O
thread"; it runs on the worker (dispatch) thread, which consumes the I/O
thread's event queue inline. The handler/close()/await() javadocs now say
so and forbid blocking calls from a handler.

Tests: QueryWorkerTest.testCloseAndAwaitFromWorkerThreadThrowInsteadOf-
Deadlocking drives close()/await()/await(timeout) on the worker thread and
asserts IllegalStateException (a method-level timeout fails the test if the
guard regresses into the old deadlock; verified it times out without the
guard). The QuestDB facade E2E suite adds an end-to-end check that a real
handler calling close()/await() throws and leaves the worker reusable.

Author: bluestreak01

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

@@ -36,15 +36,22 @@
  * {@link #await(long, TimeUnit)} returning {@code true}, or an explicit
  * {@link #cancel()} that races to terminal).
  * <p>
- * Signaling: the Completion is signaled from the I/O thread of the pooled
- * query client when the handler's terminal callback ({@code onEnd},
- * {@code onError}, or {@code onExecDone}) returns.
+ * Signaling: the Completion is signaled on the worker (dispatch) thread of the
+ * pooled query client when the handler's terminal callback ({@code onEnd},
+ * {@code onError}, or {@code onExecDone}) returns -- that callback runs inline
+ * on the worker thread, not on the I/O thread. Because of this, {@code await()}
+ * must never be called from inside a handler (it would self-deadlock on the
+ * worker thread); use {@link #cancel()} to stop a query from inside a handler.
  */
 public interface Completion {
 
     /**
      * Blocks until the query completes. Rethrows any server-reported failure
      * as a {@link QueryException}. Returns normally on success.
+     * <p>
+     * Must NOT be called from a result handler (it runs on the worker thread
+     * and would self-deadlock); calling it there throws
+     * {@link IllegalStateException}. Use {@link #cancel()} instead.
      *
      * @throws QueryException       if the server reported an error or
      *                              {@link #cancel()} won the race

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

@@ -68,14 +68,24 @@ public interface Query extends Closeable {
      * 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.
+     * <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
+     * {@link IllegalStateException}. Use {@link #cancel()} (non-blocking) to
+     * stop a query from inside a handler.
      */
     @Override
     void close();
 
     /**
-     * Sets the result-batch handler. The handler is invoked on the pooled
-     * query client's I/O thread; if it touches caller state, it is
-     * responsible for its own synchronization.
+     * Sets the result-batch handler. The handler is invoked on the worker
+     * (dispatch) thread that drives {@code execute()} -- it consumes the pooled
+     * query client's I/O-thread event queue inline, it does NOT run on the I/O
+     * thread. If it touches caller state, it is responsible for its own
+     * synchronization. A handler must not call the blocking {@link #close()} or
+     * {@link Completion#await()} (they would self-deadlock on the worker
+     * thread); use {@link #cancel()} to stop from inside a handler.
      */
     Query handler(QwpColumnBatchHandler handler);
 

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

@@ -95,6 +95,7 @@ void abandon(long gen) {
     }
 
     void await(long gen) throws InterruptedException {
+        rejectHandlerReentry("await");
         checkLive(gen);
         doneLock.lock();
         try {
@@ -108,6 +109,7 @@ void await(long gen) throws InterruptedException {
     }
 
     boolean await(long gen, long timeout, TimeUnit unit) throws InterruptedException {
+        rejectHandlerReentry("await");
         checkLive(gen);
         long remaining = unit.toNanos(timeout);
         doneLock.lock();
@@ -143,6 +145,7 @@ void cancel(long gen) {
     }
 
     void close(long gen) {
+        rejectHandlerReentry("close");
         // A stale generation means this lease was already released and the
         // worker may now be owned by another borrower. Dropping the call is
         // what keeps close() idempotent without releasing someone else's
@@ -229,6 +232,22 @@ private void checkLive(long gen) {
         }
     }
 
+    private void rejectHandlerReentry(String op) {
+        // Result handlers (onBatch/onEnd/onError) run inline on the worker's
+        // dispatch thread. A blocking lease op called from there would wait for
+        // a terminal event that only this same thread can deliver -- a
+        // permanent, uninterruptible self-deadlock plus a leaked worker. Fail
+        // loudly at the call site instead. cancel() is the non-blocking stop.
+        if (worker.isCurrentThreadWorker()) {
+            throw new IllegalStateException(
+                    op + "() must not be called from a result handler. Handlers "
+                            + "(onBatch/onEnd/onError) run on the worker thread, so " + op
+                            + "() would block forever waiting for a terminal event that only "
+                            + "this same thread can deliver. To stop a query from inside a "
+                            + "handler, call cancel() (non-blocking).");
+        }
+    }
+
     private void signalDone(byte status, String message, Throwable unexpected) {
         doneLock.lock();
         try {

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

@@ -40,7 +40,11 @@
  * The pooled query client's own I/O thread continues to drive the wire; the
  * worker thread exists only to keep {@code execute()} off the application's
  * submitting thread. Handler callbacks ({@code onBatch}, {@code onEnd},
- * {@code onError}) still run on the client's I/O thread.
+ * {@code onError}) run on this worker's own dispatch thread, which consumes the
+ * I/O thread's event queue inline -- not on the I/O thread itself. A handler
+ * must therefore never call the lease's blocking {@code close()}/{@code await()}
+ * (it would self-deadlock waiting for a terminal event only this thread can
+ * deliver); use the non-blocking {@code cancel()} to stop from inside a handler.
  */
 public final class QueryWorker {
 
@@ -101,6 +105,17 @@ long idleSinceMillis() {
         return idleSinceMillis;
     }
 
+    /**
+     * True when the calling thread is this worker's own dispatch thread -- i.e.
+     * a reentrant call from inside a result handler, which runs inline on this
+     * thread. Blocking lease operations ({@link QueryImpl#close}/
+     * {@link QueryImpl#await}) use this to fail loudly instead of
+     * self-deadlocking.
+     */
+    boolean isCurrentThreadWorker() {
+        return Thread.currentThread() == thread;
+    }
+
     void markIdleAt(long nowMillis) {
         idleSinceMillis = nowMillis;
     }

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

@@ -26,12 +26,15 @@
 
 import io.questdb.client.Completion;
 import io.questdb.client.cutlass.qwp.client.QwpQueryClient;
+import io.questdb.client.impl.QueryClientPool;
 import io.questdb.client.impl.QueryWorker;
 import org.junit.Assert;
 import org.junit.Test;
 
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
@@ -155,4 +158,88 @@ public void testShutdownRacingDispatchMustNotStrandCaller() throws Exception {
         Assert.assertNotNull("signalUnexpected must record the closed-handle error",
                 unexpectedF.get(queryImpl));
     }
+
+    /**
+     * Result handlers (onBatch/onEnd/onError) run inline on the worker's
+     * dispatch thread. The blocking lease ops -- {@code close()} and the two
+     * {@code await()} variants -- would there wait on a terminal event that
+     * only this same thread can deliver, a permanent self-deadlock. The
+     * reentrancy guard must turn that into an immediate IllegalStateException.
+     * <p>
+     * The guard compares {@code Thread.currentThread()} to the worker's
+     * dispatch thread, so this test points that field at the test thread (the
+     * worker is never started) to stand in for a reentrant in-handler call.
+     * Without the guard, {@code close()}/{@code await()} would park forever and
+     * the method-level timeout would fail the test.
+     */
+    @Test(timeout = 30_000)
+    public void testCloseAndAwaitFromWorkerThreadThrowInsteadOfDeadlocking() throws Exception {
+        Class<?> queryImplClass = Class.forName("io.questdb.client.impl.QueryImpl");
+        Field queryF = QueryWorker.class.getDeclaredField("query");
+        queryF.setAccessible(true);
+        Field threadF = QueryWorker.class.getDeclaredField("thread");
+        threadF.setAccessible(true);
+        Field doneF = queryImplClass.getDeclaredField("done");
+        doneF.setAccessible(true);
+        Method bump = QueryWorker.class.getDeclaredMethod("bumpGeneration");
+        bump.setAccessible(true);
+        Method isWorker = QueryWorker.class.getDeclaredMethod("isCurrentThreadWorker");
+        isWorker.setAccessible(true);
+        Method close = queryImplClass.getDeclaredMethod("close", long.class);
+        close.setAccessible(true);
+        Method awaitNoTimeout = queryImplClass.getDeclaredMethod("await", long.class);
+        awaitNoTimeout.setAccessible(true);
+        Method awaitTimed = queryImplClass.getDeclaredMethod("await", long.class, long.class, TimeUnit.class);
+        awaitTimed.setAccessible(true);
+
+        QueryClientPool pool = new QueryClientPool(
+                "ws::addr=localhost:9000;",
+                /*minSize*/ 0, /*maxSize*/ 2,
+                /*acquireTimeoutMillis*/ 1_000L,
+                /*idleTimeoutMillis*/ Long.MAX_VALUE,
+                /*maxLifetimeMillis*/ Long.MAX_VALUE);
+        QwpQueryClient client = QwpQueryClient.newPlainText("localhost", 9000);
+        try {
+            QueryWorker w = new QueryWorker(client, pool, 0);
+            bump.invoke(w); // generation -> 1: a live lease
+            Object impl = queryF.get(w);
+            doneF.setBoolean(impl, false); // a submit is in flight, as during a handler
+
+            // Off the worker thread the guard must NOT fire.
+            Assert.assertFalse("guard must not fire on a normal caller thread",
+                    (Boolean) isWorker.invoke(w));
+
+            // Stand in for a reentrant call from inside a result handler: the
+            // guard compares Thread.currentThread() to the worker's dispatch
+            // thread, so point that field at this thread.
+            threadF.set(w, Thread.currentThread());
+            Assert.assertTrue((Boolean) isWorker.invoke(w));
+
+            assertThrowsHandlerReentry("close", () -> close.invoke(impl, 1L));
+            assertThrowsHandlerReentry("await", () -> awaitNoTimeout.invoke(impl, 1L));
+            assertThrowsHandlerReentry("await(timeout)",
+                    () -> awaitTimed.invoke(impl, 1L, 5L, TimeUnit.SECONDS));
+        } finally {
+            client.close();
+            pool.close();
+        }
+    }
+
+    private static void assertThrowsHandlerReentry(String op, ReflectiveCall call) throws Exception {
+        try {
+            call.run();
+            Assert.fail(op + "() from the worker thread must throw, not block/deadlock");
+        } catch (InvocationTargetException e) {
+            Throwable cause = e.getCause();
+            Assert.assertTrue(op + "(): expected IllegalStateException, was " + cause,
+                    cause instanceof IllegalStateException);
+            Assert.assertTrue(op + "(): message must point at cancel(), was: " + cause.getMessage(),
+                    cause.getMessage().contains("cancel()"));
+        }
+    }
+
+    @FunctionalInterface
+    private interface ReflectiveCall {
+        void run() throws Exception;
+    }
 }