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.cc

@@ -20,15 +20,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) {
     auto res = std::async(std::launch::async, &DB::DoBgSave, db_, path, i);

src/client.h

@@ -102,7 +102,6 @@ enum ClientFlag {
   kClientFlagMaster = (1 << 3),
 };
 
-class DB;
 struct PSlaveInfo;
 
 class PClient : public std::enable_shared_from_this<PClient>, public CmdRes {
@@ -251,4 +250,4 @@ class PClient : public std::enable_shared_from_this<PClient>, public CmdRes {
 
   static thread_local PClient* s_current;
 };
-}  // namespace pikiwidb
+}  // namespace pikiwidb

src/cmd_admin.cc

@@ -6,11 +6,11 @@
  */
 
 #include "cmd_admin.h"
-
 #include "braft/raft.h"
 #include "rocksdb/version.h"
 
 #include "praft/praft.h"
+#include "store.h"
 
 namespace pikiwidb {
 

src/cmd_raft.cc

@@ -11,6 +11,7 @@
 #include <optional>
 #include <string>
 
+#include "net/event_loop.h"
 #include "praft/praft.h"
 #include "pstd/log.h"
 #include "pstd/pstd_string.h"
@@ -28,8 +29,9 @@ RaftNodeCmd::RaftNodeCmd(const std::string& name, int16_t arity)
 bool RaftNodeCmd::DoInitial(PClient* client) {
   auto cmd = client->argv_[1];
   pstd::StringToUpper(cmd);
-  if (cmd != kAddCmd && cmd != kRemoveCmd) {
-    client->SetRes(CmdRes::kErrOther, "RAFT.NODE supports ADD / REMOVE only");
+
+  if (cmd != kAddCmd && cmd != kRemoveCmd && cmd != kDoSnapshot) {
+    client->SetRes(CmdRes::kErrOther, "RAFT.NODE supports ADD / REMOVE / DOSNAPSHOT only");
     return false;
   }
   return true;
@@ -40,8 +42,10 @@ void RaftNodeCmd::DoCmd(PClient* client) {
   pstd::StringToUpper(cmd);
   if (cmd == kAddCmd) {
     DoCmdAdd(client);
-  } else {
+  } else if (cmd == kRemoveCmd) {
     DoCmdRemove(client);
+  } else {
+    DoCmdSnapshot(client);
   }
 }
 
@@ -115,6 +119,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

@@ -45,9 +45,11 @@ 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";
+  static constexpr std::string_view kDoSnapshot = "DOSNAPSHOT";
 };
 
 /* RAFT.CLUSTER INIT <id>
@@ -78,4 +80,4 @@ class RaftClusterCmd : public BaseCmd {
   static constexpr std::string_view kJoinCmd = "JOIN";
 };
 
-}  // namespace pikiwidb
+}  // namespace pikiwidb

src/cmd_table_manager.cc

@@ -7,6 +7,8 @@
 
 #include "cmd_table_manager.h"
 
+#include <memory>
+
 #include "cmd_admin.h"
 #include "cmd_hash.h"
 #include "cmd_keys.h"

src/db.cc

@@ -7,11 +7,10 @@
 
 #include "db.h"
 
-#include "praft/praft.h"
-#include "pstd/log.h"
-
 #include "checkpoint_manager.h"
 #include "config.h"
+#include "praft/praft.h"
+#include "pstd/log.h"
 
 extern pikiwidb::PConfig g_config;
 
@@ -54,7 +53,14 @@ void DB::DoBgSave(const std::string& path, int i) {
   auto status = storage_->CreateCheckpoint(path, i);
 }
 
-void DB::CreateCheckpoint(const std::string& path) { checkpoint_manager_->CreateCheckpoint(path); }
+void DB::CreateCheckpoint(const std::string& path) {
+  auto tmp_path = path + '/' + std::to_string(db_index_);
+  if (0 != pstd::CreatePath(tmp_path)) {
+    WARN("Create dir {} fail !", tmp_path);
+    return;
+  }
+  checkpoint_manager_->CreateCheckpoint(path);
+}
 
 void DB::WaitForCheckpointDone() { checkpoint_manager_->WaitForCheckpointDone(); }
 

src/pikiwidb.cc

@@ -8,6 +8,8 @@
 //
 //  PikiwiDB.cc
 
+#include "pikiwidb.h"
+
 #include <sys/fcntl.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -21,7 +23,6 @@
 #include "client.h"
 #include "config.h"
 #include "helper.h"
-#include "pikiwidb.h"
 #include "pikiwidb_logo.h"
 #include "slow_log.h"
 #include "store.h"

src/praft/praft.cc

@@ -9,6 +9,7 @@
 
 #include <cassert>
 
+#include "braft/snapshot.h"
 #include "braft/util.h"
 #include "brpc/server.h"
 
@@ -18,6 +19,8 @@
 #include "binlog.pb.h"
 #include "config.h"
 #include "pikiwidb.h"
+#include "praft.h"
+
 #include "praft_service.h"
 #include "replication.h"
 #include "store.h"
@@ -514,6 +517,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::OnClusterCmdConnectionFailed([[maybe_unused]] EventLoop* loop, const char* peer_ip, int port) {
   auto cli = cluster_cmd_ctx_.GetClient();
   if (cli) {
@@ -567,6 +580,48 @@ void PRaft::AppendLog(const Binlog& log, std::promise<rocksdb::Status>&& promise
   node_->apply(task);
 }
 
+void PRaft::AddAllFiles(const std::filesystem::path& dir, braft::SnapshotWriter* writer, const std::string& path) {
+  assert(writer);
+  for (const auto& entry : std::filesystem::directory_iterator(dir)) {
+    if (entry.is_directory()) {
+      if (entry.path() != "." && entry.path() != "..") {
+        DEBUG("dir_path = {}", entry.path().string());
+        AddAllFiles(entry.path(), writer, path);
+      }
+    } else {
+      DEBUG("file_path = {}", std::filesystem::relative(entry.path(), path).string());
+      if (writer->add_file(std::filesystem::relative(entry.path(), path)) != 0) {
+        ERROR("add file error!");
+      }
+    }
+  }
+}
+
+void PRaft::RecursiveCopy(const std::filesystem::path& source, const std::filesystem::path& destination) {
+  if (std::filesystem::is_regular_file(source)) {
+    if (source.filename() == PBRAFT_SNAPSHOT_META_FILE) {
+      return;
+    } else if (source.extension() == ".sst") {
+      // Create a hard link
+      DEBUG("hard link success! source_file = {} , destination_file = {}", source.string(), destination.string());
+      ::link(source.c_str(), destination.c_str());
+    } else {
+      // Copy the file
+      DEBUG("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)) {
+      RecursiveCopy(entry.path(), destination / entry.path().filename());
+    }
+  }
+}
+
+// @braft::StateMachine
 void PRaft::Clear() {
   if (node_) {
     node_.reset();
@@ -607,9 +662,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();
+  AddAllFiles(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();
+  RecursiveCopy(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,7 +7,12 @@
 
 #pragma once
 
+#include <filesystem>
 #include <future>
+#include <mutex>
+#include <string>
+#include <tuple>
+#include <vector>
 
 #include "braft/raft.h"
 #include "brpc/server.h"
@@ -17,6 +22,7 @@
 
 namespace pikiwidb {
 
+#define PBRAFT_SNAPSHOT_META_FILE "__raft_snapshot_meta"
 #define RAFT_GROUPID_LEN 32
 
 #define OK "+OK"
@@ -99,6 +105,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();
@@ -148,6 +155,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 AddAllFiles(const std::filesystem::path& dir, braft::SnapshotWriter* writer, const std::string& path);
+
+  void RecursiveCopy(const std::filesystem::path& source, const std::filesystem::path& destination);
+
  private:
   std::unique_ptr<brpc::Server> server_{nullptr};  // brpc
   std::unique_ptr<braft::Node> node_{nullptr};

src/storage/include/storage/storage.h

@@ -46,6 +46,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;

src/storage/src/storage.cc

@@ -105,14 +105,6 @@ Status Storage::Open(const StorageOptions& storage_options, const std::string& d
 Status Storage::CreateCheckpoint(const std::string& dump_path, int i) {
   INFO("DB{}'s RocksDB {} begin to generate a checkpoint!", db_id_, i);
   auto source_dir = AppendSubDirectory(dump_path, db_id_);
-  if (!pstd::FileExists(source_dir)) {
-    if (0 != pstd::CreatePath(source_dir)) {
-      WARN("Create Dir {} fail!", source_dir);
-      return Status::IOError("CreatePath() fail! dir_name : {} ", source_dir);
-    }
-    INFO("Create Dir {} success!", source_dir);
-  }
-
   source_dir = AppendSubDirectory(source_dir, i);
 
   auto tmp_dir = source_dir + ".tmp";

src/store.cc

@@ -36,6 +36,8 @@ void PStore::Init() {
   }
 }
 
+void PStore::Clear() { backends_.clear(); }
+
 void PStore::DoSomeThingSpecificDB(const TasksVector& tasks) {
   std::for_each(tasks.begin(), tasks.end(), [this](const auto& task) {
     switch (task.type) {