From 152c1b6802b88fbad0113adb97500d92a84e3c63 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Thu, 6 Oct 2016 17:05:38 +0200
Subject: [PATCH] Module: Ability to get context from IO context.

It was noted by @dvirsky that it is not possible to use string functions
when writing the AOF file. This sometimes is critical since the command
rewriting may need to be built in the context of the AOF callback, and
without access to the context, and the limited types that the AOF
production functions will accept, this can be an issue.

Moreover there are other needs that we can't anticipate regarding the
ability to use Redis Modules APIs using the context in order to build
representations to emit AOF / RDB.

Because of this a new API was added that allows the user to get a
temporary context from the IO context. The context is auto released
if obtained when the RDB / AOF callback returns.

Calling multiple time the function to get the context, always returns
the same one, since it is invalid to have more than a single context.
---
 src/aof.c         |  4 ++++
 src/module.c      | 15 +++++++++++++++
 src/rdb.c         |  4 ++++
 src/redismodule.h |  3 +++
 src/server.h      |  4 ++++
 5 files changed, 30 insertions(+)

diff --git a/src/aof.c b/src/aof.c
index e52b7e21..c75153cc 100644
--- a/src/aof.c
+++ b/src/aof.c
@@ -1023,6 +1023,10 @@ int rewriteModuleObject(rio *r, robj *key, robj *o) {
     moduleType *mt = mv->type;
     moduleInitIOContext(io,mt,r);
     mt->aof_rewrite(&io,key,mv->value);
+    if (io.ctx) {
+        moduleFreeContext(io.ctx);
+        zfree(io.ctx);
+    }
     return io.error ? 0 : 1;
 }
 
diff --git a/src/module.c b/src/module.c
index c339303b..742d9b97 100644
--- a/src/module.c
+++ b/src/module.c
@@ -2958,6 +2958,20 @@ void RM_EmitAOF(RedisModuleIO *io, const char *cmdname, const char *fmt, ...) {
     return;
 }
 
+/* --------------------------------------------------------------------------
+ * IO context handling
+ * -------------------------------------------------------------------------- */
+
+RedisModuleCtx *RM_GetContextFromIO(RedisModuleIO *io) {
+    if (io->ctx) return io->ctx; /* Can't have more than one... */
+    RedisModuleCtx ctxtemplate = REDISMODULE_CTX_INIT;
+    io->ctx = zmalloc(sizeof(*io));
+    *(io->ctx) = ctxtemplate;
+    io->ctx->module = io->type->module;
+    io->ctx->client = NULL;
+    return io->ctx;
+}
+
 /* --------------------------------------------------------------------------
  * Logging
  * -------------------------------------------------------------------------- */
@@ -3342,4 +3356,5 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(StringAppendBuffer);
     REGISTER_API(RetainString);
     REGISTER_API(StringCompare);
+    REGISTER_API(GetContextFromIO);
 }
diff --git a/src/rdb.c b/src/rdb.c
index 25204434..29f880da 100644
--- a/src/rdb.c
+++ b/src/rdb.c
@@ -770,6 +770,10 @@ ssize_t rdbSaveObject(rio *rdb, robj *o) {
 
         /* Then write the module-specific representation. */
         mt->rdb_save(&io,mv->value);
+        if (io.ctx) {
+            moduleFreeContext(io.ctx);
+            zfree(io.ctx);
+        }
         return io.error ? -1 : (ssize_t)io.bytes;
     } else {
         serverPanic("Unknown object type");
diff --git a/src/redismodule.h b/src/redismodule.h
index 34611009..9b3c061a 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -191,6 +191,7 @@ void REDISMODULE_API_FUNC(RedisModule_LogIOError)(RedisModuleIO *io, const char
 int REDISMODULE_API_FUNC(RedisModule_StringAppendBuffer)(RedisModuleCtx *ctx, RedisModuleString *str, const char *buf, size_t len);
 void REDISMODULE_API_FUNC(RedisModule_RetainString)(RedisModuleCtx *ctx, RedisModuleString *str);
 int REDISMODULE_API_FUNC(RedisModule_StringCompare)(RedisModuleString *a, RedisModuleString *b);
+RedisModuleCtx *REDISMODULE_API_FUNC(RM_GetContextFromIO)(RedisModuleIO *io);
 
 /* This is included inline inside each Redis module. */
 static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int apiver) __attribute__((unused));
@@ -291,6 +292,8 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(StringAppendBuffer);
     REDISMODULE_GET_API(RetainString);
     REDISMODULE_GET_API(StringCompare);
+    REDISMODULE_GET_API(GetIOContext);
+    REDISMODULE_GET_API(FreeIOContext);
 
     RedisModule_SetModuleAttribs(ctx,name,ver,apiver);
     return REDISMODULE_OK;
diff --git a/src/server.h b/src/server.h
index 8aad8f98..21e32824 100644
--- a/src/server.h
+++ b/src/server.h
@@ -463,6 +463,7 @@ typedef long long mstime_t; /* millisecond time type. */
 struct RedisModule;
 struct RedisModuleIO;
 struct RedisModuleDigest;
+struct RedisModuleCtx;
 struct redisObject;
 
 /* Each module type implementation should export a set of methods in order
@@ -516,6 +517,7 @@ typedef struct RedisModuleIO {
     rio *rio;           /* Rio stream. */
     moduleType *type;   /* Module type doing the operation. */
     int error;          /* True if error condition happened. */
+    struct RedisModuleCtx *ctx; /* Optional context, via GetIOContext() call. */
 } RedisModuleIO;
 
 #define moduleInitIOContext(iovar,mtype,rioptr) do { \
@@ -523,6 +525,7 @@ typedef struct RedisModuleIO {
     iovar.type = mtype; \
     iovar.bytes = 0; \
     iovar.error = 0; \
+    iovar.ctx = NULL; \
 } while(0);
 
 /* Objects encoding. Some kind of objects like Strings and Hashes can be
@@ -1222,6 +1225,7 @@ void moduleLoadFromQueue(void);
 int *moduleGetCommandKeysViaAPI(struct redisCommand *cmd, robj **argv, int argc, int *numkeys);
 moduleType *moduleTypeLookupModuleByID(uint64_t id);
 void moduleTypeNameByID(char *name, uint64_t moduleid);
+void moduleFreeContext(struct RedisModuleCtx *ctx);
 
 /* Utils */
 long long ustime(void);