spsc_queue_base.hpp

#pragma once
 
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <limits>
#include <new>
 
#include "libxr_def.hpp"
#include "libxr_mem.hpp"
 
namespace LibXR
{
class alignas(LibXR::CONCURRENCY_ALIGNMENT) SPSCQueueBase
{
 public:
  using IndexType = size_t;  
 
  SPSCQueueBase(size_t element_size, size_t capacity)
      : element_size_(element_size),
        capacity_(capacity),
        payload_alloc_align_(alignof(std::max_align_t)),
        payload_stride_(ComputeStride(element_size, alignof(std::byte))),
        payloads_(nullptr),
        head_(0),
        tail_(0)
  {
    InitStorage();
  }
 
  SPSCQueueBase(size_t element_size, size_t element_align, size_t capacity)
      : element_size_(element_size),
        capacity_(capacity),
        payload_alloc_align_(alignof(std::max_align_t)),
        payload_stride_(ComputeStride(element_size, element_align)),
        payloads_(nullptr),
        head_(0),
        tail_(0)
  {
    InitStorage();
  }
 
 private:
  void InitStorage()
  {
    // payload_alloc_align_ 是由 alignof(std::max_align_t) 初始化的 const 成员。
    // payload_alloc_align_ is a const member initialised from alignof(std::max_align_t).
    static_assert(alignof(std::max_align_t) > 0 &&
                  (alignof(std::max_align_t) & (alignof(std::max_align_t) - 1)) == 0);
    ASSERT(element_size_ > 0);
    ASSERT(capacity_ > 0);
    ASSERT(capacity_ <= std::numeric_limits<size_t>::max() - 1);
 
    const size_t payload_bytes = MultiplyChecked(payload_stride_, RingCapacity());
    payloads_ = static_cast<std::byte*>(
        ::operator new[](payload_bytes, std::align_val_t(payload_alloc_align_)));
  }
 
 public:
  ~SPSCQueueBase()
  {
    ::operator delete[](payloads_, std::align_val_t(payload_alloc_align_));
  }
 
  ErrorCode PushBytes(const void* value)
  {
    if (value == nullptr)
    {
      return ErrorCode::PTR_NULL;
    }
 
    const auto current_tail = tail_.load(std::memory_order_relaxed);
    const auto next_tail = Increment(current_tail);
 
    if (next_tail == head_.load(std::memory_order_acquire))
    {
      return ErrorCode::FULL;
    }
 
    LibXR::Memory::FastCopy(PayloadPtr(current_tail), value, element_size_);
    tail_.store(next_tail, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  ErrorCode PopBytes(void* value = nullptr)
  {
    const auto current_head = head_.load(std::memory_order_relaxed);
 
    if (current_head == tail_.load(std::memory_order_acquire))
    {
      return ErrorCode::EMPTY;
    }
 
    if (value != nullptr)
    {
      LibXR::Memory::FastCopy(value, PayloadPtr(current_head), element_size_);
    }
 
    head_.store(Increment(current_head), std::memory_order_release);
    return ErrorCode::OK;
  }
 
  ErrorCode PeekBytes(void* value)
  {
    if (value == nullptr)
    {
      return ErrorCode::PTR_NULL;
    }
 
    const auto current_head = head_.load(std::memory_order_relaxed);
    if (current_head == tail_.load(std::memory_order_acquire))
    {
      return ErrorCode::EMPTY;
    }
 
    LibXR::Memory::FastCopy(value, PayloadPtr(current_head), element_size_);
    return ErrorCode::OK;
  }
 
  ErrorCode PushBatchBytes(const void* data, size_t count)
  {
    if (count == 0U)
    {
      return ErrorCode::OK;
    }
    if (data == nullptr)
    {
      return ErrorCode::PTR_NULL;
    }
 
    const auto current_tail = tail_.load(std::memory_order_relaxed);
    const auto current_head = head_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t free_space = (current_tail >= current_head)
                                  ? (capacity - (current_tail - current_head) - 1)
                                  : (current_head - current_tail - 1);
 
    if (free_space < count)
    {
      return ErrorCode::FULL;
    }
 
    const auto* src = static_cast<const std::byte*>(data);
    for (size_t index = 0; index < count; ++index)
    {
      LibXR::Memory::FastCopy(PayloadPtr((current_tail + index) % capacity),
                              src + index * element_size_, element_size_);
    }
 
    tail_.store((current_tail + count) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Writer>
  ErrorCode PushBytesWithWriter(size_t count, Writer&& writer)
  {
    if (count == 0U)
    {
      return ErrorCode::OK;
    }
 
    const auto current_tail = tail_.load(std::memory_order_relaxed);
    const auto current_head = head_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t free_space = (current_tail >= current_head)
                                  ? (capacity - (current_tail - current_head) - 1)
                                  : (current_head - current_tail - 1);
 
    if (free_space < count)
    {
      return ErrorCode::FULL;
    }
 
    const size_t first_chunk = std::min(count, capacity - current_tail);
    Writer& writer_ref = writer;
    const ErrorCode first_ec = writer_ref(PayloadPtr(current_tail), first_chunk);
    if (first_ec != ErrorCode::OK)
    {
      return first_ec;
    }
 
    if (count > first_chunk)
    {
      const ErrorCode second_ec = writer_ref(PayloadPtr(0), count - first_chunk);
      if (second_ec != ErrorCode::OK)
      {
        return second_ec;
      }
    }
 
    tail_.store((current_tail + count) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Writer>
  size_t ProduceWithWriter(size_t limit, Writer&& writer)
  {
    if (limit == 0U)
    {
      return 0U;
    }
 
    const auto current_tail = tail_.load(std::memory_order_relaxed);
    const auto current_head = head_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t free_space = (current_tail >= current_head)
                                  ? (capacity - (current_tail - current_head) - 1U)
                                  : (current_head - current_tail - 1U);
    const size_t offered = std::min(limit, free_space);
    if (offered == 0U)
    {
      return 0U;
    }
 
    const size_t first_count = std::min(offered, capacity - current_tail);
    const size_t second_count = offered - first_count;
    void* const second = second_count == 0U ? nullptr : PayloadPtr(0U);
    const size_t produced =
        writer(PayloadPtr(current_tail), first_count, second, second_count);
    ASSERT(produced <= offered);
    if (produced != 0U)
    {
      tail_.store((current_tail + produced) % capacity, std::memory_order_release);
    }
    return produced;
  }
 
  ErrorCode PopBatchBytes(void* data, size_t count)
  {
    if (count == 0U)
    {
      return ErrorCode::OK;
    }
 
    const auto current_head = head_.load(std::memory_order_relaxed);
    const auto current_tail = tail_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t available = (current_tail >= current_head)
                                 ? (current_tail - current_head)
                                 : (capacity - current_head + current_tail);
 
    if (available < count)
    {
      return ErrorCode::EMPTY;
    }
 
    auto* dst = static_cast<std::byte*>(data);
    if (dst != nullptr)
    {
      for (size_t index = 0; index < count; ++index)
      {
        LibXR::Memory::FastCopy(dst + index * element_size_,
                                PayloadPtr((current_head + index) % capacity),
                                element_size_);
      }
    }
 
    head_.store((current_head + count) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Reader>
  ErrorCode PopBytesWithReader(size_t count, Reader&& reader)
  {
    if (count == 0U)
    {
      return ErrorCode::OK;
    }
 
    const auto current_head = head_.load(std::memory_order_relaxed);
    const auto current_tail = tail_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t available = (current_tail >= current_head)
                                 ? (current_tail - current_head)
                                 : (capacity - current_head + current_tail);
 
    if (available < count)
    {
      return ErrorCode::EMPTY;
    }
 
    const size_t first_chunk = std::min(count, capacity - current_head);
    Reader& reader_ref = reader;
    const ErrorCode first_ec = reader_ref(PayloadPtr(current_head), first_chunk);
    if (first_ec != ErrorCode::OK)
    {
      return first_ec;
    }
 
    if (count > first_chunk)
    {
      const ErrorCode second_ec = reader_ref(PayloadPtr(0), count - first_chunk);
      if (second_ec != ErrorCode::OK)
      {
        return second_ec;
      }
    }
 
    head_.store((current_head + count) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Reader>
  size_t ConsumeWithReader(size_t limit, Reader&& reader)
  {
    if (limit == 0U)
    {
      return 0U;
    }
 
    const auto current_head = head_.load(std::memory_order_relaxed);
    const auto current_tail = tail_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t available = (current_tail >= current_head)
                                 ? (current_tail - current_head)
                                 : (capacity - current_head + current_tail);
    const size_t offered = std::min(limit, available);
    if (offered == 0U)
    {
      return 0U;
    }
 
    const size_t first_count = std::min(offered, capacity - current_head);
    const size_t second_count = offered - first_count;
    const void* const second = second_count == 0U ? nullptr : PayloadPtr(0U);
    const size_t accepted =
        reader(PayloadPtr(current_head), first_count, second, second_count);
    ASSERT(accepted <= offered);
    if (accepted != 0U)
    {
      head_.store((current_head + accepted) % capacity, std::memory_order_release);
    }
    return accepted;
  }
 
  ErrorCode PeekBatchBytes(void* data, size_t count)
  {
    if (count == 0U)
    {
      return ErrorCode::OK;
    }
    if (data == nullptr)
    {
      return ErrorCode::PTR_NULL;
    }
 
    const auto current_head = head_.load(std::memory_order_relaxed);
    const auto current_tail = tail_.load(std::memory_order_acquire);
    const size_t capacity = RingCapacity();
    const size_t available = (current_tail >= current_head)
                                 ? (current_tail - current_head)
                                 : (capacity - current_head + current_tail);
 
    if (available < count)
    {
      return ErrorCode::EMPTY;
    }
 
    auto* dst = static_cast<std::byte*>(data);
    for (size_t index = 0; index < count; ++index)
    {
      LibXR::Memory::FastCopy(dst + index * element_size_,
                              PayloadPtr((current_head + index) % capacity),
                              element_size_);
    }
    return ErrorCode::OK;
  }
 
  void Reset()
  {
    const auto current_tail = tail_.load(std::memory_order_acquire);
    head_.store(current_tail, std::memory_order_release);
  }
 
  size_t Size() const
  {
    const auto current_head = head_.load(std::memory_order_acquire);
    const auto current_tail = tail_.load(std::memory_order_acquire);
    return (current_tail >= current_head)
               ? (current_tail - current_head)
               : (RingCapacity() - current_head + current_tail);
  }
 
  size_t EmptySize() const { return capacity_ - Size(); }
 
  size_t MaxSize() const { return capacity_; }
 
 private:
  std::byte* PayloadPtr(IndexType index) { return payloads_ + index * payload_stride_; }
 
  const std::byte* PayloadPtr(IndexType index) const
  {
    return payloads_ + index * payload_stride_;
  }
 
  size_t RingCapacity() const { return capacity_ + 1; }
 
  IndexType Increment(IndexType index) const { return (index + 1) % RingCapacity(); }
 
  SPSCQueueBase(const SPSCQueueBase&);
  SPSCQueueBase& operator=(const SPSCQueueBase&);
  SPSCQueueBase(SPSCQueueBase&&);
  SPSCQueueBase& operator=(SPSCQueueBase&&);
 
  static size_t AlignUpChecked(size_t size, size_t align)
  {
    DEV_ASSERT(align > 0);
    DEV_ASSERT((align & (align - 1)) == 0);
    ASSERT(size <= std::numeric_limits<size_t>::max() - (align - 1));
    return ((size + align - 1) / align) * align;
  }
 
  static size_t ComputeStride(size_t element_size, size_t element_align)
  {
    ASSERT(element_size > 0);
    ASSERT(element_align > 0);
    ASSERT((element_align & (element_align - 1)) == 0);
    return AlignUpChecked(element_size, element_align);
  }
 
  static size_t MultiplyChecked(size_t lhs, size_t rhs)
  {
    if (lhs == 0 || rhs == 0)
    {
      return 0;
    }
 
    ASSERT(lhs <= std::numeric_limits<size_t>::max() / rhs);
    return lhs * rhs;
  }
 
  const size_t element_size_;  
  const size_t capacity_;      
  const size_t payload_alloc_align_;  
  const size_t payload_stride_;  
  std::byte* payloads_;          
 
  alignas(LibXR::CONCURRENCY_ALIGNMENT) std::atomic<
      IndexType> head_;  
  alignas(LibXR::CONCURRENCY_ALIGNMENT) std::atomic<
      IndexType> tail_;  
};
}  // namespace LibXR