concurrent_object_arena.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <dispenso/platform.h>
 
#include <algorithm>
#include <atomic>
#include <cassert>
#include <cstring>
#include <mutex>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace detail {
 
template <class T>
constexpr T log2i(const T v) {
  T log2 = 0, val = v;
  while (val >>= 1)
    ++log2;
  return log2;
}
 
} // namespace detail
 
namespace dispenso {
template <class T, class Index = size_t, size_t alignment = dispenso::kCacheLineSize>
struct ConcurrentObjectArena {
  ConcurrentObjectArena() = delete;
  explicit ConcurrentObjectArena(const Index minBuffSize, const Index initialSize = 0)
      : kLog2BuffSize(
            ::detail::log2i(minBuffSize) +
            ((Index{1} << ::detail::log2i(minBuffSize)) == minBuffSize ? 0 : 1)),
        kBufferSize(Index{1} << kLog2BuffSize),
        kMask((Index{1} << kLog2BuffSize) - 1),
        pos_(0),
        allocatedSize_(0),
        buffers_(nullptr),
        buffersSize_(0),
        buffersPos_(0) {
    allocateBuffer();
    allocatedSize_.store(kBufferSize, std::memory_order_relaxed);
    if (initialSize > 0) {
      grow_by(initialSize);
    }
  }
 
  ConcurrentObjectArena(const ConcurrentObjectArena<T, Index, alignment>& other)
      : kLog2BuffSize(other.kLog2BuffSize),
        kBufferSize(other.kBufferSize),
        kMask(other.kMask),
        pos_(other.pos_.load(std::memory_order_relaxed)),
        allocatedSize_(other.allocatedSize_.load(std::memory_order_relaxed)),
        buffersSize_(other.buffersSize_),
        buffersPos_(other.buffersPos_),
        gaps_(other.gaps_) {
    static_assert(
        std::is_trivially_copyable<T>::value,
        "ConcurrentObjectArena copy constructor uses memcpy; T must be trivially copyable.");
    T** otherBuffers = other.buffers_.load(std::memory_order_acquire);
    T** newBuffers = new T*[buffersSize_];
    // Only the first buffersPos_ slots of the pointer table hold buffers; the rest are
    // uninitialized spare capacity, filled in by allocateBuffer() as the arena grows.
    for (Index i = 0; i < buffersPos_; ++i) {
      void* ptr = detail::alignedMallocOrThrow(kBufferSize * sizeof(T), alignment);
      std::memcpy(ptr, otherBuffers[i], kBufferSize * sizeof(T));
      newBuffers[i] = static_cast<T*>(ptr);
    }
    buffers_.store(newBuffers, std::memory_order_release);
  }
 
  ConcurrentObjectArena(ConcurrentObjectArena<T, Index, alignment>&& other) noexcept
      : kLog2BuffSize(0),
        kBufferSize(0),
        kMask(0),
        pos_(0),
        allocatedSize_(0),
        buffers_(nullptr),
        buffersSize_(0),
        buffersPos_(0) {
    swap(*this, other);
  }
 
  ~ConcurrentObjectArena() {
    T** buffers = buffers_.load(std::memory_order_acquire);
 
    destroyObjects(buffers);
 
    for (Index i = 0; i < buffersPos_; i++)
      detail::alignedFree(buffers[i]);
 
    delete[] buffers;
 
    for (T** p : deleteLater_)
      delete[] p;
  }
 
  ConcurrentObjectArena<T, Index, alignment>& operator=(
      ConcurrentObjectArena<T, Index, alignment> const& other) {
    ConcurrentObjectArena<T, Index, alignment> copy(other);
    swap(*this, copy);
    return *this;
  }
 
  ConcurrentObjectArena<T, Index, alignment>& operator=(
      ConcurrentObjectArena<T, Index, alignment>&& other) noexcept {
    swap(*this, other);
    return *this;
  }
  Index grow_by(const Index delta) {
    Index newPos;
    Index oldPos = pos_.load(std::memory_order_relaxed);
 
    do {
      Index curSize = allocatedSize_.load(std::memory_order_acquire);
 
      if (oldPos + delta >= curSize) {
        const std::lock_guard<std::mutex> guard(resizeMutex_);
        curSize = allocatedSize_.load(std::memory_order_relaxed);
        while (oldPos + delta >= curSize) {
          allocateBuffer();
          allocatedSize_.store(curSize + kBufferSize, std::memory_order_release);
          curSize = curSize + kBufferSize;
        }
      }
 
      newPos = oldPos + delta;
    } while (!std::atomic_compare_exchange_weak_explicit(
        &pos_, &oldPos, newPos, std::memory_order_release, std::memory_order_relaxed));
 
    Index built = oldPos;
#if defined(__cpp_exceptions)
    try {
      constructObjects(oldPos, newPos, built);
    } catch (...) {
      destroyRange(buffers_.load(std::memory_order_acquire), oldPos, built);
      recordGap(oldPos, newPos);
      throw;
    }
#else
    constructObjects(oldPos, newPos, built);
#endif // __cpp_exceptions
 
    return oldPos;
  }
 
  inline const T& operator[](const Index index) const {
    const Index bufIndex = index >> kLog2BuffSize;
    const Index i = index & kMask;
 
    return buffers_.load(std::memory_order_acquire)[bufIndex][i];
  }
 
  inline T& operator[](const Index index) {
    return const_cast<T&>(
        const_cast<const ConcurrentObjectArena<T, Index, alignment>&>(*this)[index]);
  }
 
  Index size() const {
    return pos_.load(std::memory_order_relaxed);
  }
 
  Index capacity() const {
    return allocatedSize_.load(std::memory_order_relaxed);
  }
 
  Index numBuffers() const {
    return buffersPos_;
  }
 
  const T* getBuffer(const Index index) const {
    return buffers_.load(std::memory_order_acquire)[index];
  }
 
  T* getBuffer(const Index index) {
    return buffers_.load(std::memory_order_acquire)[index];
  }
 
  Index getBufferSize(const Index index) const {
    const Index numBuffs = numBuffers();
    assert(index < numBuffs);
 
    if (index < numBuffs - 1)
      return kBufferSize;
    else
      return pos_.load(std::memory_order_relaxed) - (kBufferSize * (numBuffs - 1));
  }
 
  friend void swap(
      ConcurrentObjectArena<T, Index, alignment>& lhs,
      ConcurrentObjectArena<T, Index, alignment>& rhs) noexcept {
    using std::swap;
 
    swap(lhs.kLog2BuffSize, rhs.kLog2BuffSize);
    swap(lhs.kBufferSize, rhs.kBufferSize);
    swap(lhs.kMask, rhs.kMask);
 
    const Index rhs_pos = rhs.pos_.load(std::memory_order_relaxed);
    rhs.pos_.store(lhs.pos_.load(std::memory_order_relaxed), std::memory_order_relaxed);
    lhs.pos_.store(rhs_pos, std::memory_order_relaxed);
 
    const Index rhs_allocatedSize = rhs.allocatedSize_.load(std::memory_order_relaxed);
    rhs.allocatedSize_.store(
        lhs.allocatedSize_.load(std::memory_order_relaxed), std::memory_order_relaxed);
    lhs.allocatedSize_.store(rhs_allocatedSize, std::memory_order_relaxed);
 
    T** const rhs_buffers = rhs.buffers_.load(std::memory_order_acquire);
    rhs.buffers_.store(lhs.buffers_.load(std::memory_order_acquire), std::memory_order_release);
    lhs.buffers_.store(rhs_buffers, std::memory_order_release);
 
    swap(lhs.buffersSize_, rhs.buffersSize_);
    swap(lhs.buffersPos_, rhs.buffersPos_);
    swap(lhs.deleteLater_, rhs.deleteLater_);
    swap(lhs.gaps_, rhs.gaps_);
  }
 
 private:
  void allocateBuffer() {
    void* ptr = detail::alignedMallocOrThrow(kBufferSize * sizeof(T), alignment);
 
    if (buffersPos_ < buffersSize_) {
      buffers_.load(std::memory_order_acquire)[buffersPos_++] = static_cast<T*>(ptr);
    } else {
      const Index oldBuffersSize = buffersSize_;
      T** oldBuffers = buffers_.load(std::memory_order_acquire);
 
      buffersSize_ = oldBuffersSize == 0 ? 2 : oldBuffersSize * 2;
      T** newBuffers = new T*[buffersSize_];
 
      if (oldBuffers != nullptr) {
        std::memcpy(newBuffers, oldBuffers, sizeof(T*) * oldBuffersSize);
        deleteLater_.push_back(oldBuffers);
      }
 
      newBuffers[buffersPos_++] = static_cast<T*>(ptr);
      buffers_.store(newBuffers, std::memory_order_release);
    }
  }
 
  // Constructs [beginIndex, endIndex). built tracks one past the last constructed element, so a
  // caller can destroy exactly those if a constructor throws.
  void constructObjects(const Index beginIndex, const Index endIndex, Index& built) {
    const Index startBuffer = beginIndex >> kLog2BuffSize;
    const Index endBuffer = endIndex >> kLog2BuffSize;
 
    Index bufStart = beginIndex & kMask;
    for (Index b = startBuffer; b <= endBuffer; ++b) {
      T* buf = buffers_.load(std::memory_order_acquire)[b];
      const Index bufEnd = b == endBuffer ? (endIndex & kMask) : kBufferSize;
      for (Index i = bufStart; i < bufEnd; ++i) {
        new (buf + i) T();
        built = (b << kLog2BuffSize) + i + 1;
      }
 
      bufStart = 0;
    }
  }
 
  void destroyRange(T** buffers, Index beginIndex, const Index endIndex) {
    for (; beginIndex < endIndex; ++beginIndex) {
      buffers[beginIndex >> kLog2BuffSize][beginIndex & kMask].~T();
    }
  }
 
  // Destroys every element, skipping gaps left by constructors that threw. Not concurrency safe.
  void destroyObjects(T** buffers) {
    if (std::is_trivially_destructible<T>::value || buffers == nullptr) {
      return;
    }
    std::sort(gaps_.begin(), gaps_.end());
    const Index end = pos_.load(std::memory_order_relaxed);
    Index next = 0;
    for (const std::pair<Index, Index>& gap : gaps_) {
      destroyRange(buffers, next, gap.first);
      next = gap.second;
    }
    destroyRange(buffers, next, end);
  }
 
  // Records [beginIndex, endIndex), reserved by a grow_by() whose constructor threw, as holding no
  // objects.
  void recordGap(const Index beginIndex, const Index endIndex) {
    const std::lock_guard<std::mutex> guard(resizeMutex_);
    gaps_.emplace_back(beginIndex, endIndex);
  }
  //
  //    kBufferSize = 2^kLog2BuffSize
  //    mask = 0b00011111
  //                ──┬──
  //                  └─number of 1s is log2BuffSize
  //
  std::mutex resizeMutex_;
 
  Index kLog2BuffSize;
  Index kBufferSize;
  Index kMask;
 
  std::atomic<Index> pos_;
  std::atomic<Index> allocatedSize_;
 
  std::atomic<T**> buffers_;
  Index buffersSize_;
  Index buffersPos_;
  std::vector<T**> deleteLater_;
  // Index ranges reserved by a grow_by() whose constructor threw. They hold no objects.
  std::vector<std::pair<Index, Index>> gaps_;
};
 
} // namespace dispenso