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.

5 files changed, 30 insertions, 0 deletions

diff --git a/src/aof.c b/src/aof.c
index e52b7e214..c75153cc7 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 c339303b5..742d9b974 100644
--- a/src/module.c
+++ b/src/module.c
@@ -2959,6 +2959,20 @@ void RM_EmitAOF(RedisModuleIO *io, const char *cmdname, const char *fmt, ...) {
 }
 
 /* --------------------------------------------------------------------------
+ * 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 252044344..29f880dac 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 346110090..9b3c061ac 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 8aad8f983..21e32824d 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);