concurrent hash map for ghost

williamnixon20 · williamnixon20 · commit 0566f6a1353a · 2025-12-02T06:03:21.000Z
diff --git a/cachelib/allocator/CMakeLists.txt b/cachelib/allocator/CMakeLists.txt
@@ -34,6 +34,7 @@ add_library (cachelib_allocator
     CacheAllocatorLruCache.cpp
     CacheAllocatorLruCacheWithSpinBuckets.cpp
     CacheAllocatorS3FIFOCache.cpp
+    CacheAllocatorS3FIFO5BCache.cpp
     CacheAllocatorTinyLFU5BCache.cpp
     CacheAllocatorTinyLFUCache.cpp
     CacheAllocatorWTinyLFU5BCache.cpp
diff --git a/cachelib/allocator/CacheAllocator.h b/cachelib/allocator/CacheAllocator.h
@@ -6106,6 +6106,7 @@ using Lru5B2QAllocator = CacheAllocator<Lru5B2QCacheTrait>;
 // If accessed while in probation, it will eventually be promoted to the main queue.
 // Items in the tail of main queue will be reinserted if accessed.
 using S3FIFOAllocator = CacheAllocator<S3FIFOCacheTrait>;
+using S3FIFO5BAllocator = CacheAllocator<S3FIFO5BCacheTrait>;
 
 // CacheAllocator with Tiny LFU eviction policy
 // It has a window initially to gauage the frequency of accesses of newly
diff --git a/cachelib/allocator/CacheAllocatorS3FIFO5BCache.cpp b/cachelib/allocator/CacheAllocatorS3FIFO5BCache.cpp
@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cachelib/allocator/CacheAllocator.h"
+
+namespace facebook::cachelib {
+template class CacheAllocator<S3FIFO5BCacheTrait>;
+}
diff --git a/cachelib/allocator/CacheTraits.h b/cachelib/allocator/CacheTraits.h
@@ -97,6 +97,13 @@ struct Lru5B2QCacheTrait {
   using CompressedPtrType = CompressedPtr5B;
 };
 
+struct S3FIFO5BCacheTrait {
+  using MMType = MMS3FIFO;
+  using AccessType = ChainedHashTable;
+  using AccessTypeLocks = SharedMutexBuckets;
+  using CompressedPtrType = CompressedPtr5B;
+};
+
 struct TinyLFU5BCacheTrait {
   using MMType = MMTinyLFU;
   using AccessType = ChainedHashTable;
diff --git a/cachelib/allocator/MMS3FIFO.h b/cachelib/allocator/MMS3FIFO.h
@@ -33,6 +33,7 @@
 #include "cachelib/allocator/memory/serialize/gen-cpp2/objects_types.h"
 #include "cachelib/common/CompilerUtils.h"
 #include "cachelib/common/FIFOHashSet.h"
+#include "cachelib/common/FIFOConcurrentHashSet.h"
 
 namespace facebook::cachelib {
 
@@ -112,8 +113,6 @@ class MMS3FIFO {
 
     // The size of ghost queue, as a percentage of total size
     size_t ghostSizePercent{90};
-
-    size_t reserveCapacity{0};
   };
 
   // The container object which can be used to keep track of objects of type
@@ -138,9 +137,6 @@ class MMS3FIFO {
     Container(Config c, PtrCompressor compressor)
         : lru_(LruType::NumTypes, std::move(compressor)),
           config_(std::move(c)) {
-            if (config_.reserveCapacity > 0) {
-              ghostQueue_.reserve(config_.reserveCapacity);
-            }
           }
     Container(serialization::MMS3FIFOObject object, PtrCompressor compressor);
 
@@ -216,7 +212,7 @@ class MMS3FIFO {
         ++it;
         // Skip accessed items
         skipAccessed(it);
-        return *this;
+        return *this; 
       }
 
       ListIterator& skipAccessed(ListIterator& it) noexcept {
@@ -389,10 +385,15 @@ class MMS3FIFO {
     void maybeResizeGhostLocked() noexcept;
 
     // Returns the hash of node's key
-    static size_t hashNode(const T& node) noexcept {
+    static size_t hashNode64(const T& node) noexcept {
       return folly::hasher<folly::StringPiece>()(node.getKey());
     }
 
+    static uint32_t hashNode(const T& node) noexcept {
+      return static_cast<uint32_t>(
+          folly::hasher<folly::StringPiece>()(node.getKey()));
+    }
+
     void removeLocked(T& node) noexcept;
 
     static bool isTiny(const T& node) noexcept {
@@ -432,7 +433,8 @@ class MMS3FIFO {
     // the lru
     LruList lru_;
 
-    facebook::cachelib::util::FIFOHashSet ghostQueue_;
+    facebook::cachelib::util::FIFOConcurrentHashSet32 ghostQueue_;
+    // facebook::cachelib::util::FIFOHashSet32 ghostQueue_;
 
     // Config for this lru.
     // Write access to the MMS3FIFO Config is serialized.
@@ -483,7 +485,7 @@ bool MMS3FIFO::Container<T, HookPtr>::recordAccess(T& node,
   }
   const auto currTime = static_cast<Time>(util::getCurrentTimeSec());
 
-  // Remove lock from record access wince we only set bits
+  // Remove lock from record access since we only set bits
   if (node.isInMMContainer()) {
     if (!isAccessed(node)) {
       markAccessed(node);
@@ -538,12 +540,12 @@ bool MMS3FIFO::Container<T, HookPtr>::add(T& node) noexcept {
   const auto currTime = static_cast<Time>(util::getCurrentTimeSec());
 
   const auto nodeHash = hashNode(node);
-  return lruMutex_->lock_combine([this, &node, currTime, nodeHash]() {
+  const auto ghostContains = ghostQueue_.contains(nodeHash);
+  return lruMutex_->lock_combine([this, &node, currTime, ghostContains]() {
     if (node.isInMMContainer()) {
       return false;
     }
 
-    const auto ghostContains = ghostQueue_.contains(nodeHash);
     if (ghostContains) {
       // Insert to main queue
       auto& mainLru = lru_.getList(LruType::Main);
@@ -625,6 +627,7 @@ template <typename T, MMS3FIFO::Hook<T> T::* HookPtr>
 typename MMS3FIFO::Container<T, HookPtr>::LockedIterator
 MMS3FIFO::Container<T, HookPtr>::getEvictionIterator() noexcept {
   LockHolder l(*lruMutex_);
+  // This is cheap
   maybeResizeGhostLocked();
   rebalanceForEviction();
   // Cache is full now so we know it's max size
@@ -643,13 +646,12 @@ void MMS3FIFO::Container<T, HookPtr>::withContainerLock(F&& fun) {
   lruMutex_->lock_combine([&fun]() { fun(); });
 }
 
+// Ghost queue insertion done on callee, outside lock
 template <typename T, MMS3FIFO::Hook<T> T::* HookPtr>
 void MMS3FIFO::Container<T, HookPtr>::removeLocked(T& node) noexcept {
   if (isTiny(node)) {
     lru_.getList(LruType::Tiny).remove(node); 
     unmarkTiny(node);
-    // Insert into ghost queue upon eviction from tiny queue
-    ghostQueue_.insert(hashNode(node));
   } else {
     lru_.getList(LruType::Main).remove(node);
   }
@@ -661,21 +663,51 @@ void MMS3FIFO::Container<T, HookPtr>::removeLocked(T& node) noexcept {
 
 template <typename T, MMS3FIFO::Hook<T> T::* HookPtr>
 bool MMS3FIFO::Container<T, HookPtr>::remove(T& node) noexcept {
-  return lruMutex_->lock_combine([this, &node]() {
+  bool isTiny_ = isTiny(node);
+  auto result = lruMutex_->lock_combine([this, &node]() {
     if (!node.isInMMContainer()) {
       return false;
     }
     removeLocked(node);
     return true;
   });
+  if (result && isTiny_) {
+    // Insert to ghost queue
+    ghostQueue_.insert(hashNode(node));
+  }
+  return result;
 }
 
 template <typename T, MMS3FIFO::Hook<T> T::* HookPtr>
 void MMS3FIFO::Container<T, HookPtr>::remove(LockedIterator& it) noexcept {
   T& node = *it;
   XDCHECK(node.isInMMContainer());
-  ++it; 
-  removeLocked(node);
+  ++it;
+
+  bool evictedFromTiny = false;
+  if (isTiny(node)) {
+    lru_.getList(LruType::Tiny).remove(node);
+    unmarkTiny(node);
+    evictedFromTiny = true;
+    // Insert into ghost queue upon eviction from tiny queue
+  } else {
+    lru_.getList(LruType::Main).remove(node);
+  }
+
+  unmarkAccessed(node);
+  node.unmarkInMMContainer();
+  
+  if (evictedFromTiny) {
+    // We need to insert to ghost queue. 
+    // Release lock early
+    // AFAIK every call to remove(locked itr) is followed by the iterator's destruction.
+    // If this is not the case, we may need to rethink this.
+    if (it.l_.owns_lock()) {
+      it.l_.unlock();
+    }
+    ghostQueue_.insert(hashNode(node));
+  }
+  return;
 }
 
 template <typename T, MMS3FIFO::Hook<T> T::* HookPtr>
diff --git a/cachelib/common/FIFOConcurrentHashSet.h b/cachelib/common/FIFOConcurrentHashSet.h
@@ -0,0 +1,70 @@
+
+#pragma once
+#include <folly/concurrency/ConcurrentHashMap.h>
+
+#include <atomic>
+#include <deque>
+#include <mutex>
+
+namespace facebook::cachelib::util {
+namespace detail {
+
+template <typename KeyT = uint32_t>
+class FIFOConcurrentHashSetBase {
+ public:
+  using Key = KeyT; // hashed object id
+  using TS = uint64_t;  // logical time / FIFO order
+
+  explicit FIFOConcurrentHashSetBase(size_t capacity)
+      : ghost_(capacity), seq_(0), queueSize_(0) {}
+
+  FIFOConcurrentHashSetBase() : seq_(0), queueSize_(0) {}
+
+  struct PendingEntry {
+    Key k;
+    TS ts;
+  };
+
+  inline void insert(Key k) {
+    TS ts = seq_.fetch_add(1, std::memory_order_relaxed);
+    ghost_.insert_or_assign(k, ts);
+    // If it overflows and wraps around, we reset the entire structure.
+    if (UNLIKELY(ts == std::numeric_limits<TS>::max())) {
+      ghost_.clear();
+    }
+  }
+
+  inline bool contains(Key k)  {
+    auto it = ghost_.find(k);
+    if (it == ghost_.end())
+      return false;
+    // Seq - current must be larger
+    auto currDiff = seq_.load(std::memory_order_relaxed) - it->second;
+    auto isValid = currDiff <= queueSize_.load(std::memory_order_relaxed);
+
+    if (!isValid) {
+      // remove stale entry
+      ghost_.erase(k);
+    }
+
+    return isValid;
+  }
+
+  inline void resize(TS limitTS) {
+    queueSize_.store(limitTS, std::memory_order_relaxed);
+  }
+
+ private:
+  folly::ConcurrentHashMap<Key, TS> ghost_;
+
+  // Logical TS counter
+  std::atomic<TS> seq_;
+  // Maximum size of the FIFO queue
+  std::atomic<TS> queueSize_;
+};
+
+} // namespace detail
+using FIFOConcurrentHashSet32 = detail::FIFOConcurrentHashSetBase<uint32_t>;
+using FIFOConcurrentHashSet64 = detail::FIFOConcurrentHashSetBase<uint64_t>;
+  
+} // namespace facebook::cachelib::util
diff --git a/cachelib/common/FIFOHashSet.h b/cachelib/common/FIFOHashSet.h
