spsc_ring_buffer.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 <atomic>
#include <cstddef>
#include <new>
#include <type_traits>
#include <utility>
 
#include <dispenso/platform.h>
#include <dispenso/util.h>
 
namespace dispenso {
 
// TODO(bbudge): Add kDynamicCapacity specialization that heap-allocates storage
// with runtime-determined capacity.
 
template <typename T, size_t Capacity = 16, bool RoundUpToPowerOfTwo = true>
class SPSCRingBuffer {
  static_assert(Capacity >= 1, "SPSCRingBuffer capacity must be at least 1");
  static_assert(
      std::is_move_constructible<T>::value,
      "SPSCRingBuffer element type must be move-constructible");
 
 public:
  using value_type = T;
 
  using size_type = size_t;
 
  SPSCRingBuffer() = default;
 
  SPSCRingBuffer(const SPSCRingBuffer&) = delete;
 
  SPSCRingBuffer& operator=(const SPSCRingBuffer&) = delete;
 
  SPSCRingBuffer(SPSCRingBuffer&&) = delete;
 
  SPSCRingBuffer& operator=(SPSCRingBuffer&&) = delete;
 
  ~SPSCRingBuffer() {
    // Destroy any remaining elements
    size_t head = head_.load(std::memory_order_relaxed);
    size_t tail = tail_.load(std::memory_order_relaxed);
    while (head != tail) {
      elementAt(head)->~T();
      head = increment(head);
    }
  }
 
  bool try_push(T&& item) {
    const size_t currentTail = tail_.load(std::memory_order_relaxed);
    const size_t nextTail = increment(currentTail);
 
    // Check if buffer is full
    if (nextTail == head_.load(std::memory_order_acquire)) {
      return false;
    }
 
    // Construct element in-place
    new (elementAt(currentTail)) T(std::move(item));
    tail_.store(nextTail, std::memory_order_release);
    return true;
  }
 
  bool try_push(const T& item) {
    const size_t currentTail = tail_.load(std::memory_order_relaxed);
    const size_t nextTail = increment(currentTail);
 
    // Check if buffer is full
    if (nextTail == head_.load(std::memory_order_acquire)) {
      return false;
    }
 
    // Construct element in-place via copy
    new (elementAt(currentTail)) T(item);
    tail_.store(nextTail, std::memory_order_release);
    return true;
  }
 
  template <typename... Args>
  bool try_emplace(Args&&... args) {
    const size_t currentTail = tail_.load(std::memory_order_relaxed);
    const size_t nextTail = increment(currentTail);
 
    // Check if buffer is full
    if (nextTail == head_.load(std::memory_order_acquire)) {
      return false;
    }
 
    // Construct element in-place
    new (elementAt(currentTail)) T(std::forward<Args>(args)...);
    tail_.store(nextTail, std::memory_order_release);
    return true;
  }
 
  bool try_pop(T& item) {
    const size_t currentHead = head_.load(std::memory_order_relaxed);
 
    // Check if buffer is empty
    if (currentHead == tail_.load(std::memory_order_acquire)) {
      return false;
    }
 
    T* elem = elementAt(currentHead);
    item = std::move(*elem);
    elem->~T();
    head_.store(increment(currentHead), std::memory_order_release);
    return true;
  }
 
  OpResult<T> try_pop() {
    const size_t currentHead = head_.load(std::memory_order_relaxed);
 
    // Check if buffer is empty
    if (currentHead == tail_.load(std::memory_order_acquire)) {
      return {};
    }
 
    T* elem = elementAt(currentHead);
    OpResult<T> result(std::move(*elem));
    elem->~T();
    head_.store(increment(currentHead), std::memory_order_release);
    return result;
  }
 
  bool try_pop_into(T* storage) {
    const size_t currentHead = head_.load(std::memory_order_relaxed);
 
    // Check if buffer is empty
    if (currentHead == tail_.load(std::memory_order_acquire)) {
      return false;
    }
 
    T* elem = elementAt(currentHead);
    new (storage) T(std::move(*elem));
    elem->~T();
    head_.store(increment(currentHead), std::memory_order_release);
    return true;
  }
 
  template <typename InputIt>
  size_type try_push_batch(InputIt first, InputIt last) {
    const size_t currentTail = tail_.load(std::memory_order_relaxed);
    const size_t currentHead = head_.load(std::memory_order_acquire);
 
    // Calculate available space (actual capacity is kBufferSize - 1)
    size_t available;
    if (currentTail >= currentHead) {
      // Tail is ahead of or at head: available = capacity - (tail - head)
      available = (kBufferSize - 1) - (currentTail - currentHead);
    } else {
      // Tail has wrapped: available = head - tail - 1
      available = currentHead - currentTail - 1;
    }
 
    if (available == 0) {
      return 0;
    }
 
    // Push as many items as possible. The new tail is published on scope exit, so if a T
    // constructor throws, the elements already constructed are still handed to the consumer
    // (and destroyed with the buffer) rather than being stranded past the tail.
    size_t count = 0;
    IndexPublisher tail{tail_, currentTail};
    for (; first != last && count < available; ++first, ++count) {
      new (elementAt(tail.pos)) T(std::move(*first));
      tail.pos = increment(tail.pos);
    }
    return count;
  }
 
  template <typename OutputIt>
  size_type try_pop_batch(OutputIt dest, size_type maxCount) {
    const size_t currentHead = head_.load(std::memory_order_relaxed);
    const size_t currentTail = tail_.load(std::memory_order_acquire);
 
    // Calculate available items
    size_t available;
    if (currentTail >= currentHead) {
      available = currentTail - currentHead;
    } else {
      available = kBufferSize - currentHead + currentTail;
    }
 
    if (available == 0) {
      return 0;
    }
 
    // Pop as many items as requested and available
    size_t count = std::min(available, maxCount);
    // The new head is published on scope exit, so if assigning to dest throws, the elements
    // already moved out and destroyed are released, and the remaining ones stay in the buffer.
    IndexPublisher head{head_, currentHead};
    for (size_t i = 0; i < count; ++i, ++dest) {
      T* elem = elementAt(head.pos);
      *dest = std::move(*elem);
      elem->~T();
      head.pos = increment(head.pos);
    }
    return count;
  }
 
  bool empty() const {
    return head_.load(std::memory_order_acquire) == tail_.load(std::memory_order_acquire);
  }
 
  bool full() const {
    return increment(tail_.load(std::memory_order_acquire)) ==
        head_.load(std::memory_order_acquire);
  }
 
  size_type size() const {
    const size_t head = head_.load(std::memory_order_acquire);
    const size_t tail = tail_.load(std::memory_order_acquire);
    // Handle wrap-around: if tail < head, we've wrapped
    return (tail >= head) ? (tail - head) : (kBufferSize - head + tail);
  }
 
  static constexpr size_type capacity() noexcept {
    return kBufferSize - 1;
  }
 
 private:
  struct IndexPublisher {
    std::atomic<size_t>& index;
    size_t start;
    size_t pos;
    IndexPublisher(std::atomic<size_t>& idx, size_t current)
        : index(idx), start(current), pos(current) {}
    IndexPublisher(const IndexPublisher&) = delete;
    IndexPublisher& operator=(const IndexPublisher&) = delete;
    ~IndexPublisher() {
      if (pos != start) {
        index.store(pos, std::memory_order_release);
      }
    }
  };
 
  static constexpr size_t computeBufferSize() noexcept {
    return RoundUpToPowerOfTwo ? static_cast<size_t>(detail::nextPow2(Capacity + 1)) : Capacity + 1;
  }
 
  static constexpr size_t kBufferSize = computeBufferSize();
 
  static constexpr bool kIsPowerOfTwo = (kBufferSize & (kBufferSize - 1)) == 0;
 
  static constexpr size_t kMask = kBufferSize - 1;
 
  static constexpr size_t increment(size_t index) noexcept {
    return kIsPowerOfTwo ? ((index + 1) & kMask) : ((index + 1) % kBufferSize);
  }
 
  T* elementAt(size_t index) {
    return reinterpret_cast<T*>(&storage_[index * sizeof(T)]);
  }
 
  const T* elementAt(size_t index) const {
    return reinterpret_cast<const T*>(&storage_[index * sizeof(T)]);
  }
 
  alignas(kCacheLineSize) std::atomic<size_t> head_{0};
 
  alignas(kCacheLineSize) std::atomic<size_t> tail_{0};
 
  alignas(T) char storage_[sizeof(T) * kBufferSize];
};
 
} // namespace dispenso