lockfree_pool.hpp

#pragma once
 
#include <atomic>
#include <new>
 
#include "libxr_def.hpp"
 
namespace LibXR
{
 
template <typename Data>
class LockFreePool
{
 public:
  // NOLINTBEGIN
  enum class SlotState : uint32_t
  {
    FREE = 0,             
    BUSY = 1,             
    READY = 2,            
    RECYCLE = UINT32_MAX  
 
  };
  // NOLINTEND
 
  union alignas(LibXR::CACHE_LINE_SIZE) Slot
  {
    struct
    {
      std::atomic<SlotState> state;  
      Data data;                     
    } slot;
 
    uint8_t pad[LibXR::CACHE_LINE_SIZE];  
 
    Slot()
    {
      new (&slot.state) std::atomic<SlotState>(SlotState::FREE);
    }
 
    ~Slot() {}
  };
 
  LockFreePool(uint32_t slot_count)
      : SLOT_COUNT(slot_count),
        slots_(new (std::align_val_t(LibXR::CACHE_LINE_SIZE)) Slot[slot_count])
  {
    for (uint32_t index = 0; index < SLOT_COUNT; index++)
    {
      slots_[index].slot.state.store(SlotState::FREE, std::memory_order_relaxed);
    }
  }
 
  ~LockFreePool() { delete[] slots_; }
 
  ErrorCode Put(const Data& data)
  {
    uint32_t start_index = 0;
    return Put(data, start_index);
  }
 
  ErrorCode Put(const Data& data, uint32_t& start_index)
  {
    for (uint32_t index = start_index; index < SLOT_COUNT; index++)
    {
      auto expected = slots_[index].slot.state.load(std::memory_order_relaxed);
      if (expected == SlotState::FREE || expected == SlotState::RECYCLE)
      {
        if (slots_[index].slot.state.compare_exchange_strong(expected, SlotState::BUSY,
                                                             std::memory_order_release,
                                                             std::memory_order_relaxed))
        {
          slots_[index].slot.data = data;
          slots_[index].slot.state.store(SlotState::READY, std::memory_order_release);
          start_index = index;
          return ErrorCode::OK;
        }
      }
    }
    return ErrorCode::FULL;
  }
 
  ErrorCode PutToSlot(const Data& data, uint32_t index)
  {
    auto expected = slots_[index].slot.state.load(std::memory_order_relaxed);
    if (expected == SlotState::FREE || expected == SlotState::RECYCLE)
    {
      if (slots_[index].slot.state.compare_exchange_strong(expected, SlotState::BUSY,
                                                           std::memory_order_release,
                                                           std::memory_order_relaxed))
      {
        slots_[index].slot.data = data;
        slots_[index].slot.state.store(SlotState::READY, std::memory_order_release);
        return ErrorCode::OK;
      }
    }
    return ErrorCode::FULL;
  }
 
  ErrorCode Get(Data& data)
  {
    uint32_t start_index = 0;
 
    return Get(data, start_index);
  }
 
  ErrorCode Get(Data& data, uint32_t& start_index)
  {
    for (uint32_t index = start_index; index < SLOT_COUNT; index++)
    {
      auto expected = slots_[index].slot.state.load(std::memory_order_acquire);
      if (expected == SlotState::READY)
      {
        if (slots_[index].slot.state.compare_exchange_strong(expected, SlotState::BUSY,
                                                             std::memory_order_acquire,
                                                             std::memory_order_relaxed))
        {
          data = slots_[index].slot.data;
          slots_[index].slot.state.store(SlotState::RECYCLE, std::memory_order_release);
          start_index = index;
          return ErrorCode::OK;
        }
      }
      if (expected == SlotState::FREE)
      {
        start_index = 0;
        return ErrorCode::EMPTY;
      }
    }
 
    start_index = 0;
    return ErrorCode::EMPTY;
  }
 
  ErrorCode GetFromSlot(Data& data, uint32_t index)
  {
    auto expected = slots_[index].slot.state.load(std::memory_order_acquire);
    if (expected == SlotState::READY)
    {
      if (slots_[index].slot.state.compare_exchange_strong(expected, SlotState::BUSY,
                                                           std::memory_order_acquire,
                                                           std::memory_order_relaxed))
      {
        data = slots_[index].slot.data;
        slots_[index].slot.state.store(SlotState::RECYCLE, std::memory_order_release);
        return ErrorCode::OK;
      }
    }
    return ErrorCode::EMPTY;
  }
 
  ErrorCode RecycleSlot(uint32_t index)
  {
    auto expected = slots_[index].slot.state.load(std::memory_order_relaxed);
    if (expected == SlotState::READY)
    {
      if (slots_[index].slot.state.compare_exchange_strong(expected, SlotState::RECYCLE,
                                                           std::memory_order_release,
                                                           std::memory_order_relaxed))
      {
        return ErrorCode::OK;
      }
    }
    return ErrorCode::EMPTY;
  }
 
  [[nodiscard]] size_t Size() const
  {
    uint32_t size = 0;
    for (uint32_t index = 0; index < SLOT_COUNT; index++)
    {
      if (slots_[index].slot.state.load(std::memory_order_relaxed) == SlotState::READY)
      {
        size++;
      }
    }
    return size;
  }
 
  [[nodiscard]] size_t EmptySize()
  {
    uint32_t size = 0;
    for (uint32_t index = 0; index < SLOT_COUNT; index++)
    {
      auto state = slots_[index].slot.state.load(std::memory_order_relaxed);
      if (state == SlotState::FREE || state == SlotState::RECYCLE)
      {
        size++;
      }
    }
    return size;
  }
 
  uint32_t SlotCount() const { return SLOT_COUNT; }
 
 protected:
  Slot& operator[](uint32_t index)
  {
    if (index >= SLOT_COUNT)
    {
      ASSERT(false);
    }
    return slots_[index];
  }
 
 private:
  const uint32_t SLOT_COUNT;  
  Slot* slots_;               
};
}  // namespace LibXR