Reintroduce lua argument cache in luaRedisGenericCommand removed in v7.0 (#11541)

This mechanism aims to reduce calls to malloc and free when
preparing the arguments the script sends to redis commands.
This is a mechanism was originally implemented in 48c49c4
and 4f68655, and was removed in #10220 (thinking it's not needed
and that it has no impact), but it now turns out it was wrong, and it
indeed provides some 5% performance improvement.

The implementation is a little bit too simplistic, it assumes consecutive
calls use the same size in the same arg index, but that's arguably
sufficient since it's only aimed at caching very small things.

We could even consider always pre-allocating args to the full
LUA_CMD_OBJCACHE_MAX_LEN (64 bytes) rather than the right size for the argument,
that would increase the chance they'll be able to be re-used.
But in some way this is already happening since we're using
sdsalloc, which in turn uses s_malloc_usable and takes ownership
of the full side of the allocation, so we are padded to the allocator
bucket size.


Co-authored-by: Oran Agra <oran@redislabs.com>
Co-authored-by: sundb <sundbcn@gmail.com>

src/script.c

@@ -513,22 +513,18 @@ static int scriptVerifyAllowStale(client *c, sds *err) {
  * up to the engine to take and parse.
  * The err out variable is set only if error occurs and describe the error.
  * If err is set on reply is written to the run_ctx client. */
-void scriptCall(scriptRunCtx *run_ctx, robj* *argv, int argc, sds *err) {
+void scriptCall(scriptRunCtx *run_ctx, sds *err) {
     client *c = run_ctx->c;
 
     /* Setup our fake client for command execution */
-    c->argv = argv;
-    c->argc = argc;
     c->user = run_ctx->original_client->user;
 
     /* Process module hooks */
     moduleCallCommandFilters(c);
-    argv = c->argv;
-    argc = c->argc;
 
-    struct redisCommand *cmd = lookupCommand(argv, argc);
+    struct redisCommand *cmd = lookupCommand(c->argv, c->argc);
     c->cmd = c->lastcmd = c->realcmd = cmd;
-    if (scriptVerifyCommandArity(cmd, argc, err) != C_OK) {
+    if (scriptVerifyCommandArity(cmd, c->argc, err) != C_OK) {
         goto error;
     }
 

src/script.h

@@ -98,7 +98,7 @@ int scriptPrepareForRun(scriptRunCtx *r_ctx, client *engine_client, client *call
 void scriptResetRun(scriptRunCtx *r_ctx);
 int scriptSetResp(scriptRunCtx *r_ctx, int resp);
 int scriptSetRepl(scriptRunCtx *r_ctx, int repl);
-void scriptCall(scriptRunCtx *r_ctx, robj **argv, int argc, sds *err);
+void scriptCall(scriptRunCtx *r_ctx, sds *err);
 int scriptInterrupt(scriptRunCtx *r_ctx);
 void scriptKill(client *c, int is_eval);
 int scriptIsRunning();

src/script_lua.c

@@ -783,6 +783,17 @@ static void luaReplyToRedisReply(client *c, client* script_client, lua_State *lu
 /* ---------------------------------------------------------------------------
  * Lua redis.* functions implementations.
  * ------------------------------------------------------------------------- */
+void freeLuaRedisArgv(robj **argv, int argc);
+
+/* Cached argv array across calls. */
+static robj **lua_argv = NULL;
+static int lua_argv_size = 0;
+
+/* Cache of recently used small arguments to avoid malloc calls. */
+#define LUA_CMD_OBJCACHE_SIZE 32
+#define LUA_CMD_OBJCACHE_MAX_LEN 64
+static robj *lua_args_cached_objects[LUA_CMD_OBJCACHE_SIZE];
+static size_t lua_args_cached_objects_len[LUA_CMD_OBJCACHE_SIZE];
 
 static robj **luaArgsToRedisArgv(lua_State *lua, int *argc) {
     int j;
@@ -793,8 +804,11 @@ static robj **luaArgsToRedisArgv(lua_State *lua, int *argc) {
         return NULL;
     }
 
-    /* Build the arguments vector */
-    robj **argv = zcalloc(sizeof(robj*) * *argc);
+    /* Build the arguments vector (reuse a cached argv from last call) */
+    if (lua_argv_size < *argc) {
+        lua_argv = zrealloc(lua_argv,sizeof(robj*)* *argc);
+        lua_argv_size = *argc;
+    }
 
     for (j = 0; j < *argc; j++) {
         char *obj_s;
@@ -812,8 +826,18 @@ static robj **luaArgsToRedisArgv(lua_State *lua, int *argc) {
             obj_s = (char*)lua_tolstring(lua,j+1,&obj_len);
             if (obj_s == NULL) break; /* Not a string. */
         }
-
-        argv[j] = createStringObject(obj_s, obj_len);
+        /* Try to use a cached object. */
+        if (j < LUA_CMD_OBJCACHE_SIZE && lua_args_cached_objects[j] &&
+            lua_args_cached_objects_len[j] >= obj_len)
+        {
+            sds s = lua_args_cached_objects[j]->ptr;
+            lua_argv[j] = lua_args_cached_objects[j];
+            lua_args_cached_objects[j] = NULL;
+            memcpy(s,obj_s,obj_len+1);
+            sdssetlen(s, obj_len);
+        } else {
+            lua_argv[j] = createStringObject(obj_s, obj_len);
+        }
     }
 
     /* Pop all arguments from the stack, we do not need them anymore
@@ -824,17 +848,42 @@ static robj **luaArgsToRedisArgv(lua_State *lua, int *argc) {
      * is not a string or an integer (lua_isstring() return true for
      * integers as well). */
     if (j != *argc) {
-        j--;
-        while (j >= 0) {
-            decrRefCount(argv[j]);
-            j--;
-        }
-        zfree(argv);
+        freeLuaRedisArgv(lua_argv, j);
         luaPushError(lua, "Lua redis lib command arguments must be strings or integers");
         return NULL;
     }
 
-    return argv;
+    return lua_argv;
+}
+
+void freeLuaRedisArgv(robj **argv, int argc) {
+    int j;
+    for (j = 0; j < argc; j++) {
+        robj *o = argv[j];
+
+        /* Try to cache the object in the lua_args_cached_objects array.
+         * The object must be small, SDS-encoded, and with refcount = 1
+         * (we must be the only owner) for us to cache it. */
+        if (j < LUA_CMD_OBJCACHE_SIZE &&
+            o->refcount == 1 &&
+            (o->encoding == OBJ_ENCODING_RAW ||
+             o->encoding == OBJ_ENCODING_EMBSTR) &&
+            sdslen(o->ptr) <= LUA_CMD_OBJCACHE_MAX_LEN)
+        {
+            sds s = o->ptr;
+            if (lua_args_cached_objects[j]) decrRefCount(lua_args_cached_objects[j]);
+            lua_args_cached_objects[j] = o;
+            lua_args_cached_objects_len[j] = sdsalloc(s);
+        } else {
+            decrRefCount(o);
+        }
+    }
+    if (argv != lua_argv) {
+        /* The command changed argv, scrap the cache and start over. */
+        zfree(argv);
+        lua_argv = NULL;
+        lua_argv_size = 0;
+    }
 }
 
 static int luaRedisGenericCommand(lua_State *lua, int raise_error) {
@@ -845,9 +894,9 @@ static int luaRedisGenericCommand(lua_State *lua, int raise_error) {
     client* c = rctx->c;
     sds reply;
 
-    int argc;
-    robj **argv = luaArgsToRedisArgv(lua, &argc);
-    if (argv == NULL) {
+    c->argv = luaArgsToRedisArgv(lua, &c->argc);
+    c->argv_len = lua_argv_size;
+    if (c->argv == NULL) {
         return raise_error ? luaError(lua) : 1;
     }
 
@@ -883,7 +932,7 @@ static int luaRedisGenericCommand(lua_State *lua, int raise_error) {
         ldbLog(cmdlog);
     }
 
-    scriptCall(rctx, argv, argc, &err);
+    scriptCall(rctx, &err);
     if (err) {
         luaPushError(lua, err);
         sdsfree(err);
@@ -928,8 +977,11 @@ static int luaRedisGenericCommand(lua_State *lua, int raise_error) {
 cleanup:
     /* Clean up. Command code may have changed argv/argc so we use the
      * argv/argc of the client instead of the local variables. */
-    freeClientArgv(c);
+    freeLuaRedisArgv(c->argv, c->argc);
+    c->argc = c->argv_len = 0;
     c->user = NULL;
+    c->argv = NULL;
+    freeClientArgv(c);
     inuse--;
 
     if (raise_error) {
@@ -1096,8 +1148,7 @@ static int luaRedisAclCheckCmdPermissionsCommand(lua_State *lua) {
         }
     }
 
-    while (argc--) decrRefCount(argv[argc]);
-    zfree(argv);
+    freeLuaRedisArgv(argv, argc);
     if (raise_error)
         return luaError(lua);
     else

tests/unit/scripting.tcl

@@ -1956,5 +1956,20 @@ start_server {tags {"scripting"}} {
             r eval {local status, res = pcall(function() return foo end); return 'status: ' .. tostring(status) .. ' result: ' .. res} 0
         ]
     }
+
+    test "LUA test pcall with non string/integer arg" {
+        assert_error "ERR Lua redis lib command arguments must be strings or integers*" {
+            r eval {
+                local x={}
+                return redis.call("ping", x)
+            } 0
+        }
+        # run another command, to make sure the cached argv array survived
+        assert_equal [
+            r eval {
+                return redis.call("ping", "asdf")
+            } 0
+        ] {asdf}
+    }
 }