diff --git a/src/ebuckets.c b/src/ebuckets.c index 3823525a0f3..c5a615463f4 100644 --- a/src/ebuckets.c +++ b/src/ebuckets.c @@ -1434,6 +1434,70 @@ uint64_t ebCascade(ebuckets *eb, EbucketsType *type, uint64_t now, uint64_t maxC return info.itemsExpired; /* Number of items cascaded */ } +/** + * Calculate the memory usage of an ebuckets data structure. + * + * This function returns the total amount of memory (in bytes) consumed by a given + * ebuckets instance (`eb`). The calculation depends on the internal representation + * of the ebuckets, which is determined by the provided `type`. + * + * Logic Overview: + * - If the ebuckets instance is empty (`ebIsEmpty` returns true), the memory usage is 0. + * - Otherwise, the size starts with the base `sizeof(ebuckets)` structure. + * + * - If the ebuckets type is a **stack**: + * - Add the size of the `ebStack` structure. + * - Recursively include memory used by its first (`l1`) and second (`l2`) levels, + * which are themselves ebuckets. + * - If the third level (`l3`) exists, include its memory. + * + * - If the ebuckets type is a **rax**: + * - Add the size of the `rax` structure. + * - Add memory for all `raxNode` and `FirstSegHdr` entries. + * - Add additional memory for `NextSegHdr` entries, which are used + * when the number of items exceeds a certain threshold. + * + * - If the ebuckets type is a **list**, no additional memory beyond the base + * `ebuckets` structure is used. + * + * @param eb - The ebuckets instance to analyze. + * @param type - Pointer to an `EbucketsType` structure that defines the ebuckets layout. + * @return - Total size in bytes of memory used by the ebuckets instance. + */ +size_t ebMemUsage(ebuckets eb, EbucketsType *type) { + if (ebIsEmpty(eb)) + return 0; + + size_t size = 0; + if (type->isEbStack) { + ebStack *stack = (ebStack *)eb; + size += sizeof(ebStack); + EB_STACK_EXEC_L1(type, size += ebMemUsage(stack->l1, type)); + EB_STACK_EXEC_L1(type, size += ebMemUsage(stack->l2, type)); + if (stack->l3) { + size += EB_STACK_L3_SIZEOF(stack->l3->vecSize); + } + } else if(!ebIsList(eb)) { /* If not a list, then it is a rax */ + size += sizeof(rax); /* rax structure itself */ + + rax *rax = ebGetRaxPtr(eb); + raxIterator iter; + raxStart(&iter, rax); + raxSeek(&iter, "^", NULL, 0); + while (raxNext(&iter)) { + FirstSegHdr *seg = iter.data; + size += sizeof(raxNode) + iter.key_len; /* raxNode + key */ + if(seg->numSegs > 0) { + /* If the bucket has segments, then add size of FirstSegHdr */ + size += sizeof(FirstSegHdr); + size += (seg->numSegs - 1) * sizeof(NextSegHdr); /* NextSegHdr for each segment */ + } + } + raxStop(&iter); + } + return size; +} + int ebAddToRax(ebuckets *eb, EbucketsType *type, eItem item, uint64_t bucketKeyItem) { EBucketNew newBucket; /* ebAddToBucket takes care to update newBucket.segment.head */ raxIterator iter; @@ -3097,6 +3161,107 @@ int ebucketsTest(int argc, char **argv, int flags) { } } + TEST("ebMemUsage - empty ebucket returns 0") { + ebuckets eb = NULL; + size_t usage = ebMemUsage(eb, &myEbType); + assert(usage == 0); + } + + TEST("ebMemUsage - list-based ebucket returns zero memory usage") { + ebuckets eb = NULL; + MyItem *items[EB_LIST_MAX_ITEMS]; + for (int i = 0; i < EB_LIST_MAX_ITEMS; i++) { + items[i] = zmalloc(sizeof(MyItem)); + ebAdd(&eb, &myEbType, items[i], i); + } + + assert(ebIsList(eb)); + size_t usage = ebMemUsage(eb, &myEbType); + assert(usage == 0); + + ebDestroy(&eb, &myEbType, NULL); + } + + TEST("ebMemUsage - rax-based ebucket with three segments returns correct memory usage") { + ebuckets eb = NULL; + for (int i = 0; i < EB_LIST_MAX_ITEMS + 10; i++) { + MyItem *item = zmalloc(sizeof(MyItem)); + ebAdd(&eb, &myEbType, item, i + 1); + } + + assert(!ebIsList(eb)); + size_t usage = ebMemUsage(eb, &myEbType); + /* Calculation breakdown: + 3 * (16 bytes for FirstSegHdr + 4 bytes for raxNode + 6 bytes for key size) + + 24 bytes for the rax structure + Total: 102 bytes + */ + assert(usage == 102); + + ebDestroy(&eb, &myEbType, NULL); + } + + TEST("ebuckets - rax-based ebucket with extended segment returns correct memory usage") { + ebuckets eb = NULL; + for (int i = 0; i < 2*EB_SEG_MAX_ITEMS; i++) { + MyItem *item = zmalloc(sizeof(MyItem)); + ebAdd(&eb, &myEbType, item, 1); + } + + assert(!ebIsList(eb)); + size_t usage = ebMemUsage(eb, &myEbType); + /* Calculation breakdown: + 1 * (16 bytes for FirstSegHdr + 4 bytes for raxNode + 6 bytes for key size) + + 1 * 24 bytes for NextSegHdr + + 24 bytes for the rax structure + Total: 74 bytes + */ + assert(usage == 74); + + ebDestroy(&eb, &myEbType, NULL); + } + + TEST("ebMemUsage - ebStack-based ebucket returns memory usage for L1, L2 and L3") { + EbucketsType stackType = myEbType2; + stackType.isEbStack = 1; + + ebuckets eb = zmalloc(sizeof(ebStack)); + ebStack *stack = (ebStack *)eb; + + stack->l1 = ebCreate(); + stack->l2 = ebCreate(); + stack->l3 = zmalloc(EB_STACK_L3_SIZEOF(2)); + stack->l3->items = 0; + stack->l3->vecSize = 2; + + for (int i = 0; i < EB_SEG_MAX_ITEMS + 10; i++) { + MyItem *item = zmalloc(sizeof(MyItem)); + ebAdd(&stack->l1, &myEbType, item, i + 1); + } + + for (int i = 0; i < 2*EB_SEG_MAX_ITEMS; i++) { + MyItem *item = zmalloc(sizeof(MyItem)); + ebAdd(&stack->l2, &myEbType, item, 1); + } + + size_t usage = ebMemUsage(eb, &stackType); + /* Calculation breakdown: + L1: + 3 * (16 bytes for FirstSegHdr + 4 bytes for raxNode + 6 bytes for key size) + + 24 bytes for the rax structure + + 1 * (16 bytes for FirstSegHdr + 4 bytes for raxNode + 5 bytes for key size) + L2: + + 1 * 24 bytes for NextSegHdr + + 24 bytes for the rax structure + L3: + + 2 * 4 for vector item + 16 bytes for ebVector + Total: 240 bytes + */ + assert(usage == 240); + + ebDestroy(&eb, &stackType, NULL); + } + // TEST("segment - Add smaller item to full segment that all share same ebucket-key") // TEST("segment - Add item to full segment and make it extended-segment (all share same ebucket-key)") // TEST("ebuckets - Create rax tree with extended-segment and add item before") diff --git a/src/ebuckets.h b/src/ebuckets.h index 9bab743983a..beaced021aa 100644 --- a/src/ebuckets.h +++ b/src/ebuckets.h @@ -362,6 +362,8 @@ int ebAdd(ebuckets *eb, EbucketsType *type, eItem item, uint64_t expireTime); uint64_t ebCascade(ebuckets *eb, EbucketsType *type, uint64_t now, uint64_t maxCascade); +size_t ebMemUsage(ebuckets eb, EbucketsType *type); + uint64_t ebGetExpireTime(EbucketsType *type, eItem item); void ebStart(EbucketsIterator *iter, ebuckets eb, EbucketsType *type); diff --git a/src/expire.c b/src/expire.c index 9a5abf3cc7f..2180bcff1ee 100644 --- a/src/expire.c +++ b/src/expire.c @@ -413,12 +413,17 @@ ebuckets *estoreGetBucket(estore *es, int slot) { size_t estoreMemUsage(estore *es) { if (es == NULL) return 0; - // TODO_MOTI: Follow kvstoreMemUsage() to imp for estoreMemUsage() -// size_t mem = sizeof(*es); -// for (int i = 0; i < es->num_buckets; i++) { -// mem += ebMemUsage(es->buckets[i], es->bucket_type); -// } - return 0; + size_t mem = sizeof(estore); /* Size of the estore structure */ + mem += sizeof(ebuckets) * es->num_buckets; /* Buckets array */ + if (!server.cluster_enabled) { + mem += ebMemUsage(es->buckets + 0, es->bucket_type); + return mem; + } + + for(int i = 0; i < es->num_buckets; i++) + mem += ebMemUsage(es->buckets + i, es->bucket_type); + + return mem; } @@ -1107,3 +1112,97 @@ void touchCommand(client *c) { if (lookupKeyRead(c->db,c->argv[j]) != NULL) touched++; addReplyLongLong(c,touched); } + +#ifdef REDIS_TEST +#include +#include "testhelp.h" + +#define TEST(name) printf("test — %s\n", name); +typedef struct TestKVObj { + kvobj obj; + ExpireMeta mexpire; +} TestKVObj; + +ExpireMeta *getTestKVObjExpireMeta(const void *item) { + return &((TestKVObj *)item)->mexpire; +} + +void deleteTestKVObjCallback(eItem item, void *ctx) { + UNUSED(ctx); + UNUSED(item); +} + +EbucketsType testEbType = { + .getExpireMeta = getTestKVObjExpireMeta, + .onDeleteItem = deleteTestKVObjCallback, + .itemsAddrAreOdd = 0, + .ebp.precision = 0, + .ebp.keySize = EB_PRECISION2KEYSIZE(0 /*.precision*/), +}; + +/* ./redis-server test expire */ +int expireTest(int argc, char **argv, int flags) { + UNUSED(argc); + UNUSED(argv); + UNUSED(flags); + + TEST("estoreMemUsage returns 0 when estore is NULL") { + assert(estoreMemUsage(NULL) == 0); + } + + TEST("estoreMemUsage computes memory correctly in non-cluster mode") { + server.cluster_enabled = 0; + + estore *es = estoreCreate(&testEbType, 1); /* 2 buckets */ + + TestKVObj *kv1 = zmalloc(sizeof(TestKVObj)); + TestKVObj *kv2 = zmalloc(sizeof(TestKVObj)); + + estoreAdd(es, (kvobj*)kv1, 0, 1); + estoreAdd(es, (kvobj*)kv2, 0, 2); + + size_t expected = sizeof(estore) + sizeof(ebuckets) * 2; /* Size of the estore structure + buckets array */ + expected += ebMemUsage(es->buckets + 0, es->bucket_type); /* Only one bucket considered */ + + size_t actual = estoreMemUsage(es); + assert(actual == expected); + + zfree(kv1); + zfree(kv2); + estoreRelease(es); + } + + TEST("estoreMemUsage computes memory correctly in cluster mode") { + server.cluster_enabled = 1; + + estore *es = estoreCreate(&testEbType, 2); /* 4 buckets */ + + TestKVObj *kv1 = zmalloc(sizeof(TestKVObj)); + TestKVObj *kv2 = zmalloc(sizeof(TestKVObj)); + TestKVObj *kv3 = zmalloc(sizeof(TestKVObj)); + TestKVObj *kv4 = zmalloc(sizeof(TestKVObj)); + + estoreAdd(es, (kvobj*)kv1, 0, 3); + estoreAdd(es, (kvobj*)kv2, 1, 4); + estoreAdd(es, (kvobj*)kv3, 2, 5); + estoreAdd(es, (kvobj*)kv4, 3, 6); + + size_t expected = sizeof(estore) + sizeof(ebuckets) * 4; + for (int i = 0; i < 4; i++) { + expected += ebMemUsage(es->buckets + i, es->bucket_type); + } + + size_t actual = estoreMemUsage(es); + assert(actual == expected); + + zfree(kv1); + zfree(kv2); + zfree(kv3); + zfree(kv4); + estoreRelease(es); + } + + return 0; +} + +#endif \ No newline at end of file diff --git a/src/expire.h b/src/expire.h index 401e403aaf6..dd6b9ae76f2 100644 --- a/src/expire.h +++ b/src/expire.h @@ -60,4 +60,8 @@ ebuckets *estoreGetBucket(estore *es, int slot); size_t estoreMemUsage(estore *es); +#ifdef REDIS_TEST +int expireTest(int argc, char *argv[], int flags); +#endif + #endif diff --git a/src/server.c b/src/server.c index f5c7cf7558a..bbddcdd8bc7 100644 --- a/src/server.c +++ b/src/server.c @@ -7296,6 +7296,7 @@ struct redisTest { {"listpack", listpackTest}, {"kvstore", kvstoreTest}, {"ebuckets", ebucketsTest}, + {"expire", expireTest} }; redisTestProc *getTestProcByName(const char *name) { int numtests = sizeof(redisTests)/sizeof(struct redisTest);