folly/IndexedMemPool.h

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