Rax size tracking

src/rax.c

@@ -192,6 +192,7 @@ rax *raxNew(void) {
     rax->numele = 0;
     rax->numnodes = 1;
     rax->head = raxNewNode(0, 0);
+    rax->alloc_size = rax_alloc_size(rax) + rax_alloc_size(rax->head);
     if (rax->head == NULL) {
         rax_free(rax);
         return NULL;
@@ -507,11 +508,17 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
      * our key. We have just to reallocate the node and make space for the
      * data pointer. */
     if (i == len && (!h->iscompr || j == 0 /* not in the middle if j is 0 */)) {
+        size_t oldalloc = 0;
+
         debugf("### Insert: node representing key exists\n");
         /* Make space for the value pointer if needed. */
         if (!h->iskey || (h->isnull && overwrite)) {
+            oldalloc = rax_alloc_size(h);
             h = raxReallocForData(h, data);
-            if (h) memcpy(parentlink, &h, sizeof(h));
+            if (h) {
+                memcpy(parentlink, &h, sizeof(h));
+                rax->alloc_size = rax->alloc_size - oldalloc + rax_alloc_size(h);
+            }
         }
         if (h == NULL) {
             errno = ENOMEM;
@@ -706,6 +713,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
             return 0;
         }
         splitnode->data[0] = h->data[j];
+        rax->alloc_size += rax_alloc_size(splitnode);
 
         if (j == 0) {
             /* 3a: Replace the old node with the split node. */
@@ -730,6 +738,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
             memcpy(parentlink, &trimmed, sizeof(trimmed));
             parentlink = cp; /* Set parentlink to splitnode parent. */
             rax->numnodes++;
+            rax->alloc_size += rax_alloc_size(trimmed);
         }
 
         /* 4: Create the postfix node: what remains of the original
@@ -744,6 +753,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
             raxNode **cp = raxNodeLastChildPtr(postfix);
             memcpy(cp, &next, sizeof(next));
             rax->numnodes++;
+            rax->alloc_size += rax_alloc_size(postfix);
         } else {
             /* 4b: just use next as postfix node. */
             postfix = next;
@@ -756,6 +766,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
         /* 6. Continue insertion: this will cause the splitnode to
          * get a new child (the non common character at the currently
          * inserted key). */
+        rax->alloc_size -= rax_alloc_size(h);
         rax_free(h);
         h = splitnode;
     } else if (h->iscompr && i == len) {
@@ -794,6 +805,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
         raxNode **cp = raxNodeLastChildPtr(postfix);
         memcpy(cp, &next, sizeof(next));
         rax->numnodes++;
+        rax->alloc_size += rax_alloc_size(postfix);
 
         /* 3: Trim the compressed node. */
         trimmed->size = j;
@@ -806,6 +818,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
             void *aux = raxGetData(h);
             raxSetData(trimmed, aux);
         }
+        rax->alloc_size += rax_alloc_size(trimmed);
 
         /* Fix the trimmed node child pointer to point to
          * the postfix node. */
@@ -815,6 +828,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
         /* Finish! We don't need to continue with the insertion
          * algorithm for ALGO 2. The key is already inserted. */
         rax->numele++;
+        rax->alloc_size -= rax_alloc_size(h);
         rax_free(h);
         return 1; /* Key inserted. */
     }
@@ -823,6 +837,7 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
      * chars in our string. We need to insert the missing nodes. */
     while (i < len) {
         raxNode *child;
+        size_t oldalloc = rax_alloc_size(h);
 
         /* If this node is going to have a single child, and there
          * are other characters, so that that would result in a chain
@@ -848,14 +863,17 @@ int raxGenericInsert(rax *rax, unsigned char *s, size_t len, void *data, void **
             i++;
         }
         rax->numnodes++;
+        rax->alloc_size = rax->alloc_size - oldalloc + rax_alloc_size(h) + rax_alloc_size(child);
         h = child;
     }
+    size_t oldalloc = rax_alloc_size(h);
     raxNode *newh = raxReallocForData(h, data);
     if (newh == NULL) goto oom;
     h = newh;
     if (!h->iskey) rax->numele++;
     raxSetData(h, data);
     memcpy(parentlink, &h, sizeof(h));
+    rax->alloc_size = rax->alloc_size - oldalloc + rax_alloc_size(h);
     return 1; /* Element inserted. */
 
 oom:
@@ -1025,6 +1043,7 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
             child = h;
             debugf("Freeing child %p [%.*s] key:%d\n", (void *)child, (int)child->size, (char *)child->data,
                    child->iskey);
+            rax->alloc_size -= rax_alloc_size(child);
             rax_free(child);
             rax->numnodes--;
             h = raxStackPop(&ts);
@@ -1034,7 +1053,9 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
         }
         if (child) {
             debugf("Unlinking child %p from parent %p\n", (void *)child, (void *)h);
+            size_t oldalloc = rax_alloc_size(h);
             raxNode *new = raxRemoveChild(h, child);
+            rax->alloc_size = rax->alloc_size - oldalloc + rax_alloc_size(new);
             if (new != h) {
                 raxNode *parent = raxStackPeek(&ts);
                 raxNode **parentlink;
@@ -1151,6 +1172,7 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
             new->iscompr = 1;
             new->size = comprsize;
             rax->numnodes++;
+            rax->alloc_size += rax_alloc_size(new);
 
             /* Scan again, this time to populate the new node content and
              * to fix the new node child pointer. At the same time we free
@@ -1163,6 +1185,7 @@ int raxRemove(rax *rax, unsigned char *s, size_t len, void **old) {
                 raxNode **cp = raxNodeLastChildPtr(h);
                 raxNode *tofree = h;
                 memcpy(&h, cp, sizeof(h));
+                rax->alloc_size -= rax_alloc_size(tofree);
                 rax_free(tofree);
                 rax->numnodes--;
                 if (h->iskey || (!h->iscompr && h->size != 1)) break;
@@ -1764,6 +1787,11 @@ uint64_t raxSize(rax *rax) {
     return rax->numele;
 }
 
+/* Return the rax tree allocation size in bytes */
+size_t raxAllocSize(rax *rax) {
+    return rax->alloc_size;
+}
+
 /* ----------------------------- Introspection ------------------------------ */
 
 /* This function is mostly used for debugging and learning purposes.

src/rax.h

@@ -131,9 +131,10 @@ typedef struct raxNode {
 } raxNode;
 
 typedef struct rax {
-    raxNode *head;
-    uint64_t numele;
-    uint64_t numnodes;
+    raxNode *head;         /* Pointer to root node of tree */
+    uint64_t numele;       /* Number of keys in the tree */
+    uint64_t numnodes;     /* Number of rax nodes in the tree */
+    size_t alloc_size;     /* Total allocation size of the tree in bytes */
 } rax;
 
 /* Stack data structure used by raxLowWalk() in order to, optionally, return
@@ -205,6 +206,7 @@ void raxStop(raxIterator *it);
 int raxEOF(raxIterator *it);
 void raxShow(rax *rax);
 uint64_t raxSize(rax *rax);
+size_t raxAllocSize(rax *rax);
 unsigned long raxTouch(raxNode *n);
 void raxSetDebugMsg(int onoff);
 

src/rax_malloc.h

@@ -41,4 +41,5 @@
 #define rax_malloc zmalloc
 #define rax_realloc zrealloc
 #define rax_free zfree
+#define rax_alloc_size zmalloc_size
 #endif

src/unit/test_files.h

@@ -23,6 +23,18 @@ int test_kvstoreIteratorRemoveAllKeysNoDeleteEmptyDict(int argc, char **argv, in
 int test_kvstoreIteratorRemoveAllKeysDeleteEmptyDict(int argc, char **argv, int flags);
 int test_kvstoreDictIteratorRemoveAllKeysNoDeleteEmptyDict(int argc, char **argv, int flags);
 int test_kvstoreDictIteratorRemoveAllKeysDeleteEmptyDict(int argc, char **argv, int flags);
+int test_randomWalkTest(int argc, char **argv, int flags);
+int test_iteratorUnitTests(int argc, char **argv, int flags);
+int test_tryInsertUnitTests(int argc, char **argv, int flags);
+int test_regtest1(int argc, char **argv, int flags);
+int test_regtest2(int argc, char **argv, int flags);
+int test_regtest3(int argc, char **argv, int flags);
+int test_regtest4(int argc, char **argv, int flags);
+int test_regtest5(int argc, char **argv, int flags);
+int test_regtest6(int argc, char **argv, int flags);
+int test_benchmark(int argc, char **argv, int flags);
+int test_hugeKey(int argc, char **argv, int flags);
+int test_fuzz(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);
@@ -79,6 +91,7 @@ unitTest __test_crc64combine_c[] = {{"test_crc64combine", test_crc64combine}, {N
 unitTest __test_endianconv_c[] = {{"test_endianconv", test_endianconv}, {NULL, NULL}};
 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_rax_c[] = {{"test_randomWalkTest", test_randomWalkTest}, {"test_iteratorUnitTests", test_iteratorUnitTests}, {"test_tryInsertUnitTests", test_tryInsertUnitTests}, {"test_regtest1", test_regtest1}, {"test_regtest2", test_regtest2}, {"test_regtest3", test_regtest3}, {"test_regtest4", test_regtest4}, {"test_regtest5", test_regtest5}, {"test_regtest6", test_regtest6}, {"test_benchmark", test_benchmark}, {"test_hugeKey", test_hugeKey}, {"test_fuzz", test_fuzz}, {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}};
@@ -94,6 +107,7 @@ struct unitTestSuite {
     {"test_endianconv.c", __test_endianconv_c},
     {"test_intset.c", __test_intset_c},
     {"test_kvstore.c", __test_kvstore_c},
+    {"test_rax.c", __test_rax_c},
     {"test_sds.c", __test_sds_c},
     {"test_sha1.c", __test_sha1_c},
     {"test_util.c", __test_util_c},