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) {