Try to trim strings only when applicable (#11817)

As `sdsRemoveFreeSpace` have an impact on performance even if it is a no-op (see details at #11508). 
Only call the function when there is a possibility that the string contains free space.
* For strings coming from the network, it's only if they're bigger than PROTO_MBULK_BIG_ARG
* For strings coming from scripts, it's only if they're smaller than LUA_CMD_OBJCACHE_MAX_LEN
* For strings coming from modules, it could be anything.

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

7 files changed, 84 insertions, 23 deletions

diff --git a/src/module.c b/src/module.c
index b42a838eb..a6c1b552c 100644
--- a/src/module.c
+++ b/src/module.c
@@ -846,7 +846,7 @@ void RedisModuleCommandDispatcher(client *c) {
         /* Only do the work if the module took ownership of the object:
          * in that case the refcount is no longer 1. */
         if (c->argv[i]->refcount > 1)
-            trimStringObjectIfNeeded(c->argv[i]);
+            trimStringObjectIfNeeded(c->argv[i], 0);
     }
 }
 
@@ -2738,7 +2738,7 @@ int RM_StringAppendBuffer(RedisModuleCtx *ctx, RedisModuleString *str, const cha
  */
 void RM_TrimStringAllocation(RedisModuleString *str) {
     if (!str) return;
-    trimStringObjectIfNeeded(str);
+    trimStringObjectIfNeeded(str, 1);
 }
 
 /* --------------------------------------------------------------------------
diff --git a/src/object.c b/src/object.c
index f64695061..c89f5aa85 100644
--- a/src/object.c
+++ b/src/object.c
@@ -609,14 +609,20 @@ int isObjectRepresentableAsLongLong(robj *o, long long *llval) {
 }
 
 /* Optimize the SDS string inside the string object to require little space,
- * in case there is more than 10% of free space at the end of the SDS
- * string. This happens because SDS strings tend to overallocate to avoid
- * wasting too much time in allocations when appending to the string. */
-void trimStringObjectIfNeeded(robj *o) {
-    if (o->encoding == OBJ_ENCODING_RAW &&
-        sdsavail(o->ptr) > sdslen(o->ptr)/10)
-    {
-        o->ptr = sdsRemoveFreeSpace(o->ptr, 0);
+ * in case there is more than 10% of free space at the end of the SDS. */
+void trimStringObjectIfNeeded(robj *o, int trim_small_values) {
+    if (o->encoding != OBJ_ENCODING_RAW) return;
+    /* A string may have free space in the following cases:
+     * 1. When an arg len is greater than PROTO_MBULK_BIG_ARG the query buffer may be used directly as the SDS string.
+     * 2. When utilizing the argument caching mechanism in Lua. 
+     * 3. When calling from RM_TrimStringAllocation (trim_small_values is true). */
+    size_t len = sdslen(o->ptr);
+    if (len >= PROTO_MBULK_BIG_ARG ||
+        trim_small_values||
+        (server.executing_client && server.executing_client->flags & CLIENT_SCRIPT && len < LUA_CMD_OBJCACHE_MAX_LEN)) {
+        if (sdsavail(o->ptr) > len/10) {
+            o->ptr = sdsRemoveFreeSpace(o->ptr, 0);
+        }
     }
 }
 
@@ -685,15 +691,8 @@ robj *tryObjectEncoding(robj *o) {
     }
 
     /* We can't encode the object...
-     *
-     * Do the last try, and at least optimize the SDS string inside
-     * the string object to require little space, in case there
-     * is more than 10% of free space at the end of the SDS string.
-     *
-     * We do that only for relatively large strings as this branch
-     * is only entered if the length of the string is greater than
-     * OBJ_ENCODING_EMBSTR_SIZE_LIMIT. */
-    trimStringObjectIfNeeded(o);
+     * Do the last try, and at least optimize the SDS string inside */
+    trimStringObjectIfNeeded(o, 0);
 
     /* Return the original object. */
     return o;
diff --git a/src/script_lua.c b/src/script_lua.c
index b44dbc0d5..cf0755e15 100644
--- a/src/script_lua.c
+++ b/src/script_lua.c
@@ -790,8 +790,6 @@ 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];
 
diff --git a/src/server.h b/src/server.h
index 7d8e65b36..36cd8c760 100644
--- a/src/server.h
+++ b/src/server.h
@@ -2709,7 +2709,7 @@ int compareStringObjects(const robj *a, const robj *b);
 int collateStringObjects(const robj *a, const robj *b);
 int equalStringObjects(robj *a, robj *b);
 unsigned long long estimateObjectIdleTime(robj *o);
-void trimStringObjectIfNeeded(robj *o);
+void trimStringObjectIfNeeded(robj *o, int trim_small_values);
 #define sdsEncodedObject(objptr) (objptr->encoding == OBJ_ENCODING_RAW || objptr->encoding == OBJ_ENCODING_EMBSTR)
 
 /* Synchronous I/O with timeout */
@@ -3279,6 +3279,9 @@ typedef struct luaScript {
     uint64_t flags;
     robj *body;
 } luaScript;
+/* Cache of recently used small arguments to avoid malloc calls. */
+#define LUA_CMD_OBJCACHE_SIZE 32
+#define LUA_CMD_OBJCACHE_MAX_LEN 64
 
 /* Blocked clients API */
 void processUnblockedClients(void);
diff --git a/tests/modules/basics.c b/tests/modules/basics.c
index b6f9766c6..579d8f629 100644
--- a/tests/modules/basics.c
+++ b/tests/modules/basics.c
@@ -407,6 +407,26 @@ int TestStringAppendAM(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     return REDISMODULE_OK;
 }
 
+/* TEST.STRING.TRIM -- Test we trim a string with free space. */
+int TestTrimString(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argv);
+    REDISMODULE_NOT_USED(argc);
+    RedisModuleString *s = RedisModule_CreateString(ctx,"foo",3);
+    char *tmp = RedisModule_Alloc(1024);
+    RedisModule_StringAppendBuffer(ctx,s,tmp,1024);
+    size_t string_len = RedisModule_MallocSizeString(s);
+    RedisModule_TrimStringAllocation(s);
+    size_t len_after_trim = RedisModule_MallocSizeString(s);
+    if (len_after_trim < string_len) {
+        RedisModule_ReplyWithSimpleString(ctx, "OK");
+    } else {
+        RedisModule_ReplyWithError(ctx, "String was not trimmed as expected.");
+    }
+    RedisModule_Free(tmp);
+    RedisModule_FreeString(ctx,s);
+    return REDISMODULE_OK;
+}
+
 /* TEST.STRING.PRINTF -- Test string formatting. */
 int TestStringPrintf(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     RedisModule_AutoMemory(ctx);
@@ -852,6 +872,9 @@ int TestBasics(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
 
     T("test.string.append.am","");
     if (!TestAssertStringReply(ctx,reply,"foobar",6)) goto fail;
+    
+    T("test.string.trim","");
+    if (!TestAssertStringReply(ctx,reply,"OK",2)) goto fail;
 
     T("test.string.printf", "cc", "foo", "bar");
     if (!TestAssertStringReply(ctx,reply,"Got 3 args. argv[1]: foo, argv[2]: bar",38)) goto fail;
@@ -963,6 +986,10 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
         TestStringAppend,"write deny-oom",1,1,1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
+    if (RedisModule_CreateCommand(ctx,"test.string.trim",
+        TestTrimString,"write deny-oom",1,1,1) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
     if (RedisModule_CreateCommand(ctx,"test.string.append.am",
         TestStringAppendAM,"write deny-oom",1,1,1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
diff --git a/tests/unit/moduleapi/misc.tcl b/tests/unit/moduleapi/misc.tcl
index 8d63ce5b6..4bf0d58ea 100644
--- a/tests/unit/moduleapi/misc.tcl
+++ b/tests/unit/moduleapi/misc.tcl
@@ -525,6 +525,15 @@ start_server {tags {"modules"}} {
         assert_equal {0} [r test.get_n_events]
     }
 
+    test {test RM_Call with large arg for SET command} {
+        # set a big value to trigger increasing the query buf
+        r set foo [string repeat A 100000]
+        # set a smaller value but > PROTO_MBULK_BIG_ARG (32*1024) Redis will try to save the query buf itself on the DB.
+        r test.call_generic set bar [string repeat A 33000]
+        # asset the value was trimmed
+        assert {[r memory usage bar] < 42000}; # 42K to count for Jemalloc's additional memory overhead.
+    }
+
     test "Unload the module - misc" {
         assert_equal {OK} [r module unload misc]
     }
diff --git a/tests/unit/scripting.tcl b/tests/unit/scripting.tcl
index 440b9e2e1..e3b1a620b 100644
--- a/tests/unit/scripting.tcl
+++ b/tests/unit/scripting.tcl
@@ -1995,5 +1995,30 @@ start_server {tags {"scripting"}} {
             } 0
         ] {asdf}
     }
-}
 
+    test "LUA test trim string as expected" {
+        # this test may fail if we use different memory allocator than jemalloc, as libc for example may keep the old size on realloc.
+        if {[string match {*jemalloc*} [s mem_allocator]]} {
+            # test that when using LUA cache mechanism, if there is free space in the argv array, the string is trimmed.
+            r set foo [string repeat "a" 45]
+            set expected_memory [r memory usage foo]
+
+            # Jemalloc will allocate for the requested 63 bytes, 80 bytes.
+            # We can't test for larger sizes because LUA_CMD_OBJCACHE_MAX_LEN is 64.
+            # This value will be recycled to be used in the next argument.
+            # We use SETNX to avoid saving the string which will prevent us to reuse it in the next command.
+            r eval {
+                return redis.call("SETNX", "foo", string.rep("a", 63))
+            } 0
+
+            # Jemalloc will allocate for the request 45 bytes, 56 bytes.
+            # we can't test for smaller sizes because OBJ_ENCODING_EMBSTR_SIZE_LIMIT is 44 where no trim is done.
+            r eval {
+                return redis.call("SET", "foo", string.rep("a", 45))
+            } 0
+
+            # Assert the string has been trimmed and the 80 bytes from the previous alloc were not kept.
+            assert { [r memory usage foo] <= $expected_memory};
+        }
+    }
+}