Support custom env in sst_dump (facebook#5845)

Summary:
This PR allows for the creation of custom env when using sst_dump. If
the user does not set options.env or set options.env to nullptr, then sst_dump
will automatically try to create a custom env depending on the path to the sst
file or db directory. In order to use this feature, the user must call
ObjectRegistry::Register() beforehand.

Test Plan (on devserver):
```
$make all && make check
```
All tests must pass to ensure this change does not break anything.
Pull Request resolved: facebook#5845

Differential Revision: D17678038

Pulled By: riversand963

fbshipit-source-id: 58ecb4b3f75246d52b07c4c924a63ee61c1ee626

HISTORY.md

@@ -22,6 +22,7 @@
 * Deprecate `snap_refresh_nanos` option.
 * Added DisableManualCompaction/EnableManualCompaction to stop and resume manual compaction.
 * Add TryCatchUpWithPrimary() to StackableDB in non-LITE mode.
+* Add a new Env::LoadEnv() overloaded function to return a shared_ptr to Env.
 ### Performance Improvements
 * Improve the speed of the MemTable Bloom filter, reducing the write overhead of enabling it by 1/3 to 1/2, with similar benefit to read performance.
 

env/env.cc

@@ -43,6 +43,36 @@ Status Env::LoadEnv(const std::string& value, Env** result) {
   return s;
 }
 
+Status Env::LoadEnv(const std::string& value, Env** result,
+                    std::shared_ptr<Env>* guard) {
+  assert(result);
+  Status s;
+#ifndef ROCKSDB_LITE
+  Env* env = nullptr;
+  std::unique_ptr<Env> uniq_guard;
+  std::string err_msg;
+  assert(guard != nullptr);
+  env = ObjectRegistry::NewInstance()->NewObject<Env>(value, &uniq_guard,
+                                                      &err_msg);
+  if (!env) {
+    s = Status::NotFound(std::string("Cannot load ") + Env::Type() + ": " +
+                         value);
+    env = Env::Default();
+  }
+  if (s.ok() && uniq_guard) {
+    guard->reset(uniq_guard.release());
+    *result = guard->get();
+  } else {
+    *result = env;
+  }
+#else
+  (void)result;
+  (void)guard;
+  s = Status::NotSupported("Cannot load environment in LITE mode: ", value);
+#endif
+  return s;
+}
+
 std::string Env::PriorityToString(Env::Priority priority) {
   switch (priority) {
     case Env::Priority::BOTTOM:

include/rocksdb/env.h

@@ -152,6 +152,10 @@ class Env {
   // Loads the environment specified by the input value into the result
   static Status LoadEnv(const std::string& value, Env** result);
 
+  // Loads the environment specified by the input value into the result
+  static Status LoadEnv(const std::string& value, Env** result,
+                        std::shared_ptr<Env>* guard);
+
   // Return a default environment suitable for the current operating
   // system.  Sophisticated users may wish to provide their own Env
   // implementation instead of relying on this default environment.

include/rocksdb/utilities/ldb_cmd.h

@@ -29,6 +29,7 @@ namespace rocksdb {
 class LDBCommand {
  public:
   // Command-line arguments
+  static const std::string ARG_ENV_URI;
   static const std::string ARG_DB;
   static const std::string ARG_PATH;
   static const std::string ARG_SECONDARY_PATH;
@@ -128,6 +129,7 @@ class LDBCommand {
 
  protected:
   LDBCommandExecuteResult exec_state_;
+  std::string env_uri_;
   std::string db_path_;
   // If empty, open DB as primary. If non-empty, open the DB as secondary
   // with this secondary path. When running against a database opened by
@@ -176,6 +178,9 @@ class LDBCommand {
   /** List of command-line options valid for this command */
   const std::vector<std::string> valid_cmd_line_options_;
 
+  /** Shared pointer to underlying environment if applicable **/
+  std::shared_ptr<Env> env_guard_;
+
   bool ParseKeyValue(const std::string& line, std::string* key,
                      std::string* value, bool is_key_hex, bool is_value_hex);
 

tools/ldb_cmd.cc

@@ -45,6 +45,7 @@
 
 namespace rocksdb {
 
+const std::string LDBCommand::ARG_ENV_URI = "env_uri";
 const std::string LDBCommand::ARG_DB = "db";
 const std::string LDBCommand::ARG_PATH = "path";
 const std::string LDBCommand::ARG_SECONDARY_PATH = "secondary_path";
@@ -274,6 +275,17 @@ void LDBCommand::Run() {
     return;
   }
 
+  if (!options_.env || options_.env == Env::Default()) {
+    Env* env = Env::Default();
+    Status s = Env::LoadEnv(env_uri_, &env, &env_guard_);
+    if (!s.ok() && !s.IsNotFound()) {
+      fprintf(stderr, "LoadEnv: %s\n", s.ToString().c_str());
+      exec_state_ = LDBCommandExecuteResult::Failed(s.ToString());
+      return;
+    }
+    options_.env = env;
+  }
+
   if (db_ == nullptr && !NoDBOpen()) {
     OpenDB();
     if (exec_state_.IsFailed() && try_load_options_) {
@@ -318,6 +330,11 @@ LDBCommand::LDBCommand(const std::map<std::string, std::string>& options,
     db_path_ = itr->second;
   }
 
+  itr = options.find(ARG_ENV_URI);
+  if (itr != options.end()) {
+    env_uri_ = itr->second;
+  }
+
   itr = options.find(ARG_CF_NAME);
   if (itr != options.end()) {
     column_family_name_ = itr->second;
@@ -341,7 +358,7 @@ LDBCommand::LDBCommand(const std::map<std::string, std::string>& options,
 
 void LDBCommand::OpenDB() {
   if (!create_if_missing_ && try_load_options_) {
-    Status s = LoadLatestOptions(db_path_, Env::Default(), &options_,
+    Status s = LoadLatestOptions(db_path_, options_.env, &options_,
                                  &column_families_, ignore_unknown_options_);
     if (!s.ok() && !s.IsNotFound()) {
       // Option file exists but load option file error.
@@ -397,7 +414,7 @@ void LDBCommand::OpenDB() {
     if (column_families_.empty()) {
       // Try to figure out column family lists
       std::vector<std::string> cf_list;
-      st = DB::ListColumnFamilies(DBOptions(), db_path_, &cf_list);
+      st = DB::ListColumnFamilies(options_, db_path_, &cf_list);
       // There is possible the DB doesn't exist yet, for "create if not
       // "existing case". The failure is ignored here. We rely on DB::Open()
       // to give us the correct error message for problem with opening
@@ -487,7 +504,8 @@ ColumnFamilyHandle* LDBCommand::GetCfHandle() {
 
 std::vector<std::string> LDBCommand::BuildCmdLineOptions(
     std::vector<std::string> options) {
-  std::vector<std::string> ret = {ARG_DB,
+  std::vector<std::string> ret = {ARG_ENV_URI,
+                                  ARG_DB,
                                   ARG_SECONDARY_PATH,
                                   ARG_BLOOM_BITS,
                                   ARG_BLOCK_SIZE,
@@ -1095,31 +1113,23 @@ void ManifestDumpCommand::DoCommand() {
 void ListColumnFamiliesCommand::Help(std::string& ret) {
   ret.append("  ");
   ret.append(ListColumnFamiliesCommand::Name());
-  ret.append(" full_path_to_db_directory ");
   ret.append("\n");
 }
 
 ListColumnFamiliesCommand::ListColumnFamiliesCommand(
-    const std::vector<std::string>& params,
+    const std::vector<std::string>& /*params*/,
     const std::map<std::string, std::string>& options,
     const std::vector<std::string>& flags)
-    : LDBCommand(options, flags, false, {}) {
-  if (params.size() != 1) {
-    exec_state_ = LDBCommandExecuteResult::Failed(
-        "dbname must be specified for the list_column_families command");
-  } else {
-    dbname_ = params[0];
-  }
-}
+    : LDBCommand(options, flags, false, BuildCmdLineOptions({})) {}
 
 void ListColumnFamiliesCommand::DoCommand() {
   std::vector<std::string> column_families;
-  Status s = DB::ListColumnFamilies(DBOptions(), dbname_, &column_families);
+  Status s = DB::ListColumnFamilies(options_, db_path_, &column_families);
   if (!s.ok()) {
-    printf("Error in processing db %s %s\n", dbname_.c_str(),
+    printf("Error in processing db %s %s\n", db_path_.c_str(),
            s.ToString().c_str());
   } else {
-    printf("Column families in %s: \n{", dbname_.c_str());
+    printf("Column families in %s: \n{", db_path_.c_str());
     bool first = true;
     for (auto cf : column_families) {
       if (!first) {
@@ -2857,13 +2867,14 @@ void BackupCommand::DoCommand() {
   }
   printf("open db OK\n");
   Env* custom_env = nullptr;
-  Env::LoadEnv(backup_env_uri_, &custom_env);
+  Env::LoadEnv(backup_env_uri_, &custom_env, &backup_env_guard_);
+  assert(custom_env != nullptr);
 
   BackupableDBOptions backup_options =
       BackupableDBOptions(backup_dir_, custom_env);
   backup_options.info_log = logger_.get();
   backup_options.max_background_operations = num_threads_;
-  status = BackupEngine::Open(Env::Default(), backup_options, &backup_engine);
+  status = BackupEngine::Open(custom_env, backup_options, &backup_engine);
   if (status.ok()) {
     printf("open backup engine OK\n");
   } else {
@@ -2893,7 +2904,8 @@ void RestoreCommand::Help(std::string& ret) {
 
 void RestoreCommand::DoCommand() {
   Env* custom_env = nullptr;
-  Env::LoadEnv(backup_env_uri_, &custom_env);
+  Env::LoadEnv(backup_env_uri_, &custom_env, &backup_env_guard_);
+  assert(custom_env != nullptr);
 
   std::unique_ptr<BackupEngineReadOnly> restore_engine;
   Status status;
@@ -2902,8 +2914,8 @@ void RestoreCommand::DoCommand() {
     opts.info_log = logger_.get();
     opts.max_background_operations = num_threads_;
     BackupEngineReadOnly* raw_restore_engine_ptr;
-    status = BackupEngineReadOnly::Open(Env::Default(), opts,
-                                        &raw_restore_engine_ptr);
+    status =
+        BackupEngineReadOnly::Open(custom_env, opts, &raw_restore_engine_ptr);
     if (status.ok()) {
       restore_engine.reset(raw_restore_engine_ptr);
     }

tools/ldb_cmd_impl.h

@@ -183,9 +183,6 @@ class ListColumnFamiliesCommand : public LDBCommand {
   virtual void DoCommand() override;
 
   virtual bool NoDBOpen() override { return true; }
-
- private:
-  std::string dbname_;
 };
 
 class CreateColumnFamilyCommand : public LDBCommand {
@@ -510,6 +507,7 @@ class BackupableCommand : public LDBCommand {
   std::string backup_dir_;
   int num_threads_;
   std::unique_ptr<Logger> logger_;
+  std::shared_ptr<Env> backup_env_guard_;
 
  private:
   static const std::string ARG_BACKUP_DIR;

tools/ldb_cmd_test.cc

@@ -6,6 +6,7 @@
 #ifndef ROCKSDB_LITE
 
 #include "rocksdb/utilities/ldb_cmd.h"
+#include "port/stack_trace.h"
 #include "test_util/sync_point.h"
 #include "test_util/testharness.h"
 
@@ -15,7 +16,23 @@ using std::map;
 
 namespace rocksdb {
 
-class LdbCmdTest : public testing::Test {};
+class LdbCmdTest : public testing::Test {
+ public:
+  LdbCmdTest() : testing::Test() {}
+
+  Env* TryLoadCustomOrDefaultEnv() {
+    const char* test_env_uri = getenv("TEST_ENV_URI");
+    if (!test_env_uri) {
+      return Env::Default();
+    }
+    Env* env = Env::Default();
+    Env::LoadEnv(test_env_uri, &env, &env_guard_);
+    return env;
+  }
+
+ private:
+  std::shared_ptr<Env> env_guard_;
+};
 
 TEST_F(LdbCmdTest, HexToString) {
   // map input to expected outputs.
@@ -51,7 +68,8 @@ TEST_F(LdbCmdTest, HexToStringBadInputs) {
 }
 
 TEST_F(LdbCmdTest, MemEnv) {
-  std::unique_ptr<Env> env(NewMemEnv(Env::Default()));
+  Env* base_env = TryLoadCustomOrDefaultEnv();
+  std::unique_ptr<Env> env(NewMemEnv(base_env));
   Options opts;
   opts.env = env.get();
   opts.create_if_missing = true;
@@ -84,13 +102,15 @@ TEST_F(LdbCmdTest, MemEnv) {
 
 TEST_F(LdbCmdTest, OptionParsing) {
   // test parsing flags
+  Options opts;
+  opts.env = TryLoadCustomOrDefaultEnv();
   {
     std::vector<std::string> args;
     args.push_back("scan");
     args.push_back("--ttl");
     args.push_back("--timestamp");
     LDBCommand* command = rocksdb::LDBCommand::InitFromCmdLineArgs(
-        args, Options(), LDBOptions(), nullptr);
+        args, opts, LDBOptions(), nullptr);
     const std::vector<std::string> flags = command->TEST_GetFlags();
     EXPECT_EQ(flags.size(), 2);
     EXPECT_EQ(flags[0], "ttl");
@@ -107,7 +127,7 @@ TEST_F(LdbCmdTest, OptionParsing) {
         "opq:__rst.uvw.xyz?a=3+4+bcd+efghi&jk=lm_no&pq=rst-0&uv=wx-8&yz=a&bcd_"
         "ef=gh.ijk'");
     LDBCommand* command = rocksdb::LDBCommand::InitFromCmdLineArgs(
-        args, Options(), LDBOptions(), nullptr);
+        args, opts, LDBOptions(), nullptr);
     const std::map<std::string, std::string> option_map =
         command->TEST_GetOptionMap();
     EXPECT_EQ(option_map.at("db"), "/dev/shm/ldbtest/");
@@ -120,7 +140,8 @@ TEST_F(LdbCmdTest, OptionParsing) {
 }
 
 TEST_F(LdbCmdTest, ListFileTombstone) {
-  std::unique_ptr<Env> env(NewMemEnv(Env::Default()));
+  Env* base_env = TryLoadCustomOrDefaultEnv();
+  std::unique_ptr<Env> env(NewMemEnv(base_env));
   Options opts;
   opts.env = env.get();
   opts.create_if_missing = true;
@@ -209,8 +230,18 @@ TEST_F(LdbCmdTest, ListFileTombstone) {
 }
 } // namespace rocksdb
 
+#ifdef ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS
+extern "C" {
+void RegisterCustomObjects(int argc, char** argv);
+}
+#else
+void RegisterCustomObjects(int /*argc*/, char** /*argv*/) {}
+#endif  // !ROCKSDB_UNITTESTS_WITH_CUSTOM_OBJECTS_FROM_STATIC_LIBS
+
 int main(int argc, char** argv) {
+  rocksdb::port::InstallStackTraceHandler();
   ::testing::InitGoogleTest(&argc, argv);
+  RegisterCustomObjects(argc, argv);
   return RUN_ALL_TESTS();
 }
 #else

tools/ldb_test.py

@@ -516,13 +516,12 @@ def testWALDump(self):
 
     def testListColumnFamilies(self):
         print "Running testListColumnFamilies..."
-        dbPath = os.path.join(self.TMP_DIR, self.DB_NAME)
         self.assertRunOK("put x1 y1 --create_if_missing", "OK")
-        cmd = "list_column_families %s | grep -v \"Column families\""
+        cmd = "list_column_families | grep -v \"Column families\""
         # Test on valid dbPath.
-        self.assertRunOKFull(cmd % dbPath, "{default}")
+        self.assertRunOK(cmd, "{default}")
         # Test on empty path.
-        self.assertRunFAILFull(cmd % "")
+        self.assertRunFAIL(cmd)
 
     def testColumnFamilies(self):
         print "Running testColumnFamilies..."

tools/ldb_tool.cc

@@ -21,6 +21,8 @@ void LDBCommandRunner::PrintHelp(const LDBOptions& ldb_options,
   ret.append("commands MUST specify --" + LDBCommand::ARG_DB +
              "=<full_path_to_db_directory> when necessary\n");
   ret.append("\n");
+  ret.append("commands can optionally specify --" + LDBCommand::ARG_ENV_URI +
+             "=<uri_of_environment> if necessary\n\n");
   ret.append(
       "The following optional parameters control if keys/values are "
       "input/output as hex or as plain strings:\n");