address review comments

Author: glasstiger

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

@@ -614,6 +614,7 @@ public int getTimeout() {
         private PrivateKey privateKey;
         private int protocol = PARAMETER_NOT_SET_EXPLICITLY;
         private int protocolVersion = PARAMETER_NOT_SET_EXPLICITLY;
+        private boolean requestDurableAck;
         private int retryTimeoutMillis = PARAMETER_NOT_SET_EXPLICITLY;
         private boolean shouldDestroyPrivKey;
         private boolean tlsEnabled;
@@ -927,7 +928,8 @@ public Sender build() {
                         actualAutoFlushIntervalNanos,
                         actualInFlightWindowSize,
                         wsAuthHeader,
-                        actualMaxSchemasPerConnection
+                        actualMaxSchemasPerConnection,
+                        requestDurableAck
                 );
             }
 
@@ -1478,6 +1480,27 @@ public LineSenderBuilder protocolVersion(int protocolVersion) {
             return this;
         }
 
+        /**
+         * Opts the connection in for STATUS_DURABLE_ACK frames. When enabled,
+         * servers with primary replication will emit per-table durable-upload
+         * watermarks as WAL data reaches the object store.
+         * <p>
+         * This setting is only supported for WebSocket transport.
+         * <p>
+         * Observe durable progress via
+         * {@link QwpWebSocketSender#getHighestDurableSeqTxn(CharSequence)}.
+         *
+         * @param enabled true to request durable ACKs
+         * @return this instance for method chaining
+         */
+        public LineSenderBuilder requestDurableAck(boolean enabled) {
+            if (protocol != PARAMETER_NOT_SET_EXPLICITLY && protocol != PROTOCOL_WEBSOCKET) {
+                throw new LineSenderException("request_durable_ack is only supported for WebSocket transport");
+            }
+            this.requestDurableAck = enabled;
+            return this;
+        }
+
         /**
          * Configures the maximum time the Sender will spend retrying upon receiving a recoverable error from the server.
          * <br>
@@ -1870,6 +1893,18 @@ private LineSenderBuilder fromConfig(CharSequence configurationString) {
                     pos = getValue(configurationString, pos, sink, "in_flight_window");
                     int windowSize = parseIntValue(sink, "in_flight_window");
                     inFlightWindowSize(windowSize);
+                } else if (Chars.equals("request_durable_ack", sink)) {
+                    if (protocol != PROTOCOL_WEBSOCKET) {
+                        throw new LineSenderException("request_durable_ack is only supported for WebSocket transport");
+                    }
+                    pos = getValue(configurationString, pos, sink, "request_durable_ack");
+                    if (Chars.equalsIgnoreCase("on", sink)) {
+                        requestDurableAck(true);
+                    } else if (Chars.equalsIgnoreCase("off", sink)) {
+                        requestDurableAck(false);
+                    } else {
+                        throw new LineSenderException("invalid request_durable_ack [value=").put(sink).put(", allowed-values=[on, off]]");
+                    }
                 } else if (Chars.equals("max_schemas_per_connection", sink)) {
                     if (protocol != PROTOCOL_WEBSOCKET) {
                         throw new LineSenderException("max_schemas_per_connection is only supported for WebSocket transport");

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

@@ -1022,21 +1022,25 @@ public QwpWebSocketSender longColumn(CharSequence columnName, long value) {
     }
 
     /**
-     * Sends a WebSocket PING and reads frames until the PONG comes back,
-     * processing any STATUS_DURABLE_ACK or STATUS_OK frames along the way.
-     * After this method returns, {@link #getHighestDurableSeqTxn(CharSequence)} reflects
-     * the latest durable watermark reported by the server.
+     * Sends a WebSocket PING and blocks until the PONG arrives, processing
+     * any STATUS_DURABLE_ACK or STATUS_OK frames along the way.
      * <p>
-     * In async mode the PING is queued for the I/O thread; this method
-     * returns once the I/O thread has sent it and processed the response.
+     * The server flushes pending durable ACKs before sending the PONG, so
+     * after this method returns, {@link #getHighestDurableSeqTxn(CharSequence)}
+     * reflects all durable progress up to the moment the server processed
+     * the PING.
+     * <p>
+     * In async mode the PING is sent by the I/O thread; the I/O loop
+     * continues its normal work (sending batches, draining ACKs) while
+     * waiting for the PONG.
      *
      * @throws LineSenderException if the connection is closed or the ping times out
      */
     public void ping() {
         checkNotClosed();
         ensureConnected();
         if (inFlightWindowSize > 1) {
-            sendQueue.pingAndDrain();
+            sendQueue.ping();
         } else {
             syncPing();
         }

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

@@ -110,8 +110,10 @@ public class WebSocketSendQueue implements QuietCloseable {
     // Single pending buffer slot (double-buffering means at most 1 item in queue)
     // Zero allocation - just a volatile reference handoff
     private volatile MicrobatchBuffer pendingBuffer;
-    private volatile boolean pingRequested;
     private volatile boolean pingComplete;
+    private volatile boolean pingRequested;
+    private volatile boolean pongReceived;
+    private long pingDeadlineNanos;
     // Running state
     private volatile boolean running;
     private volatile boolean shuttingDown;
@@ -347,10 +349,16 @@ public long getDurableSeqTxn(String tableName) {
     }
 
     /**
-     * Requests the I/O thread to send a PING and drain any pending frames
-     * (including STATUS_DURABLE_ACK). Blocks until the I/O thread completes.
+     * Requests the I/O thread to send a WebSocket PING and blocks until
+     * the PONG arrives. The I/O loop continues its normal work (sending
+     * batches, draining ACKs) while waiting for the PONG.
+     * <p>
+     * The server flushes pending durable ACKs before sending the PONG,
+     * so after this method returns {@code getDurableSeqTxn()} reflects
+     * all durable progress up to the moment the server processed the PING.
      */
-    public void pingAndDrain() {
+    public void ping() {
+        checkError();
         synchronized (processingLock) {
             pingComplete = false;
             pingRequested = true;
@@ -369,6 +377,7 @@ public void pingAndDrain() {
                 }
             }
         }
+        checkError();
     }
 
     /**
@@ -396,12 +405,12 @@ private void checkError() {
     }
 
     /**
-     * Computes the current I/O state based on queue and in-flight status.
+     * Computes the current I/O state based on queue, in-flight, and ping status.
      */
     private IoState computeState(boolean hasInFlight) {
         if (!isPendingEmpty()) {
             return IoState.ACTIVE;
-        } else if (hasInFlight) {
+        } else if (hasInFlight || pingDeadlineNanos > 0) {
             return IoState.DRAINING;
         } else {
             return IoState.IDLE;
@@ -463,18 +472,17 @@ private void ioLoop() {
         try {
             int drainIdleCycles = 0;
             while (running || !isPendingEmpty()) {
+                // Send a pending PING if requested
                 if (pingRequested) {
                     pingRequested = false;
+                    pongReceived = false;
+                    pingDeadlineNanos = System.nanoTime() + InFlightWindow.DEFAULT_TIMEOUT_MS * 1_000_000L;
                     try {
                         client.sendPing(1000);
-                        tryReceiveAcks();
                     } catch (Exception e) {
+                        pingDeadlineNanos = 0;
                         failTransport(new LineSenderException("Ping failed", e));
-                    } finally {
-                        synchronized (processingLock) {
-                            pingComplete = true;
-                            processingLock.notifyAll();
-                        }
+                        completePing();
                     }
                 }
 
@@ -489,7 +497,7 @@ private void ioLoop() {
                         // Nothing to do - wait for work under lock
                         synchronized (processingLock) {
                             // Re-check under lock to avoid missed wakeup
-                            if (isPendingEmpty() && running) {
+                            if (isPendingEmpty() && running && !pingRequested) {
                                 try {
                                     processingLock.wait(100);
                                 } catch (InterruptedException e) {
@@ -502,10 +510,22 @@ private void ioLoop() {
                     case ACTIVE:
                     case DRAINING:
                         // Try to receive any pending ACKs (non-blocking)
-                        if (hasInFlight && client.isConnected()) {
+                        if (client.isConnected()) {
                             receivedAcks = tryReceiveAcks();
                         }
 
+                        // Check if a pending PING has been answered
+                        if (pingDeadlineNanos > 0) {
+                            if (pongReceived) {
+                                pingDeadlineNanos = 0;
+                                completePing();
+                            } else if (System.nanoTime() >= pingDeadlineNanos) {
+                                pingDeadlineNanos = 0;
+                                failTransport(new LineSenderException("Ping timed out waiting for PONG"));
+                                completePing();
+                            }
+                        }
+
                         // Try to dequeue and send a batch
                         boolean hasWindowSpace = (inFlightWindow == null || inFlightWindow.hasWindowSpace());
                         if (hasWindowSpace) {
@@ -550,6 +570,13 @@ private void ioLoop() {
         }
     }
 
+    private void completePing() {
+        synchronized (processingLock) {
+            pingComplete = true;
+            processingLock.notifyAll();
+        }
+    }
+
     private boolean isPendingEmpty() {
         return pendingBuffer == null;
     }
@@ -749,6 +776,11 @@ public void onClose(int code, String reason) {
             LOG.info("WebSocket closed by server [code={}, reason={}]", code, reason);
             failTransport(new LineSenderException("WebSocket closed by server [code=" + code + ", reason=" + reason + ']'));
         }
+
+        @Override
+        public void onPong(long payloadPtr, int payloadLen) {
+            pongReceived = true;
+        }
     }
 
     static final class SeqTxn {

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

@@ -451,24 +451,31 @@ public void testDurableAckInterleavedWithStatusOk() throws Exception {
     }
 
     @Test
-    public void testPingAndDrainCompletesSuccessfully() throws Exception {
+    public void testPingBlocksUntilPong() throws Exception {
         assertMemoryLeak(() -> {
             InFlightWindow window = new InFlightWindow(8, 5_000);
             WebSocketSendQueue queue = null;
             try (FakeWebSocketClient client = new FakeWebSocketClient()) {
-                AtomicBoolean ackDelivered = new AtomicBoolean(false);
+                AtomicInteger callCount = new AtomicInteger();
                 client.setTryReceiveBehavior(handler -> {
-                    if (ackDelivered.compareAndSet(false, true)) {
-                        emitDurableAck(handler, "t", 7);
-                        return true;
+                    int n = callCount.getAndIncrement();
+                    switch (n) {
+                        case 0:
+                            emitDurableAck(handler, "t", 7);
+                            return true;
+                        case 1:
+                            handler.onPong(0, 0);
+                            return true;
+                        default:
+                            return false;
                     }
-                    return false;
                 });
 
                 queue = new WebSocketSendQueue(client, window, 1_000, 500);
 
-                queue.pingAndDrain();
+                queue.ping();
 
+                // After ping() returns, durable ACK must already be processed
                 assertEquals(7, queue.getDurableSeqTxn("t"));
             } finally {
                 closeQuietly(queue);
@@ -477,7 +484,7 @@ public void testPingAndDrainCompletesSuccessfully() throws Exception {
     }
 
     @Test
-    public void testPingAndDrainWithInFlightBatches() throws Exception {
+    public void testPingWithInFlightBatches() throws Exception {
         assertMemoryLeak(() -> {
             InFlightWindow window = new InFlightWindow(8, 5_000);
             WebSocketSendQueue queue = null;
@@ -495,14 +502,17 @@ public void testPingAndDrainWithInFlightBatches() throws Exception {
                         case 1:
                             emitDurableAck(handler, "t", 5);
                             return true;
+                        case 2:
+                            handler.onPong(0, 0);
+                            return true;
                         default:
                             return false;
                     }
                 });
 
                 queue = new WebSocketSendQueue(client, window, 1_000, 500);
 
-                queue.pingAndDrain();
+                queue.ping();
 
                 assertEquals(0, window.getInFlightCount());
                 assertEquals(5, queue.getDurableSeqTxn("t"));
@@ -512,6 +522,54 @@ public void testPingAndDrainWithInFlightBatches() throws Exception {
         });
     }
 
+    @Test
+    public void testPingTimesOutWhenNoPong() throws Exception {
+        assertMemoryLeak(() -> {
+            InFlightWindow window = new InFlightWindow(8, 5_000);
+            WebSocketSendQueue queue = null;
+            try (FakeWebSocketClient client = new FakeWebSocketClient()) {
+                // Never emit a PONG
+                client.setTryReceiveBehavior(handler -> false);
+
+                queue = new WebSocketSendQueue(client, window, 1_000, 500);
+
+                try {
+                    queue.ping();
+                    fail("Expected ping timeout");
+                } catch (LineSenderException e) {
+                    assertTrue(e.getMessage().contains("Ping timed out"));
+                }
+            } finally {
+                closeQuietly(queue);
+            }
+        });
+    }
+
+    @Test
+    public void testPingSurfacesTransportError() throws Exception {
+        assertMemoryLeak(() -> {
+            InFlightWindow window = new InFlightWindow(8, 5_000);
+            WebSocketSendQueue queue = null;
+            try (FakeWebSocketClient client = new FakeWebSocketClient()) {
+                client.setPingSendBehavior(() -> {
+                    throw new RuntimeException("ping-send-fail");
+                });
+
+                queue = new WebSocketSendQueue(client, window, 1_000, 500);
+
+                try {
+                    queue.ping();
+                    fail("Expected error from ping");
+                } catch (LineSenderException e) {
+                    assertTrue(e.getMessage().contains("Ping failed")
+                            || e.getMessage().contains("Error in send queue"));
+                }
+            } finally {
+                closeQuietly(queue);
+            }
+        });
+    }
+
     @Test
     public void testDurableSeqTxnInitiallyMinusOne() throws Exception {
         assertMemoryLeak(() -> {
@@ -603,6 +661,7 @@ private interface ReceiveBehavior {
     private static class FakeWebSocketClient extends WebSocketClient {
         private volatile TryReceiveBehavior behavior = handler -> false;
         private volatile boolean connected = true;
+        private volatile Runnable pingSendBehavior = () -> {};
         private volatile ReceiveBehavior receiveBehavior = (handler, timeout) -> false;
         private volatile SendBehavior sendBehavior = (dataPtr, length) -> {
         };
@@ -629,6 +688,11 @@ public void sendBinary(long dataPtr, int length) {
 
         @Override
         public void sendPing(int timeout) {
+            pingSendBehavior.run();
+        }
+
+        public void setPingSendBehavior(Runnable pingSendBehavior) {
+            this.pingSendBehavior = pingSendBehavior;
         }
 
         public void setSendBehavior(SendBehavior sendBehavior) {