feat: Snapshot save & load

save_load.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+killall -9 pikiwidb
+mkdir leader follower1 follower2
+
+cd leader && ulimit -n 99999  && rm -fr *  && ../bin/pikiwidb ../pikiwidb.conf --port 7777 &
+
+cd follower1 && ulimit -n 99999 && rm -fr * && ../bin/pikiwidb ../pikiwidb.conf --port 8888 &
+sleep 10
+redis-cli -p 7777 raft.cluster init 
+redis-benchmark -p 7777 -c 5 -n 10000 -r 10000000 -d 1024 -t set
+
+
+redis-cli -p 7777 raft.node DSS 
+redis-cli -p 7777 raft.node DSS
+
+redis-cli -p 8888 raft.cluster join 127.0.0.1:7777
+
+
+

src/checkpoint_manager.cpp

@@ -23,15 +23,6 @@ void CheckpointManager::Init(int instNum, DB* db) {
 
 void CheckpointManager::CreateCheckpoint(const std::string& path) {
   res_.clear();
-
-  if (!pstd::FileExists(path)) {
-    if (0 != pstd::CreatePath(path)) {
-      WARN("Create Dir {} fail!", path);
-      return;
-    }
-    INFO("Create Dir {} success!", path);
-  }
-
   std::lock_guard Lock(shared_mutex_);
   for (int i = 0; i < checkpoint_num_; ++i) {
     checkpoint_infoes_[i].checkpoint_in_process = true;

src/client.cc

@@ -12,10 +12,10 @@
 #include "config.h"
 #include "log.h"
 #include "pikiwidb.h"
+#include "praft.h"
 #include "pstd_string.h"
 #include "slow_log.h"
 #include "store.h"
-#include "praft.h"
 
 namespace pikiwidb {
 

src/cmd_admin.cc

@@ -6,9 +6,9 @@
  */
 
 #include "cmd_admin.h"
-#include "store.h"
 #include "braft/raft.h"
 #include "praft.h"
+#include "store.h"
 
 namespace pikiwidb {
 
@@ -80,7 +80,7 @@ void SelectCmd::DoCmd(PClient* client) {
   client->SetRes(CmdRes::kOK);
 }
 
-InfoCmd::InfoCmd(const std::string& name, int16_t arity) 
+InfoCmd::InfoCmd(const std::string& name, int16_t arity)
     : BaseCmd(name, arity, kCmdFlagsAdmin | kCmdFlagsReadonly, kAclCategoryAdmin) {}
 
 bool InfoCmd::DoInitial(PClient* client) { return true; }
@@ -127,7 +127,7 @@ void InfoCmd::DoCmd(PClient* client) {
       message += "raft_state:up\r\n";
     } else {
       message += "raft_state:down\r\n";
-    }    
+    }
     message += "raft_role:" + std::string(braft::state2str(node_status.state)) + "\r\n";
     // message += "raft_is_voting:" + node_status.is_voting + "\r\n";
     message += "raft_leader_id:" + node_status.leader_id.to_string() + "\r\n";
@@ -141,9 +141,10 @@ void InfoCmd::DoCmd(PClient* client) {
       if (!status.ok()) {
         return client->SetRes(CmdRes::kErrOther, status.error_str());
       }
-      
+
       for (int i = 0; i < peers.size(); i++) {
-        message += "raft_node" + std::to_string(i) + ":addr=" + butil::ip2str(peers[i].addr.ip).c_str() + ",port=" + std::to_string(peers[i].addr.port) + "\r\n";
+        message += "raft_node" + std::to_string(i) + ":addr=" + butil::ip2str(peers[i].addr.ip).c_str() +
+                   ",port=" + std::to_string(peers[i].addr.port) + "\r\n";
       }
     }
 

src/cmd_admin.h

@@ -85,14 +85,14 @@ class SelectCmd : public BaseCmd {
 };
 
 class InfoCmd : public BaseCmd {
-  public:
-    InfoCmd(const std::string& name, int16_t arity);
+ public:
+  InfoCmd(const std::string& name, int16_t arity);
 
-  protected:
-    bool DoInitial(PClient* client) override;
+ protected:
+  bool DoInitial(PClient* client) override;
 
-  private:
-    void DoCmd(PClient* client) override;
+ private:
+  void DoCmd(PClient* client) override;
 };
 
 }  // namespace pikiwidb

src/cmd_raft.cc

@@ -36,6 +36,8 @@ void RaftNodeCmd::DoCmd(PClient* client) {
     DoCmdAdd(client);
   } else if (!strcasecmp(cmd.c_str(), "REMOVE")) {
     DoCmdRemove(client);
+  } else if (!strcasecmp(cmd.c_str(), "DSS")) {
+    DoCmdSnapshot(client);
   } else {
     client->SetRes(CmdRes::kErrOther, "RAFT.NODE supports ADD / REMOVE only");
   }
@@ -70,6 +72,13 @@ void RaftNodeCmd::DoCmdRemove(PClient* client) {
   }
 }
 
+void RaftNodeCmd::DoCmdSnapshot(PClient* client) {
+  auto s = PRAFT.DoSnapshot();
+  if (s.ok()) {
+    client->SetRes(CmdRes::kOK);
+  }
+}
+
 RaftClusterCmd::RaftClusterCmd(const std::string& name, int16_t arity)
     : BaseCmd(name, arity, kCmdFlagsRaft, kAclCategoryRaft) {}
 

src/cmd_raft.h

@@ -25,7 +25,7 @@ namespace pikiwidb {
  *   :<new node id>
  *   :<dbid>
  *
- * RAFT.NODE REMOVE [id] 
+ * RAFT.NODE REMOVE [id]
  *   Remove an existing node from the cluster.
  * Reply:
  *   -NOCLUSTER ||
@@ -45,6 +45,7 @@ class RaftNodeCmd : public BaseCmd {
   void DoCmd(PClient *client) override;
   void DoCmdAdd(PClient *client);
   void DoCmdRemove(PClient *client);
+  void DoCmdSnapshot(PClient *client);
 
   static constexpr std::string_view kAddCmd = "ADD";
   static constexpr std::string_view kRemoveCmd = "REMOVE";

src/cmd_table_manager.cc

@@ -12,8 +12,8 @@
 #include "cmd_keys.h"
 #include "cmd_kv.h"
 #include "cmd_list.h"
-#include "cmd_set.h"
 #include "cmd_raft.h"
+#include "cmd_set.h"
 #include "cmd_table_manager.h"
 #include "cmd_zset.h"
 
@@ -46,7 +46,7 @@ void CmdTableManager::InitCmdTable() {
 
   // info
   ADD_COMMAND(Info, -1);
-  
+
   // raft
   ADD_COMMAND(RaftCluster, -1);
   ADD_COMMAND(RaftNode, -2);

src/db.cpp

@@ -30,6 +30,8 @@ DB::DB(int db_index, const std::string& db_path)
     ERROR("Storage open failed! {}", s.ToString());
     abort();
   }
+  checkpoint_manager_ = std::make_unique<CheckpointManager>();
+  checkpoint_manager_->Init(g_config.db_instance_num, this);
   opened_ = true;
   INFO("Open DB{} success!", db_index_);
 }
@@ -45,8 +47,14 @@ void DB::DoBgSave(CheckpointInfo& checkpoint_info, const std::string& path, int
   checkpoint_info.checkpoint_in_process = false;
 }
 
-void DB::CreateCheckpoint(const std::string& path) { checkpoint_manager_->CreateCheckpoint(path); }
+void DB::CreateCheckpoint(const std::string& path) {
+  if (0 != pstd::CreatePath(path + '/' + std::to_string(db_index_))) {
+    WARN("Create dir {} fail !", path + '/' + std::to_string(db_index_));
+    return;
+  }
+  checkpoint_manager_->CreateCheckpoint(path);
+}
 
 void DB::WaitForCheckpointDone() { checkpoint_manager_->WaitForCheckpointDone(); }
-  
+
 }  // namespace pikiwidb

src/db.h

@@ -55,7 +55,6 @@ class DB {
   bool opened_ = false;
 
   std::unique_ptr<CheckpointManager> checkpoint_manager_;
-
 };
 }  // namespace pikiwidb
 

src/pikiwidb.h

@@ -9,8 +9,8 @@
 #include "common.h"
 #include "event_loop.h"
 #include "io_thread_pool.h"
-#include "tcp_connection.h"
 #include "praft/praft.h"
+#include "tcp_connection.h"
 
 #define kPIKIWIDB_VERSION "4.0.0"
 

src/praft/praft.cc

@@ -8,17 +8,18 @@
 //
 //  praft.cc
 
-#include "praft.h"
-
 #include <cassert>
 #include <memory>
 #include <string>
 
+#include "braft/snapshot.h"
+
 #include "client.h"
 #include "config.h"
 #include "event_loop.h"
 #include "log.h"
 #include "pikiwidb.h"
+#include "praft.h"
 #include "praft.pb.h"
 #include "pstd_string.h"
 
@@ -308,6 +309,16 @@ butil::Status PRaft::RemovePeer(const std::string& peer) {
   return {0, "OK"};
 }
 
+butil::Status PRaft::DoSnapshot() {
+  if (!node_) {
+    return ERROR_LOG_AND_STATUS("Node is not initialized");
+  }
+  braft::SynchronizedClosure done;
+  node_->snapshot(&done);
+  done.wait();
+  return {0, "OK"};
+}
+
 void PRaft::OnJoinCmdConnectionFailed([[maybe_unused]] EventLoop* loop, const char* peer_ip, int port) {
   auto cli = join_ctx_.GetClient();
   if (cli) {
@@ -345,6 +356,46 @@ void PRaft::Apply(braft::Task& task) {
   }
 }
 
+void PRaft::add_all_files(const std::filesystem::path& dir, braft::SnapshotWriter* writer, const std::string& path) {
+  for (const auto& entry : std::filesystem::directory_iterator(dir)) {
+    if (entry.is_directory()) {
+      if (entry.path() != "." && entry.path() != "..") {
+        INFO("dir_path = {}", entry.path().string());
+        add_all_files(entry.path(), writer, path);
+      }
+    } else {
+      INFO("file_path = {}", std::filesystem::relative(entry.path(), path).string());
+      if (writer->add_file(std::filesystem::relative(entry.path(), path)) != 0) {
+        WARN("出出出出 错错错错错错 啦啦啦啦啦啦");
+      }
+    }
+  }
+}
+
+void PRaft::recursive_copy(const std::filesystem::path& source, const std::filesystem::path& destination) {
+  if (std::filesystem::is_regular_file(source)) {
+    if (source.filename() == "__raft_snapshot_meta") {
+      return;
+    } else if (source.extension() == ".sst") {
+      // Create a hard link
+      INFO("hard link success! source_file = {} , destination_file = {}", source.string(), destination.string());
+      ::link(source.c_str(), destination.c_str());
+    } else {
+      // Copy the file
+      INFO("copy success! source_file = {} , destination_file = {}", source.string(), destination.string());
+      std::filesystem::copy_file(source, destination, std::filesystem::copy_options::overwrite_existing);
+    }
+  } else {
+    if (!pstd::FileExists(destination)) {
+      pstd::CreateDir(destination);
+    }
+
+    for (const auto& entry : std::filesystem::directory_iterator(source)) {
+      recursive_copy(entry.path(), destination / entry.path().filename());
+    }
+  }
+}
+
 // @braft::StateMachine
 void PRaft::on_apply(braft::Iterator& iter) {
   // A batch of tasks are committed, which must be processed through
@@ -353,9 +404,33 @@ void PRaft::on_apply(braft::Iterator& iter) {
   }
 }
 
-void PRaft::on_snapshot_save(braft::SnapshotWriter* writer, braft::Closure* done) {}
+void PRaft::on_snapshot_save(braft::SnapshotWriter* writer, braft::Closure* done) {
+  brpc::ClosureGuard done_guard(done);
+  TasksVector tasks;
+  tasks.reserve(g_config.databases);
+  for (auto i = 0; i < g_config.databases; ++i) {
+    tasks.push_back({TaskType::kCheckpoint, i, {{TaskArg::kCheckpointPath, writer->get_path()}}});
+  }
+  PSTORE.DoSomeThingSpecificDB(tasks);
+  PSTORE.WaitForCheckpointDone();
+  auto writer_path = writer->get_path();
+  add_all_files(writer_path, writer, writer_path);
+}
 
-int PRaft::on_snapshot_load(braft::SnapshotReader* reader) { return 0; }
+int PRaft::on_snapshot_load(braft::SnapshotReader* reader) {
+  CHECK(!IsLeader()) << "Leader is not supposed to load snapshot";
+  auto reader_path = reader->get_path();  // xx/snapshot_0000001
+  auto db_path = g_config.dbpath;
+  PSTORE.Clear();
+  for (int i = 0; i < g_config.databases; i++) {
+    auto sub_path = db_path + std::to_string(i);
+    pstd::DeleteDirIfExist(sub_path);
+  }
+  db_path.pop_back();
+  recursive_copy(reader_path, db_path);
+  PSTORE.Init();
+  return 0;
+}
 
 void PRaft::on_leader_start(int64_t term) {
   WARN("Node {} start to be leader, term={}", node_->node_id().to_string(), term);

src/praft/praft.h

@@ -7,11 +7,12 @@
 
 #pragma once
 
+#include <filesystem>
 #include <memory>
 #include <mutex>
+#include <string>
 #include <tuple>
 #include <vector>
-#include <string>
 
 #include "braft/configuration.h"
 #include "braft/raft.h"
@@ -88,6 +89,7 @@ class PRaft : public braft::StateMachine {
   butil::Status AddPeer(const std::string& peer);
   butil::Status RemovePeer(const std::string& peer);
   butil::Status RaftRecvEntry();
+  butil::Status DoSnapshot();
 
   void ShutDown();
   void Join();
@@ -125,6 +127,11 @@ class PRaft : public braft::StateMachine {
   void on_stop_following(const ::braft::LeaderChangeContext& ctx) override;
   void on_start_following(const ::braft::LeaderChangeContext& ctx) override;
 
+ private:
+  void add_all_files(const std::filesystem::path& dir, braft::SnapshotWriter* writer, const std::string& path);
+
+  void recursive_copy(const std::filesystem::path& source, const std::filesystem::path& destination);
+
  private:
   std::unique_ptr<brpc::Server> server_;  // brpc
   std::unique_ptr<braft::Node> node_;

src/storage/include/storage/storage.h

@@ -41,6 +41,7 @@ inline constexpr size_t BATCH_DELETE_LIMIT = 100;
 inline constexpr size_t COMPACT_THRESHOLD_COUNT = 2000;
 
 inline constexpr uint64_t kNoFlush = std::numeric_limits<uint64_t>::max();
+inline constexpr uint64_t kFlush = 0;
 
 using Options = rocksdb::Options;
 using BlockBasedTableOptions = rocksdb::BlockBasedTableOptions;