From 0fb938c4482af318817b9d0737a741681b59a117 Mon Sep 17 00:00:00 2001 From: mrambacher Date: Fri, 10 Sep 2021 05:19:47 -0700 Subject: [PATCH] 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 | 14 ++ include/rocksdb/utilities/customizable_util.h | 95 ++++++++ include/rocksdb/utilities/object_registry.h | 134 ++++++++++- options/customizable.cc | 10 + options/customizable_test.cc | 131 ++++++++++ utilities/object_registry.cc | 88 ++++++- utilities/object_registry_test.cc | 226 ++++++++++++++++++ 7 files changed, 691 insertions(+), 7 deletions(-) diff --git a/include/rocksdb/customizable.h b/include/rocksdb/customizable.h index b168f472e..d6bd07087 100644 --- a/include/rocksdb/customizable.h +++ b/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 :#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. diff --git a/include/rocksdb/utilities/customizable_util.h b/include/rocksdb/utilities/customizable_util.h index 518207ac1..05b8460e5 100644 --- a/include/rocksdb/utilities/customizable_util.h +++ b/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 +static Status NewManagedObject( + const ConfigOptions& config_options, const std::string& id, + const std::unordered_map& opt_map, + std::shared_ptr* result) { + Status status; + if (!id.empty()) { +#ifndef ROCKSDB_LITE + status = config_options.registry->GetOrCreateManagedObject( + 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 +static Status LoadManagedObject(const ConfigOptions& config_options, + const std::string& value, + std::shared_ptr* result) { + std::string id; + std::unordered_map 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 diff --git a/include/rocksdb/utilities/object_registry.h b/include/rocksdb/utilities/object_registry.h index 31ed1abfb..b21868709 100644 --- a/include/rocksdb/utilities/object_registry.h +++ b/include/rocksdb/utilities/object_registry.h @@ -8,6 +8,7 @@ #ifndef ROCKSDB_LITE #include +#include #include #include #include @@ -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; +template +using ConfigureFunc = std::function; + class ObjectLibrary { public: // Base class for an Entry in the Registry. @@ -143,11 +148,12 @@ class ObjectRegistry { std::shared_ptr AddLibrary(const std::string& id) { auto library = std::make_shared(id); - libraries_.push_back(library); + AddLibrary(library); return library; } void AddLibrary(const std::shared_ptr& library) { + std::unique_lock 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 + Status SetManagedObject(const std::shared_ptr& object) { + assert(object != nullptr); + return SetManagedObject(object->GetId(), object); + } + + template + Status SetManagedObject(const std::string& id, + const std::shared_ptr& object) { + const auto c = std::static_pointer_cast(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 + std::shared_ptr GetManagedObject(const std::string& id) const { + auto c = GetManagedObject(T::Type(), id); + return std::static_pointer_cast(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 + Status ListManagedObjects(const std::string& id, + std::vector>* results) const { + std::vector> 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(c)); + } + } + return s; + } + + template + Status ListManagedObjects(std::vector>* 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 + Status GetOrCreateManagedObject(const std::string& id, + std::shared_ptr* result, + const ConfigureFunc& cfunc = nullptr) { + if (parent_ != nullptr) { + auto object = parent_->GetManagedObject(T::Type(), id); + if (object != nullptr) { + *result = std::static_pointer_cast(object); + return Status::OK(); + } + } + { + std::unique_lock 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(object); + return Status::OK(); + } + } + std::shared_ptr object; + Status s = NewSharedObject(id, &object); + if (s.ok() && cfunc != nullptr) { + s = cfunc(object.get()); + } + if (s.ok()) { + auto c = std::static_pointer_cast(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& 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 GetManagedObject(const std::string& type, + const std::string& id) const; + Status ListManagedObjects( + const std::string& type, const std::string& pattern, + std::vector>* 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& 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> libraries_; + std::map> managed_objects_; std::shared_ptr 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 diff --git a/options/customizable.cc b/options/customizable.cc index df69f5b85..88ccf1de7 100644 --- a/options/customizable.cc +++ b/options/customizable.cc @@ -5,7 +5,10 @@ #include "rocksdb/customizable.h" +#include + #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(this) << "#" + << port::GetProcessID(); + return ostr.str(); +} + #ifndef ROCKSDB_LITE Status Customizable::GetOption(const ConfigOptions& config_options, const std::string& opt_name, diff --git a/options/customizable_test.cc b/options/customizable_test.cc index 20cd3c0c5..e4afca42c 100644 --- a/options/customizable_test.cc +++ b/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 object1, object2; + ASSERT_OK(LoadManagedObject( + config_options_, "id=A_1;int=1;bool=true", &object1)); + ASSERT_OK( + LoadManagedObject(config_options_, "A_1", &object2)); + ASSERT_EQ(object1, object2); + auto* opts = object2->GetOptions("A"); + ASSERT_NE(opts, nullptr); + ASSERT_EQ(opts->i, 1); + ASSERT_EQ(opts->b, true); + ASSERT_OK( + LoadManagedObject(config_options_, "A_2", &object2)); + ASSERT_NE(object1, object2); + object1.reset(); + ASSERT_OK(LoadManagedObject( + config_options_, "id=A_1;int=2;bool=false", &object1)); + opts = object1->GetOptions("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* result) { + return LoadManagedObject(opts, value, result); + } + + private: + std::string id_; + }; + + config_options_.registry->AddLibrary("Managed") + ->Register( + "Managed(@.*)?", [](const std::string& /*name*/, + std::unique_ptr* guard, + std::string* /* msg */) { + guard->reset(new ManagedCustomizable()); + return guard->get(); + }); + + std::shared_ptr 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> 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(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(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 diff --git a/utilities/object_registry.cc b/utilities/object_registry.cc index 79e4b5683..b0a60f53e 100644 --- a/utilities/object_registry.cc +++ b/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,10 +89,13 @@ std::shared_ptr 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 { - for (auto iter = libraries_.crbegin(); iter != libraries_.crend(); ++iter) { - const auto *entry = iter->get()->FindEntry(type, name); - if (entry != nullptr) { - return entry; + { + std::unique_lock 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) { @@ -99,10 +104,81 @@ const ObjectLibrary::Entry *ObjectRegistry::FindEntry( return nullptr; } } +Status ObjectRegistry::SetManagedObject( + const std::string &type, const std::string &id, + const std::shared_ptr &object) { + std::string object_key = ToManagedObjectKey(type, id); + std::shared_ptr 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 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 ObjectRegistry::GetManagedObject( + const std::string &type, const std::string &id) const { + { + std::unique_lock 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> *results) const { + { + std::string key = ToManagedObjectKey(type, name); + std::unique_lock 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 { - for (auto iter = libraries_.crbegin(); iter != libraries_.crend(); ++iter) { - iter->get()->Dump(logger); + { + std::unique_lock lock(library_mutex_); + for (auto iter = libraries_.crbegin(); iter != libraries_.crend(); ++iter) { + iter->get()->Dump(logger); + } } if (parent_ != nullptr) { parent_->Dump(logger); diff --git a/utilities/object_registry_test.cc b/utilities/object_registry_test.cc index e09392dd8..d9a5b1526 100644 --- a/utilities/object_registry_test.cc +++ b/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("cousin", &guard)); ASSERT_NOK(uncle->NewUniqueObject("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("", "A"); + auto m_a2 = std::make_shared("", "A"); + + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_OK(registry->SetManagedObject(m_a1)); + ASSERT_EQ(registry->GetManagedObject("A"), m_a1); + + ASSERT_NOK(registry->SetManagedObject(m_a2)); + ASSERT_OK(registry->SetManagedObject(m_a1)); + m_a1.reset(); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_OK(registry->SetManagedObject(m_a2)); + ASSERT_EQ(registry->GetManagedObject("A"), m_a2); +} + +TEST_F(EnvRegistryTest, TestTwoManagedObjects) { + auto registry = ObjectRegistry::NewInstance(); + auto m_a = std::make_shared("", "A"); + auto m_b = std::make_shared("", "B"); + std::vector> objects; + + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("B"), nullptr); + ASSERT_OK(registry->ListManagedObjects(&objects)); + ASSERT_EQ(objects.size(), 0U); + ASSERT_OK(registry->SetManagedObject(m_a)); + ASSERT_EQ(registry->GetManagedObject("B"), nullptr); + ASSERT_EQ(registry->GetManagedObject("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("A"), m_a); + ASSERT_EQ(registry->GetManagedObject("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("B"), m_b); + ASSERT_EQ(registry->GetManagedObject("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("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("B"), nullptr); +} + +TEST_F(EnvRegistryTest, TestAlternateNames) { + auto registry = ObjectRegistry::NewInstance(); + auto m_a = std::make_shared("", "A"); + auto m_b = std::make_shared("", "B"); + std::vector> objects; + // Test no objects exist + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("B"), nullptr); + ASSERT_EQ(registry->GetManagedObject("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("B"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("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("B"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), m_a); + ASSERT_EQ(registry->GetManagedObject("TheOne"), m_a); + ASSERT_OK(registry->ListManagedObjects(&objects)); + ASSERT_EQ(objects.size(), 2U); + + // Delete "A". + m_a.reset(); + objects.clear(); + + ASSERT_EQ(registry->GetManagedObject("TheOne"), nullptr); + ASSERT_OK(registry->SetManagedObject("TheOne", m_b)); + ASSERT_EQ(registry->GetManagedObject("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("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("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("", "A"); + auto m_a2 = std::make_shared("", "A"); + std::vector> obj1s; + std::vector> obj2s; + + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + + ASSERT_OK(registry->SetManagedObject(m_a1)); + ASSERT_EQ(registry->GetManagedObject("A"), m_a1); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + + ASSERT_OK(registry->SetManagedObject(m_a2)); + ASSERT_EQ(registry->GetManagedObject("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("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), m_a2); + + m_a2.reset(); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); +} + +TEST_F(EnvRegistryTest, TestManagedObjectsWithParent) { + auto base = ObjectRegistry::NewInstance(); + auto registry = ObjectRegistry::NewInstance(base); + + auto m_a = std::make_shared("", "A"); + auto m_b = std::make_shared("", "A"); + + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_OK(base->SetManagedObject(m_a)); + ASSERT_EQ(registry->GetManagedObject("A"), m_a); + + ASSERT_NOK(registry->SetManagedObject(m_b)); + ASSERT_OK(registry->SetManagedObject(m_a)); + + m_a.reset(); + ASSERT_EQ(registry->GetManagedObject("A"), nullptr); + ASSERT_OK(registry->SetManagedObject(m_b)); + ASSERT_EQ(registry->GetManagedObject("A"), m_b); +} + +TEST_F(EnvRegistryTest, TestGetOrCreateManagedObject) { + auto registry = ObjectRegistry::NewInstance(); + registry->AddLibrary("test")->Register( + "MC(@.*)?", + [](const std::string& uri, std::unique_ptr* guard, + std::string* /* errmsg */) { + guard->reset(new MyCustomizable("MC", uri)); + return guard->get(); + }); + std::shared_ptr m_a, m_b, obj; + std::vector> objs; + + std::unordered_map opt_map; + + ASSERT_EQ(registry->GetManagedObject("MC@A"), nullptr); + ASSERT_EQ(registry->GetManagedObject("MC@B"), nullptr); + ASSERT_OK(registry->GetOrCreateManagedObject("MC@A", &m_a)); + ASSERT_OK(registry->GetOrCreateManagedObject("MC@B", &m_b)); + ASSERT_EQ(registry->GetManagedObject("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) {