folly/IndexedMemPool.h

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <type_traits>
  20 #include <assert.h>
  21 #include <errno.h>
  22 #include <stdint.h>
  23 #include <boost/noncopyable.hpp>
  24 #include <folly/AtomicStruct.h>
  25 #include <folly/detail/CacheLocality.h>
  26 #include <folly/portability/SysMman.h>
  27 #include <folly/portability/Unistd.h>
  28
  29 // Ignore shadowing warnings within this file, so includers can use -Wshadow.
  30 #pragma GCC diagnostic push
  31 #pragma GCC diagnostic ignored "-Wshadow"
  32
  33 namespace folly {
  34
  35 namespace detail {
  36 template <typename Pool>
  37 struct IndexedMemPoolRecycler;
  38 }
  39
  40 template <
  41     typename T,
  42     bool EagerRecycleWhenTrivial = false,
  43     bool EagerRecycleWhenNotTrivial = true>
  44 struct IndexedMemPoolTraits {
  45   static constexpr bool eagerRecycle() {
  46     return std::is_trivial<T>::value ? EagerRecycleWhenTrivial
  47                                      : EagerRecycleWhenNotTrivial;
  48   }
  49
  50   /// Called when the element pointed to by ptr is allocated for the
  51   /// first time.
  52   static void initialize(T* ptr) {
  53     if (!eagerRecycle()) {
  54       new (ptr) T();
  55     }
  56   }
  57
  58   /// Called when the element pointed to by ptr is freed at the pool
  59   /// destruction time.
  60   static void cleanup(T* ptr) {
  61     if (!eagerRecycle()) {
  62       ptr->~T();
  63     }
  64   }
  65
  66   /// Called when the element is allocated with the arguments forwarded from
  67   /// IndexedMemPool::allocElem.
  68   template <typename... Args>
  69   static void onAllocate(T* ptr, Args&&... args) {
  70     static_assert(
  71         sizeof...(Args) == 0 || eagerRecycle(),
  72         "emplace-style allocation requires eager recycle, "
  73         "which is defaulted only for non-trivial types");
  74     if (eagerRecycle()) {
  75       new (ptr) T(std::forward<Args>(args)...);
  76     }
  77   }
  78
  79   /// Called when the element is recycled.
  80   static void onRecycle(T* ptr) {
  81     if (eagerRecycle()) {
  82       ptr->~T();
  83     }
  84   }
  85 };
  86
  87 /// IndexedMemPool traits that implements the lazy lifecycle strategy. In this
  88 /// strategy elements are default-constructed the first time they are allocated,
  89 /// and destroyed when the pool itself is destroyed.
  90 template <typename T>
  91 using IndexedMemPoolTraitsLazyRecycle = IndexedMemPoolTraits<T, false, false>;
  92
  93 /// IndexedMemPool traits that implements the eager lifecycle strategy. In this
  94 /// strategy elements are constructed when they are allocated from the pool and
  95 /// destroyed when recycled.
  96 template <typename T>
  97 using IndexedMemPoolTraitsEagerRecycle = IndexedMemPoolTraits<T, true, true>;
  98
  99 /// Instances of IndexedMemPool dynamically allocate and then pool their
 100 /// element type (T), returning 4-byte integer indices that can be passed
 101 /// to the pool's operator[] method to access or obtain pointers to the
 102 /// actual elements.  The memory backing items returned from the pool
 103 /// will always be readable, even if items have been returned to the pool.
 104 /// These two features are useful for lock-free algorithms.  The indexing
 105 /// behavior makes it easy to build tagged pointer-like-things, since
 106 /// a large number of elements can be managed using fewer bits than a
 107 /// full pointer.  The access-after-free behavior makes it safe to read
 108 /// from T-s even after they have been recycled, since it is guaranteed
 109 /// that the memory won't have been returned to the OS and unmapped
 110 /// (the algorithm must still use a mechanism to validate that the read
 111 /// was correct, but it doesn't have to worry about page faults), and if
 112 /// the elements use internal sequence numbers it can be guaranteed that
 113 /// there won't be an ABA match due to the element being overwritten with
 114 /// a different type that has the same bit pattern.
 115 ///
 116 /// The object lifecycle strategy is controlled by the Traits parameter.
 117 /// One strategy, implemented by IndexedMemPoolTraitsEagerRecycle, is to
 118 /// construct objects when they are allocated from the pool and destroy
 119 /// them when they are recycled.  In this mode allocIndex and allocElem
 120 /// have emplace-like semantics.  In another strategy, implemented by
 121 /// IndexedMemPoolTraitsLazyRecycle, objects are default-constructed the
 122 /// first time they are removed from the pool, and deleted when the pool
 123 /// itself is deleted.  By default the first mode is used for non-trivial
 124 /// T, and the second is used for trivial T.  Clients can customize the
 125 /// object lifecycle by providing their own Traits implementation.
 126 /// See IndexedMemPoolTraits for a Traits example.
 127 ///
 128 /// IMPORTANT: Space for extra elements is allocated to account for those
 129 /// that are inaccessible because they are in other local lists, so the
 130 /// actual number of items that can be allocated ranges from capacity to
 131 /// capacity + (NumLocalLists_-1)*LocalListLimit_.  This is important if
 132 /// you are trying to maximize the capacity of the pool while constraining
 133 /// the bit size of the resulting pointers, because the pointers will
 134 /// actually range up to the boosted capacity.  See maxIndexForCapacity
 135 /// and capacityForMaxIndex.
 136 ///
 137 /// To avoid contention, NumLocalLists_ free lists of limited (less than
 138 /// or equal to LocalListLimit_) size are maintained, and each thread
 139 /// retrieves and returns entries from its associated local list.  If the
 140 /// local list becomes too large then elements are placed in bulk in a
 141 /// global free list.  This allows items to be efficiently recirculated
 142 /// from consumers to producers.  AccessSpreader is used to access the
 143 /// local lists, so there is no performance advantage to having more
 144 /// local lists than L1 caches.
 145 ///
 146 /// The pool mmap-s the entire necessary address space when the pool is
 147 /// constructed, but delays element construction.  This means that only
 148 /// elements that are actually returned to the caller get paged into the
 149 /// process's resident set (RSS).
 150 template <
 151     typename T,
 152     uint32_t NumLocalLists_ = 32,
 153     uint32_t LocalListLimit_ = 200,
 154     template <typename> class Atom = std::atomic,
 155     typename Traits = IndexedMemPoolTraits<T>>
 156 struct IndexedMemPool : boost::noncopyable {
 157   typedef T value_type;
 158
 159   typedef std::unique_ptr<T, detail::IndexedMemPoolRecycler<IndexedMemPool>>
 160       UniquePtr;
 161
 162   static_assert(LocalListLimit_ <= 255, "LocalListLimit must fit in 8 bits");
 163   enum {
 164     NumLocalLists = NumLocalLists_,
 165     LocalListLimit = LocalListLimit_
 166   };
 167
 168   // these are public because clients may need to reason about the number
 169   // of bits required to hold indices from a pool, given its capacity
 170
 171   static constexpr uint32_t maxIndexForCapacity(uint32_t capacity) {
 172     // index of std::numeric_limits<uint32_t>::max() is reserved for isAllocated
 173     // tracking
 174     return uint32_t(std::min(
 175         uint64_t(capacity) + (NumLocalLists - 1) * LocalListLimit,
 176         uint64_t(std::numeric_limits<uint32_t>::max() - 1)));
 177   }
 178
 179   static constexpr uint32_t capacityForMaxIndex(uint32_t maxIndex) {
 180     return maxIndex - (NumLocalLists - 1) * LocalListLimit;
 181   }
 182
 183
 184   /// Constructs a pool that can allocate at least _capacity_ elements,
 185   /// even if all the local lists are full
 186   explicit IndexedMemPool(uint32_t capacity)
 187     : actualCapacity_(maxIndexForCapacity(capacity))
 188     , size_(0)
 189     , globalHead_(TaggedPtr{})
 190   {
 191     const size_t needed = sizeof(Slot) * (actualCapacity_ + 1);
 192     size_t pagesize = size_t(sysconf(_SC_PAGESIZE));
 193     mmapLength_ = ((needed - 1) & ~(pagesize - 1)) + pagesize;
 194     assert(needed <= mmapLength_ && mmapLength_ < needed + pagesize);
 195     assert((mmapLength_ % pagesize) == 0);
 196
 197     slots_ = static_cast<Slot*>(mmap(nullptr, mmapLength_,
 198                                      PROT_READ | PROT_WRITE,
 199                                      MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
 200     if (slots_ == MAP_FAILED) {
 201       assert(errno == ENOMEM);
 202       throw std::bad_alloc();
 203     }
 204   }
 205
 206   /// Destroys all of the contained elements
 207   ~IndexedMemPool() {
 208     for (uint32_t i = maxAllocatedIndex(); i > 0; --i) {
 209       Traits::cleanup(&slots_[i].elem);
 210     }
 211     munmap(slots_, mmapLength_);
 212   }
 213
 214   /// Returns a lower bound on the number of elements that may be
 215   /// simultaneously allocated and not yet recycled.  Because of the
 216   /// local lists it is possible that more elements than this are returned
 217   /// successfully
 218   uint32_t capacity() {
 219     return capacityForMaxIndex(actualCapacity_);
 220   }
 221
 222   /// Returns the maximum index of elements ever allocated in this pool
 223   /// including elements that have been recycled.
 224   uint32_t maxAllocatedIndex() const {
 225     // Take the minimum since it is possible that size_ > actualCapacity_.
 226     // This can happen if there are multiple concurrent requests
 227     // when size_ == actualCapacity_ - 1.
 228     return std::min(uint32_t(size_), uint32_t(actualCapacity_));
 229   }
 230
 231   /// Finds a slot with a non-zero index, emplaces a T there if we're
 232   /// using the eager recycle lifecycle mode, and returns the index,
 233   /// or returns 0 if no elements are available.  Passes a pointer to
 234   /// the element to Traits::onAllocate before the slot is marked as
 235   /// allocated.
 236   template <typename ...Args>
 237   uint32_t allocIndex(Args&&... args) {
 238     auto idx = localPop(localHead());
 239     if (idx != 0) {
 240       Slot& s = slot(idx);
 241       Traits::onAllocate(&s.elem, std::forward<Args>(args)...);
 242       markAllocated(s);
 243     }
 244     return idx;
 245   }
 246
 247   /// If an element is available, returns a std::unique_ptr to it that will
 248   /// recycle the element to the pool when it is reclaimed, otherwise returns
 249   /// a null (falsy) std::unique_ptr.  Passes a pointer to the element to
 250   /// Traits::onAllocate before the slot is marked as allocated.
 251   template <typename ...Args>
 252   UniquePtr allocElem(Args&&... args) {
 253     auto idx = allocIndex(std::forward<Args>(args)...);
 254     T* ptr = idx == 0 ? nullptr : &slot(idx).elem;
 255     return UniquePtr(ptr, typename UniquePtr::deleter_type(this));
 256   }
 257
 258   /// Gives up ownership previously granted by alloc()
 259   void recycleIndex(uint32_t idx) {
 260     assert(isAllocated(idx));
 261     Traits::onRecycle(&slot(idx).elem);
 262     localPush(localHead(), idx);
 263   }
 264
 265   /// Provides access to the pooled element referenced by idx
 266   T& operator[](uint32_t idx) {
 267     return slot(idx).elem;
 268   }
 269
 270   /// Provides access to the pooled element referenced by idx
 271   const T& operator[](uint32_t idx) const {
 272     return slot(idx).elem;
 273   }
 274
 275   /// If elem == &pool[idx], then pool.locateElem(elem) == idx.  Also,
 276   /// pool.locateElem(nullptr) == 0
 277   uint32_t locateElem(const T* elem) const {
 278     if (!elem) {
 279       return 0;
 280     }
 281
 282     static_assert(std::is_standard_layout<Slot>::value, "offsetof needs POD");
 283
 284     auto slot = reinterpret_cast<const Slot*>(
 285         reinterpret_cast<const char*>(elem) - offsetof(Slot, elem));
 286     auto rv = uint32_t(slot - slots_);
 287
 288     // this assert also tests that rv is in range
 289     assert(elem == &(*this)[rv]);
 290     return rv;
 291   }
 292
 293   /// Returns true iff idx has been alloc()ed and not recycleIndex()ed
 294   bool isAllocated(uint32_t idx) const {
 295     return slot(idx).localNext.load(std::memory_order_acquire) == uint32_t(-1);
 296   }
 297
 298
 299  private:
 300   ///////////// types
 301
 302   struct Slot {
 303     T elem;
 304     Atom<uint32_t> localNext;
 305     Atom<uint32_t> globalNext;
 306
 307     Slot() : localNext{}, globalNext{} {}
 308   };
 309
 310   struct TaggedPtr {
 311     uint32_t idx;
 312
 313     // size is bottom 8 bits, tag in top 24.  g++'s code generation for
 314     // bitfields seems to depend on the phase of the moon, plus we can
 315     // do better because we can rely on other checks to avoid masking
 316     uint32_t tagAndSize;
 317
 318     enum : uint32_t {
 319         SizeBits = 8,
 320         SizeMask = (1U << SizeBits) - 1,
 321         TagIncr = 1U << SizeBits,
 322     };
 323
 324     uint32_t size() const {
 325       return tagAndSize & SizeMask;
 326     }
 327
 328     TaggedPtr withSize(uint32_t repl) const {
 329       assert(repl <= LocalListLimit);
 330       return TaggedPtr{ idx, (tagAndSize & ~SizeMask) | repl };
 331     }
 332
 333     TaggedPtr withSizeIncr() const {
 334       assert(size() < LocalListLimit);
 335       return TaggedPtr{ idx, tagAndSize + 1 };
 336     }
 337
 338     TaggedPtr withSizeDecr() const {
 339       assert(size() > 0);
 340       return TaggedPtr{ idx, tagAndSize - 1 };
 341     }
 342
 343     TaggedPtr withIdx(uint32_t repl) const {
 344       return TaggedPtr{ repl, tagAndSize + TagIncr };
 345     }
 346
 347     TaggedPtr withEmpty() const {
 348       return withIdx(0).withSize(0);
 349     }
 350   };
 351
 352   struct FOLLY_ALIGN_TO_AVOID_FALSE_SHARING LocalList {
 353     AtomicStruct<TaggedPtr,Atom> head;
 354
 355     LocalList() : head(TaggedPtr{}) {}
 356   };
 357
 358   ////////// fields
 359
 360   /// the number of bytes allocated from mmap, which is a multiple of
 361   /// the page size of the machine
 362   size_t mmapLength_;
 363
 364   /// the actual number of slots that we will allocate, to guarantee
 365   /// that we will satisfy the capacity requested at construction time.
 366   /// They will be numbered 1..actualCapacity_ (note the 1-based counting),
 367   /// and occupy slots_[1..actualCapacity_].
 368   uint32_t actualCapacity_;
 369
 370   /// this records the number of slots that have actually been constructed.
 371   /// To allow use of atomic ++ instead of CAS, we let this overflow.
 372   /// The actual number of constructed elements is min(actualCapacity_,
 373   /// size_)
 374   Atom<uint32_t> size_;
 375
 376   /// raw storage, only 1..min(size_,actualCapacity_) (inclusive) are
 377   /// actually constructed.  Note that slots_[0] is not constructed or used
 378   FOLLY_ALIGN_TO_AVOID_FALSE_SHARING Slot* slots_;
 379
 380   /// use AccessSpreader to find your list.  We use stripes instead of
 381   /// thread-local to avoid the need to grow or shrink on thread start
 382   /// or join.   These are heads of lists chained with localNext
 383   LocalList local_[NumLocalLists];
 384
 385   /// this is the head of a list of node chained by globalNext, that are
 386   /// themselves each the head of a list chained by localNext
 387   FOLLY_ALIGN_TO_AVOID_FALSE_SHARING AtomicStruct<TaggedPtr,Atom> globalHead_;
 388
 389   ///////////// private methods
 390
 391   uint32_t slotIndex(uint32_t idx) const {
 392     assert(0 < idx &&
 393            idx <= actualCapacity_ &&
 394            idx <= size_.load(std::memory_order_acquire));
 395     return idx;
 396   }
 397
 398   Slot& slot(uint32_t idx) {
 399     return slots_[slotIndex(idx)];
 400   }
 401
 402   const Slot& slot(uint32_t idx) const {
 403     return slots_[slotIndex(idx)];
 404   }
 405
 406   // localHead references a full list chained by localNext.  s should
 407   // reference slot(localHead), it is passed as a micro-optimization
 408   void globalPush(Slot& s, uint32_t localHead) {
 409     while (true) {
 410       TaggedPtr gh = globalHead_.load(std::memory_order_acquire);
 411       s.globalNext.store(gh.idx, std::memory_order_relaxed);
 412       if (globalHead_.compare_exchange_strong(gh, gh.withIdx(localHead))) {
 413         // success
 414         return;
 415       }
 416     }
 417   }
 418
 419   // idx references a single node
 420   void localPush(AtomicStruct<TaggedPtr,Atom>& head, uint32_t idx) {
 421     Slot& s = slot(idx);
 422     TaggedPtr h = head.load(std::memory_order_acquire);
 423     while (true) {
 424       s.localNext.store(h.idx, std::memory_order_relaxed);
 425
 426       if (h.size() == LocalListLimit) {
 427         // push will overflow local list, steal it instead
 428         if (head.compare_exchange_strong(h, h.withEmpty())) {
 429           // steal was successful, put everything in the global list
 430           globalPush(s, idx);
 431           return;
 432         }
 433       } else {
 434         // local list has space
 435         if (head.compare_exchange_strong(h, h.withIdx(idx).withSizeIncr())) {
 436           // success
 437           return;
 438         }
 439       }
 440       // h was updated by failing CAS
 441     }
 442   }
 443
 444   // returns 0 if empty
 445   uint32_t globalPop() {
 446     while (true) {
 447       TaggedPtr gh = globalHead_.load(std::memory_order_acquire);
 448       if (gh.idx == 0 ||
 449           globalHead_.compare_exchange_strong(
 450               gh,
 451               gh.withIdx(
 452                   slot(gh.idx).globalNext.load(std::memory_order_relaxed)))) {
 453         // global list is empty, or pop was successful
 454         return gh.idx;
 455       }
 456     }
 457   }
 458
 459   // returns 0 if allocation failed
 460   uint32_t localPop(AtomicStruct<TaggedPtr,Atom>& head) {
 461     while (true) {
 462       TaggedPtr h = head.load(std::memory_order_acquire);
 463       if (h.idx != 0) {
 464         // local list is non-empty, try to pop
 465         Slot& s = slot(h.idx);
 466         auto next = s.localNext.load(std::memory_order_relaxed);
 467         if (head.compare_exchange_strong(h, h.withIdx(next).withSizeDecr())) {
 468           // success
 469           return h.idx;
 470         }
 471         continue;
 472       }
 473
 474       uint32_t idx = globalPop();
 475       if (idx == 0) {
 476         // global list is empty, allocate and construct new slot
 477         if (size_.load(std::memory_order_relaxed) >= actualCapacity_ ||
 478             (idx = ++size_) > actualCapacity_) {
 479           // allocation failed
 480           return 0;
 481         }
 482         Traits::initialize(&slot(idx).elem);
 483         return idx;
 484       }
 485
 486       Slot& s = slot(idx);
 487       auto next = s.localNext.load(std::memory_order_relaxed);
 488       if (head.compare_exchange_strong(
 489               h, h.withIdx(next).withSize(LocalListLimit))) {
 490         // global list moved to local list, keep head for us
 491         return idx;
 492       }
 493       // local bulk push failed, return idx to the global list and try again
 494       globalPush(s, idx);
 495     }
 496   }
 497
 498   AtomicStruct<TaggedPtr,Atom>& localHead() {
 499     auto stripe = detail::AccessSpreader<Atom>::current(NumLocalLists);
 500     return local_[stripe].head;
 501   }
 502
 503   void markAllocated(Slot& slot) {
 504     slot.localNext.store(uint32_t(-1), std::memory_order_release);
 505   }
 506 };
 507
 508 namespace detail {
 509
 510 /// This is a stateful Deleter functor, which allows std::unique_ptr
 511 /// to track elements allocated from an IndexedMemPool by tracking the
 512 /// associated pool.  See IndexedMemPool::allocElem.
 513 template <typename Pool>
 514 struct IndexedMemPoolRecycler {
 515   Pool* pool;
 516
 517   explicit IndexedMemPoolRecycler(Pool* pool) : pool(pool) {}
 518
 519   IndexedMemPoolRecycler(const IndexedMemPoolRecycler<Pool>& rhs)
 520       = default;
 521   IndexedMemPoolRecycler& operator= (const IndexedMemPoolRecycler<Pool>& rhs)
 522       = default;
 523
 524   void operator()(typename Pool::value_type* elem) const {
 525     pool->recycleIndex(pool->locateElem(elem));
 526   }
 527 };
 528
 529 }
 530
 531 } // namespace folly
 532
 533 # pragma GCC diagnostic pop