Make watchdog cancel re-check lease generation under the pool lock

QueryImpl.cancel(gen) validated the lease generation with an unlocked
volatile read and then issued the wire cancel with no lock, so the two
steps were not atomic. cancel() is a cross-thread watchdog/timeout API,
so this is a TOCTOU: a watchdog can pass the generation check, get
preempted while the lease is released (close -> release bumps the
generation) and the worker is re-borrowed and re-submitted by another
caller, then resume and call cancelInFlight() -- aborting that caller's
in-flight query with a spurious STATUS_CANCELLED. The client never
resets its cancel state on release/re-borrow, so the stale cancel lands
either via the wire requestCancel(currentRequestId) or via the surviving
pendingCancel latch the next executeOnce consumes.

Route the cancel through QueryClientPool.cancelIfCurrent(worker, gen),
which re-checks worker.generation() == gen and issues the wire cancel
together under the pool lock -- the same lock acquire()/release() bump
the generation under. Once cancelIfCurrent holds the lock the generation
cannot change, so a cancel whose lease has gone stale is dropped instead
of hitting the new borrower. The cancel itself is non-blocking (a
volatile flag plus an AtomicLong set), so the lock is held only briefly.
cancel() keeps a cheap unlocked fast-path that drops an obviously-stale
or already-done cancel without taking the lock; the authoritative check
stays under the lock. close() routes its abort-on-unawaited-close through
the same path, closing the identical window on its cancelInFlight() call.

QueryLeaseGenerationTest.testStaleCancelDoesNotReachClient only covered
the already-stale case (it pre-advanced the generation, leaving no
check->act window), so it could not catch this. Add
testConcurrentCancelDoesNotReachClientAfterReborrow: it holds the pool
lock so a concurrent cancel parks inside the re-check, advances the
generation (release + re-borrow) under the lock, then releases it and
asserts the parked cancel observed the new generation and never reached
the client. The new test fails against the unlocked check-then-cancel and
passes with the locked path.

Author: bluestreak01

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

@@ -305,6 +305,31 @@ public void close() {
         }
     }
 
+    /**
+     * Cancels the in-flight query on {@code w} only while its lease generation
+     * still equals {@code gen}, holding the pool lock across both the check and
+     * the wire cancel. acquire() and release() bump the generation under this
+     * same lock, so once this method holds it the generation cannot change: a
+     * cancel whose lease has already gone stale (the worker was released and
+     * re-borrowed) is dropped instead of aborting the new borrower's query. The
+     * cancel itself is non-blocking -- a volatile flag plus an AtomicLong set --
+     * so the lock is held only briefly.
+     */
+    void cancelIfCurrent(QueryWorker w, long gen) {
+        lock.lock();
+        try {
+            if (closed) {
+                return;
+            }
+            if (w.generation() != gen) {
+                return;
+            }
+            w.cancelInFlight();
+        } finally {
+            lock.unlock();
+        }
+    }
+
     void reapIdle() {
         if (closed) {
             return;

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

@@ -126,14 +126,20 @@ boolean await(long gen, long timeout, TimeUnit unit) throws InterruptedException
     }
 
     void cancel(long gen) {
-        // No-op on a stale lease: the worker may now be running a different
-        // borrower's query, and cancelling here would abort it.
-        if (gen != worker.generation()) {
+        // Fast-path drop of an obviously-stale or already-finished cancel,
+        // without taking the pool lock. This is only a hint -- the
+        // authoritative re-check runs under the pool lock inside
+        // worker.cancelInFlight(gen).
+        if (gen != worker.generation() || done) {
             return;
         }
-        if (!done) {
-            worker.cancelInFlight();
-        }
+        // Re-check the lease generation and issue the wire cancel atomically
+        // under the pool lock (the same lock acquire()/release() bump the
+        // generation under). An unlocked check followed by an unlocked cancel
+        // is a TOCTOU: a cross-thread watchdog can pass the check, get
+        // preempted while this lease is released and the worker re-borrowed by
+        // another caller, then resume and abort that caller's in-flight query.
+        worker.cancelInFlight(gen);
     }
 
     void close(long gen) {
@@ -151,7 +157,7 @@ void close(long gen) {
         // it returns to the pool -- otherwise the next borrower would inherit a
         // running execute().
         if (!done) {
-            worker.cancelInFlight();
+            worker.cancelInFlight(gen);
             doneLock.lock();
             try {
                 while (!done) {

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

@@ -106,8 +106,11 @@ void markIdleAt(long nowMillis) {
     }
 
     /**
-     * Cancels the in-flight query on this worker's client. Safe to call from
-     * any thread; harmless if the worker is idle.
+     * Issues an unconditional wire cancel against whatever query this worker's
+     * client is currently running. Callers must already own the worker for the
+     * current lease -- in practice this runs under the pool lock via
+     * {@link QueryClientPool#cancelIfCurrent}, which validates the lease
+     * generation first. Lease code must use {@link #cancelInFlight(long)}.
      */
     void cancelInFlight() {
         try {
@@ -117,6 +120,18 @@ void cancelInFlight() {
         }
     }
 
+    /**
+     * Cancels the in-flight query only if this worker's lease generation still
+     * equals {@code gen}. Delegates to the pool so the generation re-check and
+     * the wire cancel happen together under the pool lock that
+     * {@link QueryClientPool#acquire} and {@link QueryClientPool#release} bump
+     * the generation under. That atomicity stops a stale cross-thread cancel
+     * from aborting a later borrower's query on the same worker.
+     */
+    void cancelInFlight(long gen) {
+        pool.cancelIfCurrent(this, gen);
+    }
+
     /**
      * Returns the {@link QwpQueryClient} this worker drives. Exposed for
      * introspection and tests; callers must not invoke {@code execute()} on

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

@@ -33,6 +33,10 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.ArrayDeque;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Regression tests for M1: a stale {@code Query} lease (held after close, or a
@@ -68,32 +72,134 @@ public void testStaleCancelDoesNotReachClient() throws Exception {
         Method cancel = queryImplClass.getDeclaredMethod("cancel", long.class);
         cancel.setAccessible(true);
 
-        // Live lease: generation 1 (one acquire), query in flight -> cancel(1)
-        // must reach the client.
-        try (QwpQueryClient live = QwpQueryClient.newPlainText("localhost", 9000)) {
-            QueryWorker w = new QueryWorker(live, null, 0);
-            bump.invoke(w); // generation -> 1 (acquire stamp)
-            Object impl = queryF.get(w);
-            doneF.setBoolean(impl, false); // pretend a submit is in flight
-            cancel.invoke(impl, 1L);
-            Assert.assertTrue("cancel() on the live lease must reach the client",
-                    live.isPendingCancelForTest());
+        // cancel(gen) validates the generation under the pool lock, so the
+        // worker needs a real pool to lock on (the worker thread is never
+        // started, so no network or I/O thread is involved).
+        QueryClientPool pool = new QueryClientPool(
+                "ws::addr=localhost:9000;",
+                /*minSize*/ 0, /*maxSize*/ 2,
+                /*acquireTimeoutMillis*/ 1_000L,
+                /*idleTimeoutMillis*/ Long.MAX_VALUE,
+                /*maxLifetimeMillis*/ Long.MAX_VALUE);
+        try {
+            // Live lease: generation 1 (one acquire), query in flight -> cancel(1)
+            // must reach the client.
+            try (QwpQueryClient live = QwpQueryClient.newPlainText("localhost", 9000)) {
+                QueryWorker w = new QueryWorker(live, pool, 0);
+                bump.invoke(w); // generation -> 1 (acquire stamp)
+                Object impl = queryF.get(w);
+                doneF.setBoolean(impl, false); // pretend a submit is in flight
+                cancel.invoke(impl, 1L);
+                Assert.assertTrue("cancel() on the live lease must reach the client",
+                        live.isPendingCancelForTest());
+            }
+
+            // Stale lease: the worker was borrowed (gen 1), released and re-borrowed
+            // (gen now 3). A cancel from the old lease (gen 1) must be dropped, even
+            // though the current query is in flight.
+            try (QwpQueryClient reused = QwpQueryClient.newPlainText("localhost", 9000)) {
+                QueryWorker w = new QueryWorker(reused, pool, 0);
+                bump.invoke(w); // -> 1 (first acquire)
+                bump.invoke(w); // -> 2 (release)
+                bump.invoke(w); // -> 3 (second acquire by a new borrower)
+                Object impl = queryF.get(w);
+                doneF.setBoolean(impl, false); // the new borrower's query is in flight
+                cancel.invoke(impl, 1L); // stale lease cancels
+                Assert.assertFalse("a stale lease's cancel() must NOT reach the client and "
+                                + "cancel a different borrower's in-flight query",
+                        reused.isPendingCancelForTest());
+            }
+        } finally {
+            pool.close();
         }
+    }
 
-        // Stale lease: the worker was borrowed (gen 1), released and re-borrowed
-        // (gen now 3). A cancel from the old lease (gen 1) must be dropped, even
-        // though the current query is in flight.
-        try (QwpQueryClient reused = QwpQueryClient.newPlainText("localhost", 9000)) {
-            QueryWorker w = new QueryWorker(reused, null, 0);
-            bump.invoke(w); // -> 1 (first acquire)
-            bump.invoke(w); // -> 2 (release)
-            bump.invoke(w); // -> 3 (second acquire by a new borrower)
-            Object impl = queryF.get(w);
-            doneF.setBoolean(impl, false); // the new borrower's query is in flight
-            cancel.invoke(impl, 1L); // stale lease cancels
-            Assert.assertFalse("a stale lease's cancel() must NOT reach the client and "
-                            + "cancel a different borrower's in-flight query",
-                    reused.isPendingCancelForTest());
+    /**
+     * The TOCTOU the locked cancel closes: a cross-thread watchdog calls
+     * {@code cancel(gen)} while its lease is live, but the lease goes stale (the
+     * worker is released and re-borrowed) before the wire cancel fires. The
+     * cancel must re-validate the generation atomically with the cancel, under
+     * the pool lock, or it would abort the new borrower's query.
+     * <p>
+     * Driven deterministically: the test thread holds the pool lock, so the
+     * watchdog's cancel parks inside the pool's generation re-check. We then
+     * advance the generation (release + re-borrow) under the lock and release
+     * it. The parked cancel must observe the new generation and drop. An
+     * unlocked check-then-cancel would not park, would pass its check at the
+     * still-live generation, and would fire the wire cancel.
+     */
+    @Test
+    public void testConcurrentCancelDoesNotReachClientAfterReborrow() throws Exception {
+        Method bump = QueryWorker.class.getDeclaredMethod("bumpGeneration");
+        bump.setAccessible(true);
+        Field queryF = QueryWorker.class.getDeclaredField("query");
+        queryF.setAccessible(true);
+        Class<?> queryImplClass = Class.forName("io.questdb.client.impl.QueryImpl");
+        Field doneF = queryImplClass.getDeclaredField("done");
+        doneF.setAccessible(true);
+        Method cancel = queryImplClass.getDeclaredMethod("cancel", long.class);
+        cancel.setAccessible(true);
+        Field poolLockF = QueryClientPool.class.getDeclaredField("lock");
+        poolLockF.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 {
+            final QueryWorker w = new QueryWorker(client, pool, 0);
+            bump.invoke(w); // generation -> 1; the watchdog's lease captured 1
+            final Object impl = queryF.get(w);
+            doneF.setBoolean(impl, false); // a query is in flight
+
+            ReentrantLock poolLock = (ReentrantLock) poolLockF.get(pool);
+            final CountDownLatch atCancel = new CountDownLatch(1);
+            final CountDownLatch cancelReturned = new CountDownLatch(1);
+            final AtomicReference<Throwable> err = new AtomicReference<>();
+
+            // Hold the pool lock so the watchdog's cancel cannot finish its
+            // generation re-check + wire cancel until we let go.
+            poolLock.lock();
+            Thread watchdog = new Thread(() -> {
+                atCancel.countDown();
+                try {
+                    cancel.invoke(impl, 1L); // lease generation captured at borrow = 1
+                } catch (Throwable t) {
+                    err.set(t);
+                } finally {
+                    cancelReturned.countDown();
+                }
+            }, "watchdog-cancel");
+            watchdog.start();
+            Assert.assertTrue("watchdog must start", atCancel.await(5, TimeUnit.SECONDS));
+
+            // With the locked cancel, cancel() parks on the pool lock and cannot
+            // return while we hold it. An unlocked check-then-cancel would have
+            // already fired the wire cancel and returned.
+            Assert.assertFalse("cancel() must re-check the generation under the pool "
+                            + "lock, so it cannot complete while the lock is held",
+                    cancelReturned.await(200, TimeUnit.MILLISECONDS));
+
+            // The lease goes stale underneath the parked cancel: released (-> 2)
+            // and re-borrowed by a new owner (-> 3).
+            bump.invoke(w);
+            bump.invoke(w);
+            poolLock.unlock();
+
+            Assert.assertTrue("cancel() must return once the pool lock is free",
+                    cancelReturned.await(5, TimeUnit.SECONDS));
+            if (err.get() != null) {
+                throw new AssertionError("cancel() threw", err.get());
+            }
+            Assert.assertFalse("a cancel whose lease went stale while parked on the pool "
+                            + "lock must NOT reach the client and abort the new borrower's query",
+                    client.isPendingCancelForTest());
+        } finally {
+            client.close();
+            pool.close();
         }
     }