Add combined locking support for MMContainer

through withEvictionIterator function.

Also, expose config option to enable and disable
combined locking.

withEvictionIterator is implemented as en extra function,
getEvictionIterator() is still there and it's behavior hasn't
changed.

Author: igchor

cachelib/allocator/CacheAllocator.h

@@ -1660,7 +1660,7 @@ class CacheAllocator : public CacheBase {
   // @return An evicted item or nullptr  if there is no suitable candidate.
   Item* findEviction(PoolId pid, ClassId cid);
 
-  using EvictionIterator = typename MMContainer::Iterator;
+  using EvictionIterator = typename MMContainer::LockedIterator;
 
   // Advance the current iterator and try to evict a regular item
   //

cachelib/allocator/MM2Q-inl.h

@@ -241,29 +241,17 @@ bool MM2Q::Container<T, HookPtr>::add(T& node) noexcept {
 }
 
 template <typename T, MM2Q::Hook<T> T::*HookPtr>
-typename MM2Q::Container<T, HookPtr>::Iterator
+typename MM2Q::Container<T, HookPtr>::LockedIterator
 MM2Q::Container<T, HookPtr>::getEvictionIterator() const noexcept {
-  // we cannot use combined critical sections with folly::DistributedMutex here
-  // because the lock is held for the lifetime of the eviction iterator.  In
-  // other words, the abstraction of the iterator just does not lend itself well
-  // to combinable critical sections as the user can hold the lock for an
-  // arbitrary amount of time outside a lambda-friendly piece of code (eg. they
-  // can return the iterator from functions, pass it to functions, etc)
-  //
-  // it would be theoretically possible to refactor this interface into
-  // something like the following to allow combining
-  //
-  //    mm2q.withEvictionIterator([&](auto iterator) {
-  //      // user code
-  //    });
-  //
-  // at the time of writing it is unclear if the gains from combining are
-  // reasonable justification for the codemod required to achieve combinability
-  // as we don't expect this critical section to be the hotspot in user code.
-  // This is however subject to change at some time in the future as and when
-  // this assertion becomes false.
   LockHolder l(*lruMutex_);
-  return Iterator{std::move(l), lru_.rbegin()};
+  return LockedIterator{std::move(l), lru_.rbegin()};
+}
+
+template <typename T, MM2Q::Hook<T> T::*HookPtr>
+template <typename F>
+void MM2Q::Container<T, HookPtr>::withEvictionIterator(F&& fun) {
+  LockHolder l(*lruMutex_);
+  fun(Iterator{lru_.rbegin()});
 }
 
 template <typename T, MM2Q::Hook<T> T::*HookPtr>
@@ -462,8 +450,9 @@ void MM2Q::Container<T, HookPtr>::reconfigureLocked(const Time& currTime) {
 
 // Iterator Context Implementation
 template <typename T, MM2Q::Hook<T> T::*HookPtr>
-MM2Q::Container<T, HookPtr>::Iterator::Iterator(
-    LockHolder l, const typename LruList::Iterator& iter) noexcept
-    : LruList::Iterator(iter), l_(std::move(l)) {}
+MM2Q::Container<T, HookPtr>::LockedIterator::LockedIterator(
+    LockHolder l, const Iterator& iter) noexcept
+    : Iterator(iter), l_(std::move(l)) {}
+
 } // namespace cachelib
 } // namespace facebook

cachelib/allocator/MM2Q.h

@@ -68,6 +68,7 @@ class MM2Q {
   enum LruType { Warm, WarmTail, Hot, Cold, ColdTail, NumTypes };
 
   // Config class for MM2Q
+  // TODO: implement support for useCombinedLockForIterators
   struct Config {
     // Create from serialized config
     explicit Config(SerializationConfigType configState)
@@ -347,22 +348,24 @@ class MM2Q {
     Container(const Container&) = delete;
     Container& operator=(const Container&) = delete;
 
+    using Iterator = typename LruList::Iterator;
+
     // context for iterating the MM container. At any given point of time,
     // there can be only one iterator active since we need to lock the LRU for
     // iteration. we can support multiple iterators at same time, by using a
     // shared ptr in the context for the lock holder in the future.
-    class Iterator : public LruList::Iterator {
+    class LockedIterator : public Iterator {
      public:
       // noncopyable but movable.
-      Iterator(const Iterator&) = delete;
-      Iterator& operator=(const Iterator&) = delete;
+      LockedIterator(const LockedIterator&) = delete;
+      LockedIterator& operator=(const LockedIterator&) = delete;
 
-      Iterator(Iterator&&) noexcept = default;
+      LockedIterator(LockedIterator&&) noexcept = default;
 
       // 1. Invalidate this iterator
       // 2. Unlock
       void destroy() {
-        LruList::Iterator::reset();
+        Iterator::reset();
         if (l_.owns_lock()) {
           l_.unlock();
         }
@@ -373,15 +376,15 @@ class MM2Q {
         if (!l_.owns_lock()) {
           l_.lock();
         }
-        LruList::Iterator::resetToBegin();
+        Iterator::resetToBegin();
       }
 
      private:
       // private because it's easy to misuse and cause deadlock for MM2Q
-      Iterator& operator=(Iterator&&) noexcept = default;
+      LockedIterator& operator=(LockedIterator&&) noexcept = default;
 
       // create an lru iterator with the lock being held.
-      Iterator(LockHolder l, const typename LruList::Iterator& iter) noexcept;
+      LockedIterator(LockHolder l, const Iterator& iter) noexcept;
 
       // only the container can create iterators
       friend Container<T, HookPtr>;
@@ -422,7 +425,7 @@ class MM2Q {
 
     // same as the above but uses an iterator context. The iterator is updated
     // on removal of the corresponding node to point to the next node. The
-    // iterator context holds the lock on the lru.
+    // iterator context is responsible for locking.
     //
     // iterator will be advanced to the next node after removing the node
     //
@@ -445,7 +448,12 @@ class MM2Q {
     // Obtain an iterator that start from the tail and can be used
     // to search for evictions. This iterator holds a lock to this
     // container and only one such iterator can exist at a time
-    Iterator getEvictionIterator() const noexcept;
+    LockedIterator getEvictionIterator() const noexcept;
+
+    // Execute provided function under container lock. Function gets
+    // Iterator passed as parameter.
+    template <typename F>
+    void withEvictionIterator(F&& f);
 
     // get the current config as a copy
     Config getConfig() const;

cachelib/allocator/MMLru-inl.h

@@ -212,10 +212,21 @@ bool MMLru::Container<T, HookPtr>::add(T& node) noexcept {
 }
 
 template <typename T, MMLru::Hook<T> T::*HookPtr>
-typename MMLru::Container<T, HookPtr>::Iterator
+typename MMLru::Container<T, HookPtr>::LockedIterator
 MMLru::Container<T, HookPtr>::getEvictionIterator() const noexcept {
   LockHolder l(*lruMutex_);
-  return Iterator{std::move(l), lru_.rbegin()};
+  return LockedIterator{std::move(l), lru_.rbegin()};
+}
+
+template <typename T, MMLru::Hook<T> T::*HookPtr>
+template <typename F>
+void MMLru::Container<T, HookPtr>::withEvictionIterator(F&& fun) {
+  if (config_.useCombinedLockForIterators) {
+    lruMutex_->lock_combine([this, &fun]() { fun(Iterator{lru_.rbegin()}); });
+  } else {
+    LockHolder lck{*lruMutex_};
+    fun(Iterator{lru_.rbegin()});
+  }
 }
 
 template <typename T, MMLru::Hook<T> T::*HookPtr>
@@ -360,8 +371,9 @@ void MMLru::Container<T, HookPtr>::reconfigureLocked(const Time& currTime) {
 
 // Iterator Context Implementation
 template <typename T, MMLru::Hook<T> T::*HookPtr>
-MMLru::Container<T, HookPtr>::Iterator::Iterator(
-    LockHolder l, const typename LruList::Iterator& iter) noexcept
-    : LruList::Iterator(iter), l_(std::move(l)) {}
+MMLru::Container<T, HookPtr>::LockedIterator::LockedIterator(
+    LockHolder l, const Iterator& iter) noexcept
+    : Iterator(iter), l_(std::move(l)) {}
+
 } // namespace cachelib
 } // namespace facebook

cachelib/allocator/MMLru.h

@@ -144,15 +144,17 @@ class MMLru {
            bool updateOnR,
            bool tryLockU,
            uint8_t ipSpec,
-           uint32_t mmReconfigureInterval)
+           uint32_t mmReconfigureInterval,
+           bool useCombinedLockForIterators = false)
         : defaultLruRefreshTime(time),
           lruRefreshRatio(ratio),
           updateOnWrite(updateOnW),
           updateOnRead(updateOnR),
           tryLockUpdate(tryLockU),
           lruInsertionPointSpec(ipSpec),
           mmReconfigureIntervalSecs(
-              std::chrono::seconds(mmReconfigureInterval)) {}
+              std::chrono::seconds(mmReconfigureInterval)),
+          useCombinedLockForIterators(useCombinedLockForIterators) {}
 
     Config() = default;
     Config(const Config& rhs) = default;
@@ -195,6 +197,9 @@ class MMLru {
     //     lruInsertionPointSpec = 2, we insert at a point 1/4th from tail
     uint8_t lruInsertionPointSpec{0};
 
+    // Whether to use combined locking for withEvictionIterator.
+    bool useCombinedLockForIterators{false};
+
     // Minimum interval between reconfigurations. If 0, reconfigure is never
     // called.
     std::chrono::seconds mmReconfigureIntervalSecs{};
@@ -234,22 +239,24 @@ class MMLru {
     Container(const Container&) = delete;
     Container& operator=(const Container&) = delete;
 
+    using Iterator = typename LruList::Iterator;
+
     // context for iterating the MM container. At any given point of time,
     // there can be only one iterator active since we need to lock the LRU for
     // iteration. we can support multiple iterators at same time, by using a
     // shared ptr in the context for the lock holder in the future.
-    class Iterator : public LruList::Iterator {
+    class LockedIterator : public Iterator {
      public:
       // noncopyable but movable.
-      Iterator(const Iterator&) = delete;
-      Iterator& operator=(const Iterator&) = delete;
+      LockedIterator(const LockedIterator&) = delete;
+      LockedIterator& operator=(const LockedIterator&) = delete;
 
-      Iterator(Iterator&&) noexcept = default;
+      LockedIterator(LockedIterator&&) noexcept = default;
 
       // 1. Invalidate this iterator
       // 2. Unlock
       void destroy() {
-        LruList::Iterator::reset();
+        Iterator::reset();
         if (l_.owns_lock()) {
           l_.unlock();
         }
@@ -260,15 +267,15 @@ class MMLru {
         if (!l_.owns_lock()) {
           l_.lock();
         }
-        LruList::Iterator::resetToBegin();
+        Iterator::resetToBegin();
       }
 
      private:
       // private because it's easy to misuse and cause deadlock for MMLru
-      Iterator& operator=(Iterator&&) noexcept = default;
+      LockedIterator& operator=(LockedIterator&&) noexcept = default;
 
       // create an lru iterator with the lock being held.
-      Iterator(LockHolder l, const typename LruList::Iterator& iter) noexcept;
+      LockedIterator(LockHolder l, const Iterator& iter) noexcept;
 
       // only the container can create iterators
       friend Container<T, HookPtr>;
@@ -307,10 +314,9 @@ class MMLru {
     //          state of node is unchanged.
     bool remove(T& node) noexcept;
 
-    using Iterator = Iterator;
     // same as the above but uses an iterator context. The iterator is updated
     // on removal of the corresponding node to point to the next node. The
-    // iterator context holds the lock on the lru.
+    // iterator context is responsible for locking.
     //
     // iterator will be advanced to the next node after removing the node
     //
@@ -330,7 +336,12 @@ class MMLru {
     // Obtain an iterator that start from the tail and can be used
     // to search for evictions. This iterator holds a lock to this
     // container and only one such iterator can exist at a time
-    Iterator getEvictionIterator() const noexcept;
+    LockedIterator getEvictionIterator() const noexcept;
+
+    // Execute provided function under container lock. Function gets
+    // iterator passed as parameter.
+    template <typename F>
+    void withEvictionIterator(F&& f);
 
     // get copy of current config
     Config getConfig() const;

cachelib/allocator/MMTinyLFU-inl.h

@@ -214,10 +214,17 @@ bool MMTinyLFU::Container<T, HookPtr>::add(T& node) noexcept {
 }
 
 template <typename T, MMTinyLFU::Hook<T> T::*HookPtr>
-typename MMTinyLFU::Container<T, HookPtr>::Iterator
+typename MMTinyLFU::Container<T, HookPtr>::LockedIterator
 MMTinyLFU::Container<T, HookPtr>::getEvictionIterator() const noexcept {
   LockHolder l(lruMutex_);
-  return Iterator{std::move(l), *this};
+  return LockedIterator{std::move(l), *this};
+}
+
+template <typename T, MMTinyLFU::Hook<T> T::*HookPtr>
+template <typename F>
+void MMTinyLFU::Container<T, HookPtr>::withEvictionIterator(F&& fun) {
+  LockHolder l(lruMutex_);
+  fun(Iterator{*this});
 }
 
 template <typename T, MMTinyLFU::Hook<T> T::*HookPtr>
@@ -347,13 +354,18 @@ void MMTinyLFU::Container<T, HookPtr>::reconfigureLocked(const Time& currTime) {
   lruRefreshTime_.store(lruRefreshTime, std::memory_order_relaxed);
 }
 
+// Locked Iterator Context Implementation
+template <typename T, MMTinyLFU::Hook<T> T::*HookPtr>
+MMTinyLFU::Container<T, HookPtr>::LockedIterator::LockedIterator(
+    LockHolder l, const Container<T, HookPtr>& c) noexcept
+    : Iterator(c), l_(std::move(l)) {}
+
 // Iterator Context Implementation
 template <typename T, MMTinyLFU::Hook<T> T::*HookPtr>
 MMTinyLFU::Container<T, HookPtr>::Iterator::Iterator(
-    LockHolder l, const Container<T, HookPtr>& c) noexcept
+    const Container<T, HookPtr>& c) noexcept
     : c_(c),
       tIter_(c.lru_.getList(LruType::Tiny).rbegin()),
-      mIter_(c.lru_.getList(LruType::Main).rbegin()),
-      l_(std::move(l)) {}
+      mIter_(c.lru_.getList(LruType::Main).rbegin()) {}
 } // namespace cachelib
 } // namespace facebook