Add persistence tests for CacheComponents

rlyerly · meta-codesync[bot] · commit 456cb9a24e58 · 2026-03-31T09:58:15.000-07:00
Summary: Add tests for persist/recover functionality for all CacheComponents.

Reviewed By: pbhandar2

Differential Revision: D97173718

fbshipit-source-id: f50da3b17702f40277c53851f70596a920445ec4
diff --git a/cachelib/interface/components/FlashCacheComponent.h b/cachelib/interface/components/FlashCacheComponent.h
@@ -25,6 +25,10 @@ namespace facebook::cachelib::interface {
 class FlashCacheItem;
 class ConsistentFlashCacheItem;
 
+namespace test {
+class FlashCacheFactory;
+}
+
 /**
  * A cache component that uses Cachelib's BlockCache flash cache without RAM
  * cache.
@@ -144,6 +148,7 @@ class FlashCacheComponent : public CacheComponentWithStats {
 
   friend class FlashCacheItem;
   friend class ConsistentFlashCacheItem;
+  friend class test::FlashCacheFactory;
 };
 
 /**
diff --git a/cachelib/interface/components/tests/CacheComponentFactory.cpp b/cachelib/interface/components/tests/CacheComponentFactory.cpp
@@ -23,13 +23,34 @@
 namespace facebook::cachelib::interface::test {
 
 std::unique_ptr<CacheComponent> RAMCacheFactory::create() {
-  auto config = createConfig();
-  auto poolConfig = createPoolConfig();
-  auto ramCache = ASSERT_OK(
-      RAMCacheComponent::create(std::move(config), std::move(poolConfig)));
+  auto ramCache = ASSERT_OK(createWithPersistence(
+      RAMCacheComponent::PersistenceConfig::noPersistenceOrRecovery()));
   return std::make_unique<RAMCacheComponent>(std::move(ramCache));
 }
 
+std::unique_ptr<CacheComponent> RAMCacheFactory::createPersistent() {
+  auto ramCache = ASSERT_OK(createWithPersistence(
+      RAMCacheComponent::PersistenceConfig::persistenceButNoRecovery(
+          tmpDir_.path().string())));
+  return std::make_unique<RAMCacheComponent>(std::move(ramCache));
+}
+
+Result<std::unique_ptr<CacheComponent>> RAMCacheFactory::recover() {
+  auto result = createWithPersistence(
+      RAMCacheComponent::PersistenceConfig::persistenceAndRecovery(
+          tmpDir_.path().string()));
+  if (result.hasError()) {
+    return folly::makeUnexpected(std::move(result).error());
+  }
+  return std::make_unique<RAMCacheComponent>(std::move(result).value());
+}
+
+Result<RAMCacheComponent> RAMCacheFactory::createWithPersistence(
+    RAMCacheComponent::PersistenceConfig pc) {
+  return RAMCacheComponent::create(createConfig(), createPoolConfig(),
+                                   std::move(pc));
+}
+
 /* static */ LruAllocatorConfig RAMCacheFactory::createConfig() {
   LruAllocatorConfig config;
   config.setCacheName("CacheComponentTest");
@@ -54,31 +75,72 @@ FlashCacheFactory::FlashCacheFactory()
     : hits_(/* count */ kDeviceSize / kRegionSize, /* value */ 0) {}
 
 std::unique_ptr<CacheComponent> FlashCacheFactory::create() {
-  auto flashCache = ASSERT_OK(FlashCacheComponent::create(
-      "CacheComponentTest", makeConfig(), makeDevice()));
-  return std::make_unique<FlashCacheComponent>(std::move(flashCache));
+  return ASSERT_OK(createWithPersistence(
+      makeDevice(kCacheSize), // don't need space for metadata
+      FlashCacheComponent::PersistenceConfig::noPersistenceOrRecovery()));
+}
+
+std::unique_ptr<CacheComponent> FlashCacheFactory::createPersistent() {
+  return ASSERT_OK(createWithPersistence(
+      makeDevice(kDeviceSize),
+      FlashCacheComponent::PersistenceConfig::persistenceButNoRecovery(
+          kMetadataSize)));
+}
+
+void FlashCacheFactory::onShutdown(CacheComponent& component) {
+  savedDevice_ =
+      std::move(static_cast<FlashCacheComponent&>(component).device_);
+}
+
+Result<std::unique_ptr<CacheComponent>> FlashCacheFactory::recover() {
+  auto device =
+      savedDevice_ ? std::move(savedDevice_) : makeDevice(kDeviceSize);
+  return createWithPersistence(
+      std::move(device),
+      FlashCacheComponent::PersistenceConfig::persistenceAndRecovery(
+          kMetadataSize));
+}
+
+Result<std::unique_ptr<CacheComponent>>
+FlashCacheFactory::createWithPersistence(
+    std::unique_ptr<navy::Device> device,
+    FlashCacheComponent::PersistenceConfig pc) {
+  auto result = FlashCacheComponent::create("CacheComponentTest", makeConfig(),
+                                            std::move(device), std::move(pc));
+  if (result.hasError()) {
+    return folly::makeUnexpected(std::move(result).error());
+  }
+  return std::make_unique<FlashCacheComponent>(std::move(result).value());
 }
 
-/* static */ std::unique_ptr<navy::Device> FlashCacheFactory::makeDevice() {
-  return navy::createMemoryDevice(kDeviceSize, /* encryptor */ nullptr);
+/* static */ std::unique_ptr<navy::Device> FlashCacheFactory::makeDevice(
+    size_t size) {
+  return navy::createMemoryDevice(size, /* encryptor */ nullptr, kIOAlignSize);
 }
 
 navy::BlockCache::Config FlashCacheFactory::makeConfig() {
   navy::BlockCache::Config config;
   config.regionSize = kRegionSize;
-  config.cacheSize = kDeviceSize;
+  config.cacheSize = kCacheSize;
+  config.indexConfig.setNumSparseMapBuckets(64);
   config.evictionPolicy =
       std::make_unique<::testing::NiceMock<navy::MockPolicy>>(&hits_);
   config.reinsertionConfig.enablePctBased(100);
   return config;
 }
 
-std::unique_ptr<CacheComponent> ConsistentFlashCacheFactory::create() {
-  auto consistentFlashCache = ASSERT_OK(ConsistentFlashCacheComponent::create(
-      "CacheComponentTest", makeConfig(), makeDevice(),
-      std::make_unique<MurmurHash2>(), kShardsPower));
+Result<std::unique_ptr<CacheComponent>>
+ConsistentFlashCacheFactory::createWithPersistence(
+    std::unique_ptr<navy::Device> device,
+    FlashCacheComponent::PersistenceConfig pc) {
+  auto result = ConsistentFlashCacheComponent::create(
+      "CacheComponentTest", makeConfig(), std::move(device),
+      std::make_unique<MurmurHash2>(), kShardsPower, std::move(pc));
+  if (result.hasError()) {
+    return folly::makeUnexpected(std::move(result).error());
+  }
   return std::make_unique<ConsistentFlashCacheComponent>(
-      std::move(consistentFlashCache));
+      std::move(result).value());
 }
 
 } // namespace facebook::cachelib::interface::test
diff --git a/cachelib/interface/components/tests/CacheComponentFactory.h b/cachelib/interface/components/tests/CacheComponentFactory.h
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <folly/testing/TestUtil.h>
+
 #include "cachelib/interface/CacheComponent.h"
 #include "cachelib/interface/components/FlashCacheComponent.h"
 #include "cachelib/interface/components/RAMCacheComponent.h"
@@ -26,13 +28,22 @@ class CacheFactory {
  public:
   virtual ~CacheFactory() = default;
   virtual std::unique_ptr<CacheComponent> create() = 0;
+  virtual std::unique_ptr<CacheComponent> createPersistent() = 0;
+  virtual void onShutdown(CacheComponent& /* component */) {}
+  virtual Result<std::unique_ptr<CacheComponent>> recover() = 0;
 };
 
 class RAMCacheFactory : public CacheFactory {
  public:
   std::unique_ptr<CacheComponent> create() override;
+  std::unique_ptr<CacheComponent> createPersistent() override;
+  Result<std::unique_ptr<CacheComponent>> recover() override;
 
  private:
+  folly::test::TemporaryDirectory tmpDir_;
+
+  Result<RAMCacheComponent> createWithPersistence(
+      RAMCacheComponent::PersistenceConfig pc);
   static LruAllocatorConfig createConfig();
   static RAMCacheComponent::PoolConfig createPoolConfig();
 };
@@ -41,20 +52,33 @@ class FlashCacheFactory : public CacheFactory {
  public:
   FlashCacheFactory();
   std::unique_ptr<CacheComponent> create() override;
+  std::unique_ptr<CacheComponent> createPersistent() override;
+  void onShutdown(CacheComponent& /* component */) override;
+  Result<std::unique_ptr<CacheComponent>> recover() override;
 
  protected:
-  static std::unique_ptr<navy::Device> makeDevice();
+  static constexpr size_t kMetadataSize{/* 32KB */ 32 * 1024};
+  static constexpr size_t kCacheSize{/* 256KB */ 256 * 1024};
+  static constexpr size_t kDeviceSize{kMetadataSize + kCacheSize};
+  std::unique_ptr<navy::Device> savedDevice_;
+
+  virtual Result<std::unique_ptr<CacheComponent>> createWithPersistence(
+      std::unique_ptr<navy::Device> device,
+      FlashCacheComponent::PersistenceConfig pc);
+  static std::unique_ptr<navy::Device> makeDevice(size_t size);
   navy::BlockCache::Config makeConfig();
 
  private:
   static constexpr size_t kRegionSize{/* 16KB */ 16 * 1024};
-  static constexpr size_t kDeviceSize{/* 256KB */ 256 * 1024};
+  static constexpr uint32_t kIOAlignSize{4096};
   std::vector<uint32_t> hits_;
 };
 
 class ConsistentFlashCacheFactory : public FlashCacheFactory {
- public:
-  std::unique_ptr<CacheComponent> create() override;
+ protected:
+  Result<std::unique_ptr<CacheComponent>> createWithPersistence(
+      std::unique_ptr<navy::Device> device,
+      FlashCacheComponent::PersistenceConfig pc) override;
 
  private:
   static constexpr size_t kShardsPower{4};
diff --git a/cachelib/interface/components/tests/CacheComponentTest.cpp b/cachelib/interface/components/tests/CacheComponentTest.cpp
@@ -17,6 +17,7 @@
 #include <folly/Random.h>
 #include <folly/coro/BlockingWait.h>
 #include <folly/coro/GtestHelpers.h>
+#include <folly/testing/TestUtil.h>
 #include <gtest/gtest.h>
 
 #include <atomic>
@@ -43,9 +44,9 @@ class CacheComponentTest : public ::testing::Test {
     ASSERT_NE(cache_, nullptr) << "Failed to create cache";
   }
 
-  std::unique_ptr<CacheComponent> cache_;
+  void TearDown() override { EXPECT_OK(cache_->shutdown()); }
 
- private:
+  std::unique_ptr<CacheComponent> cache_;
   std::unique_ptr<FactoryType> factory_;
 };
 
@@ -661,4 +662,88 @@ TYPED_TEST(CacheComponentTest, MultiThreadedOperations) {
   EXPECT_EQ(dataCorruptions.load(), 0);
 }
 
+// ============================================================================
+// Persist and Recover Tests
+// ============================================================================
+
+CO_TYPED_TEST(CacheComponentTest, PersistAndRecover) {
+  const std::string key = "persist_key";
+  const std::string data = "persist_data";
+
+  {
+    auto cache = this->factory_->createPersistent();
+
+    auto handle = ASSERT_OK(co_await cache->allocate(
+        key, data.size(), util::getCurrentTimeSec(), 3600));
+    std::memcpy(handle->getMemory(), data.c_str(), data.size());
+    EXPECT_OK(co_await cache->insert(std::move(handle)));
+
+    EXPECT_OK(cache->shutdown());
+    this->factory_->onShutdown(*cache);
+  }
+
+  auto result = this->factory_->recover();
+  CO_ASSERT_TRUE(result.hasValue());
+  auto recovered = std::move(result).value();
+
+  auto findResult = ASSERT_OK(co_await recovered->find(key));
+  CO_ASSERT_TRUE(findResult.has_value());
+  std::string retrievedData(
+      findResult.value()->template getMemoryAs<const char>(), data.size());
+  EXPECT_EQ(retrievedData, data);
+}
+
+CO_TYPED_TEST(CacheComponentTest, PersistAndRecoverMultipleItems) {
+  constexpr int kNumItems = 10;
+
+  {
+    auto cache = this->factory_->createPersistent();
+
+    for (int i = 0; i < kNumItems; ++i) {
+      auto key = "key_" + std::to_string(i);
+      auto val = "data_" + std::to_string(i);
+      auto handle = ASSERT_OK(co_await cache->allocate(
+          key, val.size(), util::getCurrentTimeSec(), 3600));
+      std::memcpy(handle->getMemory(), val.c_str(), val.size());
+      EXPECT_OK(co_await cache->insert(std::move(handle)));
+    }
+
+    EXPECT_OK(cache->shutdown());
+    this->factory_->onShutdown(*cache);
+  }
+
+  auto result = this->factory_->recover();
+  CO_ASSERT_TRUE(result.hasValue());
+  auto recovered = std::move(result).value();
+
+  for (int i = 0; i < kNumItems; ++i) {
+    auto key = "key_" + std::to_string(i);
+    auto expectedVal = "data_" + std::to_string(i);
+    auto findResult = ASSERT_OK(co_await recovered->find(key));
+    CO_ASSERT_TRUE(findResult.has_value());
+    std::string retrievedData(
+        findResult.value()->template getMemoryAs<const char>(),
+        expectedVal.size());
+    EXPECT_EQ(retrievedData, expectedVal);
+  }
+  EXPECT_OK(recovered->shutdown());
+}
+
+CO_TYPED_TEST(CacheComponentTest, RecoverWithoutPriorPersist) {
+  // Cache components recover gracefully with an empty cache
+  auto recovered = ASSERT_OK(this->factory_->recover());
+  auto findResult = ASSERT_OK(co_await recovered->find("nonexistent_key"));
+  EXPECT_FALSE(findResult.has_value());
+}
+
+CO_TYPED_TEST(CacheComponentTest, PersistShutdownWithPersistence) {
+  auto cache = this->factory_->createPersistent();
+
+  auto handle = ASSERT_OK(
+      co_await cache->allocate("key", 100, util::getCurrentTimeSec(), 3600));
+  EXPECT_OK(co_await cache->insert(std::move(handle)));
+
+  EXPECT_OK(cache->shutdown());
+}
+
 } // namespace
