lockfree_queue.hpp

#pragma once
 
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
#include "libxr_def.hpp"
#include "libxr_mem.hpp"
 
namespace LibXR
{
 
template <typename Data>
class alignas(LibXR::CACHE_LINE_SIZE) LockFreeQueue
{
  inline constexpr size_t AlignUp(size_t size, size_t align)
  {
    return (size / align + 1) * align;
  }
 
 public:
  LockFreeQueue(size_t length)
      : head_(0),
        tail_(0),
        LENGTH(AlignUp(length, LibXR::ALIGN_SIZE) - 1),
        queue_handle_(new Data[LENGTH + 1])
  {
  }
 
  ~LockFreeQueue() { delete[] queue_handle_; }
 
  Data* operator[](uint32_t index) { return &queue_handle_[static_cast<size_t>(index)]; }
 
  template <typename ElementData = Data>
  ErrorCode Push(ElementData&& item)
  {
    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;
    }
 
    queue_handle_[CURRENT_TAIL] = std::forward<ElementData>(item);
    tail_.store(NEXT_TAIL, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename ElementData = Data>
  ErrorCode Pop(ElementData& item)
  {
    auto current_head = head_.load(std::memory_order_relaxed);
 
    while (true)
    {
      if (current_head == tail_.load(std::memory_order_acquire))
      {
        return ErrorCode::EMPTY;
      }
 
      item = queue_handle_[current_head];
 
      if (head_.compare_exchange_weak(current_head, Increment(current_head),
                                      std::memory_order_acq_rel,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode Pop(Data& item)
  {
    auto current_head = head_.load(std::memory_order_relaxed);
 
    while (true)
    {
      if (current_head == tail_.load(std::memory_order_acquire))
      {
        return ErrorCode::EMPTY;
      }
 
      item = queue_handle_[current_head];
 
      if (head_.compare_exchange_weak(current_head, Increment(current_head),
                                      std::memory_order_acq_rel,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode Pop()
  {
    auto current_head = head_.load(std::memory_order_relaxed);
 
    while (true)
    {
      if (current_head == tail_.load(std::memory_order_acquire))
      {
        return ErrorCode::EMPTY;
      }
 
      if (head_.compare_exchange_weak(current_head, Increment(current_head),
                                      std::memory_order_acq_rel,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode Peek(Data& item)
  {
    while (true)
    {
      auto current_head = head_.load(std::memory_order_relaxed);
      if (current_head == tail_.load(std::memory_order_acquire))
      {
        return ErrorCode::EMPTY;
      }
 
      item = queue_handle_[current_head];
 
      if (head_.load(std::memory_order_acquire) == current_head)
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode PushBatch(const Data* data, size_t size)
  {
    auto current_tail = tail_.load(std::memory_order_relaxed);
    auto current_head = head_.load(std::memory_order_acquire);
 
    size_t capacity = LENGTH + 1;
    size_t free_space = (current_tail >= current_head)
                            ? (capacity - (current_tail - current_head) - 1)
                            : (current_head - current_tail - 1);
 
    if (free_space < size)
    {
      return ErrorCode::FULL;
    }
 
    size_t first_chunk = LibXR::min(size, capacity - current_tail);
    LibXR::Memory::FastCopy(reinterpret_cast<void*>(queue_handle_ + current_tail),
                            reinterpret_cast<const void*>(data),
                            first_chunk * sizeof(Data));
 
    if (size > first_chunk)
    {
      LibXR::Memory::FastCopy(reinterpret_cast<void*>(queue_handle_),
                              reinterpret_cast<const void*>(data + first_chunk),
                              (size - first_chunk) * sizeof(Data));
    }
 
    tail_.store((current_tail + size) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Writer>
  ErrorCode PushWithWriter(size_t size, Writer&& writer)
  {
    static_assert(std::is_invocable_v<Writer&, Data*, size_t>,
                  "PushWithWriter writer must be callable as "
                  "ErrorCode(Data* buffer, size_t chunk_size)");
    using WriterRet = std::invoke_result_t<Writer&, Data*, size_t>;
    static_assert(std::is_convertible_v<WriterRet, ErrorCode>,
                  "PushWithWriter writer return type must be convertible to ErrorCode");
 
    if (size == 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 = LENGTH + 1;
    const size_t free_space =
        (current_tail >= current_head) ? (capacity - (current_tail - current_head) - 1)
                                       : (current_head - current_tail - 1);
 
    if (free_space < size)
    {
      return ErrorCode::FULL;
    }
 
    const size_t first_chunk = LibXR::min(size, capacity - static_cast<size_t>(current_tail));
    Writer& writer_ref = writer;
    const ErrorCode first_ec = writer_ref(queue_handle_ + current_tail, first_chunk);
    if (first_ec != ErrorCode::OK)
    {
      return first_ec;
    }
 
    if (size > first_chunk)
    {
      const ErrorCode second_ec = writer_ref(queue_handle_, size - first_chunk);
      if (second_ec != ErrorCode::OK)
      {
        return second_ec;
      }
    }
 
    tail_.store((current_tail + size) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
 
  template <typename Reader>
  ErrorCode PopWithReader(size_t size, Reader&& reader)
  {
    static_assert(std::is_invocable_v<Reader&, const Data*, size_t>,
                  "PopWithReader reader must be callable as "
                  "ErrorCode(const Data* buffer, size_t chunk_size)");
    using ReaderRet = std::invoke_result_t<Reader&, const Data*, size_t>;
    static_assert(std::is_convertible_v<ReaderRet, ErrorCode>,
                  "PopWithReader reader return type must be convertible to ErrorCode");
 
    if (size == 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 = LENGTH + 1;
    const size_t available = (current_tail >= current_head)
                                 ? (current_tail - current_head)
                                 : (capacity - current_head + current_tail);
 
    if (available < size)
    {
      return ErrorCode::EMPTY;
    }
 
    const size_t first_chunk = LibXR::min(size, capacity - static_cast<size_t>(current_head));
    Reader& reader_ref = reader;
    const ErrorCode first_ec = reader_ref(queue_handle_ + current_head, first_chunk);
    if (first_ec != ErrorCode::OK)
    {
      return first_ec;
    }
 
    if (size > first_chunk)
    {
      const ErrorCode second_ec = reader_ref(queue_handle_, size - first_chunk);
      if (second_ec != ErrorCode::OK)
      {
        return second_ec;
      }
    }
 
    head_.store((current_head + size) % capacity, std::memory_order_release);
    return ErrorCode::OK;
  }
  ErrorCode PopBatch(Data* data, size_t size)
  {
    size_t capacity = LENGTH + 1;
 
    while (true)
    {
      auto current_head = head_.load(std::memory_order_relaxed);
      auto current_tail = tail_.load(std::memory_order_acquire);
 
      size_t available = (current_tail >= current_head)
                             ? (current_tail - current_head)
                             : (capacity - current_head + current_tail);
 
      if (available < size)
      {
        return ErrorCode::EMPTY;
      }
 
      if (data != nullptr)
      {
        size_t first_chunk = LibXR::min(size, capacity - current_head);
        LibXR::Memory::FastCopy(
            reinterpret_cast<void*>(data),
            reinterpret_cast<const void*>(queue_handle_ + current_head),
            first_chunk * sizeof(Data));
 
        if (size > first_chunk)
        {
          LibXR::Memory::FastCopy(reinterpret_cast<void*>(data + first_chunk),
                                  reinterpret_cast<const void*>(queue_handle_),
                                  (size - first_chunk) * sizeof(Data));
        }
      }
 
      size_t new_head = (current_head + size) % capacity;
 
      if (head_.compare_exchange_weak(current_head, new_head, std::memory_order_acq_rel,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode PeekBatch(Data* data, size_t size)
  {
    if (size == 0)
    {
      return ErrorCode::OK;
    }
 
    const size_t CAPACITY = LENGTH + 1;
 
    while (true)
    {
      auto current_head = head_.load(std::memory_order_relaxed);
      auto current_tail = tail_.load(std::memory_order_acquire);
 
      size_t available = (current_tail >= current_head)
                             ? (current_tail - current_head)
                             : (CAPACITY - current_head + current_tail);
 
      if (available < size)
      {
        return ErrorCode::EMPTY;
      }
 
      if (data != nullptr)
      {
        size_t first_chunk = LibXR::min(size, CAPACITY - current_head);
        LibXR::Memory::FastCopy(
            reinterpret_cast<void*>(data),
            reinterpret_cast<const void*>(queue_handle_ + current_head),
            first_chunk * sizeof(Data));
 
        if (size > first_chunk)
        {
          LibXR::Memory::FastCopy(reinterpret_cast<void*>(data + first_chunk),
                                  reinterpret_cast<const void*>(queue_handle_),
                                  (size - first_chunk) * sizeof(Data));
        }
      }
 
      if (head_.load(std::memory_order_acquire) == current_head)
      {
        return ErrorCode::OK;
      }
    }
  }
 
  void Reset()
  {
    head_.store(0, std::memory_order_relaxed);
    tail_.store(0, std::memory_order_relaxed);
  }
 
  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)
                                          : ((LENGTH + 1) - CURRENT_HEAD + CURRENT_TAIL);
  }
 
  size_t EmptySize() { return LENGTH - Size(); }
 
  size_t MaxSize() const { return LENGTH; }
 
 private:
  alignas(LibXR::CACHE_LINE_SIZE) std::atomic<uint32_t> head_;
  alignas(LibXR::CACHE_LINE_SIZE) std::atomic<uint32_t> tail_;
  const size_t LENGTH;
  Data* queue_handle_;
 
  uint32_t Increment(uint32_t index) const { return (index + 1) % (LENGTH + 1); }
};
 
}  // namespace LibXR