Add edge-case tests for AccessTimeMap shm persistence

Summary:
Added the following unit tests:

- `RecoverMissingDataSegment`: info segment exists but data segment is gone — verifies the data `attachShm` exception path cleans up the orphaned info segment and leaves the map empty.
- `RecoverDataSegmentTooSmall`: info segment claims more entries than the data segment can hold — verifies the `dataAddr.size < expectedSize` guard.
- `RecoverCleansUpSegments`: verifies that after a successful recovery, both shm segments are removed so no stale segments leak across subsequent warm restarts.
- `PersistOverwritesPriorSegments`: verifies that calling `persist()` twice replaces (not appends to) the prior segments.
- `AccessTimeMapSurvivesWarmRoll`: verifies that NVM cache is able to persist and recover ATM on warm roll.

Reviewed By: rlyerly

Differential Revision: D95332178

fbshipit-source-id: c6121c3d7e8f0b1b58c7e3ea4974bdadbd5a1055

Author: pbhandar2

cachelib/allocator/nvmcache/tests/AccessTimeMapTest.cpp

@@ -422,6 +422,202 @@ TEST(AccessTimeMapTest, RecoverNoShm) {
   ShmManager::cleanup(shmDir, true);
 }
 
+TEST(AccessTimeMapTest, RecoverMissingDataSegment) {
+  const auto shmDir =
+      "/tmp/atm_missing_data_test_" + std::to_string(::getpid());
+  constexpr size_t kNumShards = 4;
+
+  // Populate and persist, then remove the data segment
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+
+    m.set(10, 100);
+    m.set(20, 200);
+    m.persist(shm);
+
+    // Remove only the data segment, leaving the info segment intact
+    shm.removeShm(std::string(AccessTimeMap::kShmDataName));
+    shm.shutDown();
+  }
+
+  // Recovery should detect missing data segment and start empty
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+    m.recover(shm);
+    EXPECT_EQ(m.size(), 0);
+    shm.shutDown();
+  }
+
+  ShmManager::cleanup(shmDir, true);
+}
+
+TEST(AccessTimeMapTest, RecoverDataSegmentTooSmall) {
+  const auto shmDir = "/tmp/atm_small_data_test_" + std::to_string(::getpid());
+  constexpr size_t kNumShards = 4;
+
+  // Populate and persist, then overwrite the info segment to claim more
+  // entries than the data segment can hold. The data segment created by
+  // persist() for 3 entries is page-aligned to 4096 bytes. Setting
+  // entryCount to 500 makes expectedSize = 500 * sizeof(ShmEntry) = 6000,
+  // which exceeds the 4096-byte data segment.
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+
+    m.set(10, 100);
+    m.set(20, 200);
+    m.set(30, 300);
+    m.persist(shm);
+
+    // Overwrite the info segment with an inflated entryCount
+    shm.removeShm(std::string(AccessTimeMap::kShmInfoName));
+    navy::serialization::AccessTimeMapConfig cfg;
+    *cfg.version() = static_cast<int32_t>(AccessTimeMap::kVersion);
+    *cfg.numShards() = static_cast<int64_t>(kNumShards);
+    *cfg.maxSize() = 0;
+    *cfg.entryCount() = 500; // data segment is only 4096 bytes
+
+    auto ioBuf = Serializer::serializeToIOBuf(cfg);
+    auto infoAddr = shm.createShm(std::string(AccessTimeMap::kShmInfoName),
+                                  ioBuf->length());
+    Serializer serializer(
+        reinterpret_cast<uint8_t*>(infoAddr.addr),
+        reinterpret_cast<uint8_t*>(infoAddr.addr) + ioBuf->length());
+    serializer.writeToBuffer(std::move(ioBuf));
+
+    shm.shutDown();
+  }
+
+  // Recovery should detect size mismatch and start empty
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+    m.recover(shm);
+    EXPECT_EQ(m.size(), 0);
+    shm.shutDown();
+  }
+
+  ShmManager::cleanup(shmDir, true);
+}
+
+TEST(AccessTimeMapTest, RecoverCleansUpSegments) {
+  const auto shmDir = "/tmp/atm_cleanup_test_" + std::to_string(::getpid());
+  constexpr size_t kNumShards = 4;
+  constexpr size_t kMaxSize = 1000;
+
+  // Populate and persist
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards, kMaxSize);
+
+    m.set(10, 100);
+    m.set(20, 200);
+    m.persist(shm);
+    shm.shutDown();
+  }
+
+  // Recover and verify entries loaded, then check segments are gone
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards, kMaxSize);
+    m.recover(shm);
+    ASSERT_EQ(m.size(), 2);
+    EXPECT_EQ(m.get(10), 100);
+    EXPECT_EQ(m.get(20), 200);
+
+    // Segments should have been removed after successful recovery
+    EXPECT_THROW(shm.attachShm(std::string(AccessTimeMap::kShmInfoName)),
+                 std::invalid_argument);
+    EXPECT_THROW(shm.attachShm(std::string(AccessTimeMap::kShmDataName)),
+                 std::invalid_argument);
+
+    shm.shutDown();
+  }
+
+  ShmManager::cleanup(shmDir, true);
+}
+
+TEST(AccessTimeMapTest, PersistOverwritesPriorSegments) {
+  const auto shmDir = "/tmp/atm_overwrite_test_" + std::to_string(::getpid());
+  constexpr size_t kNumShards = 4;
+
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+
+    // First persist with two entries
+    m.set(1, 100);
+    m.set(2, 200);
+    m.persist(shm);
+
+    // Clear the map and add a different entry
+    m.remove(1);
+    m.remove(2);
+    m.set(3, 300);
+    ASSERT_EQ(m.size(), 1);
+
+    // Second persist should replace the prior segments
+    m.persist(shm);
+    shm.shutDown();
+  }
+
+  // Recover should only see the second persist's data
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+    m.recover(shm);
+    ASSERT_EQ(m.size(), 1);
+    EXPECT_EQ(m.get(1), std::nullopt);
+    EXPECT_EQ(m.get(2), std::nullopt);
+    EXPECT_EQ(m.get(3), 300);
+    shm.shutDown();
+  }
+
+  ShmManager::cleanup(shmDir, true);
+}
+
+TEST(AccessTimeMapTest, RecoverZeroEntryCount) {
+  const auto shmDir = "/tmp/atm_zero_count_test_" + std::to_string(::getpid());
+  constexpr size_t kNumShards = 4;
+
+  // Create an info segment with valid version/numShards but entryCount = 0.
+  // This exercises the entryCount == 0 branch in recover(), which is NOT
+  // reached by PersistEmptyMap (persist() returns early for empty maps
+  // without creating segments, so recovery hits the "no shm" path instead).
+  {
+    ShmManager shm(shmDir, true);
+
+    navy::serialization::AccessTimeMapConfig cfg;
+    *cfg.version() = static_cast<int32_t>(AccessTimeMap::kVersion);
+    *cfg.numShards() = static_cast<int64_t>(kNumShards);
+    *cfg.maxSize() = 0;
+    *cfg.entryCount() = 0;
+
+    auto ioBuf = Serializer::serializeToIOBuf(cfg);
+    auto infoAddr = shm.createShm(std::string(AccessTimeMap::kShmInfoName),
+                                  ioBuf->length());
+    Serializer serializer(
+        reinterpret_cast<uint8_t*>(infoAddr.addr),
+        reinterpret_cast<uint8_t*>(infoAddr.addr) + ioBuf->length());
+    serializer.writeToBuffer(std::move(ioBuf));
+
+    shm.shutDown();
+  }
+
+  // Recovery should hit the entryCount == 0 branch and start empty
+  {
+    ShmManager shm(shmDir, true);
+    AccessTimeMap m(kNumShards);
+    m.recover(shm);
+    EXPECT_EQ(m.size(), 0);
+    shm.shutDown();
+  }
+
+  ShmManager::cleanup(shmDir, true);
+}
+
 } // namespace tests
 } // namespace cachelib
 } // namespace facebook

cachelib/allocator/nvmcache/tests/NvmCacheTests.cpp

@@ -3193,6 +3193,76 @@ TEST_F(NvmCacheTest, AccessTimeMapCleanupTest) {
   expectAtmCleared(groups[3], "NVM-evicted item");
 }
 
+TEST_F(NvmCacheTest, AccessTimeMapSurvivesWarmRoll) {
+  // Verify that ATM entries persist across a warm restart cycle:
+  //   shutDown() persists ATM → SharedMemAttach recovers ATM
+  this->convertToShmCache();
+  auto& nvm = this->cache();
+  auto pid = this->poolId();
+  const uint32_t allocSize = 15 * 1024;
+  const uint32_t numKeysPerRegion =
+      config_.blockCache().getRegionSize() / allocSize;
+
+  // Insert a target item and push it to NVM.
+  std::string targetKey = "atm_warm_roll_target";
+  {
+    auto it = nvm.allocate(pid, targetKey, allocSize);
+    ASSERT_NE(nullptr, it);
+    cache_->insertOrReplace(it);
+  }
+  ASSERT_TRUE(this->pushToNvmCacheFromRamForTesting(targetKey));
+  nvm.flushNvmCache();
+  this->removeFromRamForTesting(targetKey);
+
+  // Promote target from NVM (making it NvmClean in DRAM).
+  {
+    auto hdl = this->fetch(targetKey, false /* ramOnly */);
+    ASSERT_NE(nullptr, hdl);
+    ASSERT_TRUE(hdl->isNvmClean());
+  }
+
+  HashedKey targetHk{targetKey};
+  auto* atm = this->getAccessTimeMap();
+  ASSERT_NE(nullptr, atm);
+
+  // ATM should be empty after promotion (entry is created on eviction).
+  EXPECT_EQ(std::nullopt, atm->get(targetHk.keyHash()));
+
+  // Insert filler items to trigger DRAM eviction of the NvmClean target.
+  auto timeBefore = util::getCurrentTimeSec();
+  auto evictBefore = this->evictionCount();
+  for (int i = 0; i < 1024; i++) {
+    auto key = folly::sformat("atm_warm_roll_filler_{}", i);
+    auto it = nvm.allocate(pid, key, allocSize);
+    ASSERT_NE(nullptr, it);
+    cache_->insertOrReplace(it);
+    if (i % numKeysPerRegion == 0) {
+      nvm.flushNvmCache();
+    }
+  }
+  nvm.flushNvmCache();
+  ASSERT_LT(evictBefore, this->evictionCount());
+
+  // The target was NvmClean when evicted, so ATM should have its entry.
+  auto accessTime = atm->get(targetHk.keyHash());
+  ASSERT_NE(std::nullopt, accessTime);
+  EXPECT_GE(*accessTime, timeBefore);
+  auto savedTime = *accessTime;
+
+  // Warm restart: shutDown() persists ATM, SharedMemAttach recovers it.
+  this->warmRoll();
+
+  // Re-acquire the ATM pointer (NvmCache instance is recreated).
+  auto* atmAfter = this->getAccessTimeMap();
+  ASSERT_NE(nullptr, atmAfter);
+
+  // Verify the entry survived with the same timestamp.
+  auto recoveredTime = atmAfter->get(targetHk.keyHash());
+  ASSERT_NE(std::nullopt, recoveredTime)
+      << "ATM entry missing after warm restart";
+  EXPECT_EQ(savedTime, *recoveredTime);
+}
+
 } // namespace tests
 } // namespace cachelib
 } // namespace facebook