fix(qwp): harden cleanup and recovery lifecycle

Author: bluestreak01

core/src/main/c/share/net.c

@@ -128,6 +128,11 @@ JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_send
     return com_questdb_network_Net_EOTHERDISCONNECT;
 }
 
+JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_shutdown
+        (JNIEnv *e, jclass cl, jint fd) {
+    return shutdown((int) fd, SHUT_RDWR);
+}
+
 JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_recv
         (JNIEnv *e, jclass cl, jint fd, jlong ptr, jint len) {
     ssize_t n;

core/src/main/c/share/net.h

@@ -230,6 +230,15 @@ JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_configureNonBlocking
     return res;
 }
 
+JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_shutdown
+        (JNIEnv *e, jclass cl, jint fd) {
+    const int result = shutdown((SOCKET) fd, SD_BOTH);
+    if (result == SOCKET_ERROR) {
+        SaveLastError();
+    }
+    return result;
+}
+
 JNIEXPORT jint JNICALL Java_io_questdb_client_network_Net_recv
         (JNIEnv *e, jclass cl, jint fd, jlong addr, jint len) {
     const int n = recv((SOCKET) fd, (char *) addr, len, 0);

core/src/main/c/windows/net.c

@@ -248,6 +248,15 @@ public void close() {
         }
     }
 
+    /**
+     * Shuts down socket traffic without releasing the descriptor or freeing
+     * buffers that a concurrent I/O worker may still access. The owner must
+     * call {@link #close()} after joining the worker to complete cleanup.
+     */
+    public void closeTraffic() {
+        socket.closeTraffic();
+    }
+
     /**
      * Connects to a WebSocket server.
      *

core/src/main/java/io/questdb/client/cutlass/http/client/WebSocketClient.java

@@ -189,6 +189,9 @@ public class QwpWebSocketSender implements Sender {
     // Null in production; set reflectively by tests.
     @TestOnly
     private volatile java.util.function.Supplier<WebSocketClient> clientFactoryOverride;
+    // Test-only lifecycle witness. drainOnClose() invokes and clears it after
+    // confirming a real unacknowledged close target, immediately before waiting.
+    private volatile Runnable closeDrainWaitingHook;
     // close() drain timeout in millis. Default applied at construction.
     // 0 or -1 means "fast close" (skip the drain); otherwise close blocks
     // up to this many millis for ackedFsn to catch up to publishedFsn.
@@ -290,6 +293,9 @@ public class QwpWebSocketSender implements Sender {
     // re-probe retired slots to recover their capacity. volatile: written on
     // the closing thread, read by pool threads.
     private volatile boolean slotLockReleased;
+    // Optional owning-pool notification, relayed after the engine confirms the
+    // flock release and slotLockReleased is visible to pool re-probes.
+    private volatile Runnable slotLockReleaseListener;
     // Engine whose close() could not complete during sender close() — its
     // cleanup is pending on a worker/I/O-thread exit path. isSlotLockReleased()
     // re-probes it so a late flock release becomes visible to the owning pool.
@@ -1368,6 +1374,17 @@ public boolean isSlotLockReleased() {
         return false;
     }
 
+    /**
+     * Registers a callback for confirmed SF slot-lock release. Pools use this
+     * to wake borrowers without waiting for a timeout or housekeeper tick.
+     */
+    public void setSlotLockReleaseListener(Runnable listener) {
+        slotLockReleaseListener = listener;
+        if (listener != null && slotLockReleased) {
+            listener.run();
+        }
+    }
+
     @Override
     public Sender decimalColumn(CharSequence name, Decimal64 value) {
         checkNotClosed();
@@ -2207,6 +2224,16 @@ public void setClientFactoryOverride(java.util.function.Supplier<WebSocketClient
         this.clientFactoryOverride = factory;
     }
 
+    /**
+     * Installs a one-shot test witness that close-time drain invokes after it
+     * observes a real unacknowledged target and before it starts waiting.
+     * Production code never sets it.
+     */
+    @TestOnly
+    public void setCloseDrainWaitingHook(Runnable hook) {
+        this.closeDrainWaitingHook = hook;
+    }
+
     /**
      * Installs a one-shot test witness that {@link #close()} invokes after it
      * publishes the closed-state transition and before it starts drain or
@@ -2302,6 +2329,9 @@ public void setCursorEngine(CursorSendEngine engine, boolean takeOwnership) {
         }
         this.cursorEngine = engine;
         this.ownsCursorEngine = takeOwnership && engine != null;
+        if (engine != null) {
+            engine.setSlotLockReleaseListener(this::onSlotLockReleased);
+        }
     }
 
     /**
@@ -3272,6 +3302,16 @@ private void drainOnClose(boolean errorOwnedByCustomHandler) {
         if (cursorEngine.ackedFsn() >= target) {
             return;
         }
+        Runnable hook = closeDrainWaitingHook;
+        closeDrainWaitingHook = null;
+        if (hook != null) {
+            try {
+                hook.run();
+            } catch (Throwable t) {
+                // A test witness must never prevent production resource cleanup.
+                LOG.error("Error in close-drain-waiting test hook: {}", String.valueOf(t));
+            }
+        }
         long deadlineNanos = System.nanoTime() + closeFlushTimeoutMillis * 1_000_000L;
         while (cursorEngine.ackedFsn() < target) {
             // Stop on a latched terminal (acks will never reach target);
@@ -3632,6 +3672,14 @@ private void flushPendingRowsSplit(ObjList<CharSequence> keys, boolean deferComm
         resetTableBuffersAfterFlush(keys);
     }
 
+    private void onSlotLockReleased() {
+        slotLockReleased = true;
+        Runnable listener = slotLockReleaseListener;
+        if (listener != null) {
+            listener.run();
+        }
+    }
+
     private void resetTableBuffersAfterFlush(ObjList<CharSequence> keys) {
         for (int i = 0, n = keys.size(); i < n; i++) {
             CharSequence tableName = keys.getQuick(i);

core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpWebSocketSender.java

@@ -30,6 +30,8 @@
 import io.questdb.client.std.QuietCloseable;
 import org.jetbrains.annotations.TestOnly;
 
+import java.util.ArrayDeque;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.LockSupport;
 
@@ -65,7 +67,15 @@ public final class CursorSendEngine implements QuietCloseable {
     public static final long DEFAULT_APPEND_DEADLINE_NANOS = 30_000_000_000L;
     private static final org.slf4j.Logger LOG =
             org.slf4j.LoggerFactory.getLogger(CursorSendEngine.class);
+    private static final ThreadFactory DEFAULT_FLOCK_RELEASE_RETRY_THREAD_FACTORY =
+            runnable -> new Thread(runnable, "qdb-sf-flock-release-retry");
+    private static final Object FLOCK_RELEASE_RETRY_LOCK = new Object();
+    private static final ArrayDeque<CursorSendEngine> FLOCK_RELEASE_RETRY_QUEUE = new ArrayDeque<>();
+    private static volatile Runnable afterFlockReleaseRetryFailureHook;
     private static volatile Runnable beforeDeferredCloseCreationHook;
+    private static Thread flockReleaseRetryThread;
+    private static volatile ThreadFactory flockReleaseRetryThreadFactory =
+            DEFAULT_FLOCK_RELEASE_RETRY_THREAD_FACTORY;
     private final long appendDeadlineNanos;
     // Number of times appendBlocking observed BACKPRESSURE_NO_SPARE on its first
     // ring.appendOrFsn attempt. One increment per blocking-call that had to wait
@@ -136,15 +146,18 @@ public final class CursorSendEngine implements QuietCloseable {
     // isCloseCompleted() from pool threads re-probing a retired slot (see the
     // getter for why it must not synchronize).
     private volatile boolean closeCompleted;
+    // Invoked after closeCompleted publishes a confirmed flock release. Used
+    // by an owning sender pool to wake capacity-starved borrowers immediately.
+    private volatile Runnable slotLockReleaseListener;
     // Test-only hook run by finishClose() between the terminal cleanup and
     // the flock release. Lets a test park the releasing thread inside the
     // cleanup/release window and assert that closeCompleted stays false —
     // i.e. that completion is never observable while the flock is still
     // held. volatile: finishClose may run on the manager worker's exit
     // thread while the hook is installed from a test thread.
     private volatile Runnable beforeFlockReleaseHook;
-    // Exactly one daemon drives a failed flock release to completion. The
-    // error path only: ordinary closes never allocate or start a thread.
+    // Ensures this engine has at most one entry in the shared flock-release
+    // retry driver. The error path only: ordinary closes never enqueue work.
     private final AtomicBoolean flockReleaseRetryStarted = new AtomicBoolean();
     // Published before deferredClose is registered. The manager lock provides
     // the callback handoff fence; volatile also covers a direct test/retry read.
@@ -781,6 +794,15 @@ private void finishClose(boolean fullyDrained) {
         }
     }
 
+    /**
+     * Installs a hook invoked after each failed shared-driver flock release.
+     * Test-only: makes persistent retry progress observable without sleeps.
+     */
+    @TestOnly
+    public static void setAfterFlockReleaseRetryFailureHook(Runnable hook) {
+        afterFlockReleaseRetryFailureHook = hook;
+    }
+
     /**
      * Installs a constructor fault hook immediately before the bound deferred
      * close callback is created. Test-only: proves callback allocation failure
@@ -791,6 +813,23 @@ public static void setBeforeDeferredCloseCreationHook(Runnable hook) {
         beforeDeferredCloseCreationHook = hook;
     }
 
+    /**
+     * Overrides creation of the single shared flock-release retry thread.
+     * Test-only: makes thread creation/start failure and persistent retry
+     * scalability deterministic without relying on scheduler timing.
+     */
+    @TestOnly
+    public static void setFlockReleaseRetryThreadFactory(ThreadFactory factory) {
+        synchronized (FLOCK_RELEASE_RETRY_LOCK) {
+            if (flockReleaseRetryThread != null || !FLOCK_RELEASE_RETRY_QUEUE.isEmpty()) {
+                throw new IllegalStateException("flock-release retry driver is active");
+            }
+            flockReleaseRetryThreadFactory = factory == null
+                    ? DEFAULT_FLOCK_RELEASE_RETRY_THREAD_FACTORY
+                    : factory;
+        }
+    }
+
     /**
      * Installs a hook that {@link #finishClose} runs between the terminal
      * cleanup and the slot-flock release. Test-only: makes the otherwise
@@ -849,42 +888,97 @@ private boolean retryFlockReleaseIfReady() {
         }
         if (released) {
             closeCompleted = true;
+            Runnable listener = slotLockReleaseListener;
+            if (listener != null) {
+                try {
+                    listener.run();
+                } catch (Throwable ignored) {
+                    // A notification failure must not invalidate a confirmed release.
+                }
+            }
             return true;
         }
         startFlockReleaseRetry();
         return false;
     }
 
+    private static void runFlockReleaseRetryDriver() {
+        while (true) {
+            final int roundSize;
+            synchronized (FLOCK_RELEASE_RETRY_LOCK) {
+                roundSize = FLOCK_RELEASE_RETRY_QUEUE.size();
+                if (roundSize == 0) {
+                    flockReleaseRetryThread = null;
+                    return;
+                }
+            }
+            boolean hasFailures = false;
+            for (int i = 0; i < roundSize; i++) {
+                final CursorSendEngine engine;
+                synchronized (FLOCK_RELEASE_RETRY_LOCK) {
+                    engine = FLOCK_RELEASE_RETRY_QUEUE.pollFirst();
+                }
+                if (engine.retryFlockReleaseIfReady()) {
+                    engine.flockReleaseRetryStarted.set(false);
+                } else {
+                    synchronized (FLOCK_RELEASE_RETRY_LOCK) {
+                        FLOCK_RELEASE_RETRY_QUEUE.addLast(engine);
+                    }
+                    hasFailures = true;
+                    Runnable hook = afterFlockReleaseRetryFailureHook;
+                    if (hook != null) {
+                        hook.run();
+                    }
+                }
+            }
+            if (hasFailures) {
+                // Interruption must not abandon a retained flock, but clear
+                // the flag so subsequent parks still throttle retries.
+                Thread.interrupted();
+                LockSupport.parkNanos(100_000_000L);
+            }
+        }
+    }
+
     private void startFlockReleaseRetry() {
         if (!flockReleaseRetryStarted.compareAndSet(false, true)) {
             return;
         }
-        try {
-            Thread retryThread = new Thread(() -> {
-                while (!retryFlockReleaseIfReady()) {
-                    try {
-                        Thread.sleep(100L);
-                    } catch (InterruptedException ignored) {
-                        // A failed flock release must remain retryable until
-                        // confirmed or process exit; interruption is not a
-                        // safe reason to abandon the slot permanently.
+        Throwable startFailure = null;
+        synchronized (FLOCK_RELEASE_RETRY_LOCK) {
+            FLOCK_RELEASE_RETRY_QUEUE.addLast(this);
+            if (flockReleaseRetryThread == null) {
+                try {
+                    Thread retryThread = flockReleaseRetryThreadFactory.newThread(
+                            CursorSendEngine::runFlockReleaseRetryDriver);
+                    if (retryThread == null) {
+                        throw new IllegalStateException("retry thread factory returned null");
+                    }
+                    retryThread.setDaemon(true);
+                    flockReleaseRetryThread = retryThread;
+                    retryThread.start();
+                } catch (Throwable t) {
+                    startFailure = t;
+                    flockReleaseRetryThread = null;
+                    CursorSendEngine queued;
+                    while ((queued = FLOCK_RELEASE_RETRY_QUEUE.pollFirst()) != null) {
+                        queued.flockReleaseRetryStarted.set(false);
                     }
                 }
-            }, "qdb-sf-flock-release-retry");
-            retryThread.setDaemon(true);
-            retryThread.start();
+            }
+        }
+        if (startFailure == null) {
             LOG.error("SF slot flock release failed during engine close; keeping "
-                            + "closeCompleted=false and retrying outside the pool lock so "
+                            + "closeCompleted=false and retrying on the shared driver so "
                             + "retired capacity recovers after the transient failure [slot={}]",
                     sfDir == null ? "<memory>" : sfDir);
-        } catch (Throwable t) {
+        } else {
             // A later explicit close() can still retry without repeating the
-            // one-time ring/watermark cleanup. The kernel remains the final
-            // process-exit backstop if thread creation and all retries fail.
-            flockReleaseRetryStarted.set(false);
+            // one-time ring/watermark cleanup. The failed queue is cleared so
+            // the driver does not retain engines it cannot service.
             LOG.error("Could not start SF flock-release retry driver; close() may be "
                             + "invoked again to retry [slot={}, error={}]",
-                    sfDir == null ? "<memory>" : sfDir, String.valueOf(t));
+                    sfDir == null ? "<memory>" : sfDir, String.valueOf(startFailure));
         }
     }
 
@@ -910,6 +1004,17 @@ public boolean isCloseCompleted() {
         return closeCompleted;
     }
 
+    /**
+     * Registers a callback for confirmed SF slot-lock release. If release
+     * already completed, invokes the callback before returning.
+     */
+    public void setSlotLockReleaseListener(Runnable listener) {
+        slotLockReleaseListener = listener;
+        if (listener != null && closeCompleted) {
+            listener.run();
+        }
+    }
+
     /**
      * Pass-through to {@link SegmentRing#findSegmentContaining(long)}.
      */