Bound the TLS handshake and drive it off the event loop

JavaTlsClientSocket.startTlsSession ran the TLS handshake with raw
delegate.recv/delegate.send on a non-blocking socket and never waited on
socket readiness. A peer that completed the TCP connect but stalled
before its half of the handshake left the engine in NEED_UNWRAP with
recv returning 0 (would-block), so the loop re-read in a tight cycle: a
100% CPU busy-spin with no deadline. connect_timeout bounded only the TCP
connect, and the WebSocket upgrade's auth/request timeout never covered
the handshake, so a stalled wss:// (e.g. QuestDB Cloud) handshake could
pin a core indefinitely and defeat a bounded connect.

Drive the handshake through the client's existing deadline-aware ioWait,
the same primitive recvOrDie/doSend already use. When the socket would
block, the handshake hands control to a SocketReadinessWaiter that parks
on epoll/kqueue/select for the remaining connect budget and throws a
timeout-flagged exception once it is spent. This removes the spin and
bounds the handshake in one change: both NEED_UNWRAP (recv == 0) and the
NEED_WRAP send loop (send == 0) now wait instead of spinning.

doConnect now calls setupIoWait() before the handshake (so the fd is
registered when the waiter parks) and bounds the handshake by
connect_timeout, falling back to the request timeout when connect_timeout
is unset, so the handshake can no longer hang or spin even with the
default config. The connect TLS block disconnects on any handshake error
(including the waiter's timeout) so the fd and native buffers do not leak.
Both HttpClient (https) and WebSocketClient (wss) connect paths share the
fix.

Extracted the handshake loop into runHandshake(waiter) so a stub
SSLEngine can exercise the wait paths. Added two tests:
testHandshakeWaitsForReadabilityInsteadOfBusySpinning drives a stalled
peer and asserts the handshake yields to the waiter exactly once (a
method-level timeout fails the test if the spin returns), and
testHandshakeCompletesWithoutWaitingWhenEngineMakesProgress guards the
happy path.

Author: bluestreak01

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

@@ -639,19 +639,36 @@ private void connect(CharSequence host, int port) {
                 throw new HttpClientException("could not configure socket to be non-blocking [fd=").put(fd).put(", errno=").put(errno).put(']');
             }
 
+            // Register the fd with the event loop before the TLS handshake so the
+            // handshake can park on socket readiness via ioWait() instead of
+            // busy-spinning on the non-blocking socket.
+            setupIoWait();
+
             if (socket.supportsTls()) {
+                // Bound the TLS handshake by the connect budget (falling back to
+                // the request timeout when connect_timeout is unset), so a peer
+                // that completes TCP but stalls mid-handshake cannot hang or pin a
+                // CPU.
+                final long tlsHandshakeStartNanos = System.nanoTime();
+                final int tlsHandshakeBudgetMillis = connectTimeout > 0 ? connectTimeout : defaultTimeout;
                 try {
-                    socket.startTlsSession(host);
+                    socket.startTlsSession(host, op -> ioWait(remainingTime(tlsHandshakeBudgetMillis, tlsHandshakeStartNanos), op));
                 } catch (TlsSessionInitFailedException e) {
                     int errno = nf.errno();
                     disconnect();
                     throw new HttpClientException("could not start TLS session [fd=").put(fd)
                             .put(", error=").put(e.getFlyweightMessage())
                             .put(", errno=").put(errno)
                             .put(']');
+                } catch (Throwable t) {
+                    // ioWait() throws a timeout-flagged HttpClientException when the
+                    // handshake budget is exhausted; any other error can also surface
+                    // mid-handshake. Disconnect so the fd and native buffers do not
+                    // leak, then propagate.
+                    disconnect();
+                    throw t;
                 }
             }
-            setupIoWait();
         }
 
         private void doSend(long lo, long hi, int timeoutMillis) {

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

@@ -962,19 +962,35 @@ private void doConnect(CharSequence host, int port) {
                     .put(", errno=").put(errno).put(']');
         }
 
+        // Register the fd with the event loop before the TLS handshake so the
+        // handshake can park on socket readiness via ioWait() instead of
+        // busy-spinning on the non-blocking socket.
+        setupIoWait();
+
         if (socket.supportsTls()) {
+            // Bound the TLS handshake by the connect budget (falling back to the
+            // request timeout when connect_timeout is unset), so a peer that
+            // completes TCP but stalls mid-handshake cannot hang or pin a CPU.
+            final long tlsHandshakeStartNanos = System.nanoTime();
+            final int tlsHandshakeBudgetMillis = connectTimeoutMillis > 0 ? connectTimeoutMillis : defaultTimeout;
             try {
-                socket.startTlsSession(host);
+                socket.startTlsSession(host, op -> ioWait(getRemainingTimeOrThrow(tlsHandshakeBudgetMillis, tlsHandshakeStartNanos), op));
             } catch (TlsSessionInitFailedException e) {
                 int errno = nf.errno();
                 disconnect();
                 throw new HttpClientException("could not start TLS session [fd=").put(fd)
                         .put(", error=").put(e.getFlyweightMessage())
                         .put(", errno=").put(errno).put(']');
+            } catch (Throwable t) {
+                // ioWait() throws a timeout-flagged HttpClientException when the
+                // handshake budget is exhausted; any other error can also surface
+                // mid-handshake. Disconnect so the fd and native buffers do not
+                // leak, then propagate.
+                disconnect();
+                throw t;
             }
         }
 
-        setupIoWait();
         if (LOG.isDebugEnabled()) {
             LOG.debug("Connected to [host={}, port={}]", host, port);
         }

core/src/main/java/io/questdb/client/network/JavaTlsClientSocket.java

@@ -307,91 +307,13 @@ public int send(long bufferPtr, int bufferLen) {
     }
 
     @Override
-    public void startTlsSession(CharSequence peerName) throws TlsSessionInitFailedException {
+    public void startTlsSession(CharSequence peerName, SocketReadinessWaiter waiter) throws TlsSessionInitFailedException {
         assert state == STATE_PLAINTEXT;
         prepareInternalBuffers();
         try {
             this.sslEngine = createSslEngine(peerName);
             this.sslEngine.beginHandshake();
-            SSLEngineResult.HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
-            while (handshakeStatus != SSLEngineResult.HandshakeStatus.FINISHED) {
-                switch (handshakeStatus) {
-                    case NEED_TASK:
-                        Runnable task;
-                        while ((task = sslEngine.getDelegatedTask()) != null) {
-                            task.run();
-                        }
-                        handshakeStatus = sslEngine.getHandshakeStatus();
-                        break;
-                    case NEED_WRAP: {
-                        SSLEngineResult result = sslEngine.wrap(wrapInputBuffer, wrapOutputBuffer);
-                        handshakeStatus = result.getHandshakeStatus();
-                        switch (result.getStatus()) {
-                            case BUFFER_UNDERFLOW:
-                                // there cannot be underflow since wrap() during handshake does not read from the input buffer at all
-                                throw new AssertionError("Buffer underflow during TLS handshake. This should not happen. please report as a bug");
-                            case BUFFER_OVERFLOW:
-                                if (wrapOutputBuffer.position() != 0) {
-                                    // wrap() left bytes behind without producing a complete record. The OK
-                                    // branch is the only place that drains and clears, so a non-empty
-                                    // buffer here means we would re-enter NEED_WRAP with identical state
-                                    // and spin forever. Fail loudly instead.
-                                    throw new AssertionError("Buffer overflow during TLS handshake with non-empty output buffer. This should not happen, please report as a bug");
-                                }
-                                // in theory, this can happen if the output buffer is too small to fit a single TLS handshake record,
-                                // but that would indicate our starting buffer is too small.
-                                growWrapOutputBuffer();
-                                break;
-                            case OK:
-                                // wrapOutputBuffer: write mode
-                                int written = 0;
-                                int bufferLimit = wrapOutputBuffer.position();
-                                while (written < bufferLimit) {
-                                    int n = delegate.send(wrapOutputBufferPtr + written, bufferLimit - written);
-                                    if (n < 0) {
-                                        throw TlsSessionInitFailedException.instance("socket write error");
-                                    }
-                                    written += n;
-                                }
-                                wrapOutputBuffer.clear();
-                                break;
-                            case CLOSED:
-                                throw TlsSessionInitFailedException.instance("server closed connection unexpectedly");
-                        }
-                        break;
-                    }
-                    case NEED_UNWRAP: {
-                        int n = readFromSocket();
-                        if (n < 0) {
-                            throw TlsSessionInitFailedException.instance("socket read error");
-                        }
-                        SSLEngineResult result = sslEngine.unwrap(unwrapInputBuffer, unwrapOutputBuffer);
-                        handshakeStatus = result.getHandshakeStatus();
-                        switch (result.getStatus()) {
-                            case BUFFER_UNDERFLOW:
-                                // we need to receive more data from a socket, let's try again
-                                break;
-                            case BUFFER_OVERFLOW:
-                                if (unwrapOutputBuffer.position() != 0) {
-                                    // unwrap() produced plaintext but signalled overflow without consuming
-                                    // the next record. Nothing in the handshake loop drains this buffer,
-                                    // so re-entering NEED_UNWRAP would spin forever. Fail loudly.
-                                    throw new AssertionError("Buffer overflow during TLS handshake with non-empty output buffer. This should not happen, please report as a bug");
-                                }
-                                // in theory, this can happen if the output buffer is too small to fit a single TLS handshake record,
-                                // but that would indicate our starting buffer is too small.
-                                growUnwrapOutputBuffer();
-                                break;
-                            case OK:
-                                // good, let's see what we need to do next
-                                break;
-                            case CLOSED:
-                                throw TlsSessionInitFailedException.instance("server closed connection unexpectedly");
-                        }
-                    }
-                    break;
-                }
-            }
+            runHandshake(waiter);
             // unwrap input buffer: read mode and empty
             unwrapInputBuffer.position(0);
             unwrapInputBuffer.limit(0);
@@ -583,6 +505,108 @@ private int readFromSocket() {
         return n;
     }
 
+    /**
+     * Drives the TLS handshake state machine to completion. When the
+     * non-blocking socket would block, hands control to {@code waiter} (which
+     * parks on the event loop bounded by the connect deadline) instead of
+     * busy-spinning on read/write. Extracted from {@link #startTlsSession} so a
+     * stub {@code sslEngine} can exercise the wait paths in isolation.
+     */
+    private void runHandshake(SocketReadinessWaiter waiter) throws SSLException, TlsSessionInitFailedException {
+        SSLEngineResult.HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
+        while (handshakeStatus != SSLEngineResult.HandshakeStatus.FINISHED) {
+            switch (handshakeStatus) {
+                case NEED_TASK:
+                    Runnable task;
+                    while ((task = sslEngine.getDelegatedTask()) != null) {
+                        task.run();
+                    }
+                    handshakeStatus = sslEngine.getHandshakeStatus();
+                    break;
+                case NEED_WRAP: {
+                    SSLEngineResult result = sslEngine.wrap(wrapInputBuffer, wrapOutputBuffer);
+                    handshakeStatus = result.getHandshakeStatus();
+                    switch (result.getStatus()) {
+                        case BUFFER_UNDERFLOW:
+                            // there cannot be underflow since wrap() during handshake does not read from the input buffer at all
+                            throw new AssertionError("Buffer underflow during TLS handshake. This should not happen. please report as a bug");
+                        case BUFFER_OVERFLOW:
+                            if (wrapOutputBuffer.position() != 0) {
+                                // wrap() left bytes behind without producing a complete record. The OK
+                                // branch is the only place that drains and clears, so a non-empty
+                                // buffer here means we would re-enter NEED_WRAP with identical state
+                                // and spin forever. Fail loudly instead.
+                                throw new AssertionError("Buffer overflow during TLS handshake with non-empty output buffer. This should not happen, please report as a bug");
+                            }
+                            // in theory, this can happen if the output buffer is too small to fit a single TLS handshake record,
+                            // but that would indicate our starting buffer is too small.
+                            growWrapOutputBuffer();
+                            break;
+                        case OK:
+                            // wrapOutputBuffer: write mode
+                            int written = 0;
+                            int bufferLimit = wrapOutputBuffer.position();
+                            while (written < bufferLimit) {
+                                int n = delegate.send(wrapOutputBufferPtr + written, bufferLimit - written);
+                                if (n < 0) {
+                                    throw TlsSessionInitFailedException.instance("socket write error");
+                                }
+                                if (n == 0) {
+                                    // The non-blocking socket's send buffer is full. Wait for it to
+                                    // become writable -- bounded by the connect deadline -- instead of
+                                    // busy-spinning on send().
+                                    waiter.awaitReady(IOOperation.WRITE);
+                                }
+                                written += n;
+                            }
+                            wrapOutputBuffer.clear();
+                            break;
+                        case CLOSED:
+                            throw TlsSessionInitFailedException.instance("server closed connection unexpectedly");
+                    }
+                    break;
+                }
+                case NEED_UNWRAP: {
+                    int n = readFromSocket();
+                    if (n < 0) {
+                        throw TlsSessionInitFailedException.instance("socket read error");
+                    }
+                    SSLEngineResult result = sslEngine.unwrap(unwrapInputBuffer, unwrapOutputBuffer);
+                    handshakeStatus = result.getHandshakeStatus();
+                    switch (result.getStatus()) {
+                        case BUFFER_UNDERFLOW:
+                            // Not enough bytes for a complete TLS record yet. If the last read
+                            // drained the socket (n == 0, would-block on the non-blocking fd), wait
+                            // for it to become readable -- bounded by the connect deadline -- instead
+                            // of busy-spinning. A positive n means we read a partial record, so loop
+                            // immediately and read the rest.
+                            if (n == 0) {
+                                waiter.awaitReady(IOOperation.READ);
+                            }
+                            break;
+                        case BUFFER_OVERFLOW:
+                            if (unwrapOutputBuffer.position() != 0) {
+                                // unwrap() produced plaintext but signalled overflow without consuming
+                                // the next record. Nothing in the handshake loop drains this buffer,
+                                // so re-entering NEED_UNWRAP would spin forever. Fail loudly.
+                                throw new AssertionError("Buffer overflow during TLS handshake with non-empty output buffer. This should not happen, please report as a bug");
+                            }
+                            // in theory, this can happen if the output buffer is too small to fit a single TLS handshake record,
+                            // but that would indicate our starting buffer is too small.
+                            growUnwrapOutputBuffer();
+                            break;
+                        case OK:
+                            // good, let's see what we need to do next
+                            break;
+                        case CLOSED:
+                            throw TlsSessionInitFailedException.instance("server closed connection unexpectedly");
+                    }
+                }
+                break;
+            }
+        }
+    }
+
     private int writeToSocket(int bytesToSend) {
         // wrapOutputBuffer is in the write mode
         int n = delegate.send(wrapOutputBufferPtr, bytesToSend);

core/src/main/java/io/questdb/client/network/PlainSocket.java

@@ -71,7 +71,7 @@ public int send(long bufferPtr, int bufferLen) {
     }
 
     @Override
-    public void startTlsSession(CharSequence peerName) {
+    public void startTlsSession(CharSequence peerName, SocketReadinessWaiter waiter) {
         throw new UnsupportedOperationException();
     }
 

core/src/main/java/io/questdb/client/network/Socket.java

@@ -84,9 +84,12 @@ public interface Socket extends QuietCloseable {
      * on server connections.
      *
      * @param peerName server name to use for SNI and certificate validation.
+     * @param waiter   blocks until the socket is ready for the next handshake
+     *                 read/write (bounded by the connect deadline), so the
+     *                 handshake does not busy-spin on the non-blocking socket.
      * @throws TlsSessionInitFailedException if the call fails.
      */
-    void startTlsSession(@Nullable CharSequence peerName) throws TlsSessionInitFailedException;
+    void startTlsSession(@Nullable CharSequence peerName, SocketReadinessWaiter waiter) throws TlsSessionInitFailedException;
 
     /**
      * @return true if the socket support TLS encryption; false otherwise.

core/src/main/java/io/questdb/client/network/SocketReadinessWaiter.java

@@ -0,0 +1,46 @@
+/*+*****************************************************************************
+ *     ___                  _   ____  ____
+ *    / _ \ _   _  ___  ___| |_|  _ \| __ )
+ *   | | | | | | |/ _ \/ __| __| | | |  _ \
+ *   | |_| | |_| |  __/\__ \ |_| |_| | |_) |
+ *    \__\_\\__,_|\___||___/\__|____/|____/
+ *
+ *  Copyright (c) 2014-2019 Appsicle
+ *  Copyright (c) 2019-2026 QuestDB
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ ******************************************************************************/
+
+package io.questdb.client.network;
+
+/**
+ * Blocks until a non-blocking socket is ready for a given I/O operation, or
+ * throws a timeout-flagged exception once the caller's deadline passes.
+ * <p>
+ * Used to drive the TLS handshake off the client's event loop: instead of
+ * busy-spinning on a non-blocking socket that returns "would block", the
+ * handshake hands control to this waiter, which parks on epoll/kqueue/select
+ * with the remaining connect budget. This bounds the handshake by the same
+ * deadline as the TCP connect and keeps a stalled peer from pinning a CPU.
+ */
+@FunctionalInterface
+public interface SocketReadinessWaiter {
+    /**
+     * Blocks until the socket is ready for {@code ioOperation}, or throws a
+     * timeout-flagged exception when the connect deadline is exceeded.
+     *
+     * @param ioOperation {@link IOOperation#READ} or {@link IOOperation#WRITE}
+     */
+    void awaitReady(int ioOperation);
+}