remove blocking waits from websocket ACK drain

Author: jerrinot

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

@@ -400,8 +400,8 @@ public void close() {
                     // Wait for all batches to be sent and acknowledged before closing
                     if (sendQueue != null) {
                         sendQueue.flush();
-                    }
-                    if (inFlightWindow != null) {
+                        sendQueue.awaitPendingAcks();
+                    } else if (inFlightWindow != null) {
                         inFlightWindow.awaitEmpty();
                     }
                 } else {
@@ -581,7 +581,11 @@ public void flush() {
             sendQueue.flush();
 
             // Wait for all in-flight batches to be acknowledged by the server
-            inFlightWindow.awaitEmpty();
+            if (sendQueue != null) {
+                sendQueue.awaitPendingAcks();
+            } else {
+                inFlightWindow.awaitEmpty();
+            }
 
             LOG.debug("Flush complete [totalBatches={}, totalBytes={}, totalAcked={}]", sendQueue.getTotalBatchesSent(), sendQueue.getTotalBytesSent(), inFlightWindow.getTotalAcked());
         } else {

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

@@ -64,6 +64,7 @@
  */
 public class WebSocketSendQueue implements QuietCloseable {
 
+    private static final int DRAIN_SPIN_TRIES = 16;
     public static final long DEFAULT_ENQUEUE_TIMEOUT_MS = 30_000;
     public static final long DEFAULT_SHUTDOWN_TIMEOUT_MS = 10_000;
     private static final Logger LOG = LoggerFactory.getLogger(WebSocketSendQueue.class);
@@ -264,7 +265,6 @@ public boolean enqueue(MicrobatchBuffer buffer) {
                 }
             }
         }
-
         LOG.debug("Enqueued batch [id={}, bytes={}, rows={}]", buffer.getBatchId(), buffer.getBufferPos(), buffer.getRowCount());
         return true;
     }
@@ -282,27 +282,27 @@ public void flush() {
 
         long startTime = System.currentTimeMillis();
 
-        // Wait under lock - I/O thread will notify when processingCount decrements
+        // Wait under lock until the queue becomes empty and no batch is being sent.
         synchronized (processingLock) {
             while (running) {
                 // Atomically check: queue empty AND not processing
                 if (isPendingEmpty() && processingCount.get() == 0) {
                     break; // All done
                 }
 
+                long remaining = enqueueTimeoutMs - (System.currentTimeMillis() - startTime);
+                if (remaining <= 0) {
+                    throw new LineSenderException("Flush timeout after " + enqueueTimeoutMs + "ms, " +
+                            "queue=" + getPendingSize() + ", processing=" + processingCount.get());
+                }
+
                 try {
-                    processingLock.wait(10);
+                    processingLock.wait(remaining);
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
                     throw new LineSenderException("Interrupted while flushing", e);
                 }
 
-                // Check timeout
-                if (System.currentTimeMillis() - startTime > enqueueTimeoutMs) {
-                    throw new LineSenderException("Flush timeout after " + enqueueTimeoutMs + "ms, " +
-                            "queue=" + getPendingSize() + ", processing=" + processingCount.get());
-                }
-
                 // Check for errors
                 checkError();
             }
@@ -314,6 +314,19 @@ public void flush() {
         LOG.debug("Flush complete");
     }
 
+    /**
+     * Waits for all in-flight batches to be acknowledged.
+     */
+    public void awaitPendingAcks() {
+        if (inFlightWindow == null) {
+            return;
+        }
+
+        checkError();
+        inFlightWindow.awaitEmpty();
+        checkError();
+    }
+
     /**
      * Returns the last error that occurred in the I/O thread, or null if no error.
      */
@@ -387,27 +400,39 @@ private int getPendingSize() {
         return pendingBuffer == null ? 0 : 1;
     }
 
+    private int idleDuringDrain(int idleCycles) {
+        if (idleCycles < DRAIN_SPIN_TRIES) {
+            Thread.onSpinWait();
+            return idleCycles + 1;
+        }
+        Thread.yield();
+        return DRAIN_SPIN_TRIES;
+    }
+
     /**
      * The main I/O loop that handles both sending batches and receiving ACKs.
      * <p>
      * Uses a state machine:
      * <ul>
      *   <li>IDLE: block on processingLock.wait() until work arrives</li>
      *   <li>ACTIVE: non-blocking poll queue, send batches, check for ACKs</li>
-     *   <li>DRAINING: no batches but ACKs pending - poll for ACKs with short wait</li>
+     *   <li>DRAINING: no batches but ACKs pending - poll for ACKs with non-blocking backoff</li>
      * </ul>
      */
     private void ioLoop() {
         LOG.info("I/O loop started");
 
         try {
+            int drainIdleCycles = 0;
             while (running || !isPendingEmpty()) {
                 MicrobatchBuffer batch = null;
                 boolean hasInFlight = (inFlightWindow != null && inFlightWindow.getInFlightCount() > 0);
                 IoState state = computeState(hasInFlight);
+                boolean receivedAcks = false;
 
                 switch (state) {
                     case IDLE:
+                        drainIdleCycles = 0;
                         // Nothing to do - wait for work under lock
                         synchronized (processingLock) {
                             // Re-check under lock to avoid missed wakeup
@@ -425,7 +450,7 @@ private void ioLoop() {
                     case DRAINING:
                         // Try to receive any pending ACKs (non-blocking)
                         if (hasInFlight && client.isConnected()) {
-                            tryReceiveAcks();
+                            receivedAcks = tryReceiveAcks();
                         }
 
                         // Try to dequeue and send a batch
@@ -452,15 +477,16 @@ private void ioLoop() {
                             }
                         }
 
-                        // In DRAINING state with no work, short wait to avoid busy loop
+                        // In DRAINING state with no work, stay non-blocking and use
+                        // a simple spin/yield backoff.
                         if (state == IoState.DRAINING && batch == null) {
-                            synchronized (processingLock) {
-                                try {
-                                    processingLock.wait(10);
-                                } catch (InterruptedException e) {
-                                    if (!running) return;
-                                }
+                            if (receivedAcks) {
+                                drainIdleCycles = 0;
+                            } else {
+                                drainIdleCycles = idleDuringDrain(drainIdleCycles);
                             }
+                        } else {
+                            drainIdleCycles = 0;
                         }
                         break;
                 }
@@ -553,9 +579,11 @@ private void sendBatch(MicrobatchBuffer batch) {
     /**
      * Tries to receive ACKs from the server (non-blocking).
      */
-    private void tryReceiveAcks() {
+    private boolean tryReceiveAcks() {
+        boolean received = false;
         try {
             while (client.tryReceiveFrame(responseHandler)) {
+                received = true;
                 // Drain all buffered ACKs before returning to the I/O loop.
             }
         } catch (Exception e) {
@@ -564,14 +592,15 @@ private void tryReceiveAcks() {
                 failTransport(new LineSenderException("Error receiving response: " + e.getMessage(), e));
             }
         }
+        return received;
     }
 
     /**
      * I/O loop states for the state machine.
      * <ul>
      *   <li>IDLE: queue empty, no in-flight batches - can block waiting for work</li>
      *   <li>ACTIVE: have batches to send - non-blocking loop</li>
-     *   <li>DRAINING: queue empty but ACKs pending - poll for ACKs, short wait</li>
+     *   <li>DRAINING: queue empty but ACKs pending - poll for ACKs with non-blocking backoff</li>
      * </ul>
      */
     private enum IoState {

core/src/test/java/io/questdb/client/test/cutlass/qwp/client/WebSocketSendQueueTest.java

@@ -30,6 +30,7 @@
 import io.questdb.client.cutlass.line.LineSenderException;
 import io.questdb.client.cutlass.qwp.client.InFlightWindow;
 import io.questdb.client.cutlass.qwp.client.MicrobatchBuffer;
+import io.questdb.client.cutlass.qwp.client.WebSocketResponse;
 import io.questdb.client.cutlass.qwp.client.WebSocketSendQueue;
 import io.questdb.client.network.PlainSocketFactory;
 import io.questdb.client.std.MemoryTag;
@@ -40,6 +41,8 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.junit.Assert.*;
@@ -266,6 +269,84 @@ public void testFlushFailsWhenServerClosesConnection() throws Exception {
         });
     }
 
+    @Test
+    public void testAwaitPendingAcksKeepsDrainNonBlocking() throws Exception {
+        assertMemoryLeak(() -> {
+            InFlightWindow window = new InFlightWindow(8, 5_000);
+            FakeWebSocketClient client = new FakeWebSocketClient();
+            WebSocketSendQueue queue = null;
+            MicrobatchBuffer batch0 = sealedBuffer((byte) 1);
+            MicrobatchBuffer batch1 = sealedBuffer((byte) 2);
+            CountDownLatch secondBatchSent = new CountDownLatch(1);
+            AtomicBoolean deliverAcks = new AtomicBoolean(false);
+            AtomicInteger tryReceivePolls = new AtomicInteger();
+            AtomicLong highestSent = new AtomicLong(-1);
+            AtomicReference<Throwable> errorRef = new AtomicReference<>();
+
+            try {
+                client.setSendBehavior((dataPtr, length) -> {
+                    long sent = highestSent.incrementAndGet();
+                    if (sent == 1) {
+                        secondBatchSent.countDown();
+                    }
+                });
+                client.setReceiveBehavior((handler, timeout) -> {
+                    throw new AssertionError("receiveFrame() must not be used while draining ACKs");
+                });
+                client.setTryReceiveBehavior(handler -> {
+                    tryReceivePolls.incrementAndGet();
+                    if (deliverAcks.get()) {
+                        long sent = highestSent.get();
+                        if (sent >= 0 && window.getInFlightCount() > 0) {
+                            emitAck(handler, sent);
+                            return true;
+                        }
+                    }
+                    return false;
+                });
+
+                queue = new WebSocketSendQueue(client, window, 1_000, 500);
+                queue.enqueue(batch0);
+                queue.flush();
+
+                CountDownLatch finished = new CountDownLatch(1);
+                WebSocketSendQueue finalQueue = queue;
+                Thread waiter = new Thread(() -> {
+                    try {
+                        finalQueue.awaitPendingAcks();
+                    } catch (Throwable t) {
+                        errorRef.set(t);
+                    } finally {
+                        finished.countDown();
+                    }
+                });
+                waiter.start();
+
+                long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
+                while (tryReceivePolls.get() == 0 && System.nanoTime() < deadline) {
+                    Thread.onSpinWait();
+                }
+                assertTrue("Expected non-blocking ACK polls while draining", tryReceivePolls.get() > 0);
+
+                queue.enqueue(batch1);
+                assertTrue("I/O thread should still send new work while ACK drain is active",
+                        secondBatchSent.await(1, TimeUnit.SECONDS));
+
+                deliverAcks.set(true);
+
+                assertTrue("awaitPendingAcks should complete once ACK arrives",
+                        finished.await(2, TimeUnit.SECONDS));
+                assertNull(errorRef.get());
+                assertEquals(0, window.getInFlightCount());
+            } finally {
+                closeQuietly(queue);
+                batch0.close();
+                batch1.close();
+                client.close();
+            }
+        });
+    }
+
     private static void awaitThreadBlocked(Thread thread) throws InterruptedException {
         long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
         while (System.nanoTime() < deadline) {
@@ -296,6 +377,18 @@ private static void emitBinary(WebSocketFrameHandler handler, byte[] payload) {
         }
     }
 
+    private static void emitAck(WebSocketFrameHandler handler, long sequence) {
+        WebSocketResponse response = WebSocketResponse.success(sequence);
+        int size = response.serializedSize();
+        long ptr = Unsafe.malloc(size, MemoryTag.NATIVE_DEFAULT);
+        try {
+            response.writeTo(ptr);
+            handler.onBinaryMessage(ptr, size);
+        } finally {
+            Unsafe.free(ptr, size, MemoryTag.NATIVE_DEFAULT);
+        }
+    }
+
     private static MicrobatchBuffer sealedBuffer(byte value) {
         MicrobatchBuffer buffer = new MicrobatchBuffer(64);
         buffer.writeByte(value);
@@ -312,9 +405,14 @@ private interface TryReceiveBehavior {
         boolean tryReceive(WebSocketFrameHandler handler);
     }
 
+    private interface ReceiveBehavior {
+        boolean receive(WebSocketFrameHandler handler, int timeout);
+    }
+
     private static class FakeWebSocketClient extends WebSocketClient {
         private volatile TryReceiveBehavior behavior = handler -> false;
         private volatile boolean connected = true;
+        private volatile ReceiveBehavior receiveBehavior = (handler, timeout) -> false;
         private volatile SendBehavior sendBehavior = (dataPtr, length) -> {
         };
 
@@ -346,6 +444,15 @@ public void setTryReceiveBehavior(TryReceiveBehavior behavior) {
             this.behavior = behavior;
         }
 
+        public void setReceiveBehavior(ReceiveBehavior receiveBehavior) {
+            this.receiveBehavior = receiveBehavior;
+        }
+
+        @Override
+        public boolean receiveFrame(WebSocketFrameHandler handler, int timeout) {
+            return receiveBehavior.receive(handler, timeout);
+        }
+
         @Override
         public boolean tryReceiveFrame(WebSocketFrameHandler handler) {
             return behavior.tryReceive(handler);