Add support to the ObjectRegistry for ManagedObjects (#8658)

Summary:
ManagedObjects are  shared pointer objects where RocksDB wants to share a single object between multiple configurations.  For example, the Cache may be shared between multiple column families/tables or the Statistics may be shared between multiple databases.

ManagedObjects are stored in the ObjectRegistry by Type (e.g. Cache) and ID.  For a given type/ID name, a single object is stored.

APIs were added to get/set/create these objects.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/8658

Reviewed By: pdillinger

Differential Revision: D30806273

Pulled By: mrambacher

fbshipit-source-id: 832ac4423b210c4c4b4a456b35897334775d3160

include/rocksdb/customizable.h

@@ -190,6 +190,20 @@ class Customizable : public Configurable {
   virtual const Customizable* Inner() const { return nullptr; }
 
  protected:
+  // Generates a ID specific for this instance of the customizable.
+  // The unique ID is of the form <name>:<addr>#pid, where:
+  // - name is the Name() of this object;
+  // - addr is the memory address of this object;
+  // - pid is the process ID of this process ID for this process.
+  // Note that if obj1 and obj2 have the same unique IDs, they must be the
+  // same.  However, if an object is deleted and recreated, it may have the
+  // same unique ID as a predecessor
+  //
+  // This method is useful for objects (especially ManagedObjects) that
+  // wish to generate an ID that is specific for this instance and wish to
+  // override the GetId() method.
+  std::string GenerateIndividualId() const;
+
   // Some classes have both a class name (e.g. PutOperator) and a nickname
   // (e.g. put). Classes can override this method to return a
   // nickname.  Nicknames can be used by InstanceOf and object creation.

include/rocksdb/utilities/customizable_util.h

@@ -81,6 +81,56 @@ static Status NewSharedObject(
   }
 }
 
+// Creates a new managed customizable instance object based on the
+// input parameters using the object registry.  Unlike "shared" objects,
+// managed objects are limited to a single instance per ID.
+//
+// The id parameter specifies the instance class of the object to create.
+// If an object with this id exists in the registry, the existing object
+// will be returned.  If the object does not exist, a new one will be created.
+//
+// The opt_map parameter specifies the configuration of the new instance.
+// If the object already exists, the existing object is returned "as is" and
+// this parameter is ignored.
+//
+// The config_options parameter controls the process and how errors are
+// returned. If ignore_unknown_options=true, unknown values are ignored during
+// the configuration. If ignore_unsupported_options=true, unknown instance types
+// are ignored. If invoke_prepare_options=true, the resulting instance will be
+// initialized (via PrepareOptions)
+//
+// @param config_options Controls how the instance is created and errors are
+// handled
+// @param id The identifier of the object.  This string
+// will be used by the object registry to locate the appropriate object to
+// create or return.
+// @param opt_map Optional name-value pairs of properties to set for the newly
+// created object
+// @param result The managed instance.
+template <typename T>
+static Status NewManagedObject(
+    const ConfigOptions& config_options, const std::string& id,
+    const std::unordered_map<std::string, std::string>& opt_map,
+    std::shared_ptr<T>* result) {
+  Status status;
+  if (!id.empty()) {
+#ifndef ROCKSDB_LITE
+    status = config_options.registry->GetOrCreateManagedObject<T>(
+        id, result, [config_options, opt_map](T* object) {
+          return object->ConfigureFromMap(config_options, opt_map);
+        });
+#else
+    status = Status::NotSupported("Cannot load object in LITE mode ", id);
+#endif  // ROCKSDB_LITE
+    if (config_options.ignore_unsupported_options && status.IsNotSupported()) {
+      return Status::OK();
+    }
+  } else {
+    status = Status::NotSupported("Cannot reset object ");
+  }
+  return status;
+}
+
 // Creates a new shared Customizable object based on the input parameters.
 // This method parses the input value to determine the type of instance to
 // create. If there is an existing instance (in result) and it is the same ID
@@ -128,6 +178,51 @@ static Status LoadSharedObject(const ConfigOptions& config_options,
   }
 }
 
+// Creates a new shared Customizable object based on the input parameters.
+//
+// The value parameter specified the instance class of the object to create.
+// If it is a simple string (e.g. BlockBasedTable), then the instance will be
+// created using the default settings.  If the value is a set of name-value
+// pairs, then the "id" value is used to determine the instance to create and
+// the remaining parameters are used to configure the object.  Id name-value
+// pairs are specified, there should be an "id=value" pairing or an error may
+// result.
+//
+// The "id" field from the value (either the whole field or "id=XX") is used
+// to determine the type/id of the object to return.  For a given id, there
+// the same instance of the object will be returned from this method (as opposed
+// to LoadSharedObject which would create different objects for the same id.
+//
+// The config_options parameter controls the process and how errors are
+// returned. If ignore_unknown_options=true, unknown values are ignored during
+// the configuration. If ignore_unsupported_options=true, unknown instance types
+// are ignored. If invoke_prepare_options=true, the resulting instance will be
+// initialized (via PrepareOptions)
+//
+// @param config_options Controls how the instance is created and errors are
+// handled
+// @param value Either the simple name of the instance to create, or a set of
+// name-value pairs to create and initailize the object
+// @param func  Optional function to call to attempt to create an instance
+// @param result The newly created instance.
+template <typename T>
+static Status LoadManagedObject(const ConfigOptions& config_options,
+                                const std::string& value,
+                                std::shared_ptr<T>* result) {
+  std::string id;
+  std::unordered_map<std::string, std::string> opt_map;
+  Status status = Customizable::GetOptionsMap(config_options, nullptr, value,
+                                              &id, &opt_map);
+  if (!status.ok()) {  // GetOptionsMap failed
+    return status;
+  } else if (value.empty()) {  // No Id and no options.  Clear the object
+    *result = nullptr;
+    return Status::OK();
+  } else {
+    return NewManagedObject(config_options, id, opt_map, result);
+  }
+}
+
 // Creates a new unique pointer customizable instance object based on the
 // input parameters using the object registry.
 // @see NewSharedObject for more information on the inner workings of this

include/rocksdb/utilities/object_registry.h

@@ -8,6 +8,7 @@
 #ifndef ROCKSDB_LITE
 
 #include <functional>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -18,6 +19,7 @@
 #include "rocksdb/utilities/regex.h"
 
 namespace ROCKSDB_NAMESPACE {
+class Customizable;
 class Logger;
 class ObjectLibrary;
 
@@ -36,6 +38,9 @@ using FactoryFunc =
 // library
 using RegistrarFunc = std::function<int(ObjectLibrary&, const std::string&)>;
 
+template <typename T>
+using ConfigureFunc = std::function<Status(T*)>;
+
 class ObjectLibrary {
  public:
   // Base class for an Entry in the Registry.
@@ -143,11 +148,12 @@ class ObjectRegistry {
 
   std::shared_ptr<ObjectLibrary> AddLibrary(const std::string& id) {
     auto library = std::make_shared<ObjectLibrary>(id);
-    libraries_.push_back(library);
+    AddLibrary(library);
     return library;
   }
 
   void AddLibrary(const std::shared_ptr<ObjectLibrary>& library) {
+    std::unique_lock<std::mutex> lock(library_mutex_);
     libraries_.push_back(library);
   }
 
@@ -249,6 +255,111 @@ class ObjectRegistry {
     }
   }
 
+  // Sets the object for the given id/type to be the input object
+  // If the registry does not contain this id/type, the object is added and OK
+  // is returned. If the registry contains a different object, an error is
+  // returned. If the registry contains the input object, OK is returned.
+  template <typename T>
+  Status SetManagedObject(const std::shared_ptr<T>& object) {
+    assert(object != nullptr);
+    return SetManagedObject(object->GetId(), object);
+  }
+
+  template <typename T>
+  Status SetManagedObject(const std::string& id,
+                          const std::shared_ptr<T>& object) {
+    const auto c = std::static_pointer_cast<Customizable>(object);
+    return SetManagedObject(T::Type(), id, c);
+  }
+
+  // Returns the object for the given id, if one exists.
+  // If the object is not found in the registry, a nullptr is returned
+  template <typename T>
+  std::shared_ptr<T> GetManagedObject(const std::string& id) const {
+    auto c = GetManagedObject(T::Type(), id);
+    return std::static_pointer_cast<T>(c);
+  }
+
+  // Returns the set of managed objects found in the registry matching
+  // the input type and ID.
+  // If the input id is not empty, then only objects of that class
+  // (IsInstanceOf(id)) will be returned (for example, only return LRUCache
+  // objects) If the input id is empty, then all objects of that type (all Cache
+  // objects)
+  template <typename T>
+  Status ListManagedObjects(const std::string& id,
+                            std::vector<std::shared_ptr<T>>* results) const {
+    std::vector<std::shared_ptr<Customizable>> customizables;
+    results->clear();
+    Status s = ListManagedObjects(T::Type(), id, &customizables);
+    if (s.ok()) {
+      for (const auto& c : customizables) {
+        results->push_back(std::static_pointer_cast<T>(c));
+      }
+    }
+    return s;
+  }
+
+  template <typename T>
+  Status ListManagedObjects(std::vector<std::shared_ptr<T>>* results) const {
+    return ListManagedObjects("", results);
+  }
+
+  // Creates a new ManagedObject in the registry for the id if one does not
+  // currently exist.  If an object with that ID already exists, the current
+  // object is returned.
+  //
+  // The ID is the identifier of the object to be returned/created and returned
+  // in result
+  // If a new object is created (using the object factories), the cfunc
+  // parameter will be invoked to configure the new object.
+  template <typename T>
+  Status GetOrCreateManagedObject(const std::string& id,
+                                  std::shared_ptr<T>* result,
+                                  const ConfigureFunc<T>& cfunc = nullptr) {
+    if (parent_ != nullptr) {
+      auto object = parent_->GetManagedObject(T::Type(), id);
+      if (object != nullptr) {
+        *result = std::static_pointer_cast<T>(object);
+        return Status::OK();
+      }
+    }
+    {
+      std::unique_lock<std::mutex> lock(objects_mutex_);
+      auto key = ToManagedObjectKey(T::Type(), id);
+      auto iter = managed_objects_.find(key);
+      if (iter != managed_objects_.end()) {
+        auto object = iter->second.lock();
+        if (object != nullptr) {
+          *result = std::static_pointer_cast<T>(object);
+          return Status::OK();
+        }
+      }
+      std::shared_ptr<T> object;
+      Status s = NewSharedObject(id, &object);
+      if (s.ok() && cfunc != nullptr) {
+        s = cfunc(object.get());
+      }
+      if (s.ok()) {
+        auto c = std::static_pointer_cast<Customizable>(object);
+        if (id != c->Name()) {
+          // If the ID is not the base name of the class, add the new
+          // object under the input ID
+          managed_objects_[key] = c;
+        }
+        if (id != c->GetId() && c->GetId() != c->Name()) {
+          // If the input and current ID do not match, and the
+          // current ID is not the base bame, add the new object under
+          // its new ID
+          key = ToManagedObjectKey(T::Type(), c->GetId());
+          managed_objects_[key] = c;
+        }
+        *result = object;
+      }
+      return s;
+    }
+  }
+
   // Dump the contents of the registry to the logger
   void Dump(Logger* logger) const;
 
@@ -256,6 +367,24 @@ class ObjectRegistry {
   explicit ObjectRegistry(const std::shared_ptr<ObjectLibrary>& library) {
     libraries_.push_back(library);
   }
+  static std::string ToManagedObjectKey(const std::string& type,
+                                        const std::string& id) {
+    return type + "://" + id;
+  }
+
+  // Returns the Customizable managed object associated with the key (Type/ID).
+  // If not found, nullptr is returned.
+  std::shared_ptr<Customizable> GetManagedObject(const std::string& type,
+                                                 const std::string& id) const;
+  Status ListManagedObjects(
+      const std::string& type, const std::string& pattern,
+      std::vector<std::shared_ptr<Customizable>>* results) const;
+  // Sets the managed object associated with the key (Type/ID) to c.
+  // If the named managed object does not exist, the object is added and OK is
+  // returned If the object exists and is the same as c, OK is returned
+  // Otherwise, an error status is returned.
+  Status SetManagedObject(const std::string& type, const std::string& id,
+                          const std::shared_ptr<Customizable>& c);
 
   const ObjectLibrary::Entry* FindEntry(const std::string& type,
                                         const std::string& name) const;
@@ -264,7 +393,10 @@ class ObjectRegistry {
   // The libraries are searched in reverse order (back to front) when
   // searching for entries.
   std::vector<std::shared_ptr<ObjectLibrary>> libraries_;
+  std::map<std::string, std::weak_ptr<Customizable>> managed_objects_;
   std::shared_ptr<ObjectRegistry> parent_;
+  mutable std::mutex objects_mutex_;  // Mutex for managed objects
+  mutable std::mutex library_mutex_;  // Mutex for managed libraries
 };
 }  // namespace ROCKSDB_NAMESPACE
 #endif  // ROCKSDB_LITE

options/customizable.cc

@@ -5,7 +5,10 @@
 
 #include "rocksdb/customizable.h"
 
+#include <sstream>
+
 #include "options/options_helper.h"
+#include "port/port.h"
 #include "rocksdb/convenience.h"
 #include "rocksdb/status.h"
 #include "rocksdb/utilities/options_type.h"
@@ -25,6 +28,13 @@ std::string Customizable::GetOptionName(const std::string& long_name) const {
   }
 }
 
+std::string Customizable::GenerateIndividualId() const {
+  std::ostringstream ostr;
+  ostr << Name() << "@" << static_cast<const void*>(this) << "#"
+       << port::GetProcessID();
+  return ostr.str();
+}
+
 #ifndef ROCKSDB_LITE
 Status Customizable::GetOption(const ConfigOptions& config_options,
                                const std::string& opt_name,

options/customizable_test.cc

@@ -17,6 +17,7 @@
 #include "db/db_test_util.h"
 #include "options/options_helper.h"
 #include "options/options_parser.h"
+#include "port/stack_trace.h"
 #include "rocksdb/convenience.h"
 #include "rocksdb/env_encryption.h"
 #include "rocksdb/flush_block_policy.h"
@@ -1064,6 +1065,135 @@ TEST_F(CustomizableTest, MutableOptionsTest) {
   ASSERT_FALSE(mc.AreEquivalent(options, &mc2, &mismatch));
   ASSERT_EQ(mismatch, "immutable");
 }
+
+TEST_F(CustomizableTest, CustomManagedObjects) {
+  std::shared_ptr<TestCustomizable> object1, object2;
+  ASSERT_OK(LoadManagedObject<TestCustomizable>(
+      config_options_, "id=A_1;int=1;bool=true", &object1));
+  ASSERT_OK(
+      LoadManagedObject<TestCustomizable>(config_options_, "A_1", &object2));
+  ASSERT_EQ(object1, object2);
+  auto* opts = object2->GetOptions<AOptions>("A");
+  ASSERT_NE(opts, nullptr);
+  ASSERT_EQ(opts->i, 1);
+  ASSERT_EQ(opts->b, true);
+  ASSERT_OK(
+      LoadManagedObject<TestCustomizable>(config_options_, "A_2", &object2));
+  ASSERT_NE(object1, object2);
+  object1.reset();
+  ASSERT_OK(LoadManagedObject<TestCustomizable>(
+      config_options_, "id=A_1;int=2;bool=false", &object1));
+  opts = object1->GetOptions<AOptions>("A");
+  ASSERT_NE(opts, nullptr);
+  ASSERT_EQ(opts->i, 2);
+  ASSERT_EQ(opts->b, false);
+}
+
+TEST_F(CustomizableTest, CreateManagedObjects) {
+  class ManagedCustomizable : public Customizable {
+   public:
+    static const char* Type() { return "ManagedCustomizable"; }
+    static const char* kClassName() { return "Managed"; }
+    const char* Name() const override { return kClassName(); }
+    std::string GetId() const override { return id_; }
+    ManagedCustomizable() { id_ = GenerateIndividualId(); }
+    static Status CreateFromString(
+        const ConfigOptions& opts, const std::string& value,
+        std::shared_ptr<ManagedCustomizable>* result) {
+      return LoadManagedObject<ManagedCustomizable>(opts, value, result);
+    }
+
+   private:
+    std::string id_;
+  };
+
+  config_options_.registry->AddLibrary("Managed")
+      ->Register<ManagedCustomizable>(
+          "Managed(@.*)?", [](const std::string& /*name*/,
+                              std::unique_ptr<ManagedCustomizable>* guard,
+                              std::string* /* msg */) {
+            guard->reset(new ManagedCustomizable());
+            return guard->get();
+          });
+
+  std::shared_ptr<ManagedCustomizable> mc1, mc2, mc3, obj;
+  // Create a "deadbeef" customizable
+  std::string deadbeef =
+      std::string(ManagedCustomizable::kClassName()) + "@0xdeadbeef#0001";
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, deadbeef, &mc1));
+  // Create an object with the base/class name
+  ASSERT_OK(ManagedCustomizable::CreateFromString(
+      config_options_, ManagedCustomizable::kClassName(), &mc2));
+  // Creating another with the base name returns a different object
+  ASSERT_OK(ManagedCustomizable::CreateFromString(
+      config_options_, ManagedCustomizable::kClassName(), &mc3));
+  // At this point, there should be 4 managed objects (deadbeef, mc1, 2, and 3)
+  std::vector<std::shared_ptr<ManagedCustomizable>> objects;
+  ASSERT_OK(config_options_.registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 4U);
+  objects.clear();
+  // Three separate object, none of them equal
+  ASSERT_NE(mc1, mc2);
+  ASSERT_NE(mc1, mc3);
+  ASSERT_NE(mc2, mc3);
+
+  // Creating another object with "deadbeef" object
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, deadbeef, &obj));
+  ASSERT_EQ(mc1, obj);
+  // Create another with the IDs of the instances
+  ASSERT_OK(ManagedCustomizable::CreateFromString(config_options_, mc1->GetId(),
+                                                  &obj));
+  ASSERT_EQ(mc1, obj);
+  ASSERT_OK(ManagedCustomizable::CreateFromString(config_options_, mc2->GetId(),
+                                                  &obj));
+  ASSERT_EQ(mc2, obj);
+  ASSERT_OK(ManagedCustomizable::CreateFromString(config_options_, mc3->GetId(),
+                                                  &obj));
+  ASSERT_EQ(mc3, obj);
+
+  // Now get rid of deadbeef.  2 Objects left (m2+m3)
+  mc1.reset();
+  ASSERT_EQ(
+      config_options_.registry->GetManagedObject<ManagedCustomizable>(deadbeef),
+      nullptr);
+  ASSERT_OK(config_options_.registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 2U);
+  objects.clear();
+
+  // Associate deadbeef with #2
+  ASSERT_OK(config_options_.registry->SetManagedObject(deadbeef, mc2));
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, deadbeef, &obj));
+  ASSERT_EQ(mc2, obj);
+  obj.reset();
+
+  // Get the ID of mc2 and then reset it.  1 Object left
+  std::string mc2id = mc2->GetId();
+  mc2.reset();
+  ASSERT_EQ(
+      config_options_.registry->GetManagedObject<ManagedCustomizable>(mc2id),
+      nullptr);
+  ASSERT_OK(config_options_.registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 1U);
+  objects.clear();
+
+  // Create another object with the old mc2id.
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, mc2id, &mc2));
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, mc2id, &obj));
+  ASSERT_EQ(mc2, obj);
+
+  // For good measure, create another deadbeef object
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, deadbeef, &mc1));
+  ASSERT_OK(
+      ManagedCustomizable::CreateFromString(config_options_, deadbeef, &obj));
+  ASSERT_EQ(mc1, obj);
+}
+
 #endif  // !ROCKSDB_LITE
 
 namespace {
@@ -1482,6 +1612,7 @@ TEST_F(LoadCustomizableTest, LoadFlushBlockPolicyFactoryTest) {
 }  // namespace ROCKSDB_NAMESPACE
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
 #ifdef GFLAGS
   ParseCommandLineFlags(&argc, &argv, true);
 #endif  // GFLAGS

utilities/object_registry.cc

@@ -6,7 +6,9 @@
 #include "rocksdb/utilities/object_registry.h"
 
 #include "logging/logging.h"
+#include "rocksdb/customizable.h"
 #include "rocksdb/env.h"
+#include "util/string_util.h"
 
 namespace ROCKSDB_NAMESPACE {
 #ifndef ROCKSDB_LITE
@@ -87,23 +89,97 @@ std::shared_ptr<ObjectRegistry> ObjectRegistry::NewInstance(
 // Returns the entry if it is found, and nullptr otherwise
 const ObjectLibrary::Entry *ObjectRegistry::FindEntry(
     const std::string &type, const std::string &name) const {
+  {
+    std::unique_lock<std::mutex> lock(library_mutex_);
     for (auto iter = libraries_.crbegin(); iter != libraries_.crend(); ++iter) {
       const auto *entry = iter->get()->FindEntry(type, name);
       if (entry != nullptr) {
         return entry;
       }
     }
+  }
   if (parent_ != nullptr) {
     return parent_->FindEntry(type, name);
   } else {
     return nullptr;
   }
 }
+Status ObjectRegistry::SetManagedObject(
+    const std::string &type, const std::string &id,
+    const std::shared_ptr<Customizable> &object) {
+  std::string object_key = ToManagedObjectKey(type, id);
+  std::shared_ptr<Customizable> curr;
+  if (parent_ != nullptr) {
+    curr = parent_->GetManagedObject(type, id);
+  }
+  if (curr == nullptr) {
+    // We did not find the object in any parent.  Update in the current
+    std::unique_lock<std::mutex> lock(objects_mutex_);
+    auto iter = managed_objects_.find(object_key);
+    if (iter != managed_objects_.end()) {  // The object exists
+      curr = iter->second.lock();
+      if (curr != nullptr && curr != object) {
+        return Status::InvalidArgument("Object already exists: ", object_key);
+      } else {
+        iter->second = object;
+      }
+    } else {
+      // The object does not exist.  Add it
+      managed_objects_[object_key] = object;
+    }
+  } else if (curr != object) {
+    return Status::InvalidArgument("Object already exists: ", object_key);
+  }
+  return Status::OK();
+}
+
+std::shared_ptr<Customizable> ObjectRegistry::GetManagedObject(
+    const std::string &type, const std::string &id) const {
+  {
+    std::unique_lock<std::mutex> lock(objects_mutex_);
+    auto iter = managed_objects_.find(ToManagedObjectKey(type, id));
+    if (iter != managed_objects_.end()) {
+      return iter->second.lock();
+    }
+  }
+  if (parent_ != nullptr) {
+    return parent_->GetManagedObject(type, id);
+  } else {
+    return nullptr;
+  }
+}
+
+Status ObjectRegistry::ListManagedObjects(
+    const std::string &type, const std::string &name,
+    std::vector<std::shared_ptr<Customizable>> *results) const {
+  {
+    std::string key = ToManagedObjectKey(type, name);
+    std::unique_lock<std::mutex> lock(objects_mutex_);
+    for (auto iter = managed_objects_.lower_bound(key);
+         iter != managed_objects_.end() && StartsWith(iter->first, key);
+         ++iter) {
+      auto shared = iter->second.lock();
+      if (shared != nullptr) {
+        if (name.empty() || shared->IsInstanceOf(name)) {
+          results->emplace_back(shared);
+        }
+      }
+    }
+  }
+  if (parent_ != nullptr) {
+    return parent_->ListManagedObjects(type, name, results);
+  } else {
+    return Status::OK();
+  }
+}
 
 void ObjectRegistry::Dump(Logger *logger) const {
+  {
+    std::unique_lock<std::mutex> lock(library_mutex_);
     for (auto iter = libraries_.crbegin(); iter != libraries_.crend(); ++iter) {
       iter->get()->Dump(logger);
     }
+  }
   if (parent_ != nullptr) {
     parent_->Dump(logger);
   }

utilities/object_registry_test.cc

@@ -6,6 +6,8 @@
 #ifndef ROCKSDB_LITE
 
 #include "rocksdb/utilities/object_registry.h"
+
+#include "rocksdb/customizable.h"
 #include "test_util/testharness.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -205,6 +207,230 @@ TEST_F(EnvRegistryTest, TestRegistryParents) {
   ASSERT_NOK(child->NewUniqueObject<Env>("cousin", &guard));
   ASSERT_NOK(uncle->NewUniqueObject<Env>("cousin", &guard));
 }
+class MyCustomizable : public Customizable {
+ public:
+  static const char* Type() { return "MyCustomizable"; }
+  MyCustomizable(const char* prefix, const std::string& id) : id_(id) {
+    name_ = id_.substr(0, strlen(prefix) - 1);
+  }
+  const char* Name() const override { return name_.c_str(); }
+  std::string GetId() const override { return id_; }
+
+ private:
+  std::string id_;
+  std::string name_;
+};
+
+TEST_F(EnvRegistryTest, TestManagedObjects) {
+  auto registry = ObjectRegistry::NewInstance();
+  auto m_a1 = std::make_shared<MyCustomizable>("", "A");
+  auto m_a2 = std::make_shared<MyCustomizable>("", "A");
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a1));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a1);
+
+  ASSERT_NOK(registry->SetManagedObject<MyCustomizable>(m_a2));
+  ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a1));
+  m_a1.reset();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_OK(registry->SetManagedObject<MyCustomizable>(m_a2));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a2);
+}
+
+TEST_F(EnvRegistryTest, TestTwoManagedObjects) {
+  auto registry = ObjectRegistry::NewInstance();
+  auto m_a = std::make_shared<MyCustomizable>("", "A");
+  auto m_b = std::make_shared<MyCustomizable>("", "B");
+  std::vector<std::shared_ptr<MyCustomizable>> objects;
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 0U);
+  ASSERT_OK(registry->SetManagedObject(m_a));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_a);
+
+  ASSERT_OK(registry->SetManagedObject(m_b));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), m_b);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 2U);
+  ASSERT_OK(registry->ListManagedObjects("A", &objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_a);
+  ASSERT_OK(registry->ListManagedObjects("B", &objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_b);
+  ASSERT_OK(registry->ListManagedObjects("C", &objects));
+  ASSERT_EQ(objects.size(), 0U);
+
+  m_a.reset();
+  objects.clear();
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), m_b);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_b);
+
+  m_b.reset();
+  objects.clear();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+}
+
+TEST_F(EnvRegistryTest, TestAlternateNames) {
+  auto registry = ObjectRegistry::NewInstance();
+  auto m_a = std::make_shared<MyCustomizable>("", "A");
+  auto m_b = std::make_shared<MyCustomizable>("", "B");
+  std::vector<std::shared_ptr<MyCustomizable>> objects;
+  // Test no objects exist
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 0U);
+
+  // Mark "TheOne" to be A
+  ASSERT_OK(registry->SetManagedObject("TheOne", m_a));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_a);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_a);
+
+  // Try to mark "TheOne" again.
+  ASSERT_NOK(registry->SetManagedObject("TheOne", m_b));
+  ASSERT_OK(registry->SetManagedObject("TheOne", m_a));
+
+  // Add "A" as a managed object.  Registered 2x
+  ASSERT_OK(registry->SetManagedObject(m_a));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("B"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_a);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 2U);
+
+  // Delete "A".
+  m_a.reset();
+  objects.clear();
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
+  ASSERT_OK(registry->SetManagedObject("TheOne", m_b));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), m_b);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 1U);
+  ASSERT_EQ(objects.front(), m_b);
+
+  m_b.reset();
+  objects.clear();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("TheOne"), nullptr);
+  ASSERT_OK(registry->ListManagedObjects(&objects));
+  ASSERT_EQ(objects.size(), 0U);
+}
+
+TEST_F(EnvRegistryTest, TestTwoManagedClasses) {
+  class MyCustomizable2 : public MyCustomizable {
+   public:
+    static const char* Type() { return "MyCustomizable2"; }
+    MyCustomizable2(const char* prefix, const std::string& id)
+        : MyCustomizable(prefix, id) {}
+  };
+
+  auto registry = ObjectRegistry::NewInstance();
+  auto m_a1 = std::make_shared<MyCustomizable>("", "A");
+  auto m_a2 = std::make_shared<MyCustomizable2>("", "A");
+  std::vector<std::shared_ptr<MyCustomizable>> obj1s;
+  std::vector<std::shared_ptr<MyCustomizable2>> obj2s;
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
+
+  ASSERT_OK(registry->SetManagedObject(m_a1));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a1);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
+
+  ASSERT_OK(registry->SetManagedObject(m_a2));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), m_a2);
+  ASSERT_OK(registry->ListManagedObjects(&obj1s));
+  ASSERT_OK(registry->ListManagedObjects(&obj2s));
+  ASSERT_EQ(obj1s.size(), 1U);
+  ASSERT_EQ(obj2s.size(), 1U);
+  ASSERT_EQ(obj1s.front(), m_a1);
+  ASSERT_EQ(obj2s.front(), m_a2);
+  m_a1.reset();
+  obj1s.clear();
+  obj2s.clear();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), m_a2);
+
+  m_a2.reset();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable2>("A"), nullptr);
+}
+
+TEST_F(EnvRegistryTest, TestManagedObjectsWithParent) {
+  auto base = ObjectRegistry::NewInstance();
+  auto registry = ObjectRegistry::NewInstance(base);
+
+  auto m_a = std::make_shared<MyCustomizable>("", "A");
+  auto m_b = std::make_shared<MyCustomizable>("", "A");
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_OK(base->SetManagedObject(m_a));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_a);
+
+  ASSERT_NOK(registry->SetManagedObject(m_b));
+  ASSERT_OK(registry->SetManagedObject(m_a));
+
+  m_a.reset();
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), nullptr);
+  ASSERT_OK(registry->SetManagedObject(m_b));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("A"), m_b);
+}
+
+TEST_F(EnvRegistryTest, TestGetOrCreateManagedObject) {
+  auto registry = ObjectRegistry::NewInstance();
+  registry->AddLibrary("test")->Register<MyCustomizable>(
+      "MC(@.*)?",
+      [](const std::string& uri, std::unique_ptr<MyCustomizable>* guard,
+         std::string* /* errmsg */) {
+        guard->reset(new MyCustomizable("MC", uri));
+        return guard->get();
+      });
+  std::shared_ptr<MyCustomizable> m_a, m_b, obj;
+  std::vector<std::shared_ptr<MyCustomizable>> objs;
+
+  std::unordered_map<std::string, std::string> opt_map;
+
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@A"), nullptr);
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@B"), nullptr);
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &m_a));
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &m_b));
+  ASSERT_EQ(registry->GetManagedObject<MyCustomizable>("MC@A"), m_a);
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &obj));
+  ASSERT_EQ(obj, m_a);
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &obj));
+  ASSERT_EQ(obj, m_b);
+  ASSERT_OK(registry->ListManagedObjects(&objs));
+  ASSERT_EQ(objs.size(), 2U);
+
+  objs.clear();
+  m_a.reset();
+  obj.reset();
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &m_a));
+  ASSERT_EQ(1, m_a.use_count());
+  ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &obj));
+  ASSERT_EQ(2, obj.use_count());
+}
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {