Add ObjectRegistry to ConfigOptions (facebook#8166)

Summary: This change enables a couple of things: - Different ConfigOptions can have different registry/factory associated with it, thereby allowing things like a "Test" ConfigOptions versus a "Production" - The ObjectRegistry is created fewer times and can be re-used The ConfigOptions can also be initialized/constructed from a DBOptions, in which case it will grab some of its settings (Env, Logger) from the DBOptions. Pull Request resolved: facebook#8166 Reviewed By: zhichao-cao Differential Revision: D27657952 Pulled By: mrambacher fbshipit-source-id: ae1d6200bb7ab127405cdeefaba43c7fe694dfdd
maztheman · May 11, 2021 · 9f2d255 · 9f2d255
1 parent ff46374
commit 9f2d255
Show file tree

Hide file tree

Showing 12 changed files with 309 additions and 41 deletions.
diff --git a/HISTORY.md b/HISTORY.md
@@ -23,6 +23,8 @@
 * Removed unused structure `CompactionFilterContext`.
 * The `skip_filters` parameter to SstFileWriter is now considered deprecated. Use `BlockBasedTableOptions::filter_policy` to control generation of filters.
 * ClockCache is known to have bugs that could lead to crash or corruption, so should not be used until fixed. Use NewLRUCache instead.
+* Added the ObjectRegistry to the ConfigOptions class.  This registry instance will be used to find any customizable loadable objects during initialization.
+* Expanded the ObjectRegistry functionality to allow nested ObjectRegistry instances.  Added methods to register a set of functions with the registry/library as a group.
 * Deprecated backupable_db.h and BackupableDBOptions in favor of new versions with appropriate names: backup_engine.h and BackupEngineOptions. Old API compatibility is preserved.
 
 ### Default Option Change

diff --git a/db/db_options_test.cc b/db/db_options_test.cc
@@ -34,7 +34,7 @@ class DBOptionsTest : public DBTestBase {
       const DBOptions& options) {
     std::string options_str;
     std::unordered_map<std::string, std::string> mutable_map;
-    ConfigOptions config_options;
+    ConfigOptions config_options(options);
     config_options.delimiter = "; ";
 
     EXPECT_OK(GetStringFromMutableDBOptions(

diff --git a/include/rocksdb/convenience.h b/include/rocksdb/convenience.h
@@ -16,6 +16,9 @@
 
 namespace ROCKSDB_NAMESPACE {
 class Env;
+class Logger;
+class ObjectRegistry;
+
 struct ColumnFamilyOptions;
 struct DBOptions;
 struct Options;
@@ -27,6 +30,15 @@ struct Options;
 // of the serialization (e.g. delimiter), and how to compare
 // options (sanity_level).
 struct ConfigOptions {
+  // Constructs a new ConfigOptions with a new object registry.
+  // This method should only be used when a DBOptions is not available,
+  // else registry settings may be lost
+  ConfigOptions();
+
+  // Constructs a new ConfigOptions using the settings from
+  // the input DBOptions.  Currently constructs a new object registry.
+  explicit ConfigOptions(const DBOptions&);
+
   // This enum defines the RocksDB options sanity level.
   enum SanityLevel : unsigned char {
     kSanityLevelNone = 0x01,  // Performs no sanity check at all.
@@ -78,6 +90,11 @@ struct ConfigOptions {
   // The environment to use for this option
   Env* env = Env::Default();
 
+#ifndef ROCKSDB_LITE
+  // The object registry to use for this options
+  std::shared_ptr<ObjectRegistry> registry;
+#endif
+
   bool IsShallow() const { return depth == Depth::kDepthShallow; }
   bool IsDetailed() const { return depth == Depth::kDepthDetailed; }
 
@@ -500,7 +517,6 @@ Status VerifySstFileChecksum(const Options& options,
                              const EnvOptions& env_options,
                              const ReadOptions& read_options,
                              const std::string& file_path);
-
 #endif  // ROCKSDB_LITE
 
 }  // namespace ROCKSDB_NAMESPACE
diff --git a/include/rocksdb/utilities/object_registry.h b/include/rocksdb/utilities/object_registry.h
@@ -17,12 +17,23 @@
 
 namespace ROCKSDB_NAMESPACE {
 class Logger;
+class ObjectLibrary;
+
 // Returns a new T when called with a string. Populates the std::unique_ptr
 // argument if granting ownership to caller.
 template <typename T>
 using FactoryFunc =
     std::function<T*(const std::string&, std::unique_ptr<T>*, std::string*)>;
 
+// The signature of the function for loading factories
+// into an object library.  This method is expected to register
+// factory functions in the supplied ObjectLibrary.
+// The ObjectLibrary is the library in which the factories will be loaded.
+// The std::string is the argument passed to the loader function.
+// The RegistrarFunc should return the number of objects loaded into this
+// library
+using RegistrarFunc = std::function<int(ObjectLibrary&, const std::string&)>;
+
 class ObjectLibrary {
  public:
   // Base class for an Entry in the Registry.
@@ -62,9 +73,18 @@ class ObjectLibrary {
     FactoryFunc<T> factory_;
   };  // End class FactoryEntry
  public:
+  explicit ObjectLibrary(const std::string& id) { id_ = id; }
+
+  const std::string& GetID() const { return id_; }
   // Finds the entry matching the input name and type
   const Entry* FindEntry(const std::string& type,
                          const std::string& name) const;
+
+  // Returns the total number of factories registered for this library.
+  // This method returns the sum of all factories registered for all types.
+  // @param num_types returns how many unique types are registered.
+  size_t GetFactoryCount(size_t* num_types) const;
+
   void Dump(Logger* logger) const;
 
   // Registers the factory with the library for the pattern.
@@ -76,6 +96,12 @@ class ObjectLibrary {
     AddEntry(T::Type(), entry);
     return factory;
   }
+
+  // Invokes the registrar function with the supplied arg for this library.
+  int Register(const RegistrarFunc& registrar, const std::string& arg) {
+    return registrar(*this, arg);
+  }
+
   // Returns the default ObjectLibrary
   static std::shared_ptr<ObjectLibrary>& Default();
 
@@ -85,6 +111,9 @@ class ObjectLibrary {
 
   // ** FactoryFunctions for this loader, organized by type
   std::unordered_map<std::string, std::vector<std::unique_ptr<Entry>>> entries_;
+
+  // The name for this library
+  std::string id_;
 };
 
 // The ObjectRegistry is used to register objects that can be created by a
@@ -93,11 +122,26 @@ class ObjectLibrary {
 class ObjectRegistry {
  public:
   static std::shared_ptr<ObjectRegistry> NewInstance();
-
-  ObjectRegistry();
+  static std::shared_ptr<ObjectRegistry> NewInstance(
+      const std::shared_ptr<ObjectRegistry>& parent);
+  static std::shared_ptr<ObjectRegistry> Default();
+  explicit ObjectRegistry(const std::shared_ptr<ObjectRegistry>& parent)
+      : parent_(parent) {}
+
+  std::shared_ptr<ObjectLibrary> AddLibrary(const std::string& id) {
+    auto library = std::make_shared<ObjectLibrary>(id);
+    libraries_.push_back(library);
+    return library;
+  }
 
   void AddLibrary(const std::shared_ptr<ObjectLibrary>& library) {
-    libraries_.emplace_back(library);
+    libraries_.push_back(library);
+  }
+
+  void AddLibrary(const std::string& id, const RegistrarFunc& registrar,
+                  const std::string& arg) {
+    auto library = AddLibrary(id);
+    library->Register(registrar, arg);
   }
 
   // Creates a new T using the factory function that was registered with a
@@ -193,13 +237,18 @@ class ObjectRegistry {
   void Dump(Logger* logger) const;
 
  private:
+  explicit ObjectRegistry(const std::shared_ptr<ObjectLibrary>& library) {
+    libraries_.push_back(library);
+  }
+
   const ObjectLibrary::Entry* FindEntry(const std::string& type,
                                         const std::string& name) const;
 
   // The set of libraries to search for factories for this registry.
   // The libraries are searched in reverse order (back to front) when
   // searching for entries.
   std::vector<std::shared_ptr<ObjectLibrary>> libraries_;
+  std::shared_ptr<ObjectRegistry> parent_;
 };
 }  // namespace ROCKSDB_NAMESPACE
 #endif  // ROCKSDB_LITE
diff --git a/options/cf_options.cc b/options/cf_options.cc
@@ -541,12 +541,12 @@ static std::unordered_map<std::string, OptionTypeInfo>
           OptionType::kComparator, OptionVerificationType::kByName,
           OptionTypeFlags::kCompareLoose,
           // Parses the string and sets the corresponding comparator
-          [](const ConfigOptions& /*opts*/, const std::string& /*name*/,
+          [](const ConfigOptions& opts, const std::string& /*name*/,
              const std::string& value, char* addr) {
             auto old_comparator = reinterpret_cast<const Comparator**>(addr);
             const Comparator* new_comparator = *old_comparator;
-            Status status = ObjectRegistry::NewInstance()->NewStaticObject(
-                value, &new_comparator);
+            Status status =
+                opts.registry->NewStaticObject(value, &new_comparator);
             if (status.ok()) {
               *old_comparator = new_comparator;
               return status;
@@ -661,12 +661,11 @@ static std::unordered_map<std::string, OptionTypeInfo>
           OptionVerificationType::kByNameAllowFromNull,
           OptionTypeFlags::kCompareLoose,
           // Parses the input value as a MergeOperator, updating the value
-          [](const ConfigOptions& /*opts*/, const std::string& /*name*/,
+          [](const ConfigOptions& opts, const std::string& /*name*/,
              const std::string& value, char* addr) {
             auto mop = reinterpret_cast<std::shared_ptr<MergeOperator>*>(addr);
             Status status =
-                ObjectRegistry::NewInstance()->NewSharedObject<MergeOperator>(
-                    value, mop);
+                opts.registry->NewSharedObject<MergeOperator>(value, mop);
             // Only support static comparator for now.
             if (status.ok()) {
               return status;

diff --git a/options/customizable_helper.h b/options/customizable_helper.h
@@ -83,12 +83,13 @@ static Status LoadSharedObject(const ConfigOptions& config_options,
       return Status::NotSupported("Cannot reset object ", id);
     } else {
 #ifndef ROCKSDB_LITE
-      status = ObjectRegistry::NewInstance()->NewSharedObject(id, result);
+      status = config_options.registry->NewSharedObject(id, result);
 #else
       status = Status::NotSupported("Cannot load object in LITE mode ", id);
 #endif
       if (!status.ok()) {
-        if (config_options.ignore_unsupported_options) {
+        if (config_options.ignore_unsupported_options &&
+            status.IsNotSupported()) {
           return Status::OK();
         } else {
           return status;
@@ -142,12 +143,13 @@ static Status LoadUniqueObject(const ConfigOptions& config_options,
       return Status::NotSupported("Cannot reset object ", id);
     } else {
 #ifndef ROCKSDB_LITE
-      status = ObjectRegistry::NewInstance()->NewUniqueObject(id, result);
+      status = config_options.registry->NewUniqueObject(id, result);
 #else
       status = Status::NotSupported("Cannot load object in LITE mode ", id);
 #endif  // ROCKSDB_LITE
       if (!status.ok()) {
-        if (config_options.ignore_unsupported_options) {
+        if (config_options.ignore_unsupported_options &&
+            status.IsNotSupported()) {
           return Status::OK();
         } else {
           return status;
@@ -199,12 +201,13 @@ static Status LoadStaticObject(const ConfigOptions& config_options,
       return Status::NotSupported("Cannot reset object ", id);
     } else {
 #ifndef ROCKSDB_LITE
-      status = ObjectRegistry::NewInstance()->NewStaticObject(id, result);
+      status = config_options.registry->NewStaticObject(id, result);
 #else
       status = Status::NotSupported("Cannot load object in LITE mode ", id);
 #endif  // ROCKSDB_LITE
       if (!status.ok()) {
-        if (config_options.ignore_unsupported_options) {
+        if (config_options.ignore_unsupported_options &&
+            status.IsNotSupported()) {
           return Status::OK();
         } else {
           return status;

diff --git a/options/customizable_test.cc b/options/customizable_test.cc
@@ -695,6 +695,68 @@ TEST_F(CustomizableTest, MutableOptionsTest) {
 }
 #endif  // !ROCKSDB_LITE
 
+#ifndef ROCKSDB_LITE
+// This method loads existing test classes into the ObjectRegistry
+static int RegisterTestObjects(ObjectLibrary& library,
+                               const std::string& /*arg*/) {
+  size_t num_types;
+  library.Register<TableFactory>(
+      "MockTable",
+      [](const std::string& /*uri*/, std::unique_ptr<TableFactory>* guard,
+         std::string* /* errmsg */) {
+        guard->reset(new mock::MockTableFactory());
+        return guard->get();
+      });
+  return static_cast<int>(library.GetFactoryCount(&num_types));
+}
+
+static int RegisterLocalObjects(ObjectLibrary& library,
+                                const std::string& /*arg*/) {
+  size_t num_types;
+  // Load any locally defined objects here
+  return static_cast<int>(library.GetFactoryCount(&num_types));
+}
+
+class LoadCustomizableTest : public testing::Test {
+ public:
+  LoadCustomizableTest() { config_options_.ignore_unsupported_options = false; }
+  bool RegisterTests(const std::string& arg) {
+#ifndef ROCKSDB_LITE
+    config_options_.registry->AddLibrary("custom-tests", RegisterTestObjects,
+                                         arg);
+    config_options_.registry->AddLibrary("local-tests", RegisterLocalObjects,
+                                         arg);
+    return true;
+#else
+    (void)arg;
+    return false;
+#endif  // !ROCKSDB_LITE
+  }
+
+ protected:
+  DBOptions db_opts_;
+  ColumnFamilyOptions cf_opts_;
+  ConfigOptions config_options_;
+};
+
+TEST_F(LoadCustomizableTest, LoadTableFactoryTest) {
+  std::shared_ptr<TableFactory> factory;
+  ASSERT_NOK(
+      TableFactory::CreateFromString(config_options_, "MockTable", &factory));
+  ASSERT_OK(TableFactory::CreateFromString(
+      config_options_, TableFactory::kBlockBasedTableName(), &factory));
+  ASSERT_NE(factory, nullptr);
+  ASSERT_STREQ(factory->Name(), TableFactory::kBlockBasedTableName());
+
+  if (RegisterTests("Test")) {
+    ASSERT_OK(
+        TableFactory::CreateFromString(config_options_, "MockTable", &factory));
+    ASSERT_NE(factory, nullptr);
+    ASSERT_STREQ(factory->Name(), "MockTable");
+  }
+}
+#endif  // !ROCKSDB_LITE
+
 }  // namespace ROCKSDB_NAMESPACE
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);

diff --git a/options/options_helper.cc b/options/options_helper.cc
@@ -28,6 +28,20 @@
 #include "util/string_util.h"
 
 namespace ROCKSDB_NAMESPACE {
+ConfigOptions::ConfigOptions()
+#ifndef ROCKSDB_LITE
+    : registry(ObjectRegistry::NewInstance())
+#endif
+{
+  env = Env::Default();
+}
+
+ConfigOptions::ConfigOptions(const DBOptions& db_opts) : env(db_opts.env) {
+#ifndef ROCKSDB_LITE
+  registry = ObjectRegistry::NewInstance();
+#endif
+}
+
 Status ValidateOptions(const DBOptions& db_opts,
                        const ColumnFamilyOptions& cf_opts) {
   Status s;
@@ -703,7 +717,7 @@ Status GetStringFromMutableDBOptions(const ConfigOptions& config_options,
 Status GetStringFromDBOptions(std::string* opt_string,
                               const DBOptions& db_options,
                               const std::string& delimiter) {
-  ConfigOptions config_options;
+  ConfigOptions config_options(db_options);
   config_options.delimiter = delimiter;
   return GetStringFromDBOptions(config_options, db_options, opt_string);
 }
@@ -816,7 +830,7 @@ Status GetDBOptionsFromMap(
     const std::unordered_map<std::string, std::string>& opts_map,
     DBOptions* new_options, bool input_strings_escaped,
     bool ignore_unknown_options) {
-  ConfigOptions config_options;
+  ConfigOptions config_options(base_options);
   config_options.input_strings_escaped = input_strings_escaped;
   config_options.ignore_unknown_options = ignore_unknown_options;
   return GetDBOptionsFromMap(config_options, base_options, opts_map,
@@ -844,7 +858,7 @@ Status GetDBOptionsFromMap(
 Status GetDBOptionsFromString(const DBOptions& base_options,
                               const std::string& opts_str,
                               DBOptions* new_options) {
-  ConfigOptions config_options;
+  ConfigOptions config_options(base_options);
   config_options.input_strings_escaped = false;
   config_options.ignore_unknown_options = false;
 
@@ -868,7 +882,7 @@ Status GetDBOptionsFromString(const ConfigOptions& config_options,
 
 Status GetOptionsFromString(const Options& base_options,
                             const std::string& opts_str, Options* new_options) {
-  ConfigOptions config_options;
+  ConfigOptions config_options(base_options);
   config_options.input_strings_escaped = false;
   config_options.ignore_unknown_options = false;
 
@@ -883,6 +897,7 @@ Status GetOptionsFromString(const ConfigOptions& config_options,
   std::unordered_map<std::string, std::string> unused_opts;
   std::unordered_map<std::string, std::string> opts_map;
 
+  assert(new_options);
   *new_options = base_options;
   Status s = StringToMap(opts_str, &opts_map);
   if (!s.ok()) {