Embed key into dict entry (valkey-io#541)

This PR incorporates changes related to key embedding described in the
redis/redis#12216
With this change there will be no `key` pointer and embedded the key
within the `dictEntry`. 1 byte is used for additional bookkeeping.
Overall the saving would be 7 bytes on average.

Key changes:

New dict entry type introduced, which is now used as an entry for the
main dictionary:

```c
typedef struct {
    union {
        void *val;
        uint64_t u64;
        int64_t s64;
        double d;
    } v;
    struct dictEntry *next;  /* Next entry in the same hash bucket. */
    uint8_t key_header_size; /* offset into key_buf where the key is located at. */
    unsigned char key_buf[]; /* buffer with embedded key. */
} embeddedDictEntry;
```

One new function has been added to the dictType:

```c
size_t (*embedKey)(unsigned char *buf, size_t buf_len, const void *key, unsigned char *header_size);
```


Change is opt-in per dict type, hence sets, hashes and other types that
are using dictionary are not impacted.
With this change main dictionary now owns the data, so copy on insert in
dbAdd is no longer needed.

### Benchmarking results

TLDR; Around 9-10% memory usage reduction in overall memory usage for
scenario with key of 16 bytes and value of 8 bytes and 16 bytes. The
throughput per second varies but is similar or greater in most of the
run(s) with the changes against unstable (ae2d421).

---------

Signed-off-by: Harkrishn Patro <[email protected]>
Signed-off-by: Madelyn Olson <[email protected]>
Co-authored-by: Madelyn Olson <[email protected]>

src/db.c

@@ -190,7 +190,11 @@ robj *lookupKeyWriteOrReply(client *c, robj *key, robj *reply) {
     return o;
 }
 
-/* Add the key to the DB. It's up to the caller to increment the reference
+/* Add the key to the DB.
+ *
+ * In this case a copy of `key` is copied in kvstore, the caller must ensure the `key` is properly freed.
+ *
+ * It's up to the caller to increment the reference
  * counter of the value if needed.
  *
  * If the update_if_existing argument is false, the program is aborted
@@ -204,7 +208,6 @@ static void dbAddInternal(serverDb *db, robj *key, robj *val, int update_if_exis
         return;
     }
     serverAssertWithInfo(NULL, key, de != NULL);
-    kvstoreDictSetKey(db->keys, slot, de, sdsdup(key->ptr));
     initObjectLRUOrLFU(val);
     kvstoreDictSetVal(db->keys, slot, de, val);
     signalKeyAsReady(db, key, val->type);
@@ -240,15 +243,16 @@ int getKeySlot(sds key) {
 
 /* This is a special version of dbAdd() that is used only when loading
  * keys from the RDB file: the key is passed as an SDS string that is
- * retained by the function (and not freed by the caller).
+ * copied by the function and freed by the caller.
  *
  * Moreover this function will not abort if the key is already busy, to
  * give more control to the caller, nor will signal the key as ready
  * since it is not useful in this context.
  *
- * The function returns 1 if the key was added to the database, taking
- * ownership of the SDS string, otherwise 0 is returned, and is up to the
- * caller to free the SDS string. */
+ * The function returns 1 if the key was added to the database, otherwise 0 is returned.
+ *
+ * In this case a copy of `key` is copied in kvstore, the caller must ensure the `key` is properly freed.
+ */
 int dbAddRDBLoad(serverDb *db, sds key, robj *val) {
     int slot = getKeySlot(key);
     dictEntry *de = kvstoreDictAddRaw(db->keys, slot, key, NULL);

src/debug.c

@@ -864,7 +864,7 @@ void debugCommand(client *c) {
         sds sizes = sdsempty();
         sizes = sdscatprintf(sizes, "bits:%d ", (sizeof(void *) == 8) ? 64 : 32);
         sizes = sdscatprintf(sizes, "robj:%d ", (int)sizeof(robj));
-        sizes = sdscatprintf(sizes, "dictentry:%d ", (int)dictEntryMemUsage());
+        sizes = sdscatprintf(sizes, "dictentry:%d ", (int)dictEntryMemUsage(NULL));
         sizes = sdscatprintf(sizes, "sdshdr5:%d ", (int)sizeof(struct sdshdr5));
         sizes = sdscatprintf(sizes, "sdshdr8:%d ", (int)sizeof(struct sdshdr8));
         sizes = sdscatprintf(sizes, "sdshdr16:%d ", (int)sizeof(struct sdshdr16));

src/defrag.c

@@ -41,6 +41,7 @@
 typedef struct defragCtx {
     void *privdata;
     int slot;
+    void *aux;
 } defragCtx;
 
 typedef struct defragPubSubCtx {
@@ -75,6 +76,36 @@ void *activeDefragAlloc(void *ptr) {
     return newptr;
 }
 
+/* This method captures the expiry db dict entry which refers to data stored in keys db dict entry. */
+void defragEntryStartCbForKeys(void *ctx, void *oldptr) {
+    defragCtx *defragctx = (defragCtx *)ctx;
+    serverDb *db = defragctx->privdata;
+    sds oldsds = (sds)dictGetKey((dictEntry *)oldptr);
+    int slot = defragctx->slot;
+    if (kvstoreDictSize(db->expires, slot)) {
+        dictEntry *expire_de = kvstoreDictFind(db->expires, slot, oldsds);
+        defragctx->aux = expire_de;
+    }
+}
+
+/* This method updates the key of expiry db dict entry. The key might be no longer valid
+ * as it could have been cleaned up during the defrag-realloc of the main dictionary. */
+void defragEntryFinishCbForKeys(void *ctx, void *newptr) {
+    defragCtx *defragctx = (defragCtx *)ctx;
+    dictEntry *expire_de = (dictEntry *)defragctx->aux;
+    /* Item doesn't have TTL associated to it. */
+    if (!expire_de) return;
+    /* No reallocation happened. */
+    if (!newptr) {
+        expire_de = NULL;
+        return;
+    }
+    serverDb *db = defragctx->privdata;
+    sds newsds = (sds)dictGetKey((dictEntry *)newptr);
+    int slot = defragctx->slot;
+    kvstoreDictSetKey(db->expires, slot, expire_de, newsds);
+}
+
 /*Defrag helper for sds strings
  *
  * returns NULL in case the allocation wasn't moved.
@@ -650,25 +681,10 @@ void defragModule(serverDb *db, dictEntry *kde) {
 /* for each key we scan in the main dict, this function will attempt to defrag
  * all the various pointers it has. */
 void defragKey(defragCtx *ctx, dictEntry *de) {
-    sds keysds = dictGetKey(de);
-    robj *newob, *ob;
-    unsigned char *newzl;
-    sds newsds;
     serverDb *db = ctx->privdata;
     int slot = ctx->slot;
-    /* Try to defrag the key name. */
-    newsds = activeDefragSds(keysds);
-    if (newsds) {
-        kvstoreDictSetKey(db->keys, slot, de, newsds);
-        if (kvstoreDictSize(db->expires, slot)) {
-            /* We can't search in db->expires for that key after we've released
-             * the pointer it holds, since it won't be able to do the string
-             * compare, but we can find the entry using key hash and pointer. */
-            uint64_t hash = kvstoreGetHash(db->expires, newsds);
-            dictEntry *expire_de = kvstoreDictFindEntryByPtrAndHash(db->expires, slot, keysds, hash);
-            if (expire_de) kvstoreDictSetKey(db->expires, slot, expire_de, newsds);
-        }
-    }
+    robj *newob, *ob;
+    unsigned char *newzl;
 
     /* Try to defrag robj and / or string value. */
     ob = dictGetVal(de);
@@ -984,7 +1000,9 @@ void activeDefragCycle(void) {
     endtime = start + timelimit;
     latencyStartMonitor(latency);
 
-    dictDefragFunctions defragfns = {.defragAlloc = activeDefragAlloc};
+    dictDefragFunctions defragfns = {.defragAlloc = activeDefragAlloc,
+                                     .defragEntryStartCb = defragEntryStartCbForKeys,
+                                     .defragEntryFinishCb = defragEntryFinishCbForKeys};
     do {
         /* if we're not continuing a scan from the last call or loop, start a new one */
         if (!defrag_stage && !defrag_cursor && (slot < 0)) {