KEP-975 Add snapshot support to rocksdb

storage/mem_storage.cpp

@@ -211,6 +211,7 @@ mem_storage::remove(const bzn::uuid_t& uuid)
     return bzn::storage_result::not_found;
 }
 
+
 bool
 mem_storage::create_snapshot()
 {
@@ -225,20 +226,23 @@ mem_storage::create_snapshot()
 
         return true;
     }
-    catch(std::exception &e)
+    catch (std::exception& ex)
     {
-        LOG(error) << "Exception creating snapshot: " << e.what();
+        LOG(error) << "Exception creating snapshot: " << ex.what();
     }
+
     return false;
 }
 
+
 std::shared_ptr<std::string>
 mem_storage::get_snapshot()
 {
     std::shared_lock<std::shared_mutex> lock(this->lock); // lock for read access
     return this->latest_snapshot;
 }
 
+
 bool
 mem_storage::load_snapshot(const std::string& data)
 {
@@ -253,9 +257,10 @@ mem_storage::load_snapshot(const std::string& data)
 
         return true;
     }
-    catch(std::exception &e)
+    catch (std::exception& ex)
     {
-        LOG(error) << "Exception loading snapshot: " << e.what();
+        LOG(error) << "Exception loading snapshot: " << ex.what();
     }
+
     return false;
 }

storage/rocksdb_storage.cpp

@@ -14,6 +14,7 @@
 
 #include <storage/rocksdb_storage.hpp>
 #include <boost/filesystem.hpp>
+#include <rocksdb/db_dump_tool.h>
 #include <thread>
 
 using namespace bzn;
@@ -28,6 +29,15 @@ namespace
 
 
 rocksdb_storage::rocksdb_storage(const std::string& state_dir, const std::string& db_name, const bzn::uuid_t& uuid)
+    : db_path(boost::filesystem::path(state_dir).append(uuid).append(db_name).string())
+    , snapshot_file(boost::filesystem::path(state_dir).append(uuid).append("SNAPSHOT." + db_name).string())
+{
+    this->open();
+}
+
+
+void
+rocksdb_storage::open()
 {
     rocksdb::Options options;
 
@@ -37,8 +47,6 @@ rocksdb_storage::rocksdb_storage(const std::string& state_dir, const std::string
 
     rocksdb::DB* rocksdb;
 
-    const std::string db_path = boost::filesystem::path(state_dir).append(uuid).append(db_name).string();
-
     boost::filesystem::create_directories(db_path);
 
     LOG(info) << "database path: " << db_path;
@@ -256,20 +264,121 @@ rocksdb_storage::has_priv(const bzn::uuid_t& uuid, const  std::string& key)
     return false;
 }
 
+
 bool
 rocksdb_storage::create_snapshot()
 {
-    return false;
+    rocksdb::DumpOptions dump_options;
+
+    dump_options.db_path = this->db_path;
+    dump_options.dump_location = this->snapshot_file;
+    dump_options.anonymous = true;
+
+    std::shared_lock<std::shared_mutex> lock(this->lock); // lock for read access
+
+    return rocksdb::DbDumpTool().Run(dump_options);
 }
 
+
 std::shared_ptr<std::string>
 rocksdb_storage::get_snapshot()
 {
-    return nullptr;
+    std::shared_lock<std::shared_mutex> lock(this->lock); // lock for read access
+
+    std::stringstream snapshot;
+
+    // check if file exists...
+    try
+    {
+        std::ifstream s(this->snapshot_file);
+        snapshot << s.rdbuf();
+    }
+    catch (std::exception& ex)
+    {
+        LOG(error) << "Exception reading snapshot: " << ex.what();
+
+        return {};
+    }
+
+    return std::make_shared<std::string>(snapshot.str());
 }
 
+
 bool
-rocksdb_storage::load_snapshot(const std::string& /*data*/)
+rocksdb_storage::load_snapshot(const std::string& data)
 {
+    if (!boost::filesystem::exists(this->snapshot_file))
+    {
+        LOG(error) << "no snapshot found";
+
+        return false;
+    }
+
+    try
+    {
+        std::shared_lock<std::shared_mutex> lock(this->lock); // lock for read access
+
+        std::ofstream snapshot(this->snapshot_file);
+        snapshot << data;
+    }
+    catch (std::exception& ex)
+    {
+        LOG(error) << "saving snapshot failed: " << ex.what();
+
+        return false;
+    }
+
+    std::lock_guard<std::shared_mutex> lock(this->lock); // lock for write access
+
+    // bring down the database...
+    this->db.reset();
+
+    // move current database out of the way...
+    std::string tmp_path(this->db_path + ".tmp");
+
+    boost::system::error_code ec;
+    boost::filesystem::rename(this->db_path, tmp_path, ec);
+
+    if (ec)
+    {
+        LOG(error) << "creating temporary db backup failed: " << ec.message();
+
+        // bring db back online...
+        this->open();
+
+        return false;
+    }
+
+    rocksdb::UndumpOptions undump_options;
+
+    undump_options.db_path = this->db_path;
+    undump_options.dump_location = this->snapshot_file;
+    undump_options.compact_db = true;
+
+    if (rocksdb::DbUndumpTool().Run(undump_options))
+    {
+        boost::system::error_code ec;
+        boost::filesystem::remove_all(tmp_path, ec);
+
+        if (ec)
+        {
+            LOG(error) << "failed to remove temporary db backup: " << ec.message();
+        }
+
+        // bring db back online...
+        this->open();
+
+        return true;
+    }
+
+    LOG(error) << "failed to load snapshot";
+
+    // any exceptions will be fatal...
+    boost::filesystem::remove_all(this->db_path);
+    boost::filesystem::rename(tmp_path, this->db_path);
+
+    // bring db back online...
+    this->open();
+
     return false;
 }

storage/rocksdb_storage.hpp

@@ -51,6 +51,11 @@ namespace bzn
         bool load_snapshot(const std::string& data) override;
 
     private:
+        void open();
+
+        const std::string db_path;
+        const std::string snapshot_file;
+
         std::unique_ptr<rocksdb::DB> db;
 
         bool has_priv(const bzn::uuid_t& uuid, const  std::string& key);

storage/test/storage_test.cpp

@@ -265,26 +265,26 @@ TYPED_TEST(storageTest, test_that_storage_can_remove_all_keys_values_associated_
     EXPECT_EQ(std::nullopt, this->storage->read(USER_UUID, KEY));
 }
 
-// TODO: refactor this as a typed test once rocksdb_storage supports checkpoints
-TEST(mem_snapshot_test, test_snapshot)
+
+TYPED_TEST(storageTest, test_snapshot)
 {
     const bzn::uuid_t user_0{"b9dc2595-15ee-435a-8af7-7cafc132f527"};
 
-    std::shared_ptr<bzn::storage_base> storage = std::make_shared<bzn::mem_storage>();
-    storage->create(user_0, "key1", "value1");
-    EXPECT_TRUE(storage->has(user_0, "key1"));
-    EXPECT_TRUE(storage->create_snapshot());
+    this->storage->create(user_0, "key1", "value1");
+    EXPECT_TRUE(this->storage->has(user_0, "key1"));
+    EXPECT_TRUE(this->storage->create_snapshot());
 
-    storage->create(user_0, "key2", "value2");
-    storage->create(user_0, "key3", "value3");
-    EXPECT_TRUE(storage->has(user_0, "key2"));
-    EXPECT_TRUE(storage->has(user_0, "key3"));
+    this->storage->create(user_0, "key2", "value2");
+    this->storage->create(user_0, "key3", "value3");
+    EXPECT_TRUE(this->storage->has(user_0, "key2"));
+    EXPECT_TRUE(this->storage->has(user_0, "key3"));
 
-    auto state = storage->get_snapshot();
+    auto state = this->storage->get_snapshot();
     EXPECT_NE(state, nullptr);
 
-    EXPECT_TRUE(storage->load_snapshot(*state));
-    EXPECT_TRUE(storage->has(user_0, "key1"));
-    EXPECT_FALSE(storage->has(user_0, "key2"));
-    EXPECT_FALSE(storage->has(user_0, "key3"));
+    EXPECT_FALSE(this->storage->load_snapshot("aslkdfkslfdk"));
+    EXPECT_TRUE(this->storage->load_snapshot(*state));
+    EXPECT_TRUE(this->storage->has(user_0, "key1"));
+    EXPECT_FALSE(this->storage->has(user_0, "key2"));
+    EXPECT_FALSE(this->storage->has(user_0, "key3"));
 }