From f16041a342d1d452975b2106fcedc6571c9ff5e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Svensson?= <bjorn.a.svensson@est.tech>
Date: Fri, 11 Oct 2024 11:42:00 +0200
Subject: [PATCH] Add OOM testing of parsing replicas
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fix memory issues when the parsing of replicas fails.

Signed-off-by: Björn Svensson <bjorn.a.svensson@est.tech>
---
 src/cluster.c                     | 17 ++++++++++++-----
 tests/ct_out_of_memory_handling.c | 10 ++++++----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/cluster.c b/src/cluster.c
index 589f4c5..66e9b2e 100644
--- a/src/cluster.c
+++ b/src/cluster.c
@@ -630,7 +630,6 @@ static int cluster_master_slave_mapping_with_name(valkeyClusterContext *cc,
             }
 
             if (node_old->slaves != NULL) {
-                node_old->slaves->free = NULL;
                 while (listLength(node_old->slaves) > 0) {
                     lnode = listFirst(node_old->slaves);
                     if (listAddNodeHead(node->slaves, lnode->value) == NULL) {
@@ -958,7 +957,6 @@ static dict *parse_cluster_nodes(valkeyClusterContext *cc, valkeyReply *reply) {
                     ret = cluster_master_slave_mapping_with_name(
                         cc, &nodes_name, master, master->name);
                     if (ret != VALKEY_OK) {
-                        freeValkeyClusterNode(master);
                         goto error;
                     }
                 }
@@ -1050,12 +1048,21 @@ static dict *parse_cluster_nodes(valkeyClusterContext *cc, valkeyReply *reply) {
 error:
     sdsfreesplitres(part, count_part);
     sdsfreesplitres(slot_start_end, count_slot_start_end);
-    if (nodes != NULL) {
-        dictRelease(nodes);
-    }
     if (nodes_name != NULL) {
+        /* Only free parsed replicas since the `nodes` dict owns primary nodes. */
+        dictIterator di;
+        dictInitIterator(&di, nodes_name);
+        dictEntry *de;
+        while ((de = dictNext(&di))) {
+            valkeyClusterNode *node = dictGetEntryVal(de);
+            if (node->role == VALKEY_ROLE_SLAVE)
+                freeValkeyClusterNode(node);
+        }
         dictRelease(nodes_name);
     }
+    if (nodes != NULL) {
+        dictRelease(nodes);
+    }
     return NULL;
 }
 
diff --git a/tests/ct_out_of_memory_handling.c b/tests/ct_out_of_memory_handling.c
index 86c023e..0a56a4f 100644
--- a/tests/ct_out_of_memory_handling.c
+++ b/tests/ct_out_of_memory_handling.c
@@ -125,6 +125,7 @@ void test_alloc_failure_handling(void) {
         cc = valkeyClusterContextInit();
         assert(cc);
     }
+    cc->flags |= VALKEYCLUSTER_FLAG_ADD_SLAVE;
 
     // Add nodes
     {
@@ -170,14 +171,14 @@ void test_alloc_failure_handling(void) {
 
     // Connect
     {
-        for (int i = 0; i < 128; ++i) {
+        for (int i = 0; i < 148; ++i) {
             prepare_allocation_test(cc, i);
             result = valkeyClusterConnect2(cc);
             assert(result == VALKEY_ERR);
             ASSERT_STR_EQ(cc->errstr, "Out of memory");
         }
 
-        prepare_allocation_test(cc, 128);
+        prepare_allocation_test(cc, 148);
         result = valkeyClusterConnect2(cc);
         assert(result == VALKEY_OK);
     }
@@ -493,6 +494,7 @@ void test_alloc_failure_handling_async(void) {
         acc = valkeyClusterAsyncContextInit();
         assert(acc);
     }
+    acc->cc->flags |= VALKEYCLUSTER_FLAG_ADD_SLAVE;
 
     // Set callbacks
     {
@@ -519,14 +521,14 @@ void test_alloc_failure_handling_async(void) {
 
     // Connect
     {
-        for (int i = 0; i < 126; ++i) {
+        for (int i = 0; i < 146; ++i) {
             prepare_allocation_test(acc->cc, i);
             result = valkeyClusterConnect2(acc->cc);
             assert(result == VALKEY_ERR);
             ASSERT_STR_EQ(acc->cc->errstr, "Out of memory");
         }
 
-        prepare_allocation_test(acc->cc, 126);
+        prepare_allocation_test(acc->cc, 146);
         result = valkeyClusterConnect2(acc->cc);
         assert(result == VALKEY_OK);
     }