Add valkeyCreate() and valkeyGetKey() functions

Signed-off-by: Viktor Söderqvist <[email protected]>

src/object.c

@@ -44,12 +44,24 @@
 /* ===================== Creation and parsing of objects ==================== */
 
 robj *createObject(int type, void *ptr) {
-    robj *o = zmalloc(sizeof(*o));
+    robj *o;
+    /* Prepare space for an 'expire' field and a 'key' pointer, so this object
+     * can be converted to a 'valkey' object (value with a key attached) without
+     * being reallocated. */
+    size_t size = sizeof(*o) + sizeof(long long) + sizeof(void *);
+    o = zmalloc(size);
     o->type = type;
     o->encoding = OBJ_ENCODING_RAW;
     o->ptr = ptr;
     o->refcount = 1;
     o->lru = 0;
+    o->hasexpire = 1;      /* There's an expire field. */
+    o->hasembkey = 0;      /* No embedded actual key contents. */
+    o->hasembkeyptr = 1;   /* There's an embedded key pointer field. */
+    unsigned char *data = (void *)(o + 1);
+    long long expire = -1; /* -1 means no expire */
+    memcpy(data, &expire, sizeof(expire)); /* expire = -1 */
+    memset(data + sizeof(expire), 0, sizeof(void *)); /* embkeyptr = NULL */
     return o;
 }
 
@@ -102,6 +114,9 @@ robj *createEmbeddedStringObject(const char *ptr, size_t len) {
     o->ptr = sh + 1;
     o->refcount = 1;
     o->lru = 0;
+    o->hasexpire = 0;
+    o->hasembkey = 0;
+    o->hasembkeyptr = 0;
 
     sh->len = len;
     size_t usable = bufsize - (sizeof(robj) + sds_hdrlen + 1);
@@ -135,6 +150,101 @@ robj *createStringObject(const char *ptr, size_t len) {
         return createRawStringObject(ptr, len);
 }
 
+sds valkeyGetKey(robj *val) {
+    unsigned char *data = (void *)(val + 1);
+    if (val->hasexpire) {
+        /* Skip expire field */
+        data += sizeof(long long);
+    }
+    if (val->hasembkeyptr) {
+        return *(sds *)data;
+    }
+    if (val->hasembkey) {
+        uint8_t hdr_size = *(uint8_t *)data;
+        data += 1 + hdr_size;
+        return (sds)data;
+    }
+    return NULL;
+}
+
+/* Creates a new object with an embedded key. */
+valkey *valkeyCreate(robj *val, const sds key) {
+    /* A key must not already be embedded. */
+    assert(valkeyGetKey(val) == NULL);
+    if (val->encoding == OBJ_ENCODING_EMBSTR) {
+        /* TODO: If there's space in val's allocation, we can embed the key
+         * there and memmove the the embedded value, without creating a new
+         * object.
+         *
+         * TODO: If key + value are too large (allocation > 64 bytes) we may not
+         * want to embed both of them. We can embed one or the other depending
+         * on sizes. */
+
+        /* Create a new object with val and key embedded. Leave 'val' intact. */
+
+        /* Calculate sizes */
+        size_t key_sds_size = sdscopytobuffer(NULL, 0, key, NULL);
+        size_t val_len = sdslen(val->ptr);
+
+        size_t min_size = sizeof(robj);
+        min_size += sizeof(long long); /* expire */
+        /* Size of embedded key, incl. 1 byte for prefixed sds hdr size. */
+        min_size += 1 + key_sds_size;
+        /* Size of embedded value (EMBSTR) including \0 term. */
+        min_size += sizeof(struct sdshdr8) + val_len + 1;
+
+        size_t bufsize = 0;
+        valkey *o = zmalloc_usable(min_size, &bufsize);
+        o->type = val->type;
+        o->encoding = val->encoding;
+        o->refcount = 1;
+        o->lru = 0;
+        o->hasexpire = 1;
+        o->hasembkey = 1;
+        o->hasembkeyptr = 0;
+
+        /* Set the embedded data. */
+        unsigned char *data = (void *)(o + 1);
+
+        /* Set the expire field. */
+        long long expire = -1;
+        memcpy(data, &expire, sizeof(long long));
+        data += sizeof(long long);
+
+        /* Copy embedded string. */
+        sdscopytobuffer(data + 1, key_sds_size, key, data);
+        data += 1 + key_sds_size;
+
+        /* Copy embedded value (EMBSTR). */
+        struct sdshdr8 *sh = (void *)data;
+        sh->flags = SDS_TYPE_8;
+        sh->len = val_len;
+        size_t capacity = bufsize - (min_size - val_len);
+        sh->alloc = capacity;
+        serverAssert(capacity == sh->alloc); /* Overflow check. */
+        memcpy(sh->buf, val->ptr, val_len);
+        sh->buf[val_len] = '\0';
+
+        o->ptr = sh->buf;
+        return o;
+    } else {
+        /* Set key pointer in val, increment the reference counter and return it
+         * as a new object. */
+        assert(val->refcount == 1 && val->hasembkeyptr && !val->hasembkey);
+        sds dup = sdsdup(key);
+
+        /* Find the correct location in val's data field. */
+        unsigned char *data = (void *)(val + 1);
+        if (val->hasexpire) {
+            /* Skip expire field */
+            data += sizeof(long long);
+        }
+        memcpy((void *)data, (void *)&dup, sizeof(void *));
+        incrRefCount(val);
+        return val;
+    }
+}
+
 /* Same as CreateRawStringObject, can return NULL if allocation fails */
 robj *tryCreateRawStringObject(const char *ptr, size_t len) {
     sds str = sdstrynewlen(ptr, len);
@@ -314,6 +424,9 @@ void freeStringObject(robj *o) {
     if (o->encoding == OBJ_ENCODING_RAW) {
         sdsfree(o->ptr);
     }
+    if (o->hasembkeyptr) {
+        sdsfree(valkeyGetKey(o));
+    }
 }
 
 void freeListObject(robj *o) {

src/sds.c

@@ -194,12 +194,12 @@ sds sdsdup(const sds s) {
 /*
  * This method returns the minimum amount of bytes required to store the sds (header + data + NULL terminator).
  */
-static inline size_t sdsminlen(sds s) {
+static inline size_t sdsminlen(const sds s) {
     return sdslen(s) + sdsHdrSize(s[-1]) + 1;
 }
 
 /* This method copies the sds `s` into `buf` which is the target character buffer. */
-size_t sdscopytobuffer(unsigned char *buf, size_t buf_len, sds s, uint8_t *hdr_size) {
+size_t sdscopytobuffer(unsigned char *buf, size_t buf_len, const sds s, uint8_t *hdr_size) {
     size_t required_keylen = sdsminlen(s);
     if (buf == NULL) {
         return required_keylen;

src/server.h

@@ -83,6 +83,7 @@ typedef long long ustime_t; /* microsecond time type. */
 
 #define VALKEYMODULE_CORE 1
 typedef struct serverObject robj;
+typedef struct serverObject valkey;
 #include "valkeymodule.h" /* Modules API defines. */
 
 /* Following includes allow test functions to be called from main() */
@@ -881,16 +882,19 @@ struct ValkeyModuleDigest {
 #define LRU_CLOCK_MAX ((1 << LRU_BITS) - 1) /* Max value of obj->lru */
 #define LRU_CLOCK_RESOLUTION 1000           /* LRU clock resolution in ms */
 
-#define OBJ_SHARED_REFCOUNT INT_MAX       /* Global object never destroyed. */
-#define OBJ_STATIC_REFCOUNT (INT_MAX - 1) /* Object allocated in the stack. */
+#define OBJ_SHARED_REFCOUNT ((1 << 30) - 1) /* Global object never destroyed. */
+#define OBJ_STATIC_REFCOUNT ((1 << 30) - 2) /* Object allocated in the stack. */
 #define OBJ_FIRST_SPECIAL_REFCOUNT OBJ_STATIC_REFCOUNT
 struct serverObject {
     unsigned type : 4;
     unsigned encoding : 4;
     unsigned lru : LRU_BITS; /* LRU time (relative to global lru_clock) or
                               * LFU data (least significant 8 bits frequency
                               * and most significant 16 bits access time). */
-    int refcount;
+    unsigned hasexpire : 1;
+    unsigned hasembkey : 1;
+    unsigned hasembkeyptr : 1;
+    unsigned refcount : 30;
     void *ptr;
 };
 
@@ -2982,7 +2986,6 @@ robj *createObject(int type, void *ptr);
 void initObjectLRUOrLFU(robj *o);
 robj *createStringObject(const char *ptr, size_t len);
 robj *createRawStringObject(const char *ptr, size_t len);
-robj *createEmbeddedStringObject(const char *ptr, size_t len);
 robj *tryCreateRawStringObject(const char *ptr, size_t len);
 robj *tryCreateStringObject(const char *ptr, size_t len);
 robj *dupStringObject(const robj *o);

src/unit/test_files.h

@@ -96,6 +96,7 @@ int test_listpackBenchmarkLpValidateIntegrity(int argc, char **argv, int flags);
 int test_listpackBenchmarkLpCompareWithString(int argc, char **argv, int flags);
 int test_listpackBenchmarkLpCompareWithNumber(int argc, char **argv, int flags);
 int test_listpackBenchmarkFree(int argc, char **argv, int flags);
+int test_valkey_from_embstr(int argc, char **argv, int flags);
 int test_sds(int argc, char **argv, int flags);
 int test_typesAndAllocSize(int argc, char **argv, int flags);
 int test_sdsHeaderSizes(int argc, char **argv, int flags);
@@ -157,6 +158,7 @@ unitTest __test_hashtab_c[] = {{"test_cursor", test_cursor}, {"test_set_hash_fun
 unitTest __test_intset_c[] = {{"test_intsetValueEncodings", test_intsetValueEncodings}, {"test_intsetBasicAdding", test_intsetBasicAdding}, {"test_intsetLargeNumberRandomAdd", test_intsetLargeNumberRandomAdd}, {"test_intsetUpgradeFromint16Toint32", test_intsetUpgradeFromint16Toint32}, {"test_intsetUpgradeFromint16Toint64", test_intsetUpgradeFromint16Toint64}, {"test_intsetUpgradeFromint32Toint64", test_intsetUpgradeFromint32Toint64}, {"test_intsetStressLookups", test_intsetStressLookups}, {"test_intsetStressAddDelete", test_intsetStressAddDelete}, {NULL, NULL}};
 unitTest __test_kvstore_c[] = {{"test_kvstoreAdd16Keys", test_kvstoreAdd16Keys}, {"test_kvstoreIteratorRemoveAllKeysNoDeleteEmptyDict", test_kvstoreIteratorRemoveAllKeysNoDeleteEmptyDict}, {"test_kvstoreIteratorRemoveAllKeysDeleteEmptyDict", test_kvstoreIteratorRemoveAllKeysDeleteEmptyDict}, {"test_kvstoreDictIteratorRemoveAllKeysNoDeleteEmptyDict", test_kvstoreDictIteratorRemoveAllKeysNoDeleteEmptyDict}, {"test_kvstoreDictIteratorRemoveAllKeysDeleteEmptyDict", test_kvstoreDictIteratorRemoveAllKeysDeleteEmptyDict}, {NULL, NULL}};
 unitTest __test_listpack_c[] = {{"test_listpackCreateIntList", test_listpackCreateIntList}, {"test_listpackCreateList", test_listpackCreateList}, {"test_listpackLpPrepend", test_listpackLpPrepend}, {"test_listpackLpPrependInteger", test_listpackLpPrependInteger}, {"test_listpackGetELementAtIndex", test_listpackGetELementAtIndex}, {"test_listpackPop", test_listpackPop}, {"test_listpackGetELementAtIndex2", test_listpackGetELementAtIndex2}, {"test_listpackIterate0toEnd", test_listpackIterate0toEnd}, {"test_listpackIterate1toEnd", test_listpackIterate1toEnd}, {"test_listpackIterate2toEnd", test_listpackIterate2toEnd}, {"test_listpackIterateBackToFront", test_listpackIterateBackToFront}, {"test_listpackIterateBackToFrontWithDelete", test_listpackIterateBackToFrontWithDelete}, {"test_listpackDeleteWhenNumIsMinusOne", test_listpackDeleteWhenNumIsMinusOne}, {"test_listpackDeleteWithNegativeIndex", test_listpackDeleteWithNegativeIndex}, {"test_listpackDeleteInclusiveRange0_0", test_listpackDeleteInclusiveRange0_0}, {"test_listpackDeleteInclusiveRange0_1", test_listpackDeleteInclusiveRange0_1}, {"test_listpackDeleteInclusiveRange1_2", test_listpackDeleteInclusiveRange1_2}, {"test_listpackDeleteWitStartIndexOutOfRange", test_listpackDeleteWitStartIndexOutOfRange}, {"test_listpackDeleteWitNumOverflow", test_listpackDeleteWitNumOverflow}, {"test_listpackBatchDelete", test_listpackBatchDelete}, {"test_listpackDeleteFooWhileIterating", test_listpackDeleteFooWhileIterating}, {"test_listpackReplaceWithSameSize", test_listpackReplaceWithSameSize}, {"test_listpackReplaceWithDifferentSize", test_listpackReplaceWithDifferentSize}, {"test_listpackRegressionGt255Bytes", test_listpackRegressionGt255Bytes}, {"test_listpackCreateLongListAndCheckIndices", test_listpackCreateLongListAndCheckIndices}, {"test_listpackCompareStrsWithLpEntries", test_listpackCompareStrsWithLpEntries}, {"test_listpackLpMergeEmptyLps", test_listpackLpMergeEmptyLps}, {"test_listpackLpMergeLp1Larger", test_listpackLpMergeLp1Larger}, {"test_listpackLpMergeLp2Larger", test_listpackLpMergeLp2Larger}, {"test_listpackLpNextRandom", test_listpackLpNextRandom}, {"test_listpackLpNextRandomCC", test_listpackLpNextRandomCC}, {"test_listpackRandomPairWithOneElement", test_listpackRandomPairWithOneElement}, {"test_listpackRandomPairWithManyElements", test_listpackRandomPairWithManyElements}, {"test_listpackRandomPairsWithOneElement", test_listpackRandomPairsWithOneElement}, {"test_listpackRandomPairsWithManyElements", test_listpackRandomPairsWithManyElements}, {"test_listpackRandomPairsUniqueWithOneElement", test_listpackRandomPairsUniqueWithOneElement}, {"test_listpackRandomPairsUniqueWithManyElements", test_listpackRandomPairsUniqueWithManyElements}, {"test_listpackPushVariousEncodings", test_listpackPushVariousEncodings}, {"test_listpackLpFind", test_listpackLpFind}, {"test_listpackLpValidateIntegrity", test_listpackLpValidateIntegrity}, {"test_listpackNumberOfElementsExceedsLP_HDR_NUMELE_UNKNOWN", test_listpackNumberOfElementsExceedsLP_HDR_NUMELE_UNKNOWN}, {"test_listpackStressWithRandom", test_listpackStressWithRandom}, {"test_listpackSTressWithVariableSize", test_listpackSTressWithVariableSize}, {"test_listpackBenchmarkInit", test_listpackBenchmarkInit}, {"test_listpackBenchmarkLpAppend", test_listpackBenchmarkLpAppend}, {"test_listpackBenchmarkLpFindString", test_listpackBenchmarkLpFindString}, {"test_listpackBenchmarkLpFindNumber", test_listpackBenchmarkLpFindNumber}, {"test_listpackBenchmarkLpSeek", test_listpackBenchmarkLpSeek}, {"test_listpackBenchmarkLpValidateIntegrity", test_listpackBenchmarkLpValidateIntegrity}, {"test_listpackBenchmarkLpCompareWithString", test_listpackBenchmarkLpCompareWithString}, {"test_listpackBenchmarkLpCompareWithNumber", test_listpackBenchmarkLpCompareWithNumber}, {"test_listpackBenchmarkFree", test_listpackBenchmarkFree}, {NULL, NULL}};
+unitTest __test_object_c[] = {{"test_valkey_from_embstr", test_valkey_from_embstr}, {NULL, NULL}};
 unitTest __test_sds_c[] = {{"test_sds", test_sds}, {"test_typesAndAllocSize", test_typesAndAllocSize}, {"test_sdsHeaderSizes", test_sdsHeaderSizes}, {NULL, NULL}};
 unitTest __test_sha1_c[] = {{"test_sha1", test_sha1}, {NULL, NULL}};
 unitTest __test_util_c[] = {{"test_string2ll", test_string2ll}, {"test_string2l", test_string2l}, {"test_ll2string", test_ll2string}, {"test_ld2string", test_ld2string}, {"test_fixedpoint_d2string", test_fixedpoint_d2string}, {"test_version2num", test_version2num}, {"test_reclaimFilePageCache", test_reclaimFilePageCache}, {NULL, NULL}};
@@ -176,6 +178,7 @@ struct unitTestSuite {
     {"test_intset.c", __test_intset_c},
     {"test_kvstore.c", __test_kvstore_c},
     {"test_listpack.c", __test_listpack_c},
+    {"test_object.c", __test_object_c},
     {"test_sds.c", __test_sds_c},
     {"test_sha1.c", __test_sha1_c},
     {"test_util.c", __test_util_c},