From 880ee363eb49f66de2cb2a6fbc25b5a775c0b09d Mon Sep 17 00:00:00 2001
From: Wanning Jiang <wanning@fb.com>
Date: Thu, 14 Jul 2016 14:09:31 -0700
Subject: [PATCH] ldb backup support

Summary: add backup support for ldb tool, and use it to run load test for backup on two HDFS envs

Test Plan: first generate some db, then compile against load test in fbcode, run load_test --db=<db path> backup --backup_env_uri=<URI of backup env> --backup_dir=<backup directory> --num_threads=<number of thread>

Reviewers: andrewkr

Reviewed By: andrewkr

Subscribers: andrewkr, dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D60633
---
 tools/ldb_cmd.cc     | 76 ++++++++++++++++++++++++++++++++++++++++++++
 tools/ldb_cmd_impl.h | 22 +++++++++++++
 tools/ldb_tool.cc    |  1 +
 3 files changed, 99 insertions(+)

diff --git a/tools/ldb_cmd.cc b/tools/ldb_cmd.cc
index 406b873ab..48a0bbbad 100644
--- a/tools/ldb_cmd.cc
+++ b/tools/ldb_cmd.cc
@@ -20,6 +20,8 @@
 #include "port/dirent.h"
 #include "rocksdb/cache.h"
 #include "rocksdb/table_properties.h"
+#include "rocksdb/utilities/backupable_db.h"
+#include "rocksdb/utilities/env_registry.h"
 #include "rocksdb/write_batch.h"
 #include "rocksdb/write_buffer_manager.h"
 #include "table/scoped_arena_iterator.h"
@@ -217,6 +219,9 @@ LDBCommand* LDBCommand::SelectCommand(const ParsedParams& parsed_params) {
   } else if (parsed_params.cmd == RepairCommand::Name()) {
     return new RepairCommand(parsed_params.cmd_params, parsed_params.option_map,
                              parsed_params.flags);
+  } else if (parsed_params.cmd == BackupCommand::Name()) {
+    return new BackupCommand(parsed_params.cmd_params, parsed_params.option_map,
+                             parsed_params.flags);
   }
   return nullptr;
 }
@@ -2520,6 +2525,77 @@ void RepairCommand::DoCommand() {
 
 // ----------------------------------------------------------------------------
 
+const std::string BackupCommand::ARG_THREAD = "num_threads";
+const std::string BackupCommand::ARG_BACKUP_ENV = "backup_env_uri";
+const std::string BackupCommand::ARG_BACKUP_DIR = "backup_dir";
+
+BackupCommand::BackupCommand(const std::vector<std::string>& params,
+                             const std::map<std::string, std::string>& options,
+                             const std::vector<std::string>& flags)
+    : LDBCommand(
+          options, flags, false,
+          BuildCmdLineOptions({ARG_BACKUP_ENV, ARG_BACKUP_DIR, ARG_THREAD})),
+      thread_num_(1) {
+  auto itr = options.find(ARG_THREAD);
+  if (itr != options.end()) {
+    thread_num_ = std::stoi(itr->second);
+  }
+  itr = options.find(ARG_BACKUP_ENV);
+  if (itr == options.end()) {
+    exec_state_ = LDBCommandExecuteResult::Failed("--" + ARG_BACKUP_ENV +
+                                                  ": missing backup env");
+  } else {
+    test_cluster_ = itr->second;
+  }
+  itr = options.find(ARG_BACKUP_DIR);
+  if (itr == options.end()) {
+    exec_state_ = LDBCommandExecuteResult::Failed("--" + ARG_BACKUP_DIR +
+                                                  ": missing backup directory");
+  } else {
+    test_path_ = itr->second;
+  }
+}
+
+void BackupCommand::Help(std::string& ret) {
+  ret.append("  ");
+  ret.append(BackupCommand::Name());
+  ret.append(" [--" + ARG_BACKUP_ENV + "] ");
+  ret.append(" [--" + ARG_BACKUP_DIR + "] ");
+  ret.append(" [--" + ARG_THREAD + "] ");
+  ret.append("\n");
+}
+
+void BackupCommand::DoCommand() {
+  BackupEngine* backup_engine;
+  Status status;
+  if (!db_) {
+    assert(GetExecuteState().IsFailed());
+    return;
+  }
+  printf("open db OK\n");
+  std::unique_ptr<Env> custom_env_guard;
+  Env* custom_env = NewEnvFromUri(test_cluster_, &custom_env_guard);
+  BackupableDBOptions backup_options =
+      BackupableDBOptions(test_path_, custom_env);
+  backup_options.max_background_operations = thread_num_;
+  status = BackupEngine::Open(Env::Default(), backup_options, &backup_engine);
+  if (status.ok()) {
+    printf("open backup engine OK\n");
+  } else {
+    exec_state_ = LDBCommandExecuteResult::Failed(status.ToString());
+    return;
+  }
+  status = backup_engine->CreateNewBackup(db_);
+  if (status.ok()) {
+    printf("create new backup OK\n");
+  } else {
+    exec_state_ = LDBCommandExecuteResult::Failed(status.ToString());
+    return;
+  }
+}
+
+// ----------------------------------------------------------------------------
+
 namespace {
 
 void DumpSstFile(std::string filename, bool output_hex, bool show_properties) {
diff --git a/tools/ldb_cmd_impl.h b/tools/ldb_cmd_impl.h
index 5100d6543..6de9534b9 100644
--- a/tools/ldb_cmd_impl.h
+++ b/tools/ldb_cmd_impl.h
@@ -446,4 +446,26 @@ class RepairCommand : public LDBCommand {
   static void Help(std::string& ret);
 };
 
+class BackupCommand : public LDBCommand {
+ public:
+  static std::string Name() { return "backup"; }
+
+  BackupCommand(const std::vector<std::string>& params,
+                const std::map<std::string, std::string>& options,
+                const std::vector<std::string>& flags);
+
+  virtual void DoCommand() override;
+
+  static void Help(std::string& ret);
+
+ private:
+  std::string test_cluster_;
+  std::string test_path_;
+  int thread_num_;
+
+  static const std::string ARG_BACKUP_DIR;
+  static const std::string ARG_BACKUP_ENV;
+  static const std::string ARG_THREAD;
+};
+
 }  // namespace rocksdb
diff --git a/tools/ldb_tool.cc b/tools/ldb_tool.cc
index 351d7cf12..2f9219c13 100644
--- a/tools/ldb_tool.cc
+++ b/tools/ldb_tool.cc
@@ -79,6 +79,7 @@ void LDBCommandRunner::PrintHelp(const char* exec_name) {
   DBFileDumperCommand::Help(ret);
   InternalDumpCommand::Help(ret);
   RepairCommand::Help(ret);
+  BackupCommand::Help(ret);
 
   fprintf(stderr, "%s\n", ret.c_str());
 }