small fixes

Author: bluestreak01

core/src/main/c/share/zstd_jni.c

@@ -43,6 +43,41 @@ JNIEXPORT jlong JNICALL Java_io_questdb_client_std_Zstd_createDCtx(
     return (jlong) (uintptr_t) ZSTD_createDCtx();
 }
 
+JNIEXPORT jlong JNICALL Java_io_questdb_client_std_Zstd_getFrameContentSize(
+        JNIEnv *env, jclass cls,
+        jlong src_addr, jlong src_len) {
+    /*
+     * Peeks the zstd frame header at src_addr to recover the declared
+     * uncompressed size. Returns:
+     *   positive  -- declared content size in bytes
+     *        -1   -- frame valid, content size not stored (ZSTD_CONTENTSIZE_UNKNOWN)
+     *        -2   -- invalid frame, truncated header, or size > INT64_MAX
+     *
+     * Lets the Java caller size the destination buffer in a single allocation
+     * instead of retrying decompress on dst-too-small. Crucially, it also lets
+     * a corrupt frame fail BEFORE any output buffer growth, eliminating a
+     * memory-amplification vector where one bad frame would have driven
+     * scratch growth all the way to the 64 MiB cap.
+     */
+    if (src_len < 0 || (src_len > 0 && src_addr == 0)) {
+        return -2;
+    }
+    unsigned long long size = ZSTD_getFrameContentSize(
+            (const void *) (uintptr_t) src_addr, (size_t) src_len);
+    if (size == ZSTD_CONTENTSIZE_UNKNOWN) {
+        return -1;
+    }
+    if (size == ZSTD_CONTENTSIZE_ERROR) {
+        return -2;
+    }
+    if (size > (unsigned long long) INT64_MAX) {
+        /* Cast to jlong would wrap to negative and look like an error code;
+         * reject upfront so the caller doesn't double-interpret. */
+        return -2;
+    }
+    return (jlong) size;
+}
+
 JNIEXPORT void JNICALL Java_io_questdb_client_std_Zstd_freeDCtx(
         JNIEnv *env, jclass cls, jlong ptr) {
     if (ptr != 0) {

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

@@ -85,9 +85,16 @@ public class QwpEgressIoThread implements Runnable, WebSocketFrameHandler {
     // Per-query credit state (accessed only from the I/O thread).
     // creditEnabled == (initialCredit > 0); controls whether we emit CREDIT
     // replenish frames after each batch release.
+    // Set true by closePool() before it drains freeBuffers / events, so any
+    // user thread still inside releaseBuffer (or one that returns from onBatch
+    // after close has run) sees closed == true and frees the buffer in place
+    // instead of stranding it in a queue nobody will iterate again.
+    // Volatile because the writer (close-caller thread) and reader (user
+    // thread inside releaseBuffer) are different threads.
+    private volatile boolean closed;
     private boolean creditEnabled;
-    private long currentRequestId = -1L;
     private boolean currentQueryDone;
+    private long currentRequestId = -1L;
     private volatile boolean shutdown;
 
     public QwpEgressIoThread(WebSocketClient wsClient, int bufferPoolSize) {
@@ -97,8 +104,21 @@ public QwpEgressIoThread(WebSocketClient wsClient, int bufferPoolSize) {
         // frame the I/O thread may emit on shutdown, so a full buffer pool
         // never stalls the producer.
         this.events = new QwpSpscQueue<>(bufferPoolSize + 2);
-        for (int i = 0; i < bufferPoolSize; i++) {
-            freeBuffers.offer(new QwpBatchBuffer(DEFAULT_BUFFER_CAPACITY));
+        // Track allocations as we go so a partial-pool failure (e.g. native
+        // OOM at iteration K) frees the K-1 buffers already in freeBuffers.
+        // Without this, the half-built QwpEgressIoThread instance escapes the
+        // failed constructor and is unreachable from QwpQueryClient.close(),
+        // leaking those buffers' native scratches for the JVM's lifetime.
+        try {
+            for (int i = 0; i < bufferPoolSize; i++) {
+                freeBuffers.offer(new QwpBatchBuffer(DEFAULT_BUFFER_CAPACITY));
+            }
+        } catch (Throwable t) {
+            for (QwpBatchBuffer b : freeBuffers) {
+                b.close();
+            }
+            freeBuffers.clear();
+            throw t;
         }
     }
 
@@ -175,6 +195,15 @@ public void onClose(int code, String reason) {
      * compact past the consumed frame.
      */
     public void releaseBuffer(QwpBatchBuffer buffer) {
+        if (closed) {
+            // closePool already drained and cleared freeBuffers; offering this
+            // buffer back into the pool would strand it (no consumer left to
+            // close it). Free its native scratch in place. The pendingRelease
+            // queue is also abandoned -- the I/O thread that would have read
+            // the token is gone.
+            buffer.close();
+            return;
+        }
         freeBuffers.offer(buffer);
         pendingRelease.offer(RELEASE_TOKEN);
     }
@@ -289,13 +318,22 @@ private void decodeAndEmitExecDone(long payload, int payloadLen) {
 
     /**
      * RESULT_END body: msg_kind(1) + requestId(8) + final_seq(varint) + total_rows(varint).
-     * We only need total_rows.
+     * We only need total_rows. Both varint loops cap at 10 bytes so a hostile
+     * server cannot drive {@code shift} past 63 (where Java's {@code <<} masks
+     * the count to 6 bits and silently wraps high bytes back into the low bits,
+     * producing a wildly wrong total).
      */
     private long decodeResultEnd(long payload, int payloadLen) {
         long p = payload + QwpConstants.HEADER_SIZE + 1 + 8;
         long limit = payload + payloadLen;
+        int seqBytes = 0;
         while (p < limit && (Unsafe.getUnsafe().getByte(p++) & 0x80) != 0) {
-            // skip final_seq continuation bytes
+            if (++seqBytes > 9) {
+                // Continuation bit set on byte 10 of the final_seq varint --
+                // malformed. Surface a 0 total rather than read into the
+                // total_rows section with a desynced cursor.
+                return 0L;
+            }
         }
         long total = 0;
         int shift = 0;
@@ -304,6 +342,12 @@ private long decodeResultEnd(long payload, int payloadLen) {
             total |= (long) (b & 0x7F) << shift;
             if ((b & 0x80) == 0) break;
             shift += 7;
+            if (shift > 63) {
+                // Same overflow guard decodeAndEmitExecDone uses; without it,
+                // byte 10 of total_rows could land in the sign bit and bytes
+                // 11+ would silently wrap.
+                return 0L;
+            }
         }
         return total;
     }
@@ -394,26 +438,26 @@ private void handleResultBatch(long payloadPtr, int payloadLen) {
     }
 
     /**
-     * Builds and transmits a CREDIT frame on the WebSocket. Wire format:
-     * {@code msg_kind(0x15) + request_id(u64) + additional_bytes(varint)}.
+     * Builds and transmits a CANCEL frame on the WebSocket. Wire format:
+     * {@code msg_kind(0x14) + request_id(u64)}.
      */
-    private void sendCredit(long requestId, long additionalBytes) {
+    private void sendCancel(long requestId) {
         sendScratch.reset();
-        sendScratch.putByte(QwpEgressMsgKind.CREDIT);
+        sendScratch.putByte(QwpEgressMsgKind.CANCEL);
         sendScratch.putLong(requestId);
-        sendScratch.putVarint(additionalBytes);
         wsClient.sendBinary(sendScratch.getBufferPtr(), sendScratch.getPosition());
         sendScratch.reset();
     }
 
     /**
-     * Builds and transmits a CANCEL frame on the WebSocket. Wire format:
-     * {@code msg_kind(0x14) + request_id(u64)}.
+     * Builds and transmits a CREDIT frame on the WebSocket. Wire format:
+     * {@code msg_kind(0x15) + request_id(u64) + additional_bytes(varint)}.
      */
-    private void sendCancel(long requestId) {
+    private void sendCredit(long requestId, long additionalBytes) {
         sendScratch.reset();
-        sendScratch.putByte(QwpEgressMsgKind.CANCEL);
+        sendScratch.putByte(QwpEgressMsgKind.CREDIT);
         sendScratch.putLong(requestId);
+        sendScratch.putVarint(additionalBytes);
         wsClient.sendBinary(sendScratch.getBufferPtr(), sendScratch.getPosition());
         sendScratch.reset();
     }
@@ -450,6 +494,12 @@ private void sendQueryRequest(QueryRequest req) {
      * forever on an empty queue.
      */
     void closePool() {
+        // Set closed BEFORE draining so a user thread that returns from
+        // onBatch concurrently and calls releaseBuffer sees closed=true and
+        // frees the buffer in place rather than offering it into a queue
+        // we're about to clear. Volatile write pairs with the volatile read
+        // in releaseBuffer.
+        closed = true;
         Misc.free(sendScratch);
         Misc.free(decoder);
         QueryEvent ev;

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

@@ -93,20 +93,20 @@ public class QwpQueryClient implements QuietCloseable {
     // payload before it parks, and the client auto-replenishes by the size of
     // each batch as the user releases it.
     private long initialCreditBytes;
-    // Client preference for server-side per-batch row cap. 0 means "unset",
-    // server uses its default. Set via {@code max_batch_rows=N} in the
-    // connection string or {@link #withMaxBatchRows}. Smaller values give
-    // streaming consumers earlier access to the first rows at the cost of
-    // more per-batch overhead; larger values amortise fixed costs over more
-    // rows. Server may clamp down to its own hard cap.
-    private int maxBatchRows;
     // Volatile so a cancel() call from a thread other than the one that ran
     // connect() sees the published reference (and a concurrent null-out from
     // close() is observed without a stale-reference race). The thread-safety
     // contract documented on cancel() relies on this.
     private volatile QwpEgressIoThread ioThread;
     private Thread ioThreadHandle;
     private boolean lastCloseTimedOut;
+    // Client preference for server-side per-batch row cap. 0 means "unset",
+    // server uses its default. Set via {@code max_batch_rows=N} in the
+    // connection string or {@link #withMaxBatchRows}. Smaller values give
+    // streaming consumers earlier access to the first rows at the cost of
+    // more per-batch overhead; larger values amortise fixed costs over more
+    // rows. Server may clamp down to its own hard cap.
+    private int maxBatchRows;
     private int negotiatedQwpVersion;
     private long nextRequestId = 1;
     // Maximum time close() will wait for the I/O thread to exit before giving up
@@ -388,6 +388,17 @@ public void cancel() {
      * resources leak for the lifetime of the process. A warning is recorded by setting
      * {@link #lastCloseTimedOut} (queryable via {@link #wasLastCloseTimedOut}) so callers
      * can detect and report the condition.
+     * <p>
+     * <strong>Threading contract:</strong> {@code close()} must be called from a thread
+     * other than the one currently inside a batch handler, AND the user must finish
+     * any in-flight {@code execute()} before calling {@code close()}. Calling
+     * {@code close()} concurrently with a {@code handler.onBatch(...)} that is still
+     * dereferencing batch column pointers can free the WebSocket recv buffer under
+     * those pointers and SIGSEGV the JVM. The interrupt-driven I/O thread shutdown
+     * does NOT detect or wait for a still-running user handler -- the timeout-based
+     * leak fallback only protects against an unresponsive I/O thread, not an
+     * unresponsive user thread. {@link #cancel()} is the right way to ask an
+     * in-flight {@code execute()} to return; close after it does.
      */
     @Override
     public void close() {
@@ -747,11 +758,9 @@ private String buildAcceptEncodingHeader() {
      * the caller doesn't inherit a half-open socket.
      */
     private void probeZstdAvailable() {
+        long dctx;
         try {
-            long dctx = Zstd.createDCtx();
-            if (dctx != 0) {
-                Zstd.freeDCtx(dctx);
-            }
+            dctx = Zstd.createDCtx();
         } catch (UnsatisfiedLinkError e) {
             LOG.error("zstd JNI symbols missing from libquestdb; aborting connect", e);
             if (webSocketClient != null) {
@@ -764,5 +773,20 @@ private void probeZstdAvailable() {
                     + "compression=raw on the connection string to skip the probe. "
                     + "[cause=" + e.getMessage() + "]");
         }
+        if (dctx == 0) {
+            // Native createDCtx returned 0 -- the underlying ZSTD_createDCtx failed
+            // (typically allocator pressure). Don't proceed: a connection that
+            // negotiated zstd but cannot decode it would surface as a confusing
+            // mid-stream decode error on the first compressed batch instead.
+            LOG.error("zstd createDCtx returned 0 (native allocation failure); aborting connect");
+            if (webSocketClient != null) {
+                webSocketClient.close();
+                webSocketClient = null;
+            }
+            throw new HttpClientException("zstd decompression context allocation failed; "
+                    + "cannot accept compressed batches. Set compression=raw on the connection "
+                    + "string to disable compression, or retry once memory pressure subsides.");
+        }
+        Zstd.freeDCtx(dctx);
     }
 }