Use sdsAllocSize instead of sdsZmallocSize (#923)

sdsAllocSize returns the correct size without consulting the
allocator. Which is much faster than consulting the allocator.
The only exception is SDS_TYPE_5, for which it has to
consult the allocator.

This PR also sets alloc field correctly for embedded string objects.
It assumes that no allocator would allocate a buffer larger
than `259 + sizeof(robj)` for embedded string. We use embedded strings
for strings up to 44 bytes. If this assumption is wrong, the whole
function would require a rewrite. In general case sds type adjustment
might be needed. Such logic should go to sds.c.

---------

Signed-off-by: Vadym Khoptynets <vadymkh@amazon.com>

src/debug.c

@@ -37,6 +37,7 @@
 #include "cluster.h"
 #include "threads_mngr.h"
 #include "io_threads.h"
+#include "sds.h"
 
 #include <arpa/inet.h>
 #include <signal.h>
@@ -682,7 +683,7 @@ void debugCommand(client *c) {
             addReplyStatusFormat(c,
                                  "key_sds_len:%lld, key_sds_avail:%lld, key_zmalloc: %lld, "
                                  "val_sds_len:%lld, val_sds_avail:%lld, val_zmalloc: %lld",
-                                 (long long)sdslen(key), (long long)sdsavail(key), (long long)sdsZmallocSize(key),
+                                 (long long)sdslen(key), (long long)sdsavail(key), (long long)sdsAllocSize(key),
                                  (long long)sdslen(val->ptr), (long long)sdsavail(val->ptr),
                                  (long long)getStringObjectSdsUsedMemory(val));
         }

src/eval.c

@@ -42,6 +42,7 @@
 #include "monotonic.h"
 #include "resp_parser.h"
 #include "script_lua.h"
+#include "sds.h"
 
 #include <lua.h>
 #include <lauxlib.h>
@@ -490,7 +491,7 @@ sds luaCreateFunction(client *c, robj *body, int evalsha) {
     l->node = luaScriptsLRUAdd(c, sha, evalsha);
     int retval = dictAdd(lctx.lua_scripts, sha, l);
     serverAssertWithInfo(c ? c : lctx.lua_client, NULL, retval == DICT_OK);
-    lctx.lua_scripts_mem += sdsZmallocSize(sha) + getStringObjectSdsUsedMemory(body);
+    lctx.lua_scripts_mem += sdsAllocSize(sha) + getStringObjectSdsUsedMemory(body);
     incrRefCount(body);
     return sha;
 }
@@ -516,7 +517,7 @@ void luaDeleteFunction(client *c, sds sha) {
     /* We only delete `EVAL` scripts, which must exist in the LRU list. */
     serverAssert(l->node);
     listDelNode(lctx.lua_scripts_lru_list, l->node);
-    lctx.lua_scripts_mem -= sdsZmallocSize(sha) + getStringObjectSdsUsedMemory(l->body);
+    lctx.lua_scripts_mem -= sdsAllocSize(sha) + getStringObjectSdsUsedMemory(l->body);
     dictFreeUnlinkedEntry(lctx.lua_scripts, de);
 }
 

src/functions.c

@@ -114,12 +114,12 @@ static dict *engines = NULL;
 static functionsLibCtx *curr_functions_lib_ctx = NULL;
 
 static size_t functionMallocSize(functionInfo *fi) {
-    return zmalloc_size(fi) + sdsZmallocSize(fi->name) + (fi->desc ? sdsZmallocSize(fi->desc) : 0) +
+    return zmalloc_size(fi) + sdsAllocSize(fi->name) + (fi->desc ? sdsAllocSize(fi->desc) : 0) +
            fi->li->ei->engine->get_function_memory_overhead(fi->function);
 }
 
 static size_t libraryMallocSize(functionLibInfo *li) {
-    return zmalloc_size(li) + sdsZmallocSize(li->name) + sdsZmallocSize(li->code);
+    return zmalloc_size(li) + sdsAllocSize(li->name) + sdsAllocSize(li->code);
 }
 
 static void engineStatsDispose(dict *d, void *obj) {
@@ -417,7 +417,7 @@ int functionsRegisterEngine(const char *engine_name, engine *engine) {
 
     dictAdd(engines, engine_name_sds, ei);
 
-    engine_cache_memory += zmalloc_size(ei) + sdsZmallocSize(ei->name) + zmalloc_size(engine) +
+    engine_cache_memory += zmalloc_size(ei) + sdsAllocSize(ei->name) + zmalloc_size(engine) +
                            engine->get_engine_memory_overhead(engine->engine_ctx);
 
     return C_OK;

src/networking.c

@@ -31,6 +31,7 @@
 #include "cluster.h"
 #include "cluster_slot_stats.h"
 #include "script.h"
+#include "sds.h"
 #include "fpconv_dtoa.h"
 #include "fmtargs.h"
 #include "io_threads.h"
@@ -51,23 +52,14 @@ __thread sds thread_shared_qb = NULL;
 
 typedef enum { PARSE_OK = 0, PARSE_ERR = -1, PARSE_NEEDMORE = -2 } parseResult;
 
-/* Return the size consumed from the allocator, for the specified SDS string,
- * including internal fragmentation. This function is used in order to compute
- * the client output buffer size. */
-size_t sdsZmallocSize(sds s) {
-    void *sh = sdsAllocPtr(s);
-    return zmalloc_size(sh);
-}
-
 /* Return the amount of memory used by the sds string at object->ptr
  * for a string object. This includes internal fragmentation. */
 size_t getStringObjectSdsUsedMemory(robj *o) {
     serverAssertWithInfo(NULL, o, o->type == OBJ_STRING);
-    switch (o->encoding) {
-    case OBJ_ENCODING_RAW: return sdsZmallocSize(o->ptr);
-    case OBJ_ENCODING_EMBSTR: return zmalloc_size(o) - sizeof(robj);
-    default: return 0; /* Just integer encoding for now. */
+    if (o->encoding != OBJ_ENCODING_INT) {
+        return sdsAllocSize(o->ptr);
     }
+    return 0;
 }
 
 /* Return the length of a string object.
@@ -4270,7 +4262,7 @@ size_t getClientMemoryUsage(client *c, size_t *output_buffer_mem_usage) {
     size_t mem = getClientOutputBufferMemoryUsage(c);
 
     if (output_buffer_mem_usage != NULL) *output_buffer_mem_usage = mem;
-    mem += c->querybuf ? sdsZmallocSize(c->querybuf) : 0;
+    mem += c->querybuf ? sdsAllocSize(c->querybuf) : 0;
     mem += zmalloc_size(c);
     mem += c->buf_usable_size;
     /* For efficiency (less work keeping track of the argv memory), it doesn't include the used memory

src/object.c

@@ -29,8 +29,11 @@
  */
 
 #include "server.h"
+#include "serverassert.h"
 #include "functions.h"
 #include "intset.h" /* Compact integer set structure */
+#include "zmalloc.h"
+#include "sds.h"
 #include <math.h>
 #include <ctype.h>
 
@@ -89,7 +92,9 @@ robj *createRawStringObject(const char *ptr, size_t len) {
  * an object where the sds string is actually an unmodifiable string
  * allocated in the same chunk as the object itself. */
 robj *createEmbeddedStringObject(const char *ptr, size_t len) {
-    robj *o = zmalloc(sizeof(robj) + sizeof(struct sdshdr8) + len + 1);
+    size_t bufsize = 0;
+    size_t sds_hdrlen = sizeof(struct sdshdr8);
+    robj *o = zmalloc_usable(sizeof(robj) + sds_hdrlen + len + 1, &bufsize);
     struct sdshdr8 *sh = (void *)(o + 1);
 
     o->type = OBJ_STRING;
@@ -99,7 +104,11 @@ robj *createEmbeddedStringObject(const char *ptr, size_t len) {
     o->lru = 0;
 
     sh->len = len;
-    sh->alloc = len;
+    size_t usable = bufsize - (sizeof(robj) + sds_hdrlen + 1);
+    sh->alloc = usable;
+    /* Overflow check. This must not happen as we use embedded strings only
+     * for sds strings that fit into SDS_TYPE_8. */
+    serverAssert(usable == sh->alloc);
     sh->flags = SDS_TYPE_8;
     if (ptr == SDS_NOINIT)
         sh->buf[len] = '\0';
@@ -982,7 +991,7 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
         if (o->encoding == OBJ_ENCODING_INT) {
             asize = sizeof(*o);
         } else if (o->encoding == OBJ_ENCODING_RAW) {
-            asize = sdsZmallocSize(o->ptr) + sizeof(*o);
+            asize = sdsAllocSize(o->ptr) + sizeof(*o);
         } else if (o->encoding == OBJ_ENCODING_EMBSTR) {
             asize = zmalloc_size((void *)o);
         } else {
@@ -1010,7 +1019,7 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
             asize = sizeof(*o) + sizeof(dict) + (sizeof(struct dictEntry *) * dictBuckets(d));
             while ((de = dictNext(di)) != NULL && samples < sample_size) {
                 ele = dictGetKey(de);
-                elesize += dictEntryMemUsage(de) + sdsZmallocSize(ele);
+                elesize += dictEntryMemUsage(de) + sdsAllocSize(ele);
                 samples++;
             }
             dictReleaseIterator(di);
@@ -1032,7 +1041,7 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
             asize = sizeof(*o) + sizeof(zset) + sizeof(zskiplist) + sizeof(dict) +
                     (sizeof(struct dictEntry *) * dictBuckets(d)) + zmalloc_size(zsl->header);
             while (znode != NULL && samples < sample_size) {
-                elesize += sdsZmallocSize(znode->ele);
+                elesize += sdsAllocSize(znode->ele);
                 elesize += dictEntryMemUsage(NULL) + zmalloc_size(znode);
                 samples++;
                 znode = znode->level[0].forward;
@@ -1051,7 +1060,7 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) {
             while ((de = dictNext(di)) != NULL && samples < sample_size) {
                 ele = dictGetKey(de);
                 ele2 = dictGetVal(de);
-                elesize += sdsZmallocSize(ele) + sdsZmallocSize(ele2);
+                elesize += sdsAllocSize(ele) + sdsAllocSize(ele2);
                 elesize += dictEntryMemUsage(de);
                 samples++;
             }
@@ -1195,7 +1204,7 @@ struct serverMemOverhead *getMemoryOverheadData(void) {
 
     mem = 0;
     if (server.aof_state != AOF_OFF) {
-        mem += sdsZmallocSize(server.aof_buf);
+        mem += sdsAllocSize(server.aof_buf);
     }
     mh->aof_buffer = mem;
     mem_total += mem;

src/sds.c

@@ -343,11 +343,7 @@ sds sdsRemoveFreeSpace(sds s, int would_regrow) {
  * if the size is smaller than currently used len, the data will be truncated.
  *
  * The when the would_regrow argument is set to 1, it prevents the use of
- * SDS_TYPE_5, which is desired when the sds is likely to be changed again.
- *
- * The sdsAlloc size will be set to the requested size regardless of the actual
- * allocation size, this is done in order to avoid repeated calls to this
- * function when the caller detects that it has excess space. */
+ * SDS_TYPE_5, which is desired when the sds is likely to be changed again. */
 sds sdsResize(sds s, size_t size, int would_regrow) {
     void *sh, *newsh = NULL;
     char type, oldtype = s[-1] & SDS_TYPE_MASK;

src/server.c

@@ -41,6 +41,7 @@
 #include "threads_mngr.h"
 #include "fmtargs.h"
 #include "io_threads.h"
+#include "sds.h"
 
 #include <time.h>
 #include <signal.h>
@@ -813,7 +814,7 @@ size_t ClientsPeakMemInput[CLIENTS_PEAK_MEM_USAGE_SLOTS] = {0};
 size_t ClientsPeakMemOutput[CLIENTS_PEAK_MEM_USAGE_SLOTS] = {0};
 
 int clientsCronTrackExpansiveClients(client *c, int time_idx) {
-    size_t qb_size = c->querybuf ? sdsZmallocSize(c->querybuf) : 0;
+    size_t qb_size = c->querybuf ? sdsAllocSize(c->querybuf) : 0;
     size_t argv_size = c->argv ? zmalloc_size(c->argv) : 0;
     size_t in_usage = qb_size + c->argv_len_sum + argv_size;
     size_t out_usage = getClientOutputBufferMemoryUsage(c);

src/server.h

@@ -2821,7 +2821,6 @@ void addReplyLoadedModules(client *c);
 void copyReplicaOutputBuffer(client *dst, client *src);
 void addListRangeReply(client *c, robj *o, long start, long end, int reverse);
 void deferredAfterErrorReply(client *c, list *errors);
-size_t sdsZmallocSize(sds s);
 size_t getStringObjectSdsUsedMemory(robj *o);
 void freeClientReplyValue(void *o);
 void *dupClientReplyValue(void *o);