Add element construction/destruction hooks to IndexedMemPool
authorVictor Zverovich <viz@fb.com>
Tue, 6 Jun 2017 20:10:49 +0000 (13:10 -0700)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Tue, 6 Jun 2017 20:20:33 +0000 (13:20 -0700)
Summary: This diff adds a Traits template parameter to IndexedMemPool that allows more control over the lifetime management of individual elements, including mixes of lazy and eager recycle semantics (or colocation of different classes of data inside a single element). It also arranges that an index is not reported as isAllocated() until it has been passed to Traits::initialize and passed to Traits::onAllocate at least once, so code that is traversing valid indexes doesn't need to deal with the post-initialize but pre-onAllocate state (it must still deal with indexes that reported isAllocated() as true but have since been passed to onRecycle). The default behavior is unchanged.

Reviewed By: nbronson

Differential Revision: D5177462

fbshipit-source-id: e7d22c860ab6bf25083977dfb5a63955641c9cfb

folly/IndexedMemPool.h
folly/experimental/flat_combining/FlatCombining.h
folly/test/IndexedMemPoolTest.cpp

index d9e84935a077cedec6acc8fd75fe729d0507fed1..5207c71d99b48ae194058d88bec38adf9703a234 100644 (file)
@@ -37,6 +37,65 @@ template <typename Pool>
 struct IndexedMemPoolRecycler;
 }
 
+template <
+    typename T,
+    bool EagerRecycleWhenTrivial = false,
+    bool EagerRecycleWhenNotTrivial = true>
+struct IndexedMemPoolTraits {
+  static constexpr bool eagerRecycle() {
+    return std::is_trivial<T>::value ? EagerRecycleWhenTrivial
+                                     : EagerRecycleWhenNotTrivial;
+  }
+
+  /// Called when the element pointed to by ptr is allocated for the
+  /// first time.
+  static void initialize(T* ptr) {
+    if (!eagerRecycle()) {
+      new (ptr) T();
+    }
+  }
+
+  /// Called when the element pointed to by ptr is freed at the pool
+  /// destruction time.
+  static void cleanup(T* ptr) {
+    if (!eagerRecycle()) {
+      ptr->~T();
+    }
+  }
+
+  /// Called when the element is allocated with the arguments forwarded from
+  /// IndexedMemPool::allocElem.
+  template <typename... Args>
+  static void onAllocate(T* ptr, Args&&... args) {
+    static_assert(
+        sizeof...(Args) == 0 || eagerRecycle(),
+        "emplace-style allocation requires eager recycle, "
+        "which is defaulted only for non-trivial types");
+    if (eagerRecycle()) {
+      new (ptr) T(std::forward<Args>(args)...);
+    }
+  }
+
+  /// Called when the element is recycled.
+  static void onRecycle(T* ptr) {
+    if (eagerRecycle()) {
+      ptr->~T();
+    }
+  }
+};
+
+/// IndexedMemPool traits that implements the lazy lifecycle strategy. In this
+/// strategy elements are default-constructed the first time they are allocated,
+/// and destroyed when the pool itself is destroyed.
+template <typename T>
+using IndexedMemPoolTraitsLazyRecycle = IndexedMemPoolTraits<T, false, false>;
+
+/// IndexedMemPool traits that implements the eager lifecycle strategy. In this
+/// strategy elements are constructed when they are allocated from the pool and
+/// destroyed when recycled.
+template <typename T>
+using IndexedMemPoolTraitsEagerRecycle = IndexedMemPoolTraits<T, true, true>;
+
 /// Instances of IndexedMemPool dynamically allocate and then pool their
 /// element type (T), returning 4-byte integer indices that can be passed
 /// to the pool's operator[] method to access or obtain pointers to the
@@ -54,13 +113,17 @@ struct IndexedMemPoolRecycler;
 /// there won't be an ABA match due to the element being overwritten with
 /// a different type that has the same bit pattern.
 ///
-/// IndexedMemPool has two object lifecycle strategies.  The first
-/// is to construct objects when they are allocated from the pool and
-/// destroy them when they are recycled.  In this mode allocIndex and
-/// allocElem have emplace-like semantics.  In the second mode, objects
-/// are default-constructed the first time they are removed from the pool,
-/// and deleted when the pool itself is deleted.  By default the first
-/// mode is used for non-trivial T, and the second is used for trivial T.
+/// The object lifecycle strategy is controlled by the Traits parameter.
+/// One strategy, implemented by IndexedMemPoolTraitsEagerRecycle, is to
+/// construct objects when they are allocated from the pool and destroy
+/// them when they are recycled.  In this mode allocIndex and allocElem
+/// have emplace-like semantics.  In another strategy, implemented by
+/// IndexedMemPoolTraitsLazyRecycle, objects are default-constructed the
+/// first time they are removed from the pool, and deleted when the pool
+/// itself is deleted.  By default the first mode is used for non-trivial
+/// T, and the second is used for trivial T.  Clients can customize the
+/// object lifecycle by providing their own Traits implementation.
+/// See IndexedMemPoolTraits for a Traits example.
 ///
 /// IMPORTANT: Space for extra elements is allocated to account for those
 /// that are inaccessible because they are in other local lists, so the
@@ -89,8 +152,7 @@ template <
     uint32_t NumLocalLists_ = 32,
     uint32_t LocalListLimit_ = 200,
     template <typename> class Atom = std::atomic,
-    bool EagerRecycleWhenTrivial = false,
-    bool EagerRecycleWhenNotTrivial = true>
+    typename Traits = IndexedMemPoolTraits<T>>
 struct IndexedMemPool : boost::noncopyable {
   typedef T value_type;
 
@@ -103,12 +165,6 @@ struct IndexedMemPool : boost::noncopyable {
     LocalListLimit = LocalListLimit_
   };
 
-
-  static constexpr bool eagerRecycle() {
-    return std::is_trivial<T>::value
-        ? EagerRecycleWhenTrivial : EagerRecycleWhenNotTrivial;
-  }
-
   // these are public because clients may need to reason about the number
   // of bits required to hold indices from a pool, given its capacity
 
@@ -149,14 +205,8 @@ struct IndexedMemPool : boost::noncopyable {
 
   /// Destroys all of the contained elements
   ~IndexedMemPool() {
-    if (!eagerRecycle()) {
-      // Take the minimum since it is possible that size_ > actualCapacity_.
-      // This can happen if there are multiple concurrent requests
-      // when size_ == actualCapacity_ - 1.
-      uint32_t last = std::min(uint32_t(size_), uint32_t(actualCapacity_));
-      for (uint32_t i = last; i > 0; --i) {
-        slots_[i].~Slot();
-      }
+    for (uint32_t i = maxAllocatedIndex(); i > 0; --i) {
+      Traits::cleanup(&slots_[i].elem);
     }
     munmap(slots_, mmapLength_);
   }
@@ -169,25 +219,35 @@ struct IndexedMemPool : boost::noncopyable {
     return capacityForMaxIndex(actualCapacity_);
   }
 
+  /// Returns the maximum index of elements ever allocated in this pool
+  /// including elements that have been recycled.
+  uint32_t maxAllocatedIndex() const {
+    // Take the minimum since it is possible that size_ > actualCapacity_.
+    // This can happen if there are multiple concurrent requests
+    // when size_ == actualCapacity_ - 1.
+    return std::min(uint32_t(size_), uint32_t(actualCapacity_));
+  }
+
   /// Finds a slot with a non-zero index, emplaces a T there if we're
   /// using the eager recycle lifecycle mode, and returns the index,
-  /// or returns 0 if no elements are available.
+  /// or returns 0 if no elements are available.  Passes a pointer to
+  /// the element to Traits::onAllocate before the slot is marked as
+  /// allocated.
   template <typename ...Args>
   uint32_t allocIndex(Args&&... args) {
-    static_assert(sizeof...(Args) == 0 || eagerRecycle(),
-        "emplace-style allocation requires eager recycle, "
-        "which is defaulted only for non-trivial types");
     auto idx = localPop(localHead());
-    if (idx != 0 && eagerRecycle()) {
-      T* ptr = &slot(idx).elem;
-      new (ptr) T(std::forward<Args>(args)...);
+    if (idx != 0) {
+      Slot& s = slot(idx);
+      Traits::onAllocate(&s.elem, std::forward<Args>(args)...);
+      markAllocated(s);
     }
     return idx;
   }
 
   /// If an element is available, returns a std::unique_ptr to it that will
   /// recycle the element to the pool when it is reclaimed, otherwise returns
-  /// a null (falsy) std::unique_ptr
+  /// a null (falsy) std::unique_ptr.  Passes a pointer to the element to
+  /// Traits::onAllocate before the slot is marked as allocated.
   template <typename ...Args>
   UniquePtr allocElem(Args&&... args) {
     auto idx = allocIndex(std::forward<Args>(args)...);
@@ -198,9 +258,7 @@ struct IndexedMemPool : boost::noncopyable {
   /// Gives up ownership previously granted by alloc()
   void recycleIndex(uint32_t idx) {
     assert(isAllocated(idx));
-    if (eagerRecycle()) {
-      slot(idx).elem.~T();
-    }
+    Traits::onRecycle(&slot(idx).elem);
     localPush(localHead(), idx);
   }
 
@@ -234,7 +292,7 @@ struct IndexedMemPool : boost::noncopyable {
 
   /// Returns true iff idx has been alloc()ed and not recycleIndex()ed
   bool isAllocated(uint32_t idx) const {
-    return slot(idx).localNext.load(std::memory_order_relaxed) == uint32_t(-1);
+    return slot(idx).localNext.load(std::memory_order_acquire) == uint32_t(-1);
   }
 
 
@@ -408,7 +466,6 @@ struct IndexedMemPool : boost::noncopyable {
         auto next = s.localNext.load(std::memory_order_relaxed);
         if (head.compare_exchange_strong(h, h.withIdx(next).withSizeDecr())) {
           // success
-          s.localNext.store(uint32_t(-1), std::memory_order_relaxed);
           return h.idx;
         }
         continue;
@@ -422,13 +479,7 @@ struct IndexedMemPool : boost::noncopyable {
           // allocation failed
           return 0;
         }
-        // default-construct it now if we aren't going to construct and
-        // destroy on each allocation
-        if (!eagerRecycle()) {
-          T* ptr = &slot(idx).elem;
-          new (ptr) T();
-        }
-        slot(idx).localNext.store(uint32_t(-1), std::memory_order_relaxed);
+        Traits::initialize(&slot(idx).elem);
         return idx;
       }
 
@@ -437,7 +488,6 @@ struct IndexedMemPool : boost::noncopyable {
       if (head.compare_exchange_strong(
               h, h.withIdx(next).withSize(LocalListLimit))) {
         // global list moved to local list, keep head for us
-        s.localNext.store(uint32_t(-1), std::memory_order_relaxed);
         return idx;
       }
       // local bulk push failed, return idx to the global list and try again
@@ -449,6 +499,10 @@ struct IndexedMemPool : boost::noncopyable {
     auto stripe = detail::AccessSpreader<Atom>::current(NumLocalLists);
     return local_[stripe].head;
   }
+
+  void markAllocated(Slot& slot) {
+    slot.localNext.store(uint32_t(-1), std::memory_order_release);
+  }
 };
 
 namespace detail {
index 3c951eccb091ae0b8316b3e70a56813d533b61e7..146f1748bd02af13789117d4c2a689f58f2ed72b 100644 (file)
@@ -225,7 +225,8 @@ class FlatCombining {
     }
   };
 
-  using Pool = folly::IndexedMemPool<Rec, 32, 4, Atom, false, false>;
+  using Pool = folly::
+      IndexedMemPool<Rec, 32, 4, Atom, IndexedMemPoolTraitsLazyRecycle<Rec>>;
 
  public:
   /// The constructor takes three optional arguments:
index daa17e02cd94b1b3472cb7ff769128ff7fe5f938..cb2e258bf1e2f388c5c43af6a6b810aeb32b01dc 100644 (file)
  */
 
 #include <folly/IndexedMemPool.h>
-#include <folly/test/DeterministicSchedule.h>
+#include <folly/portability/GMock.h>
 #include <folly/portability/GTest.h>
 #include <folly/portability/Unistd.h>
+#include <folly/test/DeterministicSchedule.h>
 
 #include <string>
 #include <thread>
@@ -25,6 +26,7 @@
 
 using namespace folly;
 using namespace folly::test;
+using namespace testing;
 
 TEST(IndexedMemPool, unique_ptr) {
   typedef IndexedMemPool<size_t> Pool;
@@ -198,8 +200,12 @@ TEST(IndexedMemPool, eager_recycle) {
 
 TEST(IndexedMemPool, late_recycle) {
   {
-    typedef IndexedMemPool<NonTrivialStruct, 8, 8, std::atomic, false, false>
-        Pool;
+    using Pool = IndexedMemPool<
+        NonTrivialStruct,
+        8,
+        8,
+        std::atomic,
+        IndexedMemPoolTraitsLazyRecycle<NonTrivialStruct>>;
     Pool pool(100);
 
     EXPECT_EQ(NonTrivialStruct::count, 0);
@@ -259,7 +265,12 @@ TEST(IndexedMemPool, construction_destruction) {
   std::atomic<bool> start{false};
   std::atomic<int> started{0};
 
-  using Pool = IndexedMemPool<Foo, 1, 1, std::atomic, false, false>;
+  using Pool = IndexedMemPool<
+      Foo,
+      1,
+      1,
+      std::atomic,
+      IndexedMemPoolTraitsLazyRecycle<Foo>>;
   int nthreads = 20;
   int count = 1000;
 
@@ -291,3 +302,97 @@ TEST(IndexedMemPool, construction_destruction) {
 
   CHECK_EQ(cnum.load(), dnum.load());
 }
+
+/// Global Traits mock. It can't be a regular (non-global) mock because we
+/// don't have access to the instance.
+struct MockTraits {
+  static MockTraits* instance;
+
+  MockTraits() {
+    instance = this;
+  }
+
+  ~MockTraits() {
+    instance = nullptr;
+  }
+
+  MOCK_METHOD2(onAllocate, void(std::string*, std::string));
+  MOCK_METHOD1(onRecycle, void(std::string*));
+
+  struct Forwarder {
+    static void initialize(std::string* ptr) {
+      new (ptr) std::string();
+    }
+
+    static void cleanup(std::string* ptr) {
+      using std::string;
+      ptr->~string();
+    }
+
+    static void onAllocate(std::string* ptr, std::string s) {
+      instance->onAllocate(ptr, s);
+    }
+
+    static void onRecycle(std::string* ptr) {
+      instance->onRecycle(ptr);
+    }
+  };
+};
+
+MockTraits* MockTraits::instance;
+
+using TraitsTestPool =
+    IndexedMemPool<std::string, 1, 1, std::atomic, MockTraits::Forwarder>;
+
+void testTraits(TraitsTestPool& pool) {
+  MockTraits traits;
+  const std::string* elem = nullptr;
+  EXPECT_CALL(traits, onAllocate(_, _))
+      .WillOnce(Invoke([&](std::string* s, auto) {
+        EXPECT_FALSE(pool.isAllocated(pool.locateElem(s)));
+        elem = s;
+      }));
+  std::string* ptr = pool.allocElem("foo").release();
+  EXPECT_EQ(ptr, elem);
+
+  elem = nullptr;
+  EXPECT_CALL(traits, onRecycle(_)).WillOnce(Invoke([&](std::string* s) {
+    EXPECT_TRUE(pool.isAllocated(pool.locateElem(s)));
+    elem = s;
+  }));
+  pool.recycleIndex(pool.locateElem(ptr));
+  EXPECT_EQ(ptr, elem);
+}
+
+// Test that Traits is used when both local and global lists are empty.
+TEST(IndexedMemPool, use_traits_empty) {
+  TraitsTestPool pool(10);
+  testTraits(pool);
+}
+
+// Test that Traits is used when allocating from a local list.
+TEST(IndexedMemPool, use_traits_local_list) {
+  TraitsTestPool pool(10);
+  MockTraits traits;
+  EXPECT_CALL(traits, onAllocate(_, _));
+  // Allocate and immediately recycle an element to populate the local list.
+  pool.allocElem("");
+  testTraits(pool);
+}
+
+// Test that Traits is used when allocating from a global list.
+TEST(IndexedMemPool, use_traits_global_list) {
+  TraitsTestPool pool(10);
+  MockTraits traits;
+  EXPECT_CALL(traits, onAllocate(_, _)).Times(2);
+  auto global = pool.allocElem("");
+  // Allocate and immediately recycle an element to fill the local list.
+  pool.allocElem("");
+  // Recycle an element to populate the global list.
+  global.reset();
+  testTraits(pool);
+}
+
+// Test that IndexedMemPool works with incomplete element types.
+struct IncompleteTestElement;
+using IncompleteTestPool = IndexedMemPool<IncompleteTestElement>;