Env registry for URI-based Env selection [pluggable Env part 1]

Summary:
This enables configurable Envs without recompiling. For example, my
next diff will make env_test test an Env created by NewEnvFromUri(). Then,
users can determine which Env is tested simply by providing the URI for
NewEnvFromUri() (e.g., through a CLI argument or environment variable).

The registration process allows us to register any Env that is linked with the
RocksDB library, so we can register our internal Envs as well.

The registration code is inspired by our internal InitRegistry.

Test Plan: new unit test

Reviewers: IslamAbdelRahman, lightmark, ldemailly, sdong

Reviewed By: sdong

Subscribers: leveldb, dhruba, andrewkr

Differential Revision: https://reviews.facebook.net/D58449

CMakeLists.txt

@@ -254,6 +254,7 @@ set(SOURCES
         utilities/document/json_document.cc
         utilities/document/json_document_builder.cc
         utilities/env_mirror.cc
+        utilities/env_registry.cc
         utilities/flashcache/flashcache.cc
         utilities/geodb/geodb_impl.cc
         utilities/leveldb_options/leveldb_options.cc
@@ -421,6 +422,7 @@ set(TESTS
         utilities/checkpoint/checkpoint_test.cc
         utilities/document/document_db_test.cc
         utilities/document/json_document_test.cc
+        utilities/env_registry_test.cc
         utilities/geodb/geodb_test.cc
         utilities/memory/memory_test.cc
         utilities/merge_operators/string_append/stringappend_test.cc

Makefile

@@ -970,6 +970,9 @@ spatial_db_test: utilities/spatialdb/spatial_db_test.o $(LIBOBJECTS) $(TESTHARNE
 env_mirror_test: utilities/env_mirror_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 
+env_registry_test: utilities/env_registry_test.o $(LIBOBJECTS) $(TESTHARNESS)
+	$(AM_LINK)
+
 ttl_test: utilities/ttl/ttl_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 

include/rocksdb/utilities/env_registry.h

@@ -0,0 +1,45 @@
+// Copyright (c) 2016-present, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+
+#pragma once
+
+#ifndef ROCKSDB_LITE
+
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "rocksdb/env.h"
+
+namespace rocksdb {
+
+// Returns a new Env when called with a URI string. Populates the unique_ptr
+// argument if granting ownership to caller.
+typedef std::function<Env*(const std::string&, std::unique_ptr<Env>*)>
+    EnvFactoryFunc;
+
+// Creates a new Env using the registered factory function corresponding to a
+// prefix of uri.
+//
+// If no prefixes match, returns nullptr. If multiple prefixes match, the
+// factory function used is unspecified.
+//
+// Populates env_guard with result pointer if caller is granted ownership.
+Env* NewEnvFromUri(const std::string& uri, std::unique_ptr<Env>* env_guard);
+
+// To register an Env factory function, initialize an EnvRegistrar object with
+// static storage duration. For example:
+//
+//   static EnvRegistrar hdfs_reg("hdfs://", &CreateHdfsEnv);
+//
+// Then, calling NewEnvFromUri("hdfs://some_path", ...) will use CreateHdfsEnv
+// to make a new Env.
+class EnvRegistrar {
+ public:
+  explicit EnvRegistrar(std::string uri_prefix, EnvFactoryFunc env_factory);
+};
+
+}  // namespace rocksdb
+#endif  // ROCKSDB_LITE

src.mk

@@ -120,6 +120,7 @@ LIB_SOURCES =                                                   \
   utilities/document/json_document_builder.cc                   \
   utilities/document/json_document.cc                           \
   utilities/env_mirror.cc                                       \
+  utilities/env_registry.cc                                     \
   utilities/flashcache/flashcache.cc                            \
   utilities/geodb/geodb_impl.cc                                 \
   utilities/leveldb_options/leveldb_options.cc                  \
@@ -269,6 +270,7 @@ TEST_BENCH_SOURCES =                                                    \
   utilities/checkpoint/checkpoint_test.cc                               \
   utilities/document/document_db_test.cc                                \
   utilities/document/json_document_test.cc                              \
+  utilities/env_registry_test.cc                                        \
   utilities/geodb/geodb_test.cc                                         \
   utilities/memory/memory_test.cc                                       \
   utilities/merge_operators/string_append/stringappend_test.cc          \

utilities/env_registry.cc

@@ -0,0 +1,47 @@
+// Copyright (c) 2016-present, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+
+#ifndef ROCKSDB_LITE
+
+#include "rocksdb/utilities/env_registry.h"
+
+#include <map>
+#include <string>
+
+namespace rocksdb {
+
+struct EnvRegistryEntry {
+  std::string prefix;
+  EnvFactoryFunc env_factory;
+};
+
+struct EnvRegistry {
+  static EnvRegistry* Get() {
+    static EnvRegistry instance;
+    return &instance;
+  }
+  std::vector<EnvRegistryEntry> entries;
+
+ private:
+  EnvRegistry() = default;
+};
+
+Env* NewEnvFromUri(const std::string& uri, std::unique_ptr<Env>* env_guard) {
+  env_guard->reset();
+  for (const auto& entry : EnvRegistry::Get()->entries) {
+    if (uri.compare(0, entry.prefix.size(), entry.prefix) == 0) {
+      return entry.env_factory(uri, env_guard);
+    }
+  }
+  return nullptr;
+}
+
+EnvRegistrar::EnvRegistrar(std::string uri_prefix, EnvFactoryFunc env_factory) {
+  EnvRegistry::Get()->entries.emplace_back(
+      EnvRegistryEntry{std::move(uri_prefix), std::move(env_factory)});
+}
+
+}  // namespace rocksdb
+#endif  // ROCKSDB_LITE

utilities/env_registry_test.cc

@@ -0,0 +1,72 @@
+// Copyright (c) 2016-present, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+
+#ifndef ROCKSDB_LITE
+
+#include "rocksdb/utilities/env_registry.h"
+#include "util/testharness.h"
+
+namespace rocksdb {
+
+class EnvRegistryTest : public testing::Test {
+ public:
+  static int num_a, num_b;
+};
+
+int EnvRegistryTest::num_a = 0;
+int EnvRegistryTest::num_b = 0;
+
+static EnvRegistrar test_reg_a("a://", [](const std::string& uri,
+                                          std::unique_ptr<Env>* env_guard) {
+  ++EnvRegistryTest::num_a;
+  return Env::Default();
+});
+
+static EnvRegistrar test_reg_b("b://", [](const std::string& uri,
+                                          std::unique_ptr<Env>* env_guard) {
+  ++EnvRegistryTest::num_b;
+  // Env::Default() is a singleton so we can't grant ownership directly to the
+  // caller - we must wrap it first.
+  env_guard->reset(new EnvWrapper(Env::Default()));
+  return env_guard->get();
+});
+
+TEST_F(EnvRegistryTest, Basics) {
+  std::unique_ptr<Env> env_guard;
+  auto res = NewEnvFromUri("a://test", &env_guard);
+  ASSERT_NE(res, nullptr);
+  ASSERT_EQ(env_guard, nullptr);
+  ASSERT_EQ(1, num_a);
+  ASSERT_EQ(0, num_b);
+
+  res = NewEnvFromUri("b://test", &env_guard);
+  ASSERT_NE(res, nullptr);
+  ASSERT_NE(env_guard, nullptr);
+  ASSERT_EQ(1, num_a);
+  ASSERT_EQ(1, num_b);
+
+  res = NewEnvFromUri("c://test", &env_guard);
+  ASSERT_EQ(res, nullptr);
+  ASSERT_EQ(env_guard, nullptr);
+  ASSERT_EQ(1, num_a);
+  ASSERT_EQ(1, num_b);
+}
+
+}  // namespace rocksdb
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
+
+#else  // ROCKSDB_LITE
+#include <stdio.h>
+
+int main(int argc, char** argv) {
+  fprintf(stderr, "SKIPPED as EnvRegistry is not supported in ROCKSDB_LITE\n");
+  return 0;
+}
+
+#endif  // ROCKSDB_LITE