BoundedSPSCQueue.h

#pragma once

#include "quill/core/Attributes.h"
#include "quill/core/Common.h"
#include "quill/core/MathUtilities.h"
#include "quill/core/QuillError.h"

#include <atomic>
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <string>

#if defined(_WIN32)
  #include <malloc.h>
#else
  #include <sys/mman.h>
#endif

#if defined(QUILL_X86ARCH)
  #if defined(_WIN32)
    #include <intrin.h>
  #else
    #if __has_include(<x86gprintrin.h>)
      #if defined(__GNUC__) && __GNUC__ > 10
        #include <emmintrin.h>
        #include <x86gprintrin.h>
      #elif defined(__clang_major__)
      // clang needs immintrin for _mm_clflushopt
        #include <immintrin.h>
      #endif
    #else
      #include <immintrin.h>
      #include <x86intrin.h>
    #endif
  #endif
#endif

QUILL_BEGIN_NAMESPACE

namespace detail
{

#if defined(_WIN32) && defined(_MSC_VER) && !defined(__GNUC__)
  #pragma warning(push)
  #pragma warning(disable : 4324)
#endif

template <typename T>
class BoundedSPSCQueueImpl
{
public:
  using integer_type = T;

  static_assert(integer_type{0} < static_cast<integer_type>(~integer_type{0}),
                "BoundedSPSCQueueImpl integer_type must be unsigned");

  struct WriteReservation
  {
    std::byte* write_buffer{nullptr};
    integer_type writer_pos{0};
    BoundedSPSCQueueImpl* bounded_queue{nullptr};
  };

  QUILL_ATTRIBUTE_HOT explicit BoundedSPSCQueueImpl(integer_type capacity,
                                                    HugePagesPolicy huge_pages_policy = HugePagesPolicy::Never,
                                                    integer_type reader_store_percent = 5)
    : _capacity(_validate_capacity(capacity)),
      _mask(_capacity - 1),
      _bytes_per_batch(static_cast<integer_type>(
        (static_cast<double>(_capacity) * static_cast<double>(reader_store_percent)) / 100.0)),
      _storage(static_cast<std::byte*>(_alloc_aligned(
        2u * static_cast<size_t>(_capacity), QUILL_CACHE_LINE_ALIGNED, huge_pages_policy))),
      _huge_pages_policy(huge_pages_policy)
  {
    size_t const storage_size = 2u * static_cast<size_t>(_capacity);
    std::memset(_storage, 0, storage_size);

    _atomic_writer_pos.store(0);
    _atomic_reader_pos.store(0);

    // reader_pos starts at 0, so cached value is 0 + capacity
    _reader_pos_cache_plus_capacity = _capacity;

#if defined(QUILL_X86ARCH)
    // remove log memory from cache
    for (size_t i = 0; i < storage_size; i += QUILL_CACHE_LINE_SIZE)
    {
      _mm_clflush(_storage + i);
    }

    uint64_t const cache_lines = (_capacity >= 2048) ? 32 : 16;

    for (uint64_t i = 0; i < cache_lines; ++i)
    {
      _mm_prefetch(reinterpret_cast<char const*>(_storage + (QUILL_CACHE_LINE_SIZE * i)), _MM_HINT_T0);
    }
#endif
  }

  ~BoundedSPSCQueueImpl() { _free_aligned(_storage); }

  BoundedSPSCQueueImpl(BoundedSPSCQueueImpl const&) = delete;
  BoundedSPSCQueueImpl& operator=(BoundedSPSCQueueImpl const&) = delete;

  QUILL_NODISCARD QUILL_ATTRIBUTE_HOT std::byte* prepare_write(integer_type n) noexcept
  {
    if (static_cast<integer_type>(_reader_pos_cache_plus_capacity - _writer_pos) < n)
    {
      // not enough space, we need to load reader and re-check
      _reader_pos_cache_plus_capacity = _atomic_reader_pos.load(std::memory_order_acquire) + _capacity;

      if (static_cast<integer_type>(_reader_pos_cache_plus_capacity - _writer_pos) < n)
      {
        return nullptr;
      }
    }

    return _storage + (_writer_pos & _mask);
  }

  QUILL_NODISCARD QUILL_ATTRIBUTE_HOT WriteReservation prepare_write_reserve_cached(integer_type n) noexcept
  {
    integer_type const writer_pos = _writer_pos;

    if (QUILL_UNLIKELY(static_cast<integer_type>(_reader_pos_cache_plus_capacity - writer_pos) < n))
    {
      return WriteReservation{nullptr, writer_pos, this};
    }

    // _storage is never null after construction; hint lets the compiler
    // eliminate a redundant null-check on the computed write pointer.
    QUILL_ASSUME(_storage != nullptr);

    return WriteReservation{_storage + (writer_pos & _mask), writer_pos, this};
  }

  QUILL_ATTRIBUTE_HOT void finish_write(integer_type n) noexcept { _writer_pos += n; }

  QUILL_ATTRIBUTE_HOT void commit_write() noexcept
  {
    // set the atomic flag so the reader can see write
    _atomic_writer_pos.store(_writer_pos, std::memory_order_release);

#if defined(QUILL_X86ARCH)
    // flush writen cache lines
    _flush_cachelines(_last_flushed_writer_pos, _writer_pos);

    // prefetch a future cache line
    _mm_prefetch(
      reinterpret_cast<char const*>(_storage + ((_writer_pos + QUILL_CACHE_LINE_SIZE * 10) & _mask)), _MM_HINT_T0);
#endif
  }

  QUILL_ATTRIBUTE_HOT void finish_and_commit_write(integer_type n) noexcept
  {
    finish_and_commit_write_reservation(_writer_pos + n);
  }

  QUILL_ATTRIBUTE_HOT void finish_and_commit_write_reservation(integer_type new_writer_pos) noexcept
  {
    _writer_pos = new_writer_pos;

    // set the atomic flag so the reader can see write
    _atomic_writer_pos.store(new_writer_pos, std::memory_order_release);

#if defined(QUILL_X86ARCH)
    // flush writen cache lines
    _flush_cachelines(_last_flushed_writer_pos, new_writer_pos);

    // prefetch a future cache line
    _mm_prefetch(
      reinterpret_cast<char const*>(_storage + ((new_writer_pos + QUILL_CACHE_LINE_SIZE * 10) & _mask)), _MM_HINT_T0);
#endif
  }

  QUILL_NODISCARD QUILL_ATTRIBUTE_HOT std::byte* prepare_read() noexcept
  {
    if (empty())
    {
      return nullptr;
    }

    return _storage + (_reader_pos & _mask);
  }

  QUILL_ATTRIBUTE_HOT void finish_read(integer_type n) noexcept { _reader_pos += n; }

  QUILL_ATTRIBUTE_HOT void commit_read() noexcept
  {
    if (static_cast<integer_type>(_reader_pos - _atomic_reader_pos.load(std::memory_order_relaxed)) >= _bytes_per_batch)
    {
      _atomic_reader_pos.store(_reader_pos, std::memory_order_release);

#if defined(QUILL_X86ARCH)
      _flush_cachelines(_last_flushed_reader_pos, _reader_pos);
#endif
    }
  }

  QUILL_NODISCARD QUILL_ATTRIBUTE_HOT bool empty() const noexcept
  {
    if (_writer_pos_cache == _reader_pos)
    {
      // if we think the queue is empty we also load the atomic variable to check further
      _writer_pos_cache = _atomic_writer_pos.load(std::memory_order_acquire);

      if (_writer_pos_cache == _reader_pos)
      {
        return true;
      }
    }

    return false;
  }

  QUILL_NODISCARD integer_type capacity() const noexcept
  {
    return static_cast<integer_type>(_capacity);
  }

  QUILL_NODISCARD HugePagesPolicy huge_pages_policy() const noexcept { return _huge_pages_policy; }

private:
#if defined(QUILL_X86ARCH)
  QUILL_ATTRIBUTE_HOT void _flush_cachelines(integer_type& last, integer_type offset)
  {
    integer_type last_diff = last - (last & QUILL_CACHE_LINE_MASK);
    integer_type const cur_diff = offset - (offset & QUILL_CACHE_LINE_MASK);

    if (cur_diff > last_diff)
    {
      do
      {
        _mm_clflushopt(_storage + (last_diff & _mask));
        last_diff += QUILL_CACHE_LINE_SIZE;
      } while (cur_diff > last_diff);

      last = last_diff;
    }
  }
#endif

  QUILL_NODISCARD static std::byte* _align_pointer(void* pointer, size_t alignment) noexcept
  {
    QUILL_ASSERT(is_power_of_two(alignment),
                 "alignment must be a power of two in BoundedSPSCQueue::_align_pointer()");
    return reinterpret_cast<std::byte*>((reinterpret_cast<uintptr_t>(pointer) + (alignment - 1ul)) &
                                        ~(alignment - 1ul));
  }

  static integer_type _validate_capacity(QUILL_MAYBE_UNUSED integer_type capacity)
  {
#if defined(QUILL_X86ARCH)
    if (capacity < 1024)
    {
      QUILL_THROW(QuillError{"Capacity must be at least 1024"});
    }
#endif

    size_t const rounded_capacity = next_power_of_two(static_cast<size_t>(capacity));

    constexpr auto max_integer_capacity = []() constexpr
    {
      // avoid including limits
      if constexpr (sizeof(integer_type) >= sizeof(size_t))
      {
        return SIZE_MAX;
      }
      else if constexpr (sizeof(integer_type) <= sizeof(uint8_t))
      {
        return static_cast<size_t>(UINT8_MAX);
      }
      else if constexpr (sizeof(integer_type) <= sizeof(uint16_t))
      {
        return static_cast<size_t>(UINT16_MAX);
      }
      else if constexpr (sizeof(integer_type) <= sizeof(uint32_t))
      {
        return static_cast<size_t>(UINT32_MAX);
      }
      else
      {
        return static_cast<size_t>(UINT64_MAX);
      }
    }();

    if (QUILL_UNLIKELY(rounded_capacity > max_integer_capacity))
    {
      QUILL_THROW(QuillError{"BoundedSPSCQueue capacity exceeds the queue position type"});
    }

    // The logical capacity is tracked as integer_type, but the backing allocation is byte-sized:
    // two mirrored queue regions plus POSIX alignment metadata. Validate the allocation size in
    // size_t before multiplying by two, otherwise an extreme capacity can wrap to a tiny mmap.
    constexpr size_t metadata_size{2u * sizeof(size_t)};
    constexpr size_t max_rounded_capacity = (SIZE_MAX - metadata_size - QUILL_CACHE_LINE_ALIGNED) / 2u;

    if (QUILL_UNLIKELY(rounded_capacity > max_rounded_capacity))
    {
      QUILL_THROW(QuillError{"BoundedSPSCQueue capacity is too large"});
    }

    return static_cast<integer_type>(rounded_capacity);
  }

  QUILL_NODISCARD static void* _alloc_aligned(size_t size, size_t alignment,
                                              QUILL_MAYBE_UNUSED HugePagesPolicy huge_pages_policy)
  {
#if defined(_WIN32)
    void* p = _aligned_malloc(size, alignment);

    if (!p)
    {
      QUILL_THROW(QuillError{std::string{"_alloc_aligned failed with error message errno: "} +
                             std::to_string(errno)});
    }

    return p;
#else
    // Calculate the total size including the metadata and alignment
    constexpr size_t metadata_size{2u * sizeof(size_t)};
    size_t const total_size{size + metadata_size + alignment};

    // Allocate the memory
    int flags = MAP_PRIVATE | MAP_ANONYMOUS;

  #if defined(__linux__)
    if (huge_pages_policy != HugePagesPolicy::Never)
    {
      flags |= MAP_HUGETLB;
    }
  #endif

    void* mem = ::mmap(nullptr, total_size, PROT_READ | PROT_WRITE, flags, -1, 0);

  #if defined(__linux__)
    if ((mem == MAP_FAILED) && (huge_pages_policy == HugePagesPolicy::Try))
    {
      // we tried but failed allocating huge pages, try normal pages instead
      flags &= ~MAP_HUGETLB;
      mem = ::mmap(nullptr, total_size, PROT_READ | PROT_WRITE, flags, -1, 0);
    }
  #endif

    if (mem == MAP_FAILED)
    {
      QUILL_THROW(QuillError{std::string{"mmap failed. errno: "} + std::to_string(errno) +
                             " error: " + std::strerror(errno)});
    }

    // Calculate the aligned address after the metadata
    std::byte* aligned_address = _align_pointer(static_cast<std::byte*>(mem) + metadata_size, alignment);

    // Calculate the offset from the original memory location
    auto const offset = static_cast<size_t>(aligned_address - static_cast<std::byte*>(mem));

    // Store the size and offset information in the metadata
    std::memcpy(aligned_address - sizeof(size_t), &total_size, sizeof(total_size));
    std::memcpy(aligned_address - (2u * sizeof(size_t)), &offset, sizeof(offset));

    return aligned_address;
#endif
  }

  void static _free_aligned(void* ptr) noexcept
  {
#if defined(_WIN32)
    _aligned_free(ptr);
#else
    // Retrieve the size and offset information from the metadata
    size_t offset;
    std::memcpy(&offset, static_cast<std::byte*>(ptr) - (2u * sizeof(size_t)), sizeof(offset));

    size_t total_size;
    std::memcpy(&total_size, static_cast<std::byte*>(ptr) - sizeof(size_t), sizeof(total_size));

    // Calculate the original memory block address
    void* mem = static_cast<std::byte*>(ptr) - offset;

    ::munmap(mem, total_size);
#endif
  }

private:
  static constexpr integer_type QUILL_CACHE_LINE_MASK{QUILL_CACHE_LINE_SIZE - 1};

  integer_type const _capacity;
  integer_type const _mask;
  integer_type const _bytes_per_batch;
  std::byte* const _storage{nullptr};
  HugePagesPolicy const _huge_pages_policy;

  alignas(QUILL_CACHE_LINE_ALIGNED) std::atomic<integer_type> _atomic_writer_pos{0};
  alignas(QUILL_CACHE_LINE_ALIGNED) integer_type _writer_pos{0};
  integer_type _reader_pos_cache_plus_capacity{0};
  integer_type _last_flushed_writer_pos{0};

  alignas(QUILL_CACHE_LINE_ALIGNED) std::atomic<integer_type> _atomic_reader_pos{0};
  alignas(QUILL_CACHE_LINE_ALIGNED) integer_type _reader_pos{0};
  mutable integer_type _writer_pos_cache{0};
  integer_type _last_flushed_reader_pos{0};
};

using BoundedSPSCQueue = BoundedSPSCQueueImpl<size_t>;

#if defined(_WIN32) && defined(_MSC_VER) && !defined(__GNUC__)
  #pragma warning(pop)
#endif

} // namespace detail

QUILL_END_NAMESPACE