From 54cf110f7b3303e602da044171aa2054578f8c85 Mon Sep 17 00:00:00 2001
From: Liang Tang <tangliang@google.com>
Date: Tue, 28 May 2024 17:57:14 +0000
Subject: [PATCH] Persistence - Remove Unowned Keys

For cluster, after startup loading, remove keys
that shouldn't be served by this server based on slot
assignment of a cluster.

Also added stat fields in server to count the total
removed keys and skipped slots from last loading.
---
 src/cluster_legacy.c | 26 ++++++++++++++++++++++++++
 src/server.c         |  9 ++++++++-
 src/server.h         |  2 ++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/cluster_legacy.c b/src/cluster_legacy.c
index 63364d3596..cd83af54da 100644
--- a/src/cluster_legacy.c
+++ b/src/cluster_legacy.c
@@ -6552,3 +6552,29 @@ void clusterReplicateOpenSlots(void) {
 
     zfree(argv);
 }
+
+del_stat delUnOwnedKeys() {
+    del_stat result = {0, 0};
+    if (!server.cluster_enabled) {
+        return result;
+    }
+
+    clusterNode *master;
+    if (clusterNodeIsMaster(server.cluster->myself)) {
+        master = myself;
+    } else {
+        master = server.cluster->myself->slaveof;
+    }
+    for (int i = 0; i < CLUSTER_SLOTS; i++) {
+        if (server.cluster->slots[i] != master) {
+            unsigned int deleted = delKeysInSlot(i);
+            result.deleted_keys += deleted;
+            if (deleted > 0) {
+                result.unowned_slots++;
+                serverLog(LL_NOTICE, "Deleted %u keys from unowned slot: %d after loading RDB", deleted, i);
+            }
+        }
+    }
+    serverLog(LL_NOTICE, "Deleted totoal: %d unowned keys from total: %d unowned slots after loading RDB", result.deleted_keys, result.unowned_slots);
+    return result;
+}
\ No newline at end of file
diff --git a/src/server.c b/src/server.c
index bdfb27ec16..94484e3295 100644
--- a/src/server.c
+++ b/src/server.c
@@ -2656,6 +2656,8 @@ void initServer(void) {
     server.rdb_save_time_start = -1;
     server.rdb_last_load_keys_expired = 0;
     server.rdb_last_load_keys_loaded = 0;
+    server.stat_persistence_startup_load_keys_deleted = 0;
+    server.stat_persistence_startup_load_unowned_slots = 0;
     server.dirty = 0;
     resetServerStats();
     /* A few stats we don't want to reset: server startup time, and peak mem. */
@@ -5582,7 +5584,9 @@ sds genValkeyInfoString(dict *section_dict, int all_sections, int everything) {
                                              aof_bio_fsync_status == C_OK) ? "ok" : "err",
             "aof_last_cow_size:%zu\r\n", server.stat_aof_cow_bytes,
             "module_fork_in_progress:%d\r\n", server.child_type == CHILD_TYPE_MODULE,
-            "module_fork_last_cow_size:%zu\r\n", server.stat_module_cow_bytes));
+            "module_fork_last_cow_size:%zu\r\n", server.stat_module_cow_bytes,
+            "total_startup_load_deleted_keys:%lld\r\n", server.stat_persistence_startup_load_keys_deleted,
+            "total_startup_load_unowned_slots:%lld\r\n", server.stat_persistence_startup_load_unowned_slots));
         /* clang-format on */
 
         if (server.aof_enabled) {
@@ -6941,6 +6945,9 @@ int main(int argc, char **argv) {
         serverLog(LL_NOTICE, "Server initialized");
         aofLoadManifestFromDisk();
         loadDataFromDisk();
+        del_stat stat = delUnOwnedKeys();
+        server.stat_persistence_startup_load_keys_deleted = stat.deleted_keys;
+        server.stat_persistence_startup_load_unowned_slots = stat.unowned_slots;
         aofOpenIfNeededOnServerStart();
         aofDelHistoryFiles();
         if (server.cluster_enabled) {
diff --git a/src/server.h b/src/server.h
index a179328bd9..164188c27d 100644
--- a/src/server.h
+++ b/src/server.h
@@ -1814,6 +1814,8 @@ struct valkeyServer {
                                                  invocation of the event loop. */
     unsigned int max_new_conns_per_cycle;     /* The maximum number of tcp connections that will be accepted during each
                                                  invocation of the event loop. */
+    long long stat_persistence_startup_load_keys_deleted;
+    long long stat_persistence_startup_load_unowned_slots;
     /* AOF persistence */
     int aof_enabled;                    /* AOF configuration */
     int aof_state;                      /* AOF_(ON|OFF|WAIT_REWRITE) */