Modules: initial pool allocator and a LEFTPAD usage example.

src/module.c

@@ -35,6 +35,29 @@ struct AutoMemEntry {
 #define REDISMODULE_AM_REPLY 2
 #define REDISMODULE_AM_FREED 3 /* Explicitly freed by user already. */
 
+/* The pool allocator block. Redis Modules can allocate memory via this special
+ * allocator that will automatically release it all once the callback returns.
+ * This means that it can only be used for ephemeral allocations. However
+ * there are two advantages for modules to use this API:
+ *
+ * 1) The memory is automatically released when the callback returns.
+ * 2) This allocator is faster for many small allocations since whole blocks
+ *    are allocated, and small pieces returned to the caller just advancing
+ *    the index of the allocation.
+ *
+ * Allocations are always rounded to the size of the void pointer in order
+ * to always return aligned memory chunks. */
+
+#define REDISMODULE_POOL_ALLOC_MIN_SIZE (1024*8)
+#define REDISMODULE_POOL_ALLOC_ALIGN (sizeof(void*))
+
+typedef struct RedisModulePoolAllocBlock {
+    uint32_t size;
+    uint32_t used;
+    struct RedisModulePoolAllocBlock *next;
+    char memory[];
+} RedisModulePoolAllocBlock;
+
 /* This structure represents the context in which Redis modules operate.
  * Most APIs module can access, get a pointer to the context, so that the API
  * implementation can hold state across calls, or remember what to free after
@@ -56,10 +79,12 @@ struct RedisModuleCtx {
     /* Used if there is the REDISMODULE_CTX_KEYS_POS_REQUEST flag set. */
     int *keys_pos;
     int keys_count;
+
+    struct RedisModulePoolAllocBlock *pa_head;
 };
 typedef struct RedisModuleCtx RedisModuleCtx;
 
-#define REDISMODULE_CTX_INIT {(void*)(unsigned long)&RM_GetApi, NULL, NULL, NULL, 0, 0, 0, NULL, 0, NULL, 0}
+#define REDISMODULE_CTX_INIT {(void*)(unsigned long)&RM_GetApi, NULL, NULL, NULL, 0, 0, 0, NULL, 0, NULL, 0, NULL}
 #define REDISMODULE_CTX_MULTI_EMITTED (1<<0)
 #define REDISMODULE_CTX_AUTO_MEMORY (1<<1)
 #define REDISMODULE_CTX_KEYS_POS_REQUEST (1<<2)
@@ -111,7 +136,7 @@ typedef struct RedisModuleCommandProxy RedisModuleCommandProxy;
 /* Reply of RM_Call() function. The function is filled in a lazy
  * way depending on the function called on the reply structure. By default
  * only the type, proto and protolen are filled. */
-struct RedisModuleCallReply {
+typedef struct RedisModuleCallReply {
     RedisModuleCtx *ctx;
     int type;       /* REDISMODULE_REPLY_... */
     int flags;      /* REDISMODULE_REPLYFLAG_...  */
@@ -126,8 +151,7 @@ struct RedisModuleCallReply {
         long long ll;    /* Reply value for integer reply. */
         struct RedisModuleCallReply *array; /* Array of sub-reply elements. */
     } val;
-};
+} RedisModuleCallReply;
-typedef struct RedisModuleCallReply RedisModuleCallReply;
 
 /* --------------------------------------------------------------------------
  * Prototypes
@@ -140,6 +164,64 @@ robj **moduleCreateArgvFromUserFormat(const char *cmdname, const char *fmt, int
 void moduleReplicateMultiIfNeeded(RedisModuleCtx *ctx);
 void RM_ZsetRangeStop(RedisModuleKey *key);
 
+/* --------------------------------------------------------------------------
+ * Pool allocator
+ * -------------------------------------------------------------------------- */
+
+/* Release the chain of blocks used for pool allocations. */
+void poolAllocRelease(RedisModuleCtx *ctx) {
+    RedisModulePoolAllocBlock *head = ctx->pa_head, *next;
+
+    while(head != NULL) {
+        next = head->next;
+        zfree(head);
+        head = next;
+    }
+    ctx->pa_head = NULL;
+}
+
+/* Return heap allocated memory that will be freed automatically when the
+ * module callback function returns. Mostly suitable for small allocations
+ * that are short living and must be released when the callback returns
+ * anyway. The returned memory is aligned to the architecture word size
+ * if at least word size bytes are requested, otherwise it is just
+ * aligned to the next power of two, so for example a 3 bytes request is
+ * 4 bytes aligned while a 2 bytes request is 2 bytes aligned.
+ *
+ * There is no realloc style function since when this is needed to use the
+ * pool allocator is not a good idea.
+ *
+ * The function returns NULL if `bytes` is 0. */
+void *RM_PoolAlloc(RedisModuleCtx *ctx, size_t bytes) {
+    if (bytes == 0) return NULL;
+    RedisModulePoolAllocBlock *b = ctx->pa_head;
+    size_t left = b ? b->size - b->used : 0;
+
+    /* Fix alignment. */
+    if (left >= bytes) {
+        size_t alignment = REDISMODULE_POOL_ALLOC_ALIGN;
+        while (bytes < alignment && alignment/2 >= bytes) alignment /= 2;
+        if (b->used % alignment)
+            b->used += alignment - (b->used % alignment);
+        left = (b->used > b->size) ? 0 : b->size - b->used;
+    }
+
+    /* Create a new block if needed. */
+    if (left < bytes) {
+        size_t blocksize = REDISMODULE_POOL_ALLOC_MIN_SIZE;
+        if (blocksize < bytes) blocksize = bytes;
+        b = zmalloc(sizeof(*b) + blocksize);
+        b->size = blocksize;
+        b->used = 0;
+        b->next = ctx->pa_head;
+        ctx->pa_head = b;
+    }
+
+    char *retval = b->memory + b->used;
+    b->used += bytes;
+    return retval;
+}
+
 /* --------------------------------------------------------------------------
  * Helpers for modules API implementation
  * -------------------------------------------------------------------------- */
@@ -240,6 +322,7 @@ int RM_GetApi(const char *funcname, void **targetPtrPtr) {
 /* Free the context after the user function was called. */
 void moduleFreeContext(RedisModuleCtx *ctx) {
     autoMemoryCollect(ctx);
+    poolAllocRelease(ctx);
     if (ctx->postponed_arrays) {
         zfree(ctx->postponed_arrays);
         ctx->postponed_arrays_count = 0;
@@ -2292,6 +2375,7 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(IsKeysPositionRequest);
     REGISTER_API(KeyAtPos);
     REGISTER_API(GetClientId);
+    REGISTER_API(PoolAlloc);
 }
 
 /* Global initialization at Redis startup. */

src/modules/helloworld.c

@@ -2,6 +2,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <ctype.h>
+#include <string.h>
 
 /* HELLO.SIMPLE is among the simplest commands you can implement.
  * It just returns the currently selected DB id, a functionality which is
@@ -448,6 +449,58 @@ int HelloHCopy_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int a
     return REDISMODULE_OK;
 }
 
+/* HELLO.LEFTPAD str len ch
+ * This is an implementation of the infamous LEFTPAD function, that
+ * was at the center of an issue with the npm modules system in March 2016.
+ *
+ * LEFTPAD is a good example of using a Redis Modules API called
+ * "pool allocator", that was a famous way to allocate memory in yet another
+ * open source project, the Apache web server.
+ *
+ * The concept is very simple: there is memory that is useful to allocate
+ * only in the context of serving a request, and must be freed anyway when
+ * the callback implementing the command returns. So in that case the module
+ * does not need to retain a reference to these allocations, it is just
+ * required to free the memory before returning. When this is the case the
+ * module can call RedisModule_PoolAlloc() instead, that works like malloc()
+ * but will automatically free the memory when the module callback returns.
+ *
+ * Note that PoolAlloc() does not necessarily require AutoMemory to be
+ * active. */
+int HelloLeftPad_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    RedisModule_AutoMemory(ctx); /* Use automatic memory management. */
+    long long padlen;
+
+    if (argc != 4) return RedisModule_WrongArity(ctx);
+
+    if ((RedisModule_StringToLongLong(argv[2],&padlen) != REDISMODULE_OK) ||
+        (padlen< 0)) {
+        return RedisModule_ReplyWithError(ctx,"ERR invalid padding length");
+    }
+    size_t strlen, chlen;
+    const char *str = RedisModule_StringPtrLen(argv[1], &strlen);
+    const char *ch = RedisModule_StringPtrLen(argv[3], &chlen);
+
+    /* If the string is already larger than the target len, just return
+     * the string itself. */
+    if (strlen >= padlen)
+        return RedisModule_ReplyWithString(ctx,argv[1]);
+
+    /* Padding must be a single character in this simple implementation. */
+    if (chlen != 1)
+        return RedisModule_ReplyWithError(ctx,
+            "ERR padding must be a single char");
+
+    /* Here we use our pool allocator, for our throw-away allocation. */
+    padlen -= strlen;
+    char *buf = RedisModule_PoolAlloc(ctx,padlen+strlen);
+    for (size_t j = 0; j < padlen; j++) buf[j] = *ch;
+    memcpy(buf+padlen,str,strlen);
+
+    RedisModule_ReplyWithStringBuffer(ctx,buf,padlen+strlen);
+    return REDISMODULE_OK;
+}
+
 /* This function must be present on each Redis module. It is used in order to
  * register the commands into the Redis server. */
 int RedisModule_OnLoad(RedisModuleCtx *ctx) {
@@ -515,5 +568,9 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx) {
         HelloHCopy_RedisCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
+    if (RedisModule_CreateCommand(ctx,"hello.leftpad",
+        HelloLeftPad_RedisCommand,"",1,1,1) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
     return REDISMODULE_OK;
 }

src/redismodule.h

@@ -150,6 +150,7 @@ int REDISMODULE_API_FUNC(RedisModule_HashGet)(RedisModuleKey *key, int flags, ..
 int REDISMODULE_API_FUNC(RedisModule_IsKeysPositionRequest)(RedisModuleCtx *ctx);
 void REDISMODULE_API_FUNC(RedisModule_KeyAtPos)(RedisModuleCtx *ctx, int pos);
 unsigned long long REDISMODULE_API_FUNC(RedisModule_GetClientId)(RedisModuleCtx *ctx);
+void *REDISMODULE_API_FUNC(RedisModule_PoolAlloc)(RedisModuleCtx *ctx, size_t bytes);
 
 /* This is included inline inside each Redis module. */
 static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int apiver) {
@@ -219,6 +220,7 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(IsKeysPositionRequest);
     REDISMODULE_GET_API(KeyAtPos);
     REDISMODULE_GET_API(GetClientId);
+    REDISMODULE_GET_API(PoolAlloc);
 
     RedisModule_SetModuleAttribs(ctx,name,ver,apiver);
     return REDISMODULE_OK;