Make Comparator into a Customizable Object (facebook#8336)

Summary: Makes the Comparator class into a Customizable object. Added/Updated the CreateFromString method to create Comparators. Added test for using the ObjectRegistry to create one. Pull Request resolved: facebook#8336 Reviewed By: jay-zhuang Differential Revision: D28999612 Pulled By: mrambacher fbshipit-source-id: bff2cb2814eeb9fef6a00fddc61d6e34b6fbcf2e
jay-zhuang · Jun 11, 2021 · 6ad0810 · 6ad0810
1 parent 3897ce3
commit 6ad0810
Show file tree

Hide file tree

Showing 8 changed files with 152 additions and 44 deletions.
diff --git a/include/rocksdb/comparator.h b/include/rocksdb/comparator.h
@@ -10,6 +10,7 @@
 
 #include <string>
 
+#include "rocksdb/customizable.h"
 #include "rocksdb/rocksdb_namespace.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -20,7 +21,7 @@ class Slice;
 // used as keys in an sstable or a database.  A Comparator implementation
 // must be thread-safe since rocksdb may invoke its methods concurrently
 // from multiple threads.
-class Comparator {
+class Comparator : public Customizable {
  public:
   Comparator() : timestamp_size_(0) {}
 
@@ -37,7 +38,11 @@ class Comparator {
 
   virtual ~Comparator() {}
 
+  static Status CreateFromString(const ConfigOptions& opts,
+                                 const std::string& id,
+                                 const Comparator** comp);
   static const char* Type() { return "Comparator"; }
+
   // Three-way comparison.  Returns value:
   //   < 0 iff "a" < "b",
   //   == 0 iff "a" == "b",

diff --git a/include/rocksdb/configurable.h b/include/rocksdb/configurable.h
@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -269,7 +270,7 @@ class Configurable {
  protected:
   // True once the object is prepared.  Once the object is prepared, only
   // mutable options can be configured.
-  bool prepared_;
+  std::atomic<bool> prepared_;
 
   // Returns the raw pointer for the associated named option.
   // The name is typically the name of an option registered via the

diff --git a/include/rocksdb/utilities/options_type.h b/include/rocksdb/utilities/options_type.h
@@ -35,7 +35,6 @@ enum class OptionType {
   kCompactionPri,
   kSliceTransform,
   kCompressionType,
-  kComparator,
   kCompactionFilter,
   kCompactionFilterFactory,
   kCompactionStopStyle,

diff --git a/options/cf_options.cc b/options/cf_options.cc
@@ -535,22 +535,30 @@ static std::unordered_map<std::string, OptionTypeInfo>
              OptionVerificationType::kNormal, OptionTypeFlags::kNone,
              {0, OptionType::kCompressionType})},
         {"comparator",
-         {offset_of(&ImmutableCFOptions::user_comparator),
-          OptionType::kComparator, OptionVerificationType::kByName,
-          OptionTypeFlags::kCompareLoose,
-          // Parses the string and sets the corresponding comparator
-          [](const ConfigOptions& opts, const std::string& /*name*/,
-             const std::string& value, void* addr) {
-            auto old_comparator = static_cast<const Comparator**>(addr);
-            const Comparator* new_comparator = *old_comparator;
-            Status status =
-                opts.registry->NewStaticObject(value, &new_comparator);
-            if (status.ok()) {
-              *old_comparator = new_comparator;
-              return status;
-            }
-            return Status::OK();
-          }}},
+         OptionTypeInfo::AsCustomRawPtr<const Comparator>(
+             offset_of(&ImmutableCFOptions::user_comparator),
+             OptionVerificationType::kByName, OptionTypeFlags::kCompareLoose,
+             // Serializes a Comparator
+             [](const ConfigOptions& /*opts*/, const std::string&,
+                const void* addr, std::string* value) {
+               // it's a const pointer of const Comparator*
+               const auto* ptr = static_cast<const Comparator* const*>(addr);
+
+               // Since the user-specified comparator will be wrapped by
+               // InternalKeyComparator, we should persist the user-specified
+               // one instead of InternalKeyComparator.
+               if (*ptr == nullptr) {
+                 *value = kNullptrString;
+               } else {
+                 const Comparator* root_comp = (*ptr)->GetRootComparator();
+                 if (root_comp == nullptr) {
+                   root_comp = (*ptr);
+                 }
+                 *value = root_comp->Name();
+               }
+               return Status::OK();
+             },
+             /* Use the default match function*/ nullptr)},
         {"memtable_insert_with_hint_prefix_extractor",
          {offset_of(
               &ImmutableCFOptions::memtable_insert_with_hint_prefix_extractor),

diff --git a/options/customizable_test.cc b/options/customizable_test.cc
@@ -707,6 +707,15 @@ static int RegisterTestObjects(ObjectLibrary& library,
         guard->reset(new mock::MockTableFactory());
         return guard->get();
       });
+  library.Register<const Comparator>(
+      test::SimpleSuffixReverseComparator::kClassName(),
+      [](const std::string& /*uri*/,
+         std::unique_ptr<const Comparator>* /*guard*/,
+         std::string* /* errmsg */) {
+        static test::SimpleSuffixReverseComparator ssrc;
+        return &ssrc;
+      });
+
   return static_cast<int>(library.GetFactoryCount(&num_types));
 }
 
@@ -716,6 +725,7 @@ static int RegisterLocalObjects(ObjectLibrary& library,
   // Load any locally defined objects here
   return static_cast<int>(library.GetFactoryCount(&num_types));
 }
+#endif  // !ROCKSDB_LITE
 
 class LoadCustomizableTest : public testing::Test {
  public:
@@ -755,7 +765,31 @@ TEST_F(LoadCustomizableTest, LoadTableFactoryTest) {
     ASSERT_STREQ(factory->Name(), "MockTable");
   }
 }
-#endif  // !ROCKSDB_LITE
+
+TEST_F(LoadCustomizableTest, LoadComparatorTest) {
+  const Comparator* bytewise = BytewiseComparator();
+  const Comparator* reverse = ReverseBytewiseComparator();
+
+  const Comparator* result = nullptr;
+  ASSERT_NOK(Comparator::CreateFromString(
+      config_options_, test::SimpleSuffixReverseComparator::kClassName(),
+      &result));
+  ASSERT_OK(
+      Comparator::CreateFromString(config_options_, bytewise->Name(), &result));
+  ASSERT_EQ(result, bytewise);
+  ASSERT_OK(
+      Comparator::CreateFromString(config_options_, reverse->Name(), &result));
+  ASSERT_EQ(result, reverse);
+
+  if (RegisterTests("Test")) {
+    ASSERT_OK(Comparator::CreateFromString(
+        config_options_, test::SimpleSuffixReverseComparator::kClassName(),
+        &result));
+    ASSERT_NE(result, nullptr);
+    ASSERT_STREQ(result->Name(),
+                 test::SimpleSuffixReverseComparator::kClassName());
+  }
+}
 
 }  // namespace ROCKSDB_NAMESPACE
 int main(int argc, char** argv) {

diff --git a/options/options_helper.cc b/options/options_helper.cc
@@ -562,23 +562,6 @@ bool SerializeSingleOptionHelper(const void* opt_address,
                                           : kNullptrString;
       break;
     }
-    case OptionType::kComparator: {
-      // it's a const pointer of const Comparator*
-      const auto* ptr = static_cast<const Comparator* const*>(opt_address);
-      // Since the user-specified comparator will be wrapped by
-      // InternalKeyComparator, we should persist the user-specified one
-      // instead of InternalKeyComparator.
-      if (*ptr == nullptr) {
-        *value = kNullptrString;
-      } else {
-        const Comparator* root_comp = (*ptr)->GetRootComparator();
-        if (root_comp == nullptr) {
-          root_comp = (*ptr);
-        }
-        *value = root_comp->Name();
-      }
-      break;
-    }
     case OptionType::kCompactionFilter: {
       // it's a const pointer of const CompactionFilter*
       const auto* ptr =

diff --git a/test_util/testutil.h b/test_util/testutil.h
@@ -98,10 +98,8 @@ class PlainInternalKeyComparator : public InternalKeyComparator {
 class SimpleSuffixReverseComparator : public Comparator {
  public:
   SimpleSuffixReverseComparator() {}
-
-  virtual const char* Name() const override {
-    return "SimpleSuffixReverseComparator";
-  }
+  static const char* kClassName() { return "SimpleSuffixReverseComparator"; }
+  virtual const char* Name() const override { return kClassName(); }
 
   virtual int Compare(const Slice& a, const Slice& b) const override {
     Slice prefix_a = Slice(a.data(), 8);

diff --git a/util/comparator.cc b/util/comparator.cc
@@ -8,20 +8,26 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 
 #include "rocksdb/comparator.h"
+
 #include <stdint.h>
+
 #include <algorithm>
 #include <memory>
+#include <mutex>
+
+#include "options/configurable_helper.h"
 #include "port/port.h"
 #include "rocksdb/slice.h"
+#include "rocksdb/utilities/object_registry.h"
 
 namespace ROCKSDB_NAMESPACE {
 
 namespace {
 class BytewiseComparatorImpl : public Comparator {
  public:
   BytewiseComparatorImpl() { }
-
-  const char* Name() const override { return "leveldb.BytewiseComparator"; }
+  static const char* kClassName() { return "leveldb.BytewiseComparator"; }
+  const char* Name() const override { return kClassName(); }
 
   int Compare(const Slice& a, const Slice& b) const override {
     return a.compare(b);
@@ -139,9 +145,10 @@ class ReverseBytewiseComparatorImpl : public BytewiseComparatorImpl {
  public:
   ReverseBytewiseComparatorImpl() { }
 
-  const char* Name() const override {
+  static const char* kClassName() {
     return "rocksdb.ReverseBytewiseComparator";
   }
+  const char* Name() const override { return kClassName(); }
 
   int Compare(const Slice& a, const Slice& b) const override {
     return -a.compare(b);
@@ -220,4 +227,77 @@ const Comparator* ReverseBytewiseComparator() {
   return &rbytewise;
 }
 
+#ifndef ROCKSDB_LITE
+static int RegisterBuiltinComparators(ObjectLibrary& library,
+                                      const std::string& /*arg*/) {
+  library.Register<const Comparator>(
+      BytewiseComparatorImpl::kClassName(),
+      [](const std::string& /*uri*/,
+         std::unique_ptr<const Comparator>* /*guard */,
+         std::string* /* errmsg */) { return BytewiseComparator(); });
+  library.Register<const Comparator>(
+      ReverseBytewiseComparatorImpl::kClassName(),
+      [](const std::string& /*uri*/,
+         std::unique_ptr<const Comparator>* /*guard */,
+         std::string* /* errmsg */) { return ReverseBytewiseComparator(); });
+  return 2;
+}
+#endif  // ROCKSDB_LITE
+
+Status Comparator::CreateFromString(const ConfigOptions& config_options,
+                                    const std::string& value,
+                                    const Comparator** result) {
+#ifndef ROCKSDB_LITE
+  static std::once_flag once;
+  std::call_once(once, [&]() {
+    RegisterBuiltinComparators(*(ObjectLibrary::Default().get()), "");
+  });
+#endif  // ROCKSDB_LITE
+  std::string id;
+  std::unordered_map<std::string, std::string> opt_map;
+  Status status =
+      ConfigurableHelper::GetOptionsMap(value, *result, &id, &opt_map);
+  if (!status.ok()) {  // GetOptionsMap failed
+    return status;
+  }
+  std::string curr_opts;
+#ifndef ROCKSDB_LITE
+  if (*result != nullptr && (*result)->GetId() == id) {
+    // Try to get the existing options, ignoring any errors
+    ConfigOptions embedded = config_options;
+    embedded.delimiter = ";";
+    (*result)->GetOptionString(embedded, &curr_opts).PermitUncheckedError();
+  }
+#endif
+  if (id == BytewiseComparatorImpl::kClassName()) {
+    *result = BytewiseComparator();
+  } else if (id == ReverseBytewiseComparatorImpl::kClassName()) {
+    *result = ReverseBytewiseComparator();
+  } else if (value.empty()) {
+    // No Id and no options.  Clear the object
+    *result = nullptr;
+    return Status::OK();
+  } else if (id.empty()) {  // We have no Id but have options.  Not good
+    return Status::NotSupported("Cannot reset object ", id);
+  } else {
+#ifndef ROCKSDB_LITE
+    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 &&
+          status.IsNotSupported()) {
+        return Status::OK();
+      } else {
+        return status;
+      }
+    } else if (!curr_opts.empty() || !opt_map.empty()) {
+      Comparator* comparator = const_cast<Comparator*>(*result);
+      status = ConfigurableHelper::ConfigureNewObject(
+          config_options, comparator, id, curr_opts, opt_map);
+    }
+  }
+  return status;
+}
 }  // namespace ROCKSDB_NAMESPACE