address review comments

Author: glasstiger

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

@@ -1757,10 +1757,10 @@ private void syncPing() {
             if (received) {
                 if (sawBinaryAck) {
                     if (ackResponse.isDurableAck()) {
-                        updateSyncDurableSeqTxns();
+                        updateSyncSeqTxns(syncDurableSeqTxns);
                     } else if (ackResponse.isSuccess()) {
                         inFlightWindow.acknowledgeUpTo(ackResponse.getSequence());
-                        updateSyncCommittedSeqTxns();
+                        updateSyncSeqTxns(syncCommittedSeqTxns);
                     }
                 }
                 if (sawPong) {
@@ -1771,22 +1771,12 @@ private void syncPing() {
         throw new LineSenderException("Ping timed out");
     }
 
-    private void updateSyncCommittedSeqTxns() {
+    private void updateSyncSeqTxns(CharSequenceLongHashMap seqTxns) {
         for (int i = 0, n = ackResponse.getTableEntryCount(); i < n; i++) {
             String name = ackResponse.getTableName(i);
             long seqTxn = ackResponse.getTableSeqTxn(i);
-            if (seqTxn > syncCommittedSeqTxns.get(name)) {
-                syncCommittedSeqTxns.put(name, seqTxn);
-            }
-        }
-    }
-
-    private void updateSyncDurableSeqTxns() {
-        for (int i = 0, n = ackResponse.getTableEntryCount(); i < n; i++) {
-            String name = ackResponse.getTableName(i);
-            long seqTxn = ackResponse.getTableSeqTxn(i);
-            if (seqTxn > syncDurableSeqTxns.get(name)) {
-                syncDurableSeqTxns.put(name, seqTxn);
+            if (seqTxn > seqTxns.get(name)) {
+                seqTxns.put(name, seqTxn);
             }
         }
     }
@@ -1842,12 +1832,12 @@ private void waitForAck(long expectedSequence) {
                     if (ackResponse.isSuccess()) {
                         long sequence = ackResponse.getSequence();
                         inFlightWindow.acknowledgeUpTo(sequence);
-                        updateSyncCommittedSeqTxns();
+                        updateSyncSeqTxns(syncCommittedSeqTxns);
                         if (sequence >= expectedSequence) {
                             return;
                         }
                     } else if (ackResponse.isDurableAck()) {
-                        updateSyncDurableSeqTxns();
+                        updateSyncSeqTxns(syncDurableSeqTxns);
                     } else {
                         long sequence = ackResponse.getSequence();
                         String errorMessage = ackResponse.getErrorMessage();

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

@@ -809,11 +809,8 @@ private static void advanceSeqTxn(ConcurrentHashMap<String, SeqTxn> map, String
         SeqTxn existing = map.get(tableName);
         if (existing != null) {
             existing.advance(seqTxn);
-        } else {
-            existing = map.putIfAbsent(tableName, new SeqTxn(seqTxn));
-            if (existing != null) {
-                existing.advance(seqTxn);
-            }
+            return;
         }
+        map.computeIfAbsent(tableName, k -> new SeqTxn(seqTxn)).advance(seqTxn);
     }
 }

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

@@ -24,6 +24,7 @@
 
 package io.questdb.client.test.cutlass.qwp.client;
 
+import io.questdb.client.cutlass.line.LineSenderException;
 import io.questdb.client.cutlass.qwp.client.QwpWebSocketSender;
 import io.questdb.client.cutlass.qwp.client.WebSocketResponse;
 import io.questdb.client.cutlass.qwp.websocket.WebSocketCloseCode;
@@ -34,8 +35,13 @@
 import org.junit.Test;
 
 import java.io.IOException;
+import java.io.InputStream;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Integration tests for QWP v1 WebSocket ACK delivery mechanism.
@@ -195,6 +201,151 @@ public void testSyncFlushIgnoresPingAndWaitsForAck() throws Exception {
         }
     }
 
+    @Test
+    public void testDurableAckUpgradeHeaderNotSentByDefault() throws Exception {
+        int port = TEST_PORT + 31;
+        AtomicReference<String> capturedRequest = new AtomicReference<>();
+
+        try (ServerSocket serverSocket = new ServerSocket(port)) {
+            serverSocket.setSoTimeout(5000);
+
+            Thread serverThread = new Thread(() -> {
+                try {
+                    Socket client = serverSocket.accept();
+                    InputStream in = client.getInputStream();
+                    StringBuilder request = new StringBuilder();
+                    byte[] buf = new byte[1];
+                    while (true) {
+                        int read = in.read(buf);
+                        if (read <= 0) {
+                            break;
+                        }
+                        request.append((char) buf[0]);
+                        if (request.toString().endsWith("\r\n\r\n")) {
+                            break;
+                        }
+                    }
+                    capturedRequest.set(request.toString());
+                    client.close();
+                } catch (Exception e) {
+                    // expected
+                }
+            });
+            serverThread.start();
+
+            try {
+                QwpWebSocketSender.connect("localhost", port, null,
+                        0, 0, 0, 1, null).close();
+            } catch (LineSenderException e) {
+                // expected - server doesn't complete handshake
+            }
+
+            serverThread.join(5000);
+
+            String request = capturedRequest.get();
+            Assert.assertNotNull("Server should have received upgrade request", request);
+            Assert.assertFalse("Request should NOT contain X-QWP-Request-Durable-Ack header",
+                    request.contains("X-QWP-Request-Durable-Ack"));
+        }
+    }
+
+    @Test
+    public void testDurableAckUpgradeHeaderSent() throws Exception {
+        int port = TEST_PORT + 30;
+        AtomicReference<String> capturedRequest = new AtomicReference<>();
+
+        try (ServerSocket serverSocket = new ServerSocket(port)) {
+            serverSocket.setSoTimeout(5000);
+
+            Thread serverThread = new Thread(() -> {
+                try {
+                    Socket client = serverSocket.accept();
+                    InputStream in = client.getInputStream();
+                    StringBuilder request = new StringBuilder();
+                    byte[] buf = new byte[1];
+                    while (true) {
+                        int read = in.read(buf);
+                        if (read <= 0) {
+                            break;
+                        }
+                        request.append((char) buf[0]);
+                        if (request.toString().endsWith("\r\n\r\n")) {
+                            break;
+                        }
+                    }
+                    capturedRequest.set(request.toString());
+                    client.close();
+                } catch (Exception e) {
+                    // expected
+                }
+            });
+            serverThread.start();
+
+            try {
+                QwpWebSocketSender.connect("localhost", port, null,
+                        0, 0, 0, 1, null,
+                        QwpWebSocketSender.DEFAULT_MAX_SCHEMAS_PER_CONNECTION,
+                        true).close();
+            } catch (LineSenderException e) {
+                // expected - server doesn't complete handshake
+            }
+
+            serverThread.join(5000);
+
+            String request = capturedRequest.get();
+            Assert.assertNotNull("Server should have received upgrade request", request);
+            Assert.assertTrue("Request should contain X-QWP-Request-Durable-Ack header",
+                    request.contains("X-QWP-Request-Durable-Ack: true"));
+        }
+    }
+
+    @Test
+    public void testSyncDurableAckDuringWaitForAck() throws Exception {
+        int port = TEST_PORT + 25;
+        DurableAckThenStatusOkHandler handler = new DurableAckThenStatusOkHandler();
+
+        try (TestWebSocketServer server = new TestWebSocketServer(port, handler)) {
+            server.start();
+            Assert.assertTrue("Server failed to start", server.awaitStart(5, TimeUnit.SECONDS));
+
+            // window=1 for sync mode
+            try (QwpWebSocketSender sender = QwpWebSocketSender.connect(
+                    "localhost", port, null, 0, 0, 0, 1, null)) {
+                sender.table("trades")
+                        .longColumn("price", 100)
+                        .atNow();
+                sender.flush();
+
+                Assert.assertEquals(42L, sender.getHighestDurableSeqTxn("trades"));
+                Assert.assertEquals(10L, sender.getHighestAckedSeqTxn("trades"));
+            }
+        }
+    }
+
+    @Test
+    public void testSyncFlushUpdatesCommittedSeqTxnsWithTableEntries() throws Exception {
+        int port = TEST_PORT + 24;
+        AckWithTableEntriesHandler handler = new AckWithTableEntriesHandler();
+
+        try (TestWebSocketServer server = new TestWebSocketServer(port, handler)) {
+            server.start();
+            Assert.assertTrue("Server failed to start", server.awaitStart(5, TimeUnit.SECONDS));
+
+            // window=1 for sync mode
+            try (QwpWebSocketSender sender = QwpWebSocketSender.connect(
+                    "localhost", port, null, 0, 0, 0, 1, null)) {
+                sender.table("trades")
+                        .longColumn("price", 100)
+                        .atNow();
+                sender.flush();
+
+                Assert.assertEquals(10L, sender.getHighestAckedSeqTxn("trades"));
+                Assert.assertEquals(20L, sender.getHighestAckedSeqTxn("orders"));
+                Assert.assertEquals(-1L, sender.getHighestAckedSeqTxn("other"));
+            }
+        }
+    }
+
     /**
      * Creates a binary ACK response using WebSocketResponse format.
      * Format: status (1) + sequence (8) + tableCount (2, zero entries)
@@ -220,6 +371,75 @@ private static byte[] createAckResponse(long sequence) {
         return response;
     }
 
+    private static byte[] createAckResponseWithTables(long sequence, String[] tableNames, long[] seqTxns) {
+        byte[][] nameBytes = new byte[tableNames.length][];
+        int size = 1 + 8 + 2;
+        for (int i = 0; i < tableNames.length; i++) {
+            nameBytes[i] = tableNames[i].getBytes(StandardCharsets.UTF_8);
+            size += 2 + nameBytes[i].length + 8;
+        }
+
+        byte[] response = new byte[size];
+        int offset = 0;
+        response[offset++] = WebSocketResponse.STATUS_OK;
+        for (int i = 0; i < 8; i++) {
+            response[offset++] = (byte) ((sequence >> (i * 8)) & 0xFF);
+        }
+        response[offset++] = (byte) (tableNames.length & 0xFF);
+        response[offset++] = (byte) ((tableNames.length >> 8) & 0xFF);
+        for (int i = 0; i < tableNames.length; i++) {
+            response[offset++] = (byte) (nameBytes[i].length & 0xFF);
+            response[offset++] = (byte) ((nameBytes[i].length >> 8) & 0xFF);
+            System.arraycopy(nameBytes[i], 0, response, offset, nameBytes[i].length);
+            offset += nameBytes[i].length;
+            for (int j = 0; j < 8; j++) {
+                response[offset++] = (byte) ((seqTxns[i] >> (j * 8)) & 0xFF);
+            }
+        }
+        return response;
+    }
+
+    private static byte[] createDurableAckResponse(String[] tableNames, long[] seqTxns) {
+        byte[][] nameBytes = new byte[tableNames.length][];
+        int size = 1 + 2;
+        for (int i = 0; i < tableNames.length; i++) {
+            nameBytes[i] = tableNames[i].getBytes(StandardCharsets.UTF_8);
+            size += 2 + nameBytes[i].length + 8;
+        }
+
+        byte[] response = new byte[size];
+        int offset = 0;
+        response[offset++] = WebSocketResponse.STATUS_DURABLE_ACK;
+        response[offset++] = (byte) (tableNames.length & 0xFF);
+        response[offset++] = (byte) ((tableNames.length >> 8) & 0xFF);
+        for (int i = 0; i < tableNames.length; i++) {
+            response[offset++] = (byte) (nameBytes[i].length & 0xFF);
+            response[offset++] = (byte) ((nameBytes[i].length >> 8) & 0xFF);
+            System.arraycopy(nameBytes[i], 0, response, offset, nameBytes[i].length);
+            offset += nameBytes[i].length;
+            for (int j = 0; j < 8; j++) {
+                response[offset++] = (byte) ((seqTxns[i] >> (j * 8)) & 0xFF);
+            }
+        }
+        return response;
+    }
+
+    private static class AckWithTableEntriesHandler implements TestWebSocketServer.WebSocketServerHandler {
+        private final AtomicLong nextSequence = new AtomicLong(0);
+
+        @Override
+        public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) {
+            long sequence = nextSequence.getAndIncrement();
+            try {
+                client.sendBinary(createAckResponseWithTables(sequence,
+                        new String[]{"trades", "orders"},
+                        new long[]{10L, 20L}));
+            } catch (IOException e) {
+                LOG.error("Failed to send ACK with tables", e);
+            }
+        }
+    }
+
     private static class ClosingServerHandler implements TestWebSocketServer.WebSocketServerHandler {
         @Override
         public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) {
@@ -261,6 +481,27 @@ public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] dat
         }
     }
 
+    private static class DurableAckThenStatusOkHandler implements TestWebSocketServer.WebSocketServerHandler {
+        private final AtomicLong nextSequence = new AtomicLong(0);
+
+        @Override
+        public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) {
+            long sequence = nextSequence.getAndIncrement();
+            try {
+                // Send durable ACK first
+                client.sendBinary(createDurableAckResponse(
+                        new String[]{"trades"},
+                        new long[]{42L}));
+                // Then send STATUS_OK with committed seqTxns
+                client.sendBinary(createAckResponseWithTables(sequence,
+                        new String[]{"trades"},
+                        new long[]{10L}));
+            } catch (IOException e) {
+                LOG.error("Failed to send ACK frames", e);
+            }
+        }
+    }
+
     private static class InvalidAckPayloadHandler implements TestWebSocketServer.WebSocketServerHandler {
         @Override
         public void onBinaryMessage(TestWebSocketServer.ClientHandler client, byte[] data) {