lockfree_queue.hpp

#pragma once
 
#include <atomic>
 
#include "libxr_def.hpp"
 
namespace LibXR
{
 
template <typename Data>
class alignas(LIBXR_CACHE_LINE_SIZE) LockFreeQueue
{
 public:
  LockFreeQueue(size_t length)
      : head_(0), tail_(0), queue_handle_(new Data[length + 1]), LENGTH(length)
  {
  }
 
  ~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);
    memcpy(reinterpret_cast<void *>(queue_handle_ + current_tail),
           reinterpret_cast<const void *>(data), first_chunk * sizeof(Data));
 
    if (size > first_chunk)
    {
      memcpy(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);
        memcpy(reinterpret_cast<void *>(data),
               reinterpret_cast<const void *>(queue_handle_ + current_head),
               first_chunk * sizeof(Data));
 
        if (size > first_chunk)
        {
          memcpy(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);
      memcpy(reinterpret_cast<void *>(data),
             reinterpret_cast<const void *>(queue_handle_ + current_head),
             first_chunk * sizeof(Data));
 
      if (size > first_chunk)
      {
        memcpy(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_;
  Data *queue_handle_;
  const size_t LENGTH;
 
  uint32_t Increment(uint32_t index) const { return (index + 1) % (LENGTH + 1); }
};
 
}  // namespace LibXR