Generalize Env registration framework

Summary:
The Env registration framework supports registering client Envs and selecting which one to instantiate according to a text field. This enabled things like adding the -env_uri argument to db_bench, so the same binary could be reused with different Envs just by changing CLI config.

Now this problem has come up again in a non-Env context, as I want to instantiate a client Statistics implementation from db_bench, which is configured entirely via text parameters. Also, in the future we may wish to use it for deserializing client objects when loading OPTIONS file.

This diff generalizes the Env registration logic to work with arbitrary types.

- Generalized registration and instantiation code by templating them
- The entire implementation is in a header file as that's Google style guide's recommendation for template definitions
- Pattern match with std::regex_match rather than checking prefix, which was the previous behavior
- Rename functions/files to be non-Env-specific
Closes https://github.com/facebook/rocksdb/pull/1776

Differential Revision: D4421933

Pulled By: ajkr

fbshipit-source-id: 34647d1

CMakeLists.txt

@@ -418,7 +418,6 @@ set(SOURCES
         utilities/document/json_document.cc
         utilities/document/json_document_builder.cc
         utilities/env_mirror.cc
-        utilities/env_registry.cc
         utilities/geodb/geodb_impl.cc
         utilities/leveldb_options/leveldb_options.cc
         utilities/lua/rocks_lua_compaction_filter.cc
@@ -629,7 +628,7 @@ set(TESTS
         utilities/date_tiered/date_tiered_test.cc
         utilities/document/document_db_test.cc
         utilities/document/json_document_test.cc
-        utilities/env_registry_test.cc
+        utilities/object_registry_test.cc
         utilities/geodb/geodb_test.cc
         utilities/memory/memory_test.cc
         utilities/merge_operators/string_append/stringappend_test.cc

Makefile

@@ -421,6 +421,7 @@ TESTS = \
 	lua_test \
 	range_del_aggregator_test \
 	lru_cache_test \
+	object_registry_test \
 
 PARALLEL_TEST = \
 	backupable_db_test \
@@ -1085,7 +1086,7 @@ env_librados_test: utilities/env_librados_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_V_CCLD)$(CXX) $^ $(EXEC_LDFLAGS) -o $@ $(LDFLAGS) $(COVERAGEFLAGS)
 endif
 
-env_registry_test: utilities/env_registry_test.o $(LIBOBJECTS) $(TESTHARNESS)
+object_registry_test: utilities/object_registry_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 
 ttl_test: utilities/ttl/ttl_test.o $(LIBOBJECTS) $(TESTHARNESS)

include/rocksdb/utilities/env_registry.h

-// 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

include/rocksdb/utilities/object_registry.h

+// 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 <regex>
+#include <string>
+#include <vector>
+
+#include "rocksdb/env.h"
+
+namespace rocksdb {
+
+// Creates a new T using the factory function that was registered with a pattern
+// that matches the provided "target" string according to std::regex_match.
+//
+// If no registered functions match, returns nullptr. If multiple functions
+// match, the factory function used is unspecified.
+//
+// Populates res_guard with result pointer if caller is granted ownership.
+template <typename T>
+T* NewCustomObject(const std::string& target, std::unique_ptr<T>* res_guard);
+
+// Returns a new T when called with a string. Populates the unique_ptr argument
+// if granting ownership to caller.
+template <typename T>
+using FactoryFunc = std::function<T*(const std::string&, std::unique_ptr<T>*)>;
+
+// To register a factory function for a type T, initialize a Registrar<T> object
+// with static storage duration. For example:
+//
+//   static Registrar<Env> hdfs_reg("hdfs://.*", &CreateHdfsEnv);
+//
+// Then, calling NewCustomObject<Env>("hdfs://some_path", ...) will match the
+// regex provided above, so it returns the result of invoking CreateHdfsEnv.
+template <typename T>
+class Registrar {
+ public:
+  explicit Registrar(std::string pattern, FactoryFunc<T> factory);
+};
+
+// Implementation details follow.
+
+namespace internal {
+
+template <typename T>
+struct RegistryEntry {
+  std::regex pattern;
+  FactoryFunc<T> factory;
+};
+
+template <typename T>
+struct Registry {
+  static Registry* Get() {
+    static Registry<T> instance;
+    return &instance;
+  }
+  std::vector<RegistryEntry<T>> entries;
+
+ private:
+  Registry() = default;
+};
+
+}  // namespace internal
+
+template <typename T>
+T* NewCustomObject(const std::string& target, std::unique_ptr<T>* res_guard) {
+  res_guard->reset();
+  for (const auto& entry : internal::Registry<T>::Get()->entries) {
+    if (std::regex_match(target, entry.pattern)) {
+      return entry.factory(target, res_guard);
+    }
+  }
+  return nullptr;
+}
+
+template <typename T>
+Registrar<T>::Registrar(std::string pattern, FactoryFunc<T> factory) {
+  internal::Registry<T>::Get()->entries.emplace_back(internal::RegistryEntry<T>{
+      std::regex(std::move(pattern)), std::move(factory)});
+}
+
+}  // namespace rocksdb
+#endif  // ROCKSDB_LITE

src.mk

@@ -153,7 +153,6 @@ LIB_SOURCES =                                                   \
   utilities/document/json_document_builder.cc                   \
   utilities/document/json_document.cc                           \
   utilities/env_mirror.cc                                       \
-  utilities/env_registry.cc                                     \
   utilities/geodb/geodb_impl.cc                                 \
   utilities/leveldb_options/leveldb_options.cc                  \
   utilities/lua/rocks_lua_compaction_filter.cc                  \
@@ -299,7 +298,7 @@ MAIN_SOURCES =                                                    \
   utilities/checkpoint/checkpoint_test.cc                               \
   utilities/document/document_db_test.cc                                \
   utilities/document/json_document_test.cc                              \
-  utilities/env_registry_test.cc                                        \
+  utilities/object_registry_test.cc                                     \
   utilities/geodb/geodb_test.cc                                         \
   utilities/memory/memory_test.cc                                       \
   utilities/merge_operators/string_append/stringappend_test.cc          \

tools/db_bench_tool.cc

@@ -47,7 +47,7 @@
 #include "rocksdb/rate_limiter.h"
 #include "rocksdb/slice.h"
 #include "rocksdb/slice_transform.h"
-#include "rocksdb/utilities/env_registry.h"
+#include "rocksdb/utilities/object_registry.h"
 #include "rocksdb/utilities/optimistic_transaction_db.h"
 #include "rocksdb/utilities/options_util.h"
 #include "rocksdb/utilities/sim_cache.h"
@@ -4903,7 +4903,7 @@ int db_bench_tool(int argc, char** argv) {
     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);
+    FLAGS_env = NewCustomObject<Env>(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);

tools/ldb_cmd.cc

@@ -21,7 +21,7 @@
 #include "rocksdb/cache.h"
 #include "rocksdb/table_properties.h"
 #include "rocksdb/utilities/backupable_db.h"
-#include "rocksdb/utilities/env_registry.h"
+#include "rocksdb/utilities/object_registry.h"
 #include "rocksdb/write_batch.h"
 #include "rocksdb/write_buffer_manager.h"
 #include "table/scoped_arena_iterator.h"
@@ -2634,7 +2634,7 @@ void BackupCommand::DoCommand() {
   }
   printf("open db OK\n");
   std::unique_ptr<Env> custom_env_guard;
-  Env* custom_env = NewEnvFromUri(test_cluster_, &custom_env_guard);
+  Env* custom_env = NewCustomObject<Env>(test_cluster_, &custom_env_guard);
   BackupableDBOptions backup_options =
       BackupableDBOptions(test_path_, custom_env);
   backup_options.max_background_operations = thread_num_;
@@ -2696,7 +2696,7 @@ void RestoreCommand::Help(std::string& ret) {
 
 void RestoreCommand::DoCommand() {
   std::unique_ptr<Env> custom_env_guard;
-  Env* custom_env = NewEnvFromUri(backup_env_uri_, &custom_env_guard);
+  Env* custom_env = NewCustomObject<Env>(backup_env_uri_, &custom_env_guard);
   std::unique_ptr<BackupEngineReadOnly> restore_engine;
   Status status;
   {

util/env_basic_test.cc

@@ -8,7 +8,7 @@
 #include <algorithm>
 
 #include "rocksdb/env.h"
-#include "rocksdb/utilities/env_registry.h"
+#include "rocksdb/utilities/object_registry.h"
 #include "util/mock_env.h"
 #include "util/testharness.h"
 
@@ -110,7 +110,7 @@ std::vector<Env*> GetCustomEnvs() {
     init = true;
     const char* uri = getenv("TEST_ENV_URI");
     if (uri != nullptr) {
-      custom_env = NewEnvFromUri(uri, &custom_env_guard);
+      custom_env = NewCustomObject<Env>(uri, &custom_env_guard);
     }
   }
 

utilities/env_registry.cc

-// 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 → utilities/object_registry_test.cc

@@ -5,7 +5,7 @@
 
 #ifndef ROCKSDB_LITE
 
-#include "rocksdb/utilities/env_registry.h"
+#include "rocksdb/utilities/object_registry.h"
 #include "util/testharness.h"
 
 namespace rocksdb {
@@ -18,14 +18,14 @@ class EnvRegistryTest : public testing::Test {
 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) {
+static Registrar<Env> 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) {
+static Registrar<Env> 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.
@@ -35,19 +35,19 @@ static EnvRegistrar test_reg_b("b://", [](const std::string& uri,
 
 TEST_F(EnvRegistryTest, Basics) {
   std::unique_ptr<Env> env_guard;
-  auto res = NewEnvFromUri("a://test", &env_guard);
+  auto res = NewCustomObject<Env>("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);
+  res = NewCustomObject<Env>("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);
+  res = NewCustomObject<Env>("c://test", &env_guard);
   ASSERT_EQ(res, nullptr);
   ASSERT_EQ(env_guard, nullptr);
   ASSERT_EQ(1, num_a);