bugfixes and tests

Author: bluestreak01

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

@@ -70,6 +70,10 @@ public abstract class WebSocketClient implements QuietCloseable {
 
     private static final int DEFAULT_RECV_BUFFER_SIZE = 65536;
     private static final int DEFAULT_SEND_BUFFER_SIZE = 65536;
+    // tryParseFrame() return values
+    private static final int PARSE_INCOMPLETE = 0;
+    private static final int PARSE_NEED_MORE = -1;
+    private static final int PARSE_OK = 1;
     private static final Logger LOG = LoggerFactory.getLogger(WebSocketClient.class);
     private static final String QWP_VERSION_HEADER_NAME = "X-QWP-Version:";
     private static final ThreadLocal<MessageDigest> SHA1_DIGEST = ThreadLocal.withInitial(() -> {
@@ -129,21 +133,26 @@ public WebSocketClient(HttpClientConfiguration configuration, SocketFactory sock
             this.recvBufSize = Math.max(configuration.getResponseBufferSize(), DEFAULT_RECV_BUFFER_SIZE);
             this.maxRecvBufSize = Math.max(configuration.getMaximumResponseBufferSize(), recvBufSize);
             this.recvBufPtr = Unsafe.malloc(recvBufSize, MemoryTag.NATIVE_DEFAULT);
+
+            this.sendBuffer = sendBuf;
+            this.controlFrameBuffer = controlBuf;
+            this.recvPos = 0;
+            this.recvReadPos = 0;
+
+            this.frameParser = new WebSocketFrameParser();
+            this.rnd = new SecureRnd();
+            this.upgraded = false;
+            this.closed = false;
         } catch (Throwable t) {
+            if (recvBufPtr != 0) {
+                Unsafe.free(recvBufPtr, recvBufSize, MemoryTag.NATIVE_DEFAULT);
+                recvBufPtr = 0;
+            }
             Misc.free(controlBuf);
             Misc.free(sendBuf);
             Misc.free(socket);
             throw t;
         }
-        this.sendBuffer = sendBuf;
-        this.controlFrameBuffer = controlBuf;
-        this.recvPos = 0;
-        this.recvReadPos = 0;
-
-        this.frameParser = new WebSocketFrameParser();
-        this.rnd = new SecureRnd();
-        this.upgraded = false;
-        this.closed = false;
     }
 
     @Override
@@ -211,6 +220,9 @@ public void connect(CharSequence host, int port) {
     /**
      * Disconnects the socket without closing the client.
      * The client can be reconnected by calling connect() again.
+     * <p>
+     * This method is NOT thread-safe. Only call it when no other thread
+     * is using this client (e.g., after the I/O thread has stopped).
      */
     public void disconnect() {
         Misc.free(socket);
@@ -222,6 +234,17 @@ public void disconnect() {
         resetFragmentState();
     }
 
+    /**
+     * Closes the socket to force-unblock a thread blocked in send/recv.
+     * <p>
+     * Unlike {@link #disconnect()}, this method only closes the socket
+     * and does not reset client state. It is safe to call from a different
+     * thread than the one performing I/O.
+     */
+    public void forceDisconnect() {
+        Misc.free(socket);
+    }
+
     /**
      * Returns the connected host.
      */
@@ -278,9 +301,9 @@ public boolean receiveFrame(WebSocketFrameHandler handler, int timeout) {
         checkConnected();
 
         // First, try to parse any data already in buffer
-        Boolean result = tryParseFrame(handler);
-        if (result != null) {
-            return result;
+        int result = tryParseFrame(handler);
+        if (result != PARSE_NEED_MORE) {
+            return result == PARSE_OK;
         }
 
         // Need more data
@@ -303,8 +326,8 @@ public boolean receiveFrame(WebSocketFrameHandler handler, int timeout) {
             recvPos += bytesRead;
 
             result = tryParseFrame(handler);
-            if (result != null) {
-                return result;
+            if (result != PARSE_NEED_MORE) {
+                return result == PARSE_OK;
             }
         }
     }
@@ -382,9 +405,9 @@ public boolean tryReceiveFrame(WebSocketFrameHandler handler) {
         checkConnected();
 
         // First, try to parse any data already in buffer
-        Boolean result = tryParseFrame(handler);
-        if (result != null) {
-            return result;
+        int result = tryParseFrame(handler);
+        if (result != PARSE_NEED_MORE) {
+            return result == PARSE_OK;
         }
 
         // Try one non-blocking recv
@@ -403,7 +426,7 @@ public boolean tryReceiveFrame(WebSocketFrameHandler handler) {
 
         // Try to parse again
         result = tryParseFrame(handler);
-        return result != null && result;
+        return result == PARSE_OK;
     }
 
     /**
@@ -800,17 +823,17 @@ private void sendPongFrame(long payloadPtr, int payloadLen) {
         }
     }
 
-    private Boolean tryParseFrame(WebSocketFrameHandler handler) {
+    private int tryParseFrame(WebSocketFrameHandler handler) {
         if (recvPos <= recvReadPos) {
-            return null; // No data
+            return PARSE_NEED_MORE;
         }
 
         frameParser.reset();
         int consumed = frameParser.parse(recvBufPtr + recvReadPos, recvBufPtr + recvPos);
 
         if (frameParser.getState() == WebSocketFrameParser.STATE_NEED_MORE ||
                 frameParser.getState() == WebSocketFrameParser.STATE_NEED_PAYLOAD) {
-            return null; // Need more data
+            return PARSE_NEED_MORE;
         }
 
         if (frameParser.getState() == WebSocketFrameParser.STATE_ERROR) {
@@ -909,10 +932,10 @@ private Boolean tryParseFrame(WebSocketFrameHandler handler) {
             // Compact buffer if needed
             compactRecvBuffer();
 
-            return true;
+            return PARSE_OK;
         }
 
-        return false;
+        return PARSE_INCOMPLETE;
     }
 
     private void validateUpgradeResponse(int headerEnd) {

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

@@ -165,13 +165,6 @@ public FrameInfo endFrame(int opcode) {
         return frameInfo.set(actualFrameStart, actualHeaderSize + payloadLen);
     }
 
-    /**
-     * Finishes the current text frame, writing the header and applying masking.
-     */
-    public FrameInfo endTextFrame() {
-        return endFrame(WebSocketOpcode.TEXT);
-    }
-
     /**
      * Ensures the buffer has capacity for the specified number of additional bytes.
      * May reallocate the buffer if necessary.
@@ -200,13 +193,6 @@ public int getCapacity() {
         return bufCapacity;
     }
 
-    /**
-     * Gets the payload length of the current frame being built.
-     */
-    public int getCurrentPayloadLength() {
-        return writePos - payloadStartOffset;
-    }
-
     /**
      * Gets the current write position (number of bytes written).
      */
@@ -215,6 +201,16 @@ public int getPosition() {
         return writePos;
     }
 
+    @Override
+    public long getWriteAddress() {
+        return bufPtr + writePos;
+    }
+
+    @Override
+    public int getWritableBytes() {
+        return bufCapacity - writePos;
+    }
+
     /**
      * Gets the current write position (total bytes written since last reset).
      */

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

@@ -107,10 +107,16 @@ public void close() {
      */
     @Override
     public void ensureCapacity(int needed) {
-        if (position + needed > capacity) {
-            int newCapacity = Math.max(capacity * 2, position + needed);
-            bufferPtr = Unsafe.realloc(bufferPtr, capacity, newCapacity, MemoryTag.NATIVE_DEFAULT);
-            capacity = newCapacity;
+        if ((long) position + needed > capacity) {
+            long required = (long) position + needed;
+            long doubled = (long) capacity * 2;
+            long newCapacity = Math.max(doubled, required);
+            if (newCapacity > Integer.MAX_VALUE) {
+                throw new OutOfMemoryError("NativeBufferWriter capacity overflow: " + newCapacity);
+            }
+            int cap = (int) newCapacity;
+            bufferPtr = Unsafe.realloc(bufferPtr, capacity, cap, MemoryTag.NATIVE_DEFAULT);
+            capacity = cap;
         }
     }
 
@@ -138,6 +144,16 @@ public int getPosition() {
         return position;
     }
 
+    @Override
+    public long getWriteAddress() {
+        return bufferPtr + position;
+    }
+
+    @Override
+    public int getWritableBytes() {
+        return capacity - position;
+    }
+
     /**
      * Patches an int value at the specified offset.
      * Used for updating length fields after writing content.
@@ -231,10 +247,35 @@ public void putString(String value) {
             return;
         }
 
-        int utf8Len = utf8Length(value);
-        putVarint(utf8Len);
-        ensureCapacity(utf8Len);
-        encodeUtf8(value);
+        int charLen = value.length();
+        // Optimistic: assume ASCII (utf8Len == charLen).
+        // Reserve varint(charLen) + charLen bytes.
+        int varintLen = varintSize(charLen);
+        ensureCapacity(varintLen + charLen);
+
+        // Single-pass: write ASCII bytes directly after varint space
+        long varintAddr = bufferPtr + position;
+        long addr = varintAddr + varintLen;
+        int i = 0;
+        for (; i < charLen; i++) {
+            char c = value.charAt(i);
+            if (c >= 0x80) {
+                break;
+            }
+            Unsafe.getUnsafe().putByte(addr++, (byte) c);
+        }
+
+        if (i == charLen) {
+            // All ASCII — write varint prefix, done in a single pass
+            writeVarintDirect(varintAddr, charLen);
+            position += varintLen + charLen;
+        } else {
+            // Non-ASCII — fall back to two-pass
+            int utf8Len = utf8Length(value);
+            putVarint(utf8Len);
+            ensureCapacity(utf8Len);
+            encodeUtf8(value);
+        }
     }
 
     /**
@@ -245,21 +286,45 @@ public void putUtf8(String value) {
         if (value == null || value.isEmpty()) {
             return;
         }
-        int utf8Len = utf8Length(value);
-        ensureCapacity(utf8Len);
-        encodeUtf8(value);
+
+        int charLen = value.length();
+        ensureCapacity(charLen);
+
+        // Single-pass: try ASCII encoding
+        long addr = bufferPtr + position;
+        int i = 0;
+        for (; i < charLen; i++) {
+            char c = value.charAt(i);
+            if (c >= 0x80) {
+                break;
+            }
+            Unsafe.getUnsafe().putByte(addr++, (byte) c);
+        }
+
+        if (i == charLen) {
+            // All ASCII — done in a single pass
+            position += charLen;
+        } else {
+            // Non-ASCII — fall back to two-pass (re-encodes from start)
+            int utf8Len = utf8Length(value);
+            ensureCapacity(utf8Len);
+            encodeUtf8(value);
+        }
     }
 
     /**
      * Writes a varint (unsigned LEB128).
      */
     @Override
     public void putVarint(long value) {
+        ensureCapacity(10); // max varint bytes
+        long addr = bufferPtr + position;
         while (value > 0x7F) {
-            putByte((byte) ((value & 0x7F) | 0x80));
+            Unsafe.getUnsafe().putByte(addr++, (byte) ((value & 0x7F) | 0x80));
             value >>>= 7;
         }
-        putByte((byte) value);
+        Unsafe.getUnsafe().putByte(addr++, (byte) value);
+        position = (int) (addr - bufferPtr);
     }
 
     /**
@@ -282,6 +347,14 @@ public void skip(int bytes) {
         position += bytes;
     }
 
+    private static void writeVarintDirect(long addr, long value) {
+        while (value > 0x7F) {
+            Unsafe.getUnsafe().putByte(addr++, (byte) ((value & 0x7F) | 0x80));
+            value >>>= 7;
+        }
+        Unsafe.getUnsafe().putByte(addr, (byte) value);
+    }
+
     private void encodeUtf8(String value) {
         long addr = bufferPtr + position;
         for (int i = 0, n = value.length(); i < n; i++) {

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

@@ -69,6 +69,23 @@ public interface QwpBufferWriter extends ArrayBufferAppender {
      */
     int getPosition();
 
+    /**
+     * Returns the native address where the next write will go.
+     * <p>
+     * Unlike {@code getBufferPtr() + getPosition()}, this method returns
+     * the correct write address for all buffer implementations, including
+     * segmented buffers where {@link #getPosition()} is a global offset.
+     * <p>
+     * The returned pointer is valid until the next buffer growth or flush.
+     */
+    long getWriteAddress();
+
+    /**
+     * Returns the number of bytes available for writing at
+     * {@link #getWriteAddress()}.
+     */
+    int getWritableBytes();
+
     /**
      * Patches an int value at the specified offset in the buffer.
      * <p>