feat(cluster): Allow appending RDB to existing store

src/server/cluster/cluster_family_test.cc

@@ -597,7 +597,7 @@ TEST_F(ClusterFamilyTest, ClusterFirstConfigCallDropsEntriesNotOwnedByNode) {
   EXPECT_EQ(Run({"save", "df"}), "OK");
 
   auto save_info = service_->server_family().GetLastSaveInfo();
-  EXPECT_EQ(Run({"debug", "load", save_info.file_name}), "OK");
+  EXPECT_EQ(Run({"dfly", "load", save_info.file_name}), "OK");
   EXPECT_EQ(CheckedInt({"dbsize"}), 50000);
 
   ConfigSingleNodeCluster("abcd1234");

src/server/debugcmd.cc

@@ -373,7 +373,6 @@ void DebugCmd::Run(CmdArgList args) {
         "    arguments. Each descriptor is prefixed by its frequency count",
         "OBJECT <key> [COMPRESS]",
         "    Show low-level info about `key` and associated value.",
-        "LOAD <filename>",
         "RELOAD [option ...]",
         "    Save the RDB on disk and reload it back to memory. Valid <option> values:",
         "    * NOSAVE: the database will be loaded from an existing RDB file.",
@@ -431,10 +430,6 @@ void DebugCmd::Run(CmdArgList args) {
     return Watched();
   }
 
-  if (subcmd == "LOAD" && args.size() == 2) {
-    return Load(ArgS(args, 1));
-  }
-
   if (subcmd == "OBJECT" && args.size() >= 2) {
     string_view key = ArgS(args, 1);
     args.remove_prefix(2);
@@ -500,7 +495,19 @@ void DebugCmd::Reload(CmdArgList args) {
   }
 
   string last_save_file = sf_.GetLastSaveInfo().file_name;
-  Load(last_save_file);
+
+  sf_.FlushAll(cntx_);
+
+  if (auto fut_ec = sf_.Load(last_save_file, ServerFamily::LoadExistingKeys::kFail); fut_ec) {
+    GenericError ec = fut_ec->Get();
+    if (ec) {
+      string msg = ec.Format();
+      LOG(WARNING) << "Could not load file " << msg;
+      return cntx_->SendError(msg);
+    }
+  }
+
+  cntx_->SendOk();
 }
 
 void DebugCmd::Replica(CmdArgList args) {
@@ -529,52 +536,6 @@ void DebugCmd::Replica(CmdArgList args) {
   return cntx_->SendError(UnknownSubCmd("replica", "DEBUG"));
 }
 
-void DebugCmd::Load(string_view filename) {
-  if (!ServerState::tlocal()->is_master) {
-    return cntx_->SendError("Replica cannot load data");
-  }
-
-  auto new_state = sf_.service().SwitchState(GlobalState::ACTIVE, GlobalState::LOADING);
-
-  if (new_state != GlobalState::LOADING) {
-    LOG(WARNING) << new_state << " in progress, ignored";
-    return cntx_->SendError("Could not load file");
-  }
-
-  absl::Cleanup rev_state = [this] {
-    sf_.service().SwitchState(GlobalState::LOADING, GlobalState::ACTIVE);
-  };
-
-  const CommandId* cid = sf_.service().FindCmd("FLUSHALL");
-  intrusive_ptr<Transaction> flush_trans(new Transaction{cid});
-  flush_trans->InitByArgs(cntx_->ns, 0, {});
-  VLOG(1) << "Performing flush";
-  error_code ec = sf_.Drakarys(flush_trans.get(), DbSlice::kDbAll);
-  if (ec) {
-    LOG(ERROR) << "Error flushing db " << ec.message();
-  }
-
-  fs::path path(filename);
-
-  if (filename.empty()) {
-    fs::path dir_path(GetFlag(FLAGS_dir));
-    string filename = GetFlag(FLAGS_dbfilename);
-    dir_path.append(filename);
-    path = dir_path;
-  }
-
-  if (auto fut_ec = sf_.Load(path.generic_string()); fut_ec) {
-    GenericError ec = fut_ec->Get();
-    if (ec) {
-      string msg = ec.Format();
-      LOG(WARNING) << "Could not load file " << msg;
-      return cntx_->SendError(msg);
-    }
-  }
-
-  cntx_->SendOk();
-}
-
 optional<DebugCmd::PopulateOptions> DebugCmd::ParsePopulateArgs(CmdArgList args) {
   if (args.size() < 2) {
     cntx_->SendError(UnknownSubCmd("populate", "DEBUG"));

src/server/debugcmd.h

@@ -30,9 +30,6 @@ class DebugCmd {
 
   void Run(CmdArgList args);
 
-  // A public function that loads a snapshot.
-  void Load(std::string_view filename);
-
   static void Shutdown();
 
  private:

src/server/dflycmd.cc

@@ -163,11 +163,33 @@ void DflyCmd::Run(CmdArgList args, ConnectionContext* cntx) {
     return ReplicaOffset(args, cntx);
   }
 
-  if (sub_cmd == "LOAD" && args.size() == 2) {
-    DebugCmd debug_cmd{sf_, cntx};
-    debug_cmd.Load(ArgS(args, 1));
-    return;
+  if (sub_cmd == "LOAD") {
+    return Load(args, cntx);
+  }
+
+  if (sub_cmd == "HELP") {
+    string_view help_arr[] = {
+        "DFLY <subcommand> [<arg> [value] [opt] ...]. Subcommands are:",
+        "THREAD",
+        "    Returns connection thread index and number of threads",
+        "THREAD <thread-id>",
+        "    Migrates connection to thread <thread-id>",
+        "EXPIRE",
+        "    Collects all expired items.",
+        "REPLICAOFFSET",
+        "    Returns LSN (log sequence number) per shard. These are the sequential ids of the ",
+        "    journal entry.",
+        "LOAD <filename> [APPEND]",
+        "    Loads <filename> RDB/DFS file into the data store.",
+        "    * APPEND: Existing keys are NOT removed before loading the file, conflicting ",
+        "      keys (that exist in both data store and in file) are overridden.",
+        "HELP",
+        "    Prints this help.",
+    };
+    auto* rb = static_cast<RedisReplyBuilder*>(cntx->reply_builder());
+    return rb->SendSimpleStrArr(help_arr);
   }
+
   cntx->SendError(kSyntaxErr);
 }
 
@@ -500,6 +522,41 @@ void DflyCmd::ReplicaOffset(CmdArgList args, ConnectionContext* cntx) {
   }
 }
 
+void DflyCmd::Load(CmdArgList args, ConnectionContext* cntx) {
+  CmdArgParser parser{args};
+  parser.ExpectTag("LOAD");
+  string_view filename = parser.Next();
+  ServerFamily::LoadExistingKeys existing_keys = ServerFamily::LoadExistingKeys::kFail;
+
+  if (parser.HasNext()) {
+    parser.ExpectTag("APPEND");
+    existing_keys = ServerFamily::LoadExistingKeys::kOverride;
+  }
+
+  if (parser.HasNext()) {
+    parser.Error();
+  }
+
+  if (parser.HasError()) {
+    return cntx->SendError(kSyntaxErr);
+  }
+
+  if (existing_keys == ServerFamily::LoadExistingKeys::kFail) {
+    sf_->FlushAll(cntx);
+  }
+
+  if (auto fut_ec = sf_->Load(filename, existing_keys); fut_ec) {
+    GenericError ec = fut_ec->Get();
+    if (ec) {
+      string msg = ec.Format();
+      LOG(WARNING) << "Could not load file " << msg;
+      return cntx->SendError(msg);
+    }
+  }
+
+  cntx->SendOk();
+}
+
 OpStatus DflyCmd::StartFullSyncInThread(FlowInfo* flow, Context* cntx, EngineShard* shard) {
   DCHECK(!flow->full_sync_fb.IsJoinable());
   DCHECK(shard);

src/server/dflycmd.h

@@ -199,6 +199,8 @@ class DflyCmd {
   // Return journal records num sent for each flow of replication.
   void ReplicaOffset(CmdArgList args, ConnectionContext* cntx);
 
+  void Load(CmdArgList args, ConnectionContext* cntx);
+
   // Start full sync in thread. Start FullSyncFb. Called for each flow.
   facade::OpStatus StartFullSyncInThread(FlowInfo* flow, Context* cntx, EngineShard* shard);
 

src/server/rdb_load.cc

@@ -35,6 +35,8 @@ extern "C" {
 #include "core/sorted_map.h"
 #include "core/string_map.h"
 #include "core/string_set.h"
+#include "server/cluster/cluster_defs.h"
+#include "server/cluster/cluster_family.h"
 #include "server/engine_shard_set.h"
 #include "server/error.h"
 #include "server/hset_family.h"
@@ -2481,7 +2483,7 @@ void RdbLoader::LoadItemsBuffer(DbIndex db_ind, const ItemsBuf& ib) {
 
     auto& res = *op_res;
     res.it->first.SetSticky(item->is_sticky);
-    if (!res.is_new) {
+    if (!override_existing_keys_ && !res.is_new) {
       LOG(WARNING) << "RDB has duplicated key '" << item->key << "' in DB " << db_ind;
     }
 
@@ -2520,6 +2522,13 @@ error_code RdbLoader::LoadKeyValPair(int type, ObjSettings* settings) {
     return ec;
   }
 
+  if (!load_unowned_slots_ && cluster::IsClusterEnabled()) {
+    const cluster::ClusterConfig* cluster_config = cluster::ClusterFamily::cluster_config();
+    if (cluster_config != nullptr && !cluster_config->IsMySlot(item->key)) {
+      return kOk;  // Ignoring item
+    }
+  }
+
   /* Check if the key already expired. This function is used when loading
    * an RDB file from disk, either at startup, or when an RDB was
    * received from the master. In the latter case, the master is

src/server/rdb_load.h

@@ -181,7 +181,16 @@ class RdbLoader : protected RdbLoaderBase {
 
   ~RdbLoader();
 
+  void SetOverrideExistingKeys(bool override) {
+    override_existing_keys_ = override;
+  }
+
+  void SetLoadUnownedSlots(bool load_unowned) {
+    load_unowned_slots_ = load_unowned;
+  }
+
   std::error_code Load(::io::Source* src);
+
   void set_source_limit(size_t n) {
     source_limit_ = n;
   }
@@ -273,6 +282,8 @@ class RdbLoader : protected RdbLoaderBase {
 
  private:
   Service* service_;
+  bool override_existing_keys_ = false;
+  bool load_unowned_slots_ = false;
   ScriptMgr* script_mgr_;
   std::vector<ItemsBuf> shard_buf_;
 

src/server/rdb_test.cc

@@ -167,7 +167,7 @@ TEST_F(RdbTest, ComressionModeSaveDragonflyAndReload) {
     ASSERT_EQ(resp, "OK");
 
     auto save_info = service_->server_family().GetLastSaveInfo();
-    resp = Run({"debug", "load", save_info.file_name});
+    resp = Run({"dfly", "load", save_info.file_name});
     ASSERT_EQ(resp, "OK");
     ASSERT_EQ(50000, CheckedInt({"dbsize"}));
   }
@@ -182,7 +182,7 @@ TEST_F(RdbTest, RdbLoaderOnReadCompressedDataShouldNotEnterEnsureReadFlow) {
   ASSERT_EQ(resp, "OK");
 
   auto save_info = service_->server_family().GetLastSaveInfo();
-  resp = Run({"debug", "load", save_info.file_name});
+  resp = Run({"dfly", "load", save_info.file_name});
   ASSERT_EQ(resp, "OK");
 }
 
@@ -265,7 +265,7 @@ TEST_F(RdbTest, ReloadExpired) {
   ASSERT_EQ(resp, "OK");
   auto save_info = service_->server_family().GetLastSaveInfo();
   AdvanceTime(2000);
-  resp = Run({"debug", "load", save_info.file_name});
+  resp = Run({"dfly", "load", save_info.file_name});
   ASSERT_EQ(resp, "OK");
   resp = Run({"get", "key"});
   ASSERT_THAT(resp, ArgType(RespExpr::NIL));
@@ -543,4 +543,26 @@ TEST_F(RdbTest, SBF) {
   EXPECT_THAT(Run({"BF.EXISTS", "k", "1"}), IntArg(1));
 }
 
+TEST_F(RdbTest, DflyLoadAppend) {
+  // Create an RDB with (k1,1) value in it saved as `filename`
+  EXPECT_EQ(Run({"set", "k1", "1"}), "OK");
+  EXPECT_EQ(Run({"save", "df"}), "OK");
+  string filename = service_->server_family().GetLastSaveInfo().file_name;
+
+  // Without APPEND option - db should be flushed
+  EXPECT_EQ(Run({"set", "k1", "TO-BE-FLUSHED"}), "OK");
+  EXPECT_EQ(Run({"set", "k2", "TO-BE-FLUSHED"}), "OK");
+  EXPECT_EQ(Run({"dfly", "load", filename}), "OK");
+  EXPECT_THAT(Run({"dbsize"}), IntArg(1));
+  EXPECT_EQ(Run({"get", "k1"}), "1");
+
+  // With APPEND option - db shouldn't be flushed, but k1 should be overridden
+  EXPECT_EQ(Run({"set", "k1", "TO-BE-OVERRIDDEN"}), "OK");
+  EXPECT_EQ(Run({"set", "k2", "2"}), "OK");
+  EXPECT_EQ(Run({"dfly", "load", filename, "append"}), "OK");
+  EXPECT_THAT(Run({"dbsize"}), IntArg(2));
+  EXPECT_EQ(Run({"get", "k1"}), "1");
+  EXPECT_EQ(Run({"get", "k2"}), "2");
+}
+
 }  // namespace dfly

src/server/replica.cc

@@ -425,6 +425,7 @@ error_code Replica::InitiatePSync() {
     }
 
     RdbLoader loader(NULL);
+    loader.SetLoadUnownedSlots(true);
     loader.set_source_limit(snapshot_size);
     // TODO: to allow registering callbacks within loader to send '\n' pings back to master.
     // Also to allow updating last_io_time_.
@@ -935,6 +936,7 @@ DflyShardReplica::DflyShardReplica(ServerContext server_context, MasterContext m
       flow_id_(flow_id) {
   executor_ = std::make_unique<JournalExecutor>(service);
   rdb_loader_ = std::make_unique<RdbLoader>(&service_);
+  rdb_loader_->SetLoadUnownedSlots(true);
 }
 
 DflyShardReplica::~DflyShardReplica() {