DRAM LockGroupIterator

Summary: Lock once, iterate over all buckets belonging to the lock.

Reviewed By: rlyerly

Differential Revision: D100740472

fbshipit-source-id: acb3a5713f851c16e1ec07c639956ec0e1dd9a48

Author: AlnisM

cachelib/allocator/CacheAllocator.h

@@ -652,6 +652,34 @@ class CacheAllocator : public CacheBase {
 
   AccessIterator end() { return accessContainer_->end(); }
 
+  // Alternative iterator that batches by hash table lock group instead of
+  // per-bucket.
+  //
+  // Differences from AccessIterator:
+  //   - Fewer lock acquisitions: O(numLocks) vs O(numBuckets).
+  //   - Larger pinning window: each lock group covers numBuckets/numLocks
+  //     buckets. All items across those buckets are snapshotted as Handles
+  //     at once, blocking eviction until the caller advances past them.
+  //     AccessIterator snapshots one bucket at a time.
+  using LockGroupAccessIterator = typename AccessContainer::LockGroupIterator;
+
+  LockGroupAccessIterator beginLockGroup() {
+    return accessContainer_->beginLockGroup(
+        [this](Item* it) { return tryAcquire(it); },
+        [this](Key key) -> WriteHandle { return findInternal(key); });
+  }
+
+  LockGroupAccessIterator beginLockGroup(util::Throttler::Config config) {
+    return accessContainer_->beginLockGroup(
+        [this](Item* it) { return tryAcquire(it); },
+        [this](Key key) -> WriteHandle { return findInternal(key); },
+        config);
+  }
+
+  LockGroupAccessIterator endLockGroup() {
+    return accessContainer_->endLockGroup();
+  }
+
   enum class RemoveRes : uint8_t {
     kSuccess,
     kNotFoundInRam,

cachelib/allocator/ChainedHashTable.h

@@ -624,6 +624,110 @@ class ChainedHashTable {
     Iterator begin() { return Iterator(*this); }
     Iterator end() { return Iterator(*this, Iterator::EndIter); }
 
+    // Like Iterator, but batches by lock group instead of per-bucket.
+    // Acquires each lock once and snapshots handles for all buckets under
+    // that lock, reducing the total number of lock acquisitions from
+    // O(numBuckets) to O(numLocks). Uses non-blocking handle creation
+    // under the lock with retry outside, avoiding the deadlock where
+    // blocking handleMaker_ waits for item moves while holding the lock.
+    //
+    // Trade-off: larger snapshot per lock group means more items are pinned
+    // (not evictable) at once compared to the per-bucket Iterator.
+    class LockGroupIterator {
+     public:
+      using TryHandleMakerFn =
+          std::function<std::pair<Handle, TryAcquireResult>(T*)>;
+      using FindByKeyFn = std::function<Handle(folly::StringPiece)>;
+
+      ~LockGroupIterator() {
+        XDCHECK_GT(container_->numIterators_.load(), 0u);
+        --container_->numIterators_;
+      }
+      LockGroupIterator(const LockGroupIterator&) = delete;
+      LockGroupIterator& operator=(const LockGroupIterator&) = delete;
+
+      LockGroupIterator(LockGroupIterator&&) noexcept;
+      LockGroupIterator& operator=(LockGroupIterator&&) noexcept;
+      enum EndIterT { EndIter };
+
+      LockGroupIterator& operator++();
+
+      T& operator*();
+      T* operator->() { return &(*(*this)); }
+      const T& operator*() const;
+      const T* operator->() const { return &(*(*this)); }
+
+      bool operator==(const LockGroupIterator& other) const noexcept {
+        return container_ == other.container_ && currLock_ == other.currLock_ &&
+               cursor_ == other.cursor_;
+      }
+
+      bool operator!=(const LockGroupIterator& other) const noexcept {
+        return !(*this == other);
+      }
+
+      const Handle& asHandle() { return curr(); }
+
+      void reset();
+
+     private:
+      using C = Container<T, HookPtr, LockT>;
+
+      friend C;
+      explicit LockGroupIterator(C& ht,
+                                 TryHandleMakerFn tryHandleMaker,
+                                 FindByKeyFn findByKey,
+                                 folly::Optional<util::Throttler::Config>
+                                     throttlerConfig = folly::none);
+
+      LockGroupIterator(C& ht, EndIterT);
+
+      mutable C* container_;
+      TryHandleMakerFn tryHandleMaker_;
+      FindByKeyFn findByKey_;
+
+      // current lock group that the iterator is pointing to.
+      mutable size_t currLock_{0};
+      // cursor into the current lock group's snapshot.
+      mutable unsigned int cursor_{0};
+      // snapshot of handles for all items in the current lock group.
+      mutable std::vector<Handle> lockGroupElems_;
+
+      folly::Optional<util::Throttler> throttler_ = folly::none;
+
+      Handle& curr() const {
+        if (cursor_ < lockGroupElems_.size()) {
+          return lockGroupElems_[cursor_];
+        }
+        throw std::logic_error(
+            "LockGroupIterator in invalid state with cursor_: " +
+            folly::to<std::string>(cursor_) + ", currLock_: " +
+            folly::to<std::string>(currLock_) + ", total locks: " +
+            folly::to<std::string>(container_->config_.getNumLocks()));
+      }
+
+      // Returns handles for all items in the given lock group. Uses
+      // non-blocking tryHandleMaker_ under the lock, then retries moving
+      // items via findByKey_ outside it.
+      std::vector<Handle> getLockGroupElems(size_t lockIdx);
+    };
+
+    LockGroupIterator beginLockGroup(
+        typename LockGroupIterator::TryHandleMakerFn tryHandleMaker,
+        typename LockGroupIterator::FindByKeyFn findByKey,
+        folly::Optional<util::Throttler::Config> throttlerConfig);
+
+    LockGroupIterator beginLockGroup(
+        typename LockGroupIterator::TryHandleMakerFn tryHandleMaker,
+        typename LockGroupIterator::FindByKeyFn findByKey) {
+      return LockGroupIterator(*this, std::move(tryHandleMaker),
+                               std::move(findByKey));
+    }
+
+    LockGroupIterator endLockGroup() {
+      return LockGroupIterator(*this, LockGroupIterator::EndIter);
+    }
+
     // Stats describing the distribution of items (keys) in the hash table
     struct DistributionStats {
       uint64_t numKeys{0};
@@ -1320,4 +1424,201 @@ void ChainedHashTable::Container<T, HookPtr, LockT>::Iterator::reset() {
   }
   XDCHECK_EQ(0u, curSor_);
 }
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+std::vector<typename ChainedHashTable::Container<T, HookPtr, LockT>::Handle>
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+    getLockGroupElems(size_t lockIdx) {
+  std::vector<Handle> elems;
+  std::vector<std::string> retryKeys;
+
+  {
+    auto guard = container_->locks_.lockShared(lockIdx);
+    const auto numBuckets = container_->config_.getNumBuckets();
+
+    container_->locks_.forEachBucketForLock(
+        lockIdx, numBuckets, [this, &elems, &retryKeys](size_t bucket) {
+          container_->ht_.forEachBucketElem(
+              bucket, [this, &elems, &retryKeys](T* elem) {
+                try {
+                  auto [h, tryRes] = tryHandleMaker_(elem);
+                  if (tryRes == TryAcquireResult::kSuccess) {
+                    elems.emplace_back(std::move(h));
+                  } else if (tryRes == TryAcquireResult::kMoving) {
+                    // Can't retry under the lock — findByKey_ may block on the
+                    // move which needs exclusive access to this same lock
+                    // group. Save the key and retry after releasing the lock.
+                    auto key = elem->getKey();
+                    retryKeys.emplace_back(key.data(), key.size());
+                  }
+                  // kSkip: handle not acquirable, skip it.
+                } catch (const std::exception&) {
+                  // if we are not able to acquire a handle, skip over them.
+                }
+              });
+        });
+  }
+
+  // Retry items that were being moved. Now that we don't hold any lock,
+  // findByKey_ can safely block waiting for the move to complete.
+  for (auto& key : retryKeys) {
+    try {
+      auto h = findByKey_(folly::StringPiece(key));
+      if (h) {
+        elems.emplace_back(std::move(h));
+      }
+    } catch (const std::exception&) {
+      // if we are not able to acquire a handle, skip over them.
+    }
+  }
+
+  return elems;
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+typename ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator&
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+operator++() {
+  if (throttler_) {
+    throttler_->throttle();
+  }
+
+  // Release the handle we're advancing past so it no longer pins the item.
+  // This unblocks evictions and slab rebalances on iterated-past items
+  // before we move to the next lock group.
+  if (cursor_ < lockGroupElems_.size()) {
+    lockGroupElems_[cursor_].reset();
+  }
+
+  ++cursor_;
+  if (cursor_ < lockGroupElems_.size()) {
+    return *this;
+  }
+
+  ++currLock_;
+  for (; currLock_ < container_->config_.getNumLocks(); ++currLock_) {
+    lockGroupElems_ = getLockGroupElems(currLock_);
+    if (!lockGroupElems_.empty()) {
+      cursor_ = 0;
+      return *this;
+    } else if (throttler_) {
+      throttler_->throttle();
+    }
+  }
+
+  // reached the end
+  lockGroupElems_.clear();
+  cursor_ = 0;
+  return *this;
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+T& ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+operator*() {
+  return *curr();
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+const T&
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::operator*()
+    const {
+  return *curr();
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+    LockGroupIterator(Container<T, HookPtr, LockT>& container,
+                      TryHandleMakerFn tryHandleMaker,
+                      FindByKeyFn findByKey,
+                      folly::Optional<util::Throttler::Config> throttlerConfig)
+    : container_(&container),
+      tryHandleMaker_(std::move(tryHandleMaker)),
+      findByKey_(std::move(findByKey)) {
+  if (throttlerConfig) {
+    throttler_.assign(util::Throttler(*throttlerConfig));
+  }
+
+  ++container_->numIterators_;
+
+  reset();
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+    LockGroupIterator(Container<T, HookPtr, LockT>& container, EndIterT)
+    : container_(&container), currLock_{container_->config_.getNumLocks()} {
+  ++container_->numIterators_;
+  XDCHECK_EQ(0u, cursor_);
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+    LockGroupIterator(LockGroupIterator&& other) noexcept
+    : container_{other.container_},
+      tryHandleMaker_(std::move(other.tryHandleMaker_)),
+      findByKey_(std::move(other.findByKey_)),
+      currLock_{other.currLock_},
+      cursor_{other.cursor_},
+      lockGroupElems_(std::move(other.lockGroupElems_)),
+      throttler_(std::move(other.throttler_)) {
+  ++container_->numIterators_;
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+typename ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator&
+ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::operator=(
+    LockGroupIterator&& other) noexcept {
+  if (this != &other) {
+    this->~LockGroupIterator();
+    new (this) LockGroupIterator(std::move(other));
+  }
+  return *this;
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+typename ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator
+ChainedHashTable::Container<T, HookPtr, LockT>::beginLockGroup(
+    typename LockGroupIterator::TryHandleMakerFn tryHandleMaker,
+    typename LockGroupIterator::FindByKeyFn findByKey,
+    folly::Optional<util::Throttler::Config> throttlerConfig) {
+  return LockGroupIterator(*this, std::move(tryHandleMaker),
+                           std::move(findByKey), throttlerConfig);
+}
+
+template <typename T,
+          typename ChainedHashTable::Hook<T> T::* HookPtr,
+          typename LockT>
+void ChainedHashTable::Container<T, HookPtr, LockT>::LockGroupIterator::
+    reset() {
+  cursor_ = 0;
+  currLock_ = 0;
+  lockGroupElems_ = getLockGroupElems(currLock_);
+  while (lockGroupElems_.empty() &&
+         ++currLock_ < container_->config_.getNumLocks()) {
+    if (throttler_) {
+      throttler_->throttle();
+    }
+    lockGroupElems_ = getLockGroupElems(currLock_);
+  }
+  XDCHECK_EQ(0u, cursor_);
+}
+
 } // namespace facebook::cachelib

cachelib/allocator/tests/AllocatorTypeTest.cpp

@@ -149,6 +149,10 @@ TYPED_TEST(BaseAllocatorTest, IterateAndRemoveWithIter) {
   this->testIterateAndRemoveWithIter();
 }
 
+TYPED_TEST(BaseAllocatorTest, IterateLockGroup) {
+  this->testIterateLockGroup();
+}
+
 TYPED_TEST(BaseAllocatorTest, IterateWithEvictions) {
   this->testIterateWithEvictions();
 }

cachelib/allocator/tests/BaseAllocatorTest.h

@@ -2081,6 +2081,35 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
     }
   }
 
+  // Verify that LockGroupAccessIterator visits every accessible item exactly
+  // once. Exercises CacheAllocator::beginLockGroup() / endLockGroup() and the
+  // tryHandleMaker / findByKey wiring inside CacheAllocator.
+  void testIterateLockGroup() {
+    typename AllocatorT::Config config;
+    config.setCacheSize(10 * Slab::kSize);
+    AllocatorT alloc(config);
+    const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
+    std::set<uint32_t> allocSizes{1024};
+    auto poolId = alloc.addPool("foobar", numBytes, allocSizes);
+
+    const unsigned int numItems = 100;
+    const uint32_t itemSize = 100;
+    std::set<std::string> expectedKeys;
+    for (unsigned int i = 0; i < numItems; ++i) {
+      const std::string key = "key_" + folly::to<std::string>(i);
+      auto handle = util::allocateAccessible(alloc, poolId, key, itemSize);
+      ASSERT_NE(nullptr, handle);
+      expectedKeys.insert(key);
+    }
+
+    std::set<std::string> visitedKeys;
+    for (auto it = alloc.beginLockGroup(); it != alloc.endLockGroup(); ++it) {
+      const bool inserted = visitedKeys.insert(it->getKey().str()).second;
+      ASSERT_TRUE(inserted);
+    }
+    ASSERT_EQ(expectedKeys, visitedKeys);
+  }
+
   void testIterateWithEvictions() {
     std::set<std::string> evictedKeys;
     auto removeCb =