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