remove client id method

Signed-off-by: Ada-Church-Closure <[email protected]>

Author: Ada-Church-Closure

src/connection.h

@@ -511,13 +511,10 @@ static inline int connIsTLS(connection *conn) {
     return conn && conn->type == connectionTypeTls();
 }
 
-/* Return 1 on state updated, 0 if otherwise. */
-static inline int connUpdateState(connection *conn) {
+static inline void connUpdateState(connection *conn) {
     if (conn->type->update_state) {
         conn->type->update_state(conn);
-        return 1;
     }
-    return 0;
 }
 
 static inline void connSetPostponeUpdateState(connection *conn, int on) {

src/memory_prefetch.c

@@ -9,8 +9,10 @@
  */
 
 #include "memory_prefetch.h"
+#include "adlist.h"
 #include "server.h"
 #include "io_threads.h"
+#include <stdlib.h>
 
 typedef enum {
     PREFETCH_ENTRY, /* Initial state, prefetch entries associated with the given key's hash */
@@ -214,7 +216,7 @@ static void prefetchCommands(void) {
 }
 
 /* Processes all the prefetched commands in the current batch. */
-void processClientsCommandsBatch(void) {
+void processClientsCommandsBatch(list *handled_clients) {
     if (!batch || batch->client_count == 0) return;
 
     /* If executed_commands is not 0,
@@ -231,7 +233,10 @@ void processClientsCommandsBatch(void) {
         /* Set the client to null immediately to avoid accessing it again recursively when ProcessingEventsWhileBlocked */
         batch->clients[i] = NULL;
         batch->executed_commands++;
-        if (processPendingCommandAndInputBuffer(c) != C_ERR) beforeNextClient(c);
+        if (processPendingCommandAndInputBuffer(c) != C_ERR) {
+            beforeNextClient(c);
+            if (handled_clients) listAddNodeTail(handled_clients, c);
+        }
     }
 
     resetCommandsBatch();
@@ -260,7 +265,7 @@ static void addCommandToBatch(struct serverCommand *cmd, robj **argv, int argc,
  * if it becomes full.
  *
  * Returns C_OK if the command was added successfully, C_ERR otherwise. */
-int addCommandToBatchAndProcessIfFull(client *c) {
+int addCommandToBatchAndProcessIfFull(client *c, list *handled_clients) {
     if (!batch) return C_ERR;
 
     batch->clients[batch->client_count++] = c;
@@ -283,7 +288,7 @@ int addCommandToBatchAndProcessIfFull(client *c) {
      * We also check the client count to handle cases where
      * no keys exist for the clients' commands. */
     if (batch->client_count == batch->max_prefetch_size || batch->key_count == batch->max_prefetch_size) {
-        processClientsCommandsBatch();
+        processClientsCommandsBatch(handled_clients);
     }
 
     return C_OK;

src/memory_prefetch.h

@@ -1,11 +1,13 @@
 #ifndef MEMORY_PREFETCH_H
 #define MEMORY_PREFETCH_H
 
+#include "adlist.h"
+
 struct client;
 
 void prefetchCommandsBatchInit(void);
-void processClientsCommandsBatch(void);
-int addCommandToBatchAndProcessIfFull(struct client *c);
+void processClientsCommandsBatch(list *handled_clients);
+int addCommandToBatchAndProcessIfFull(struct client *c, list *handled_clients);
 void removeClientFromPendingCommandsBatch(struct client *c);
 int onMaxBatchSizeChange(const char **err);
 

src/networking.c

@@ -27,6 +27,8 @@
  * POSSIBILITY OF SUCH DAMAGE.
  */
 
+#include "adlist.h"
+#include "list.h"
 #include "server.h"
 #include "cluster.h"
 #include "cluster_slot_stats.h"
@@ -6373,16 +6375,15 @@ int processIOThreadsReadDone(void) {
     if (ProcessingEventsWhileBlocked) {
         /* When ProcessingEventsWhileBlocked we may call processIOThreadsReadDone recursively.
          * In this case, there may be some clients left in the batch waiting to be processed. */
-        processClientsCommandsBatch();
+        processClientsCommandsBatch(NULL);
     }
 
     if (listLength(server.clients_pending_io_read) == 0) return 0;
     int processed = 0;
     /* Collect clients handled in this iteration to continue parsing any
      * residual buffered data synchronously after batch execution (important
-     * for edge-triggered transports like RDMA). Use IDs instead of raw client
-     * pointers to avoid UAF if a client is freed while executing the batch. */
-    list *handled_clients_ids = listCreate();
+     * for edge-triggered transports like RDMA). */
+    list *handled_clients = listCreate();
     listNode *ln;
 
     listNode *next = listFirst(server.clients_pending_io_read);
@@ -6444,39 +6445,30 @@ int processIOThreadsReadDone(void) {
 
         size_t list_length_before_command_execute = listLength(server.clients_pending_io_read);
         /* try to add the command to the batch */
-        int ret = addCommandToBatchAndProcessIfFull(c);
+        int ret = addCommandToBatchAndProcessIfFull(c, handled_clients);
         /* If the command was not added to the commands batch, process it immediately */
         if (ret == C_ERR) {
-            if (processPendingCommandAndInputBuffer(c) == C_OK) beforeNextClient(c);
+            if (processPendingCommandAndInputBuffer(c) == C_OK) {
+                beforeNextClient(c);
+                listAddNodeTail(handled_clients,c);
+            }
         }
         if (list_length_before_command_execute != listLength(server.clients_pending_io_read)) {
             /* A client was unlink from the list possibly making the next node invalid */
             next = listFirst(server.clients_pending_io_read);
         }
-
-        /* Track the client by ID so we can drain any leftover buffered data even if no
-         * further network events arrive. */
-        listAddNodeTail(handled_clients_ids, (void *)(uintptr_t)c->id);
     }
 
-    processClientsCommandsBatch();
+    processClientsCommandsBatch(handled_clients);
 
-    /* For clients handled in this iteration, if residual data remains in
-     * querybuf,  handle it now to avoid edge-triggered lost wakeups. Clients may
-     * have been freed while executing the batch, so re-lookup by ID. */
     listIter handled_li;
     listNode *handled_ln;
-    listRewind(handled_clients_ids, &handled_li);
+    listRewind(handled_clients, &handled_li);
     while ((handled_ln = listNext(&handled_li))) {
-        uint64_t id = (uint64_t)(uintptr_t)listNodeValue(handled_ln);
-        client *c = lookupClientByID(id);
-        if (!c) continue;
-
-        if (connUpdateState(c->conn)) {
-            processPendingCommandAndInputBuffer(c); /* try to handle new arrival data if possible */
-        }
+        client *c = listNodeValue(handled_ln);
+        connUpdateState(c->conn);
     }
-    listRelease(handled_clients_ids);
+    listRelease(handled_clients);
 
     return processed;
 }