Fix race condition issues between the main thread and module threads

src/blocked.c

@@ -239,7 +239,7 @@ void replyToBlockedClientTimedOut(client *c) {
         addReplyLongLong(c,server.fsynced_reploff >= c->bstate.reploffset);
         addReplyLongLong(c,replicationCountAOFAcksByOffset(c->bstate.reploffset));
     } else if (c->bstate.btype == BLOCKED_MODULE) {
-        moduleBlockedClientTimedOut(c);
+        moduleBlockedClientTimedOut(c, 0);
     } else {
         serverPanic("Unknown btype in replyToBlockedClientTimedOut().");
     }

src/module.c

@@ -304,8 +304,11 @@ static size_t moduleTempClientMinCount = 0; /* Min client count in pool since
 
 /* We need a mutex that is unlocked / relocked in beforeSleep() in order to
  * allow thread safe contexts to execute commands at a safe moment. */
-static pthread_mutex_t moduleGIL = PTHREAD_MUTEX_INITIALIZER;
-
+typedef struct {
+    pthread_mutex_t mutex;
+    redisAtomic pthread_t owner; /* The thread that currently holds the lock */
+} ModuleGIL;
+ModuleGIL moduleGIL = { PTHREAD_MUTEX_INITIALIZER, (pthread_t)0 };
 
 /* Function pointer type for keyspace event notification subscriptions from modules. */
 typedef int (*RedisModuleNotificationFunc) (RedisModuleCtx *ctx, int type, const char *event, RedisModuleString *key);
@@ -8202,7 +8205,7 @@ int RM_UnblockClient(RedisModuleBlockedClient *bc, void *privdata) {
          * argument, but better to be safe than sorry. */
         if (bc->timeout_callback == NULL) return REDISMODULE_ERR;
         if (bc->unblocked) return REDISMODULE_OK;
-        if (bc->client) moduleBlockedClientTimedOut(bc->client);
+        if (bc->client) moduleBlockedClientTimedOut(bc->client, 1);
     }
     moduleUnblockClientByHandle(bc,privdata);
     return REDISMODULE_OK;
@@ -8301,8 +8304,10 @@ void moduleHandleBlockedClients(void) {
          * This needs to be out of the reply callback above given that a
          * module might not define any callback and still do blocking ops.
          */
-        if (c && !clientHasModuleAuthInProgress(c) && !bc->blocked_on_keys) {
-            updateStatsOnUnblock(c, bc->background_duration, reply_us, server.stat_total_error_replies != prev_error_replies);
+        if (c && !clientHasModuleAuthInProgress(c)) {
+            int had_errors = c->deferred_reply_errors ? !!listLength(c->deferred_reply_errors) :
+                (server.stat_total_error_replies != prev_error_replies);
+            updateStatsOnUnblock(c, bc->background_duration, reply_us, had_errors);
         }
 
         if (c != NULL) {
@@ -8355,8 +8360,15 @@ int moduleBlockedClientMayTimeout(client *c) {
 /* Called when our client timed out. After this function unblockClient()
  * is called, and it will invalidate the blocked client. So this function
  * does not need to do any cleanup. Eventually the module will call the
- * API to unblock the client and the memory will be released. */
-void moduleBlockedClientTimedOut(client *c) {
+ * API to unblock the client and the memory will be released. 
+ *
+ * If this function is called from a module, we handle the timeout callback
+ * and the update of the unblock status in a thread-safe manner to avoid race
+ * conditions with the main thread.
+ * If this function is called from the main thread, we must handle the unblocking
+ * of the client synchronously. This ensures that we can reply to the client before
+ * resetClient() is called. */
+void moduleBlockedClientTimedOut(client *c, int from_module) {
     RedisModuleBlockedClient *bc = c->bstate.module_blocked_handle;
 
     /* Protect against re-processing: don't serve clients that are already
@@ -8365,14 +8377,22 @@ void moduleBlockedClientTimedOut(client *c) {
     if (bc->unblocked) return;
 
     RedisModuleCtx ctx;
-    moduleCreateContext(&ctx, bc->module, REDISMODULE_CTX_BLOCKED_TIMEOUT);
+    int flags = REDISMODULE_CTX_BLOCKED_TIMEOUT;
+    if (from_module) flags |= REDISMODULE_CTX_THREAD_SAFE;
+    moduleCreateContext(&ctx, bc->module, flags);
     ctx.client = bc->client;
     ctx.blocked_client = bc;
     ctx.blocked_privdata = bc->privdata;
-    long long prev_error_replies = server.stat_total_error_replies;
+
+    long long prev_error_replies;
+    if (!from_module)
+        prev_error_replies = server.stat_total_error_replies;
+
     bc->timeout_callback(&ctx,(void**)c->argv,c->argc);
     moduleFreeContext(&ctx);
-    updateStatsOnUnblock(c, bc->background_duration, 0, server.stat_total_error_replies != prev_error_replies);
+
+    if (!from_module)
+        updateStatsOnUnblock(c, bc->background_duration, 0, server.stat_total_error_replies != prev_error_replies);
 
     /* For timeout events, we do not want to call the disconnect callback,
      * because the blocked client will be automatically disconnected in
@@ -8551,17 +8571,26 @@ void RM_ThreadSafeContextUnlock(RedisModuleCtx *ctx) {
 }
 
 void moduleAcquireGIL(void) {
-    pthread_mutex_lock(&moduleGIL);
+    pthread_mutex_lock(&moduleGIL.mutex);
+    atomicSet(moduleGIL.owner, pthread_self());
 }
 
 int moduleTryAcquireGIL(void) {
-    return pthread_mutex_trylock(&moduleGIL);
+    int ret = pthread_mutex_trylock(&moduleGIL.mutex);
+    if (ret == 0) atomicSet(moduleGIL.owner, pthread_self());
+    return ret;
 }
 
 void moduleReleaseGIL(void) {
-    pthread_mutex_unlock(&moduleGIL);
+    atomicSet(moduleGIL.owner, (pthread_t)0);
+    pthread_mutex_unlock(&moduleGIL.mutex);
 }
 
+int moduleOwnsGIL(void) {
+    pthread_t owner_thread;
+    atomicGet(moduleGIL.owner, owner_thread);
+    return pthread_equal(pthread_self(), owner_thread);
+}
 
 /* --------------------------------------------------------------------------
  * ## Module Keyspace Notifications API
@@ -11928,7 +11957,7 @@ void moduleInitModulesSystem(void) {
 
     /* Our thread-safe contexts GIL must start with already locked:
      * it is just unlocked when it's safe. */
-    pthread_mutex_lock(&moduleGIL);
+    moduleAcquireGIL();
 }
 
 void modulesCron(void) {

src/networking.c

@@ -415,7 +415,8 @@ void _addReplyToBufferOrList(client *c, const char *s, size_t len) {
      * after the command's reply (specifically important during multi-exec). the exception is
      * the SUBSCRIBE command family, which (currently) have a push message instead of a proper reply.
      * The check for executing_client also avoids affecting push messages that are part of eviction. */
-    if (c == server.current_client && (c->flags & CLIENT_PUSHING) &&
+    // TODO: should forbid the ReplyWith* module family api from being called outside the lock?
+    if ((c->flags & CLIENT_PUSHING) && c == server.current_client &&
         server.executing_client && !cmdHasPushAsReply(server.executing_client->cmd))
     {
         _addReplyProtoToList(c,server.pending_push_messages,s,len);
@@ -1434,7 +1435,7 @@ void unlinkClient(client *c) {
     listNode *ln;
 
     /* If this is marked as current client unset it. */
-    if (server.current_client == c) server.current_client = NULL;
+    if (c->conn && server.current_client == c) server.current_client = NULL;
 
     /* Certain operations must be done only if the client has an active connection.
      * If the client was already unlinked or if it's a "fake client" the
@@ -1478,7 +1479,7 @@ void unlinkClient(client *c) {
     }
 
     /* Remove from the list of pending reads if needed. */
-    serverAssert(io_threads_op == IO_THREADS_OP_IDLE);
+    serverAssert(!c->conn || io_threads_op == IO_THREADS_OP_IDLE);
     if (c->pending_read_list_node != NULL) {
         listDelNode(server.clients_pending_read,c->pending_read_list_node);
         c->pending_read_list_node = NULL;
@@ -1631,6 +1632,12 @@ void freeClient(client *c) {
     reqresReset(c, 1);
 #endif
 
+    /* Remove the contribution that this client gave to our
+     * incrementally computed memory usage. */
+    if (c->conn)
+        server.stat_clients_type_memory[c->last_memory_type] -=
+            c->last_memory_usage;
+
     /* Unlink the client: this will close the socket, remove the I/O
      * handlers, and remove references of the client from different
      * places where active clients may be referenced. */
@@ -1679,10 +1686,6 @@ void freeClient(client *c) {
      * we lost the connection with the master. */
     if (c->flags & CLIENT_MASTER) replicationHandleMasterDisconnection();
 
-    /* Remove the contribution that this client gave to our
-     * incrementally computed memory usage. */
-    server.stat_clients_type_memory[c->last_memory_type] -=
-        c->last_memory_usage;
     /* Remove client from memory usage buckets */
     if (c->mem_usage_bucket) {
         c->mem_usage_bucket->mem_usage_sum -= c->last_memory_usage;

src/server.c

@@ -949,6 +949,7 @@ static inline clientMemUsageBucket *getMemUsageBucket(size_t mem) {
  * usage bucket.
  */
 void updateClientMemoryUsage(client *c) {
+    if (!c->conn) return;
     size_t mem = getClientMemoryUsage(c, NULL);
     int type = getClientType(c);
     /* Now that we have the memory used by the client, remove the old
@@ -961,7 +962,7 @@ void updateClientMemoryUsage(client *c) {
 }
 
 int clientEvictionAllowed(client *c) {
-    if (server.maxmemory_clients == 0 || c->flags & CLIENT_NO_EVICT) {
+    if (server.maxmemory_clients == 0 || c->flags & CLIENT_NO_EVICT || !c->conn) {
         return 0;
     }
     int type = getClientType(c);
@@ -1864,19 +1865,20 @@ void afterSleep(struct aeEventLoop *eventLoop) {
     /********************* WARNING ********************
      * Do NOT add anything above moduleAcquireGIL !!! *
      ***************************** ********************/
+    /* Acquire the modules GIL so that their threads won't touch anything. */
+    if (moduleCount() && !moduleOwnsGIL()) {
+        mstime_t latency;
+        latencyStartMonitor(latency);
+
+        moduleAcquireGIL();
+        moduleFireServerEvent(REDISMODULE_EVENT_EVENTLOOP,
+                                REDISMODULE_SUBEVENT_EVENTLOOP_AFTER_SLEEP,
+                                NULL);
+        latencyEndMonitor(latency);
+        latencyAddSampleIfNeeded("module-acquire-GIL",latency);
+    }
+
     if (!ProcessingEventsWhileBlocked) {
-        /* Acquire the modules GIL so that their threads won't touch anything. */
-        if (moduleCount()) {
-            mstime_t latency;
-            latencyStartMonitor(latency);
-
-            moduleAcquireGIL();
-            moduleFireServerEvent(REDISMODULE_EVENT_EVENTLOOP,
-                                  REDISMODULE_SUBEVENT_EVENTLOOP_AFTER_SLEEP,
-                                  NULL);
-            latencyEndMonitor(latency);
-            latencyAddSampleIfNeeded("module-acquire-GIL",latency);
-        }
         /* Set the eventloop start time. */
         server.el_start = getMonotonicUs();
         /* Set the eventloop command count at start. */

src/server.h

@@ -2517,12 +2517,13 @@ void moduleFreeContext(struct RedisModuleCtx *ctx);
 void moduleCallCommandUnblockedHandler(client *c);
 void unblockClientFromModule(client *c);
 void moduleHandleBlockedClients(void);
-void moduleBlockedClientTimedOut(client *c);
+void moduleBlockedClientTimedOut(client *c, int from_module);
 void modulePipeReadable(aeEventLoop *el, int fd, void *privdata, int mask);
 size_t moduleCount(void);
 void moduleAcquireGIL(void);
 int moduleTryAcquireGIL(void);
 void moduleReleaseGIL(void);
+int moduleOwnsGIL(void);
 void moduleNotifyKeyspaceEvent(int type, const char *event, robj *key, int dbid);
 void firePostExecutionUnitJobs(void);
 void moduleCallCommandFilters(client *c);

tests/modules/blockedclient.c

@@ -135,8 +135,6 @@ void *bg_call_worker(void *arg) {
     RedisModuleCallReply *rep = RedisModule_Call(ctx, cmd, format, bg->argv + cmd_pos + 1, bg->argc - cmd_pos - 1);
     RedisModule_FreeString(NULL, format_redis_str);
 
-    // Release GIL
-    RedisModule_ThreadSafeContextUnlock(ctx);
 
     // Reply to client
     if (!rep) {
@@ -155,6 +153,9 @@ void *bg_call_worker(void *arg) {
     RedisModule_Free(bg->argv);
     RedisModule_Free(bg);
 
+    // Release GIL
+    RedisModule_ThreadSafeContextUnlock(ctx);
+
     // Free the Redis module context
     RedisModule_FreeThreadSafeContext(ctx);
 

tests/modules/usercall.c

@@ -136,8 +136,6 @@ void *bg_call_worker(void *arg) {
     RedisModuleCallReply *rep = RedisModule_Call(ctx, cmd, format, bg->argv + 3, bg->argc - 3);
     RedisModule_FreeString(NULL, format_redis_str);
 
-    // Release GIL
-    RedisModule_ThreadSafeContextUnlock(ctx);
 
     // Reply to client
     if (!rep) {
@@ -156,6 +154,9 @@ void *bg_call_worker(void *arg) {
     RedisModule_Free(bg->argv);
     RedisModule_Free(bg);
 
+    // Release GIL
+    RedisModule_ThreadSafeContextUnlock(ctx);
+
     // Free the Redis module context
     RedisModule_FreeThreadSafeContext(ctx);