IndexedMemPool.h

Go to the documentation of this file.

/*
 * Copyright 2014-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <assert.h>
#include <errno.h>
#include <stdint.h>

#include <type_traits>

#include <boost/noncopyable.hpp>
#include <folly/Portability.h>
#include <folly/concurrency/CacheLocality.h>
#include <folly/portability/SysMman.h>
#include <folly/portability/Unistd.h>
#include <folly/synchronization/AtomicStruct.h>

// Ignore shadowing warnings within this file, so includers can use -Wshadow.
FOLLY_PUSH_WARNING
FOLLY_GNU_DISABLE_WARNING("-Wshadow")

namespace folly {

namespace detail {
template <typename Pool>
struct IndexedMemPoolRecycler;
} // namespace detail

template <
    typename T,
    bool EagerRecycleWhenTrivial = false,
    bool EagerRecycleWhenNotTrivial = true>
struct IndexedMemPoolTraits {
  static constexpr bool eagerRecycle() {
    return std::is_trivial<T>::value ? EagerRecycleWhenTrivial
                                     : EagerRecycleWhenNotTrivial;
  }

  static void initialize(T* ptr) {
    if (!eagerRecycle()) {
      new (ptr) T();
    }
  }

  static void cleanup(T* ptr) {
    if (!eagerRecycle()) {
      ptr->~T();
    }
  }

  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)...);
    }
  }

  static void onRecycle(T* ptr) {
    if (eagerRecycle()) {
      ptr->~T();
    }
  }
};

template <typename T>
using IndexedMemPoolTraitsLazyRecycle = IndexedMemPoolTraits<T, false, false>;

template <typename T>
using IndexedMemPoolTraitsEagerRecycle = IndexedMemPoolTraits<T, true, true>;

template <
    typename T,
    uint32_t NumLocalLists_ = 32,
    uint32_t LocalListLimit_ = 200,
    template <typename> class Atom = std::atomic,
    typename Traits = IndexedMemPoolTraits<T>>
struct IndexedMemPool : boost::noncopyable {
  typedef T value_type;

  typedef std::unique_ptr<T, detail::IndexedMemPoolRecycler<IndexedMemPool>>
      UniquePtr;

  static_assert(LocalListLimit_ <= 255, "LocalListLimit must fit in 8 bits");
  enum {
    NumLocalLists = NumLocalLists_,
    LocalListLimit = LocalListLimit_,
  };

  static_assert(
      std::is_nothrow_default_constructible<Atom<uint32_t>>::value,
      "Atom must be nothrow default constructible");

  // these are public because clients may need to reason about the number
  // of bits required to hold indices from a pool, given its capacity

  static constexpr uint32_t maxIndexForCapacity(uint32_t capacity) {
    // index of std::numeric_limits<uint32_t>::max() is reserved for isAllocated
    // tracking
    return uint32_t(std::min(
        uint64_t(capacity) + (NumLocalLists - 1) * LocalListLimit,
        uint64_t(std::numeric_limits<uint32_t>::max() - 1)));
  }

  static constexpr uint32_t capacityForMaxIndex(uint32_t maxIndex) {
    return maxIndex - (NumLocalLists - 1) * LocalListLimit;
  }

  explicit IndexedMemPool(uint32_t capacity)
      : actualCapacity_(maxIndexForCapacity(capacity)),
        size_(0),
        globalHead_(TaggedPtr{}) {
    const size_t needed = sizeof(Slot) * (actualCapacity_ + 1);
    size_t pagesize = size_t(sysconf(_SC_PAGESIZE));
    mmapLength_ = ((needed - 1) & ~(pagesize - 1)) + pagesize;
    assert(needed <= mmapLength_ && mmapLength_ < needed + pagesize);
    assert((mmapLength_ % pagesize) == 0);

    slots_ = static_cast<Slot*>(mmap(
        nullptr,
        mmapLength_,
        PROT_READ | PROT_WRITE,
        MAP_PRIVATE | MAP_ANONYMOUS,
        -1,
        0));
    if (slots_ == MAP_FAILED) {
      assert(errno == ENOMEM);
      throw std::bad_alloc();
    }
    // Atom is expected to be std::atomic, which is trivially
    // constructible, so we can avoid explicitly initializing the
    // slots which would cause the whole mapped area to be paged in.
    //
    // TODO: Switch to is_trivially_constructible once support
    // for GCC 4.9 is dropped for this file.
    if /* constexpr */ (!std::is_trivial<Atom<uint32_t>>::value) {
      for (size_t i = 1; i < actualCapacity_ + 1; i++) {
        // Atom is enforced above to be nothrow-default-constructible
        new (&slots_[i].localNext) Atom<uint32_t>;
        new (&slots_[i].globalNext) Atom<uint32_t>;
      }
    }
  }

  ~IndexedMemPool() {
    using A = Atom<uint32_t>;
    for (uint32_t i = maxAllocatedIndex(); i > 0; --i) {
      Traits::cleanup(&slots_[i].elem);
    }
    for (size_t i = 1; i < actualCapacity_ + 1; i++) {
      slots_[i].localNext.~A();
      slots_[i].globalNext.~A();
    }
    munmap(slots_, mmapLength_);
  }

  uint32_t capacity() {
    return capacityForMaxIndex(actualCapacity_);
  }

  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_));
  }

  template <typename... Args>
  uint32_t allocIndex(Args&&... args) {
    auto idx = localPop(localHead());
    if (idx != 0) {
      Slot& s = slot(idx);
      Traits::onAllocate(&s.elem, std::forward<Args>(args)...);
      markAllocated(s);
    }
    return idx;
  }

  template <typename... Args>
  UniquePtr allocElem(Args&&... args) {
    auto idx = allocIndex(std::forward<Args>(args)...);
    T* ptr = idx == 0 ? nullptr : &slot(idx).elem;
    return UniquePtr(ptr, typename UniquePtr::deleter_type(this));
  }

  void recycleIndex(uint32_t idx) {
    assert(isAllocated(idx));
    localPush(localHead(), idx);
  }

  T& operator[](uint32_t idx) {
    return slot(idx).elem;
  }

  const T& operator[](uint32_t idx) const {
    return slot(idx).elem;
  }

  uint32_t locateElem(const T* elem) const {
    if (!elem) {
      return 0;
    }

    static_assert(std::is_standard_layout<Slot>::value, "offsetof needs POD");

    auto slot = reinterpret_cast<const Slot*>(
        reinterpret_cast<const char*>(elem) - offsetof(Slot, elem));
    auto rv = uint32_t(slot - slots_);

    // this assert also tests that rv is in range
    assert(elem == &(*this)[rv]);
    return rv;
  }

  bool isAllocated(uint32_t idx) const {
    return slot(idx).localNext.load(std::memory_order_acquire) == uint32_t(-1);
  }

 private:

  struct Slot {
    T elem;
    Atom<uint32_t> localNext;
    Atom<uint32_t> globalNext;

    Slot() : localNext{}, globalNext{} {}
  };

  struct TaggedPtr {
    uint32_t idx;

    // size is bottom 8 bits, tag in top 24.  g++'s code generation for
    // bitfields seems to depend on the phase of the moon, plus we can
    // do better because we can rely on other checks to avoid masking
    uint32_t tagAndSize;

    enum : uint32_t {
      SizeBits = 8,
      SizeMask = (1U << SizeBits) - 1,
      TagIncr = 1U << SizeBits,
    };

    uint32_t size() const {
      return tagAndSize & SizeMask;
    }

    TaggedPtr withSize(uint32_t repl) const {
      assert(repl <= LocalListLimit);
      return TaggedPtr{idx, (tagAndSize & ~SizeMask) | repl};
    }

    TaggedPtr withSizeIncr() const {
      assert(size() < LocalListLimit);
      return TaggedPtr{idx, tagAndSize + 1};
    }

    TaggedPtr withSizeDecr() const {
      assert(size() > 0);
      return TaggedPtr{idx, tagAndSize - 1};
    }

    TaggedPtr withIdx(uint32_t repl) const {
      return TaggedPtr{repl, tagAndSize + TagIncr};
    }

    TaggedPtr withEmpty() const {
      return withIdx(0).withSize(0);
    }
  };

  struct alignas(hardware_destructive_interference_size) LocalList {
    AtomicStruct<TaggedPtr, Atom> head;

    LocalList() : head(TaggedPtr{}) {}
  };


  size_t mmapLength_;

  uint32_t actualCapacity_;

  Atom<uint32_t> size_;

  alignas(hardware_destructive_interference_size) Slot* slots_;

  LocalList local_[NumLocalLists];

  alignas(hardware_destructive_interference_size)
      AtomicStruct<TaggedPtr, Atom> globalHead_;


  uint32_t slotIndex(uint32_t idx) const {
    assert(
        0 < idx && idx <= actualCapacity_ &&
        idx <= size_.load(std::memory_order_acquire));
    return idx;
  }

  Slot& slot(uint32_t idx) {
    return slots_[slotIndex(idx)];
  }

  const Slot& slot(uint32_t idx) const {
    return slots_[slotIndex(idx)];
  }

  // localHead references a full list chained by localNext.  s should
  // reference slot(localHead), it is passed as a micro-optimization
  void globalPush(Slot& s, uint32_t localHead) {
    while (true) {
      TaggedPtr gh = globalHead_.load(std::memory_order_acquire);
      s.globalNext.store(gh.idx, std::memory_order_relaxed);
      if (globalHead_.compare_exchange_strong(gh, gh.withIdx(localHead))) {
        // success
        return;
      }
    }
  }

  // idx references a single node
  void localPush(AtomicStruct<TaggedPtr, Atom>& head, uint32_t idx) {
    Slot& s = slot(idx);
    TaggedPtr h = head.load(std::memory_order_acquire);
    while (true) {
      s.localNext.store(h.idx, std::memory_order_release);
      Traits::onRecycle(&slot(idx).elem);

      if (h.size() == LocalListLimit) {
        // push will overflow local list, steal it instead
        if (head.compare_exchange_strong(h, h.withEmpty())) {
          // steal was successful, put everything in the global list
          globalPush(s, idx);
          return;
        }
      } else {
        // local list has space
        if (head.compare_exchange_strong(h, h.withIdx(idx).withSizeIncr())) {
          // success
          return;
        }
      }
      // h was updated by failing CAS
    }
  }

  // returns 0 if empty
  uint32_t globalPop() {
    while (true) {
      TaggedPtr gh = globalHead_.load(std::memory_order_acquire);
      if (gh.idx == 0 ||
          globalHead_.compare_exchange_strong(
              gh,
              gh.withIdx(
                  slot(gh.idx).globalNext.load(std::memory_order_relaxed)))) {
        // global list is empty, or pop was successful
        return gh.idx;
      }
    }
  }

  // returns 0 if allocation failed
  uint32_t localPop(AtomicStruct<TaggedPtr, Atom>& head) {
    while (true) {
      TaggedPtr h = head.load(std::memory_order_acquire);
      if (h.idx != 0) {
        // local list is non-empty, try to pop
        Slot& s = slot(h.idx);
        auto next = s.localNext.load(std::memory_order_relaxed);
        if (head.compare_exchange_strong(h, h.withIdx(next).withSizeDecr())) {
          // success
          return h.idx;
        }
        continue;
      }

      uint32_t idx = globalPop();
      if (idx == 0) {
        // global list is empty, allocate and construct new slot
        if (size_.load(std::memory_order_relaxed) >= actualCapacity_ ||
            (idx = ++size_) > actualCapacity_) {
          // allocation failed
          return 0;
        }
        Traits::initialize(&slot(idx).elem);
        return idx;
      }

      Slot& s = slot(idx);
      auto next = s.localNext.load(std::memory_order_relaxed);
      if (head.compare_exchange_strong(
              h, h.withIdx(next).withSize(LocalListLimit))) {
        // global list moved to local list, keep head for us
        return idx;
      }
      // local bulk push failed, return idx to the global list and try again
      globalPush(s, idx);
    }
  }

  AtomicStruct<TaggedPtr, Atom>& localHead() {
    auto stripe = AccessSpreader<Atom>::current(NumLocalLists);
    return local_[stripe].head;
  }

  void markAllocated(Slot& slot) {
    slot.localNext.store(uint32_t(-1), std::memory_order_release);
  }

 public:
  static constexpr std::size_t kSlotSize = sizeof(Slot);
};

namespace detail {

template <typename Pool>
struct IndexedMemPoolRecycler {
  Pool* pool;

  explicit IndexedMemPoolRecycler(Pool* pool) : pool(pool) {}

  IndexedMemPoolRecycler(const IndexedMemPoolRecycler<Pool>& rhs) = default;
  IndexedMemPoolRecycler& operator=(const IndexedMemPoolRecycler<Pool>& rhs) =
      default;

  void operator()(typename Pool::value_type* elem) const {
    pool->recycleIndex(pool->locateElem(elem));
  }
};

} // namespace detail

} // namespace folly

FOLLY_POP_WARNING