Embed key into dict entry

src/db.c

@@ -201,7 +201,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);

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)) {

src/dict.c

@@ -35,6 +35,7 @@
 
 #include "fmacros.h"
 
+#include <stddef.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <stdint.h>
@@ -48,6 +49,10 @@
 #include "serverassert.h"
 #include "monotonic.h"
 
+#ifndef static_assert
+#define static_assert(expr, lit) extern char __static_assert_failure[(expr) ? 1 : -1]
+#endif
+
 /* Using dictSetResizeEnabled() we make possible to disable
  * resizing and rehashing of the hash table as needed. This is very important
  * for us, as we use copy-on-write and don't want to move too much memory
@@ -74,11 +79,36 @@ struct dictEntry {
     struct dictEntry *next; /* Next entry in the same hash bucket. */
 };
 
+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 to embed the key. */
+} embeddedDictEntry;
+
+static_assert(sizeof(embeddedDictEntry) == 24, "unexpected total size of embeddedDictEntry");
+static_assert(offsetof(embeddedDictEntry, v) == 0, "unexpected field offset");
+static_assert(offsetof(embeddedDictEntry, next) == 8, "unexpected field offset");
+static_assert(offsetof(embeddedDictEntry, key_header_size) == 16, "unexpected field offset");
+static_assert(offsetof(embeddedDictEntry, key_buf) == 17, "unexpected field offset");
+
+/* The minimum amount of bytes required for embedded dict entry. */
+static inline size_t compactSizeEmbeddedDictEntry(void) {
+    return offsetof(embeddedDictEntry, key_buf);
+}
+
 typedef struct {
     void *key;
     dictEntry *next;
 } dictEntryNoValue;
 
+/* Helper and validation for `embeddedDictEntry` */
+
 /* -------------------------- private prototypes ---------------------------- */
 
 static void _dictExpandIfNeeded(dict *d);
@@ -124,6 +154,8 @@ uint64_t dictGenCaseHashFunction(const unsigned char *buf, size_t len) {
 #define ENTRY_PTR_MASK 7     /* 111 */
 #define ENTRY_PTR_NORMAL 0   /* 000 */
 #define ENTRY_PTR_NO_VALUE 2 /* 010 */
+#define ENTRY_PTR_EMBEDDED 4 /* 100 */
+/*      ENTRY_PTR_IS_KEY        xx1 */
 
 /* Returns 1 if the entry pointer is a pointer to a key, rather than to an
  * allocated entry. Returns 0 otherwise. */
@@ -143,12 +175,9 @@ static inline int entryIsNoValue(const dictEntry *de) {
     return ((uintptr_t)(void *)de & ENTRY_PTR_MASK) == ENTRY_PTR_NO_VALUE;
 }
 
-/* Creates an entry without a value field. */
-static inline dictEntry *createEntryNoValue(void *key, dictEntry *next) {
-    dictEntryNoValue *entry = zmalloc(sizeof(*entry));
-    entry->key = key;
-    entry->next = next;
-    return (dictEntry *)(void *)((uintptr_t)(void *)entry | ENTRY_PTR_NO_VALUE);
+
+static inline int entryIsEmbedded(const dictEntry *de) {
+    return ((uintptr_t)(void *)de & ENTRY_PTR_MASK) == ENTRY_PTR_EMBEDDED;
 }
 
 static inline dictEntry *encodeMaskedPtr(const void *ptr, unsigned int bits) {
@@ -161,15 +190,40 @@ static inline void *decodeMaskedPtr(const dictEntry *de) {
     return (void *)((uintptr_t)(void *)de & ~ENTRY_PTR_MASK);
 }
 
+/* Creates an entry without a value field. */
+static inline dictEntry *createEntryNoValue(void *key, dictEntry *next) {
+    dictEntryNoValue *entry = zmalloc(sizeof(*entry));
+    entry->key = key;
+    entry->next = next;
+    return encodeMaskedPtr(entry, ENTRY_PTR_NO_VALUE);
+}
+
+static inline dictEntry *createEmbeddedEntry(void *key, dictEntry *next, dictType *dt) {
+    size_t key_len = dt->embedKey(NULL, 0, key, NULL);
+    embeddedDictEntry *entry = zmalloc(compactSizeEmbeddedDictEntry() + key_len);
+    dt->embedKey(entry->key_buf, key_len, key, &entry->key_header_size);
+    entry->next = next;
+    return encodeMaskedPtr(entry, ENTRY_PTR_EMBEDDED);
+}
+
+static inline void *getEmbeddedKey(const dictEntry *de) {
+    embeddedDictEntry *entry = (embeddedDictEntry *)decodeMaskedPtr(de);
+    return &entry->key_buf[entry->key_header_size];
+}
+
 /* Decodes the pointer to an entry without value, when you know it is an entry
  * without value. Hint: Use entryIsNoValue to check. */
 static inline dictEntryNoValue *decodeEntryNoValue(const dictEntry *de) {
     return decodeMaskedPtr(de);
 }
 
+static inline embeddedDictEntry *decodeEmbeddedEntry(const dictEntry *de) {
+    return decodeMaskedPtr(de);
+}
+
 /* Returns 1 if the entry has a value field and 0 otherwise. */
 static inline int entryHasValue(const dictEntry *de) {
-    return entryIsNormal(de);
+    return !entryIsNoValue(de);
 }
 
 /* ----------------------------- API implementation ------------------------- */
@@ -509,6 +563,8 @@ dictEntry *dictInsertAtPosition(dict *d, void *key, void *position) {
             /* Allocate an entry without value. */
             entry = createEntryNoValue(key, *bucket);
         }
+    } else if (d->type->embedded_entry) {
+        entry = createEmbeddedEntry(key, *bucket, d->type);
     } else {
         /* Allocate the memory and store the new entry.
          * Insert the element in top, with the assumption that in a database
@@ -656,6 +712,7 @@ void dictFreeUnlinkedEntry(dict *d, dictEntry *he) {
     if (he == NULL) return;
     dictFreeKey(d, he);
     dictFreeVal(d, he);
+    /* Clear the dictEntry */
     if (!entryIsKey(he)) zfree(decodeMaskedPtr(he));
 }
 
@@ -801,7 +858,12 @@ void dictSetKey(dict *d, dictEntry *de, void *key) {
 
 void dictSetVal(dict *d, dictEntry *de, void *val) {
     assert(entryHasValue(de));
-    de->v.val = d->type->valDup ? d->type->valDup(d, val) : val;
+    void *v = d->type->valDup ? d->type->valDup(d, val) : val;
+    if (entryIsEmbedded(de)) {
+        decodeEmbeddedEntry(de)->v.val = v;
+    } else {
+        de->v.val = v;
+    }
 }
 
 void dictSetSignedIntegerVal(dictEntry *de, int64_t val) {
@@ -837,11 +899,15 @@ double dictIncrDoubleVal(dictEntry *de, double val) {
 void *dictGetKey(const dictEntry *de) {
     if (entryIsKey(de)) return (void *)de;
     if (entryIsNoValue(de)) return decodeEntryNoValue(de)->key;
+    if (entryIsEmbedded(de)) return getEmbeddedKey(de);
     return de->key;
 }
 
 void *dictGetVal(const dictEntry *de) {
     assert(entryHasValue(de));
+    if (entryIsEmbedded(de)) {
+        return decodeEmbeddedEntry(de)->v.val;
+    }
     return de->v.val;
 }
 
@@ -871,6 +937,7 @@ double *dictGetDoubleValPtr(dictEntry *de) {
 static dictEntry *dictGetNext(const dictEntry *de) {
     if (entryIsKey(de)) return NULL; /* there's no next */
     if (entryIsNoValue(de)) return decodeEntryNoValue(de)->next;
+    if (entryIsEmbedded(de)) return decodeEmbeddedEntry(de)->next;
     return de->next;
 }
 
@@ -879,14 +946,16 @@ static dictEntry *dictGetNext(const dictEntry *de) {
 static dictEntry **dictGetNextRef(dictEntry *de) {
     if (entryIsKey(de)) return NULL;
     if (entryIsNoValue(de)) return &decodeEntryNoValue(de)->next;
+    if (entryIsEmbedded(de)) return &decodeEmbeddedEntry(de)->next;
     return &de->next;
 }
 
 static void dictSetNext(dictEntry *de, dictEntry *next) {
     assert(!entryIsKey(de));
     if (entryIsNoValue(de)) {
-        dictEntryNoValue *entry = decodeEntryNoValue(de);
-        entry->next = next;
+        decodeEntryNoValue(de)->next = next;
+    } else if (entryIsEmbedded(de)) {
+        decodeEmbeddedEntry(de)->next = next;
     } else {
         de->next = next;
     }
@@ -1164,7 +1233,7 @@ unsigned int dictGetSomeKeys(dict *d, dictEntry **des, unsigned int count) {
 
 /* Reallocate the dictEntry, key and value allocations in a bucket using the
  * provided allocation functions in order to defrag them. */
-static void dictDefragBucket(dictEntry **bucketref, dictDefragFunctions *defragfns) {
+static void dictDefragBucket(dictEntry **bucketref, dictDefragFunctions *defragfns, void *privdata) {
     dictDefragAllocFunction *defragalloc = defragfns->defragAlloc;
     dictDefragAllocFunction *defragkey = defragfns->defragKey;
     dictDefragAllocFunction *defragval = defragfns->defragVal;
@@ -1182,6 +1251,17 @@ static void dictDefragBucket(dictEntry **bucketref, dictDefragFunctions *defragf
                 entry = newentry;
             }
             if (newkey) entry->key = newkey;
+        } else if (entryIsEmbedded(de)) {
+            defragfns->defragEntryStartCb(privdata, de);
+            embeddedDictEntry *entry = decodeEmbeddedEntry(de), *newentry;
+            if ((newentry = defragalloc(entry))) {
+                newde = encodeMaskedPtr(newentry, ENTRY_PTR_EMBEDDED);
+                entry = newentry;
+                defragfns->defragEntryFinishCb(privdata, newde);
+            } else {
+                defragfns->defragEntryFinishCb(privdata, NULL);
+            }
+            if (newval) entry->v.val = newval;
         } else {
             assert(entryIsNormal(de));
             newde = defragalloc(de);
@@ -1345,7 +1425,7 @@ dictScanDefrag(dict *d, unsigned long v, dictScanFunction *fn, dictDefragFunctio
 
         /* Emit entries at cursor */
         if (defragfns) {
-            dictDefragBucket(&d->ht_table[htidx0][v & m0], defragfns);
+            dictDefragBucket(&d->ht_table[htidx0][v & m0], defragfns, privdata);
         }
         de = d->ht_table[htidx0][v & m0];
         while (de) {
@@ -1378,7 +1458,7 @@ dictScanDefrag(dict *d, unsigned long v, dictScanFunction *fn, dictDefragFunctio
 
         /* Emit entries at cursor */
         if (defragfns) {
-            dictDefragBucket(&d->ht_table[htidx0][v & m0], defragfns);
+            dictDefragBucket(&d->ht_table[htidx0][v & m0], defragfns, privdata);
         }
         de = d->ht_table[htidx0][v & m0];
         while (de) {
@@ -1392,7 +1472,7 @@ dictScanDefrag(dict *d, unsigned long v, dictScanFunction *fn, dictDefragFunctio
         do {
             /* Emit entries at cursor */
             if (defragfns) {
-                dictDefragBucket(&d->ht_table[htidx1][v & m1], defragfns);
+                dictDefragBucket(&d->ht_table[htidx1][v & m1], defragfns, privdata);
             }
             de = d->ht_table[htidx1][v & m1];
             while (de) {
@@ -1565,29 +1645,6 @@ uint64_t dictGetHash(dict *d, const void *key) {
     return dictHashKey(d, key);
 }
 
-/* Finds the dictEntry using pointer and pre-calculated hash.
- * oldkey is a dead pointer and should not be accessed.
- * the hash value should be provided using dictGetHash.
- * no string / key comparison is performed.
- * return value is a pointer to the dictEntry if found, or NULL if not found. */
-dictEntry *dictFindEntryByPtrAndHash(dict *d, const void *oldptr, uint64_t hash) {
-    dictEntry *he;
-    unsigned long idx, table;
-
-    if (dictSize(d) == 0) return NULL; /* dict is empty */
-    for (table = 0; table <= 1; table++) {
-        idx = hash & DICTHT_SIZE_MASK(d->ht_size_exp[table]);
-        if (table == 0 && (long)idx < d->rehashidx) continue;
-        he = d->ht_table[table][idx];
-        while (he) {
-            if (oldptr == dictGetKey(he)) return he;
-            he = dictGetNext(he);
-        }
-        if (!dictIsRehashing(d)) return NULL;
-    }
-    return NULL;
-}
-
 /* Provides the old and new ht size for a given dictionary during rehashing. This method
  * should only be invoked during initialization/rehashing. */
 void dictRehashingInfo(dict *d, unsigned long long *from_size, unsigned long long *to_size) {