Offload TLS negotiation to I/O threads

Signed-off-by: Uri Yagelnik <uriy@amazon.com>

src/connection.h

@@ -56,6 +56,7 @@ typedef enum {
 
 #define CONN_FLAG_CLOSE_SCHEDULED (1 << 0) /* Closed scheduled by a handler */
 #define CONN_FLAG_WRITE_BARRIER (1 << 1)   /* Write barrier requested */
+#define CONN_FLAG_CLIENT (1 << 2)          /* Connection is of a client - not a cluster link. */
 
 #define CONN_TYPE_SOCKET "tcp"
 #define CONN_TYPE_UNIX "unix"

src/io_threads.c

@@ -9,6 +9,7 @@
 static __thread int thread_id = 0; /* Thread local var */
 static pthread_t io_threads[IO_THREADS_MAX_NUM] = {0};
 static pthread_mutex_t io_threads_mutex[IO_THREADS_MAX_NUM];
+void (*tls_negotiation_cb)(void *);
 
 /* IO jobs queue functions - Used to send jobs from the main-thread to the IO thread. */
 typedef void (*job_handler)(void *);
@@ -554,3 +555,56 @@ void trySendPollJobToIOThreads(void) {
     aeSetPollProtect(server.el, 1);
     IOJobQueue_push(jq, IOThreadPoll, server.el);
 }
+
+void setTLSNegotiationCallback(void (*cb)(void *)) {
+    tls_negotiation_cb = cb;
+}
+
+static void ioThreadTLSNegotiation(void *data) {
+    client *c = (client *)data;
+    tls_negotiation_cb(c->conn);
+    c->io_read_state = CLIENT_COMPLETED_IO;
+}
+
+/*
+ * This function attempts to offload TLS negotiation for a client connection to an I/O thread.
+ * Returns C_OK if the TLS negotiation was successfully queued for processing by an I/O thread,
+ * or C_ERR if the client is not eligible for offloading.
+ * Parameters:
+ *   conn: The connection object for which TLS negotiation should be performed
+ */
+int trySendTLSNegotiationToIOThreads(connection *conn) {
+    if (server.io_threads_num <= 1) {
+        return C_ERR;
+    }
+
+    if (!(conn->flags & CONN_FLAG_CLIENT)) {
+        return C_ERR;
+    }
+
+    client *c = connGetPrivateData(conn);
+    if (c->io_read_state != CLIENT_IDLE) {
+        return C_OK;
+    }
+
+    if (server.active_io_threads_num <= 1) {
+        return C_ERR;
+    }
+
+    size_t thread_id = (c->id % (server.active_io_threads_num - 1)) + 1;
+    IOJobQueue *job_queue = &io_jobs[thread_id];
+
+    if (IOJobQueue_isFull(job_queue)) {
+        return C_ERR;
+    }
+
+    c->read_flags = READ_FLAGS_TLS_NEGOTIATION;
+    c->io_read_state = CLIENT_PENDING_IO;
+    c->flag.pending_read = 1;
+    listLinkNodeTail(server.clients_pending_io_read, &c->pending_read_list_node);
+    connSetPostponeUpdateState(c->conn, 1);
+    server.stat_io_tls_negotiation_offloaded++;
+    IOJobQueue_push(job_queue, ioThreadTLSNegotiation, c);
+
+    return C_OK;
+}

src/io_threads.h

@@ -13,5 +13,6 @@ int tryOffloadFreeArgvToIOThreads(client *c);
 void adjustIOThreadsByEventLoad(int numevents, int increase_only);
 void drainIOThreadsQueue(void);
 void trySendPollJobToIOThreads(void);
-
+int trySendTLSNegotiationToIOThreads(connection *conn);
+void setTLSNegotiationCallback(void (*cb)(void *));
 #endif /* IO_THREADS_H */

src/networking.c

@@ -134,6 +134,7 @@ client *createClient(connection *conn) {
         if (server.tcpkeepalive) connKeepAlive(conn, server.tcpkeepalive);
         connSetReadHandler(conn, readQueryFromClient);
         connSetPrivateData(conn, c);
+        conn->flags |= CONN_FLAG_CLIENT;
     }
     c->buf = zmalloc_usable(PROTO_REPLY_CHUNK_BYTES, &c->buf_usable_size);
     selectDb(c, 0);
@@ -4725,6 +4726,9 @@ int processIOThreadsReadDone(void) {
         connSetPostponeUpdateState(c->conn, 0);
         connUpdateState(c->conn);
 
+        /* No client's data was read only TLS handshake. */
+        if (c->read_flags & READ_FLAGS_TLS_NEGOTIATION) continue;
+
         /* On read error - stop here. */
         if (handleReadResult(c) == C_ERR) {
             continue;

src/server.c

@@ -2604,6 +2604,7 @@ void resetServerStats(void) {
     server.stat_total_reads_processed = 0;
     server.stat_io_writes_processed = 0;
     server.stat_io_freed_objects = 0;
+    server.stat_io_tls_negotiation_offloaded = 0;
     server.stat_poll_processed_by_io_threads = 0;
     server.stat_total_writes_processed = 0;
     server.stat_client_qbuf_limit_disconnections = 0;
@@ -5862,6 +5863,7 @@ sds genValkeyInfoString(dict *section_dict, int all_sections, int everything) {
                 "io_threaded_reads_processed:%lld\r\n", server.stat_io_reads_processed,
                 "io_threaded_writes_processed:%lld\r\n", server.stat_io_writes_processed,
                 "io_threaded_freed_objects:%lld\r\n", server.stat_io_freed_objects,
+                "io_threaded_tls_negotiations:%lld\r\n", server.stat_io_tls_negotiation_offloaded,
                 "io_threaded_poll_processed:%lld\r\n", server.stat_poll_processed_by_io_threads,
                 "io_threaded_total_prefetch_batches:%lld\r\n", server.stat_total_prefetch_batches,
                 "io_threaded_total_prefetch_entries:%lld\r\n", server.stat_total_prefetch_entries,

src/server.h

@@ -1841,6 +1841,7 @@ struct valkeyServer {
     long long stat_io_reads_processed;                 /* Number of read events processed by IO threads */
     long long stat_io_writes_processed;                /* Number of write events processed by IO threads */
     long long stat_io_freed_objects;                   /* Number of objects freed by IO threads */
+    long long stat_io_tls_negotiation_offloaded;       /* Number of TLS negotiation offloads */
     long long stat_poll_processed_by_io_threads;       /* Total number of poll jobs processed by IO */
     long long stat_total_reads_processed;              /* Total number of read events processed */
     long long stat_total_writes_processed;             /* Total number of write events processed */
@@ -2767,6 +2768,7 @@ void dictVanillaFree(void *val);
 #define READ_FLAGS_PRIMARY (1 << 14)
 #define READ_FLAGS_DONT_PARSE (1 << 15)
 #define READ_FLAGS_AUTH_REQUIRED (1 << 16)
+#define READ_FLAGS_TLS_NEGOTIATION (1 << 17)
 
 /* Write flags for various write errors and states */
 #define WRITE_FLAGS_WRITE_ERROR (1 << 0)