Make Cache a customizable class (#13024)

Summary: This PR allows a Cache object to be created using the object registry. Pull Request resolved: #13024 Reviewed By: pdillinger Differential Revision: D63043233 Pulled By: anand1976 fbshipit-source-id: 5bc3f7c29b35ad62638ff8205451303e2cecea9d
facebook · Sep 20, 2024 · 6549b11 · 6549b11
1 parent 71e38db
commit 6549b11
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 13 deletions.
diff --git a/cache/cache.cc b/cache/cache.cc
@@ -133,19 +133,23 @@ Status Cache::CreateFromString(const ConfigOptions& config_options,
                                std::shared_ptr<Cache>* result) {
   Status status;
   std::shared_ptr<Cache> cache;
-  if (value.find('=') == std::string::npos) {
-    cache = NewLRUCache(ParseSizeT(value));
-  } else {
-    LRUCacheOptions cache_opts;
-    status = OptionTypeInfo::ParseStruct(config_options, "",
-                                         &lru_cache_options_type_info, "",
-                                         value, &cache_opts);
+  if (value.find("://") == std::string::npos) {
+    if (value.find('=') == std::string::npos) {
+      cache = NewLRUCache(ParseSizeT(value));
+    } else {
+      LRUCacheOptions cache_opts;
+      status = OptionTypeInfo::ParseStruct(config_options, "",
+                                           &lru_cache_options_type_info, "",
+                                           value, &cache_opts);
+      if (status.ok()) {
+        cache = NewLRUCache(cache_opts);
+      }
+    }
     if (status.ok()) {
-      cache = NewLRUCache(cache_opts);
+      result->swap(cache);
     }
-  }
-  if (status.ok()) {
-    result->swap(cache);
+  } else {
+    status = LoadSharedObject<Cache>(config_options, value, result);
   }
   return status;
 }

diff --git a/include/rocksdb/advanced_cache.h b/include/rocksdb/advanced_cache.h
@@ -41,7 +41,7 @@ class Statistics;
 //
 // INTERNAL: See typed_cache.h for convenient wrappers on top of this API.
 // New virtual functions must also be added to CacheWrapper below.
-class Cache {
+class Cache : public Customizable {
  public:  // types hidden from API client
   // Opaque handle to an entry stored in the cache.
   struct Handle {};
@@ -190,6 +190,8 @@ class Cache {
   // Destroys all remaining entries by calling the associated "deleter"
   virtual ~Cache() {}
 
+  static const char* Type() { return "Cache"; }
+
   // Creates a new Cache based on the input value string and returns the result.
   // Currently, this method can be used to create LRUCaches only
   // @param config_options

diff --git a/options/customizable_test.cc b/options/customizable_test.cc
@@ -1396,6 +1396,19 @@ class MockFilterPolicy : public FilterPolicy {
   }
 };
 
+class MockCache : public CacheWrapper {
+ public:
+  static const char* kClassName() { return "MockCache"; }
+  const char* Name() const override { return kClassName(); }
+
+  MockCache()
+      : CacheWrapper(NewLRUCache(LRUCacheOptions(100, 0, false, 0.0))) {}
+
+  bool IsInstanceOf(const std::string& name) const override {
+    return name.find(Name()) == 0;
+  }
+};
+
 static int RegisterLocalObjects(ObjectLibrary& library,
                                 const std::string& /*arg*/) {
   size_t num_types;
@@ -1519,6 +1532,15 @@ static int RegisterLocalObjects(ObjectLibrary& library,
         return guard->get();
       });
 
+  library.AddFactory<Cache>(
+      ObjectLibrary::PatternEntry(MockCache::kClassName())
+          .AddSeparator("://", /*at_least_one=*/false),
+      [](const std::string& /*uri*/, std::unique_ptr<Cache>* guard,
+         std::string* /* errmsg */) {
+        guard->reset(new MockCache());
+        return guard->get();
+      });
+
   return static_cast<int>(library.GetFactoryCount(&num_types));
 }
 }  // namespace
@@ -2111,6 +2133,15 @@ TEST_F(LoadCustomizableTest, LoadFlushBlockPolicyFactoryTest) {
   }
 }
 
+TEST_F(LoadCustomizableTest, LoadCacheTest) {
+  if (RegisterTests("Test")) {
+    std::string uri(MockCache::kClassName());
+    uri.append("://");
+    auto cache = ExpectCreateShared<Cache>(uri);
+    ASSERT_TRUE(cache->IsInstanceOf(MockCache::kClassName()));
+  }
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);

diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc
@@ -1753,6 +1753,7 @@ DEFINE_bool(read_with_latest_user_timestamp, true,
             "If true, always use the current latest timestamp for read. If "
             "false, choose a random timestamp from the past.");
 
+DEFINE_string(cache_uri, "", "Full URI for creating a custom cache object");
 DEFINE_string(secondary_cache_uri, "",
               "Full URI for creating a custom secondary cache object");
 static class std::shared_ptr<ROCKSDB_NAMESPACE::SecondaryCache> secondary_cache;
@@ -3138,7 +3139,15 @@ class Benchmark {
     }
 
     std::shared_ptr<Cache> block_cache;
-    if (FLAGS_cache_type == "clock_cache") {
+    if (!FLAGS_cache_uri.empty()) {
+      Status s = Cache::CreateFromString(ConfigOptions(), FLAGS_cache_uri,
+                                         &block_cache);
+      if (block_cache == nullptr) {
+        fprintf(stderr, "No  cache registered matching string: %s status=%s\n",
+                FLAGS_cache_uri.c_str(), s.ToString().c_str());
+        exit(1);
+      }
+    } else if (FLAGS_cache_type == "clock_cache") {
       fprintf(stderr, "Old clock cache implementation has been removed.\n");
       exit(1);
     } else if (EndsWith(FLAGS_cache_type, "hyper_clock_cache")) {

diff --git a/unreleased_history/new_features/customizable_cache.md b/unreleased_history/new_features/customizable_cache.md
@@ -0,0 +1 @@
+Make Cache a customizable class that can be instantiated by the object registry.