lockfree_queue.hpp

#pragma once
 
#include <atomic>
 
#include "libxr_def.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) * 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;
      }
 
      if (head_.compare_exchange_weak(current_head, Increment(current_head),
                                      std::memory_order_acquire,
                                      std::memory_order_relaxed))
      {
        item = queue_handle_[current_head];
        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;
      }
 
      if (head_.compare_exchange_weak(current_head, Increment(current_head),
                                      std::memory_order_acquire,
                                      std::memory_order_relaxed))
      {
        std::atomic_thread_fence(std::memory_order_acquire);
        item = queue_handle_[current_head];
        return ErrorCode::OK;
      }
      current_head = head_.load(std::memory_order_relaxed);
    }
  }
 
  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_acquire,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
      current_head = head_.load(std::memory_order_relaxed);
    }
  }
 
  ErrorCode Peek(Data &item)
  {
    const auto CURRENT_HEAD = head_.load(std::memory_order_acquire);
    if (CURRENT_HEAD == tail_.load(std::memory_order_acquire))
    {
      return ErrorCode::EMPTY;
    }
 
    item = queue_handle_[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;
  }
 
  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_acquire,
                                      std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
  }
 
  ErrorCode PeekBatch(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;
      }
 
      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(); }
 
 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