From a683d4aba9c1ba230248fc985b26208ee81bbc63 Mon Sep 17 00:00:00 2001
From: Andrew Kryczka <andrewkr@fb.com>
Date: Thu, 9 Jun 2016 17:53:03 -0700
Subject: [PATCH] URI-based Env selection for db_bench

Summary:
Added an option, --env_uri. When provided, it is used as an argument to
NewEnvFromUri(), which instantiates an Env based on it.

Test Plan:
built a simple binary that registers ChrootEnv for prefix "/", then
ran:

  $ ./tmp --env_uri /tmp/ --db /abcde

/tmp/ is the chroot directory and /abcde is the db_name. Then I verified
db_bench uses /tmp/abcde

Reviewers: sdong, kradhakrishnan, lightmark

Reviewed By: lightmark

Subscribers: andrewkr, dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D59325
---
 tools/db_bench_tool.cc | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc
index f40ed8ad9..17bb38dfa 100644
--- a/tools/db_bench_tool.cc
+++ b/tools/db_bench_tool.cc
@@ -49,6 +49,7 @@
 #include "rocksdb/rate_limiter.h"
 #include "rocksdb/slice.h"
 #include "rocksdb/slice_transform.h"
+#include "rocksdb/utilities/env_registry.h"
 #include "rocksdb/utilities/flashcache.h"
 #include "rocksdb/utilities/optimistic_transaction_db.h"
 #include "rocksdb/utilities/options_util.h"
@@ -628,8 +629,12 @@ static bool ValidateTableCacheNumshardbits(const char* flagname,
 }
 DEFINE_int32(table_cache_numshardbits, 4, "");
 
-DEFINE_string(hdfs, "", "Name of hdfs environment");
-// posix or hdfs environment
+#ifndef ROCKSDB_LITE
+DEFINE_string(env_uri, "", "URI for registry Env lookup. Mutually exclusive"
+              " with --hdfs.");
+#endif  // ROCKSDB_LITE
+DEFINE_string(hdfs, "", "Name of hdfs environment. Mutually exclusive with"
+              " --env_uri.");
 static rocksdb::Env* FLAGS_env = rocksdb::Env::Default();
 
 DEFINE_int64(stats_interval, 0, "Stats are reported every N operations when "
@@ -4074,6 +4079,19 @@ int db_bench_tool(int argc, char** argv) {
   FLAGS_compression_type_e =
     StringToCompressionType(FLAGS_compression_type.c_str());
 
+#ifndef ROCKSDB_LITE
+  std::unique_ptr<Env> custom_env_guard;
+  if (!FLAGS_hdfs.empty() && !FLAGS_env_uri.empty()) {
+    fprintf(stderr, "Cannot provide both --hdfs and --env_uri.\n");
+    exit(1);
+  } else if (!FLAGS_env_uri.empty()) {
+    FLAGS_env = NewEnvFromUri(FLAGS_env_uri, &custom_env_guard);
+    if (FLAGS_env == nullptr) {
+      fprintf(stderr, "No Env registered for URI: %s\n", FLAGS_env_uri.c_str());
+      exit(1);
+    }
+  }
+#endif  // ROCKSDB_LITE
   if (!FLAGS_hdfs.empty()) {
     FLAGS_env  = new rocksdb::HdfsEnv(FLAGS_hdfs);
   }