Add shared memory persistence to AccessTimeMap

Summary: Implement shared memory persistence for AccessTimeMap (ATM).

Reviewed By: rlyerly

Differential Revision: D95409622

fbshipit-source-id: 51faec02e08d6ed86675324ca6079e62a0de98aa

Author: pbhandar2

cachelib/allocator/nvmcache/AccessTimeMap.h

@@ -16,17 +16,24 @@
 
 #pragma once
 
+#include <folly/ScopeGuard.h>
 #include <folly/container/F14Map.h>
 #include <folly/lang/Align.h>
+#include <folly/logging/xlog.h>
 
 #include <atomic>
 #include <mutex>
 #include <optional>
 #include <stdexcept>
+#include <string_view>
 #include <vector>
 
 #include "cachelib/common/AtomicCounter.h"
+#include "cachelib/common/Serialization.h"
+#include "cachelib/common/Time.h"
 #include "cachelib/common/Utils.h"
+#include "cachelib/navy/serialization/gen-cpp2/objects_types.h"
+#include "cachelib/shm/ShmManager.h"
 
 namespace facebook::cachelib {
 
@@ -42,10 +49,22 @@ namespace facebook::cachelib {
 // mutex overhead on the read path.
 class AccessTimeMap {
  public:
+  static constexpr std::string_view kShmInfoName = "shm_atm_info";
+  static constexpr std::string_view kShmDataName = "shm_atm_data";
+  static constexpr uint32_t kVersion = 1;
+
+  // Flat entry layout for serialization into shared memory.
+  struct ShmEntry {
+    uint64_t keyHash;
+    uint32_t accessTimeSecs;
+  };
+
   // maxSize: approximate upper bound on total entries across all shards.
   // Enforced per-shard as ceil(maxSize / numShards). 0 means unbounded.
   explicit AccessTimeMap(size_t numShards, size_t maxSize = 0)
-      : shards_(numShards),
+      : numShards_(numShards),
+        maxSize_(maxSize),
+        shards_(numShards),
         maxEntriesPerShard_(computeMaxEntriesPerShard(numShards, maxSize)) {}
 
   // Store a timestamp for the given key hash.
@@ -120,8 +139,162 @@ class AccessTimeMap {
   // memory ordering, but avoids the contention of locking all shards.
   size_t size() const { return size_.load(std::memory_order_relaxed); }
 
-  // Counters are read without a global barrier, so a snapshot may be
-  // internally inconsistent (e.g. hits + misses > gets).
+  void persist(ShmManager& shm) {
+    auto startTimeSecs = util::getCurrentTimeSec();
+
+    size_t entryCount = 0;
+    for (auto& shard : shards_) {
+      std::lock_guard l(shard.mutex_);
+      entryCount += shard.map_.size();
+    }
+
+    // Remove any stale segments to avoid recovering old state on warm roll.
+    shm.removeShm(std::string(kShmDataName));
+    shm.removeShm(std::string(kShmInfoName));
+
+    if (entryCount == 0) {
+      XLOG(INFO) << "AccessTimeMap is empty, nothing to persist";
+      return;
+    }
+
+    const size_t dataSize = entryCount * sizeof(ShmEntry);
+    auto dataAddr = shm.createShm(std::string(kShmDataName), dataSize);
+    auto* entries = reinterpret_cast<ShmEntry*>(dataAddr.addr);
+
+    size_t idx = 0;
+    for (auto& shard : shards_) {
+      std::lock_guard l(shard.mutex_);
+      for (const auto& [keyHash, accessTimeSecs] : shard.map_) {
+        if (idx >= entryCount) {
+          break;
+        }
+        entries[idx++] = ShmEntry{keyHash, accessTimeSecs};
+      }
+      if (idx >= entryCount) {
+        break;
+      }
+    }
+
+    navy::serialization::AccessTimeMapConfig cfg;
+    *cfg.version() = static_cast<int32_t>(kVersion);
+    *cfg.numShards() = static_cast<int64_t>(numShards_);
+    *cfg.maxSize() = static_cast<int64_t>(maxSize_);
+    *cfg.entryCount() = static_cast<int64_t>(idx);
+
+    auto ioBuf = Serializer::serializeToIOBuf(cfg);
+    auto infoAddr = shm.createShm(std::string(kShmInfoName), ioBuf->length());
+    Serializer serializer(
+        reinterpret_cast<uint8_t*>(infoAddr.addr),
+        reinterpret_cast<uint8_t*>(infoAddr.addr) + ioBuf->length());
+    serializer.writeToBuffer(std::move(ioBuf));
+
+    XLOGF(INFO,
+          "Persisted AccessTimeMap: {} entries, {} bytes to shared memory "
+          "in {} secs",
+          idx, dataSize, util::getCurrentTimeSec() - startTimeSecs);
+  }
+
+  // Restore entries from shared memory. On failure, starts empty.
+  void recover(ShmManager& shm) {
+    auto startTimeSecs = util::getCurrentTimeSec();
+
+    ShmAddr infoAddr;
+    try {
+      infoAddr = shm.attachShm(std::string(kShmInfoName));
+    } catch (const std::exception&) {
+      XLOG(INFO) << "No AccessTimeMap shm segments found, starting empty";
+      return;
+    }
+
+    SCOPE_EXIT {
+      shm.removeShm(std::string(kShmInfoName));
+      shm.removeShm(std::string(kShmDataName));
+    };
+
+    Deserializer deserializer(
+        reinterpret_cast<uint8_t*>(infoAddr.addr),
+        reinterpret_cast<uint8_t*>(infoAddr.addr) + infoAddr.size);
+    auto cfg =
+        deserializer.deserialize<navy::serialization::AccessTimeMapConfig>();
+
+    if (static_cast<uint32_t>(*cfg.version()) != kVersion) {
+      XLOGF(WARN,
+            "AccessTimeMap version mismatch: stored={}, current={}. "
+            "Starting empty.",
+            *cfg.version(), kVersion);
+      return;
+    }
+
+    const auto entryCount = static_cast<size_t>(*cfg.entryCount());
+    if (entryCount == 0) {
+      // persist() skips empty maps, so this may indicate corruption.
+      XLOG(WARN) << "AccessTimeMap shm segment has zero entries but info "
+                    "segment exists, possible corruption. Starting empty.";
+      return;
+    }
+
+    ShmAddr dataAddr;
+    try {
+      dataAddr = shm.attachShm(std::string(kShmDataName));
+    } catch (const std::exception& e) {
+      XLOGF(WARN,
+            "Failed to attach AccessTimeMap data segment: {}. "
+            "Starting empty.",
+            e.what());
+      return;
+    }
+
+    const size_t expectedSize = entryCount * sizeof(ShmEntry);
+    if (dataAddr.size < expectedSize) {
+      XLOGF(WARN,
+            "AccessTimeMap data segment size mismatch: got={}, expected={}. "
+            "Starting empty.",
+            dataAddr.size, expectedSize);
+      return;
+    }
+
+    const auto* entries = reinterpret_cast<const ShmEntry*>(dataAddr.addr);
+
+    // Group entries by shard so we only grab each lock once.
+    std::vector<folly::F14FastMap<uint64_t, uint32_t>> perShard(shards_.size());
+    for (size_t i = 0; i < entryCount; ++i) {
+      auto shardIdx = entries[i].keyHash % shards_.size();
+      perShard[shardIdx][entries[i].keyHash] = entries[i].accessTimeSecs;
+    }
+    size_t totalInserted = 0;
+    for (size_t s = 0; s < shards_.size(); ++s) {
+      if (perShard[s].empty()) {
+        continue;
+      }
+      auto& shard = shards_[s];
+      std::lock_guard l(shard.mutex_);
+      size_t inserted = 0;
+      for (const auto& [keyHash, ts] : perShard[s]) {
+        if (maxEntriesPerShard_ > 0 &&
+            shard.map_.size() >= maxEntriesPerShard_) {
+          break;
+        }
+        shard.map_.insert_or_assign(keyHash, ts);
+        ++inserted;
+      }
+      size_.fetch_add(inserted, std::memory_order_relaxed);
+      totalInserted += inserted;
+    }
+
+    if (totalInserted < entryCount) {
+      XLOGF(WARN,
+            "AccessTimeMap recovery dropped {} entries due to per-shard "
+            "size cap (maxEntriesPerShard={}, persisted={})",
+            entryCount - totalInserted, maxEntriesPerShard_, entryCount);
+    }
+
+    XLOGF(INFO,
+          "Recovered AccessTimeMap: {}/{} entries, {} bytes from shared memory "
+          "in {} secs",
+          totalInserted, entryCount, expectedSize,
+          util::getCurrentTimeSec() - startTimeSecs);
+  }
+
   void getCounters(const util::CounterVisitor& visitor) const {
     visitor("navy_atm_size", size_.load(std::memory_order_relaxed),
             util::CounterVisitor::CounterType::COUNT);
@@ -164,6 +337,8 @@ class AccessTimeMap {
     return shards_[keyHash % shards_.size()];
   }
 
+  const size_t numShards_;
+  const size_t maxSize_;
   std::vector<Shard> shards_;
   const size_t maxEntriesPerShard_{0};
   std::atomic<size_t> size_{0};