From e07af6a2b782de809c1908a236332e48af5867b3 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Fri, 13 Apr 2018 13:48:11 +0200
Subject: Modules API: Add call to get the blocked client handle from the
 context.

This is useful in the reply and timeout callback, if the module wants to
do some cleanup of the blocked client handle that may be stored around
in the module-private data structures.
---
 src/module.c      | 14 +++++++++++++-
 src/redismodule.h |  2 ++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/module.c b/src/module.c
index d255926d0..8c7544608 100644
--- a/src/module.c
+++ b/src/module.c
@@ -3542,9 +3542,10 @@ int RM_UnblockClient(RedisModuleBlockedClient *bc, void *privdata) {
 }
 
 /* Abort a blocked client blocking operation: the client will be unblocked
- * without firing the reply callback. */
+ * without firing any callback. */
 int RM_AbortBlock(RedisModuleBlockedClient *bc) {
     bc->reply_callback = NULL;
+    bc->disconnect_callback = NULL;
     return RM_UnblockClient(bc,NULL);
 }
 
@@ -3603,6 +3604,7 @@ void moduleHandleBlockedClients(void) {
             ctx.blocked_privdata = bc->privdata;
             ctx.module = bc->module;
             ctx.client = bc->client;
+            ctx.blocked_client = bc;
             bc->reply_callback(&ctx,(void**)c->argv,c->argc);
             moduleHandlePropagationAfterCommandCallback(&ctx);
             moduleFreeContext(&ctx);
@@ -3672,6 +3674,7 @@ void moduleBlockedClientTimedOut(client *c) {
     ctx.flags |= REDISMODULE_CTX_BLOCKED_TIMEOUT;
     ctx.module = bc->module;
     ctx.client = bc->client;
+    ctx.blocked_client = bc;
     bc->timeout_callback(&ctx,(void**)c->argv,c->argc);
     moduleFreeContext(&ctx);
     /* For timeout events, we do not want to call the disconnect callback,
@@ -3697,6 +3700,14 @@ void *RM_GetBlockedClientPrivateData(RedisModuleCtx *ctx) {
     return ctx->blocked_privdata;
 }
 
+/* Get the blocked client associated with a given context.
+ * This is useful in the reply and timeout callbacks of blocked clients,
+ * before sometimes the module has the blocked client handle references
+ * around, and wants to cleanup it. */
+RedisModuleBlockedClient *RM_GetBlockedClientHandle(RedisModuleCtx *ctx) {
+    return ctx->blocked_client;
+}
+
 /* Return true if when the free callback of a blocked client is called,
  * the reason for the client to be unblocked is that it disconnected
  * while it was blocked. */
@@ -4677,4 +4688,5 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(GetRandomHexChars);
     REGISTER_API(BlockedClientDisconnected);
     REGISTER_API(SetDisconnectCallback);
+    REGISTER_API(GetBlockedClientHandle);
 }
diff --git a/src/redismodule.h b/src/redismodule.h
index 9c52012dc..47ecb1308 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -282,6 +282,7 @@ int REDISMODULE_API_FUNC(RedisModule_UnblockClient)(RedisModuleBlockedClient *bc
 int REDISMODULE_API_FUNC(RedisModule_IsBlockedReplyRequest)(RedisModuleCtx *ctx);
 int REDISMODULE_API_FUNC(RedisModule_IsBlockedTimeoutRequest)(RedisModuleCtx *ctx);
 void *REDISMODULE_API_FUNC(RedisModule_GetBlockedClientPrivateData)(RedisModuleCtx *ctx);
+RedisModuleBlockedClient *REDISMODULE_API_FUNC(RedisModule_GetBlockedClientHandle)(RedisModuleCtx *ctx);
 int REDISMODULE_API_FUNC(RedisModule_AbortBlock)(RedisModuleBlockedClient *bc);
 RedisModuleCtx *REDISMODULE_API_FUNC(RedisModule_GetThreadSafeContext)(RedisModuleBlockedClient *bc);
 void REDISMODULE_API_FUNC(RedisModule_FreeThreadSafeContext)(RedisModuleCtx *ctx);
@@ -422,6 +423,7 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(IsBlockedReplyRequest);
     REDISMODULE_GET_API(IsBlockedTimeoutRequest);
     REDISMODULE_GET_API(GetBlockedClientPrivateData);
+    REDISMODULE_GET_API(GetBlockedClientHandle);
     REDISMODULE_GET_API(AbortBlock);
     REDISMODULE_GET_API(SetDisconnectCallback);
     REDISMODULE_GET_API(SubscribeToKeyspaceEvents);
-- 
cgit v1.2.1