vyukov_bounded_queue.hpp

//
// Copyright (c) 2018-2020 Manuel Pöter.
// Licensed under the MIT License. See LICENSE file in the project root for full license information.
//
 
#ifndef XENIUM_VYUKOV_BOUNDED_QUEUE_HPP
#define XENIUM_VYUKOV_BOUNDED_QUEUE_HPP
 
#include <xenium/parameter.hpp>
#include <xenium/utils.hpp>
 
#include <atomic>
#include <cassert>
#include <cstdint>
#include <memory>
 
#ifdef _MSC_VER
  #pragma warning(push)
  #pragma warning(disable : 4324) // structure was padded due to alignment specifier
#endif
 
namespace xenium {
 
namespace policy {
  template <bool Value>
  struct default_to_weak;
 
} // namespace policy
template <class T, class... Policies>
struct vyukov_bounded_queue {
public:
  using value_type = T;
 
  static constexpr bool default_to_weak =
    parameter::value_param_t<bool, policy::default_to_weak, false, Policies...>::value;
  ;
 
  explicit vyukov_bounded_queue(std::size_t size) : cells(new cell[size]), index_mask(size - 1) {
    assert(size >= 2 && utils::is_power_of_two(size));
    for (std::size_t i = 0; i < size; ++i) {
      cells[i].sequence.store(i, std::memory_order_relaxed);
    }
    enqueue_pos.store(0, std::memory_order_relaxed);
    dequeue_pos.store(0, std::memory_order_relaxed);
  }
 
  vyukov_bounded_queue(const vyukov_bounded_queue&) = delete;
  vyukov_bounded_queue(vyukov_bounded_queue&&) = delete;
 
  vyukov_bounded_queue& operator=(const vyukov_bounded_queue&) = delete;
  vyukov_bounded_queue& operator=(vyukov_bounded_queue&&) = delete;
 
  template <class... Args>
  bool try_push(Args&&... args) {
    return do_try_push<default_to_weak>(std::forward<Args>(args)...);
  }
 
  template <class... Args>
  bool try_push_strong(Args&&... args) {
    return do_try_push<false>(std::forward<Args>(args)...);
  }
 
  template <class... Args>
  bool try_push_weak(Args&&... args) {
    return do_try_push<true>(std::forward<Args>(args)...);
  }
 
  [[nodiscard]] bool try_pop(T& result) { return do_try_pop<default_to_weak>(result); }
 
  [[nodiscard]] bool try_pop_strong(T& result) { return do_try_pop<false>(result); }
 
  [[nodiscard]] bool try_pop_weak(T& result) { return do_try_pop<true>(result); }
 
private:
  template <bool Weak, class... Args>
  bool do_try_push(Args&&... args) {
    cell* c;
    std::size_t pos = enqueue_pos.load(std::memory_order_relaxed);
    for (;;) {
      c = &cells[pos & index_mask];
      // (3) - this acquire-load synchronizes-with the release-store (2)
      std::size_t seq = c->sequence.load(std::memory_order_acquire);
      if (seq == pos) {
        if (enqueue_pos.compare_exchange_weak(pos, pos + 1, std::memory_order_relaxed)) {
          break;
        }
      } else {
        if (Weak) {
          if (seq < pos) {
            return false;
          }
          pos = enqueue_pos.load(std::memory_order_relaxed);
        } else {
          auto pos2 = enqueue_pos.load(std::memory_order_relaxed);
          if (pos2 == pos && dequeue_pos.load(std::memory_order_relaxed) + index_mask + 1 == pos) {
            return false;
          }
          pos = pos2;
        }
      }
    }
    assign_value(c->value, std::forward<Args>(args)...);
    // (4) - this release-store synchronizes-with the acquire-load (1)
    c->sequence.store(pos + 1, std::memory_order_release);
    return true;
  }
 
  template <bool Weak>
  bool do_try_pop(T& result) {
    cell* c;
    std::size_t pos = dequeue_pos.load(std::memory_order_relaxed);
    for (;;) {
      c = &cells[pos & index_mask];
      // (1) - this acquire-load synchronizes-with the release-store (4)
      std::size_t seq = c->sequence.load(std::memory_order_acquire);
      auto new_pos = pos + 1;
      if (seq == new_pos) {
        if (dequeue_pos.compare_exchange_weak(pos, new_pos, std::memory_order_relaxed)) {
          break;
        }
      } else {
        if (Weak) {
          if (seq < new_pos) {
            return false;
          }
          pos = dequeue_pos.load(std::memory_order_relaxed);
        } else {
          auto pos2 = dequeue_pos.load(std::memory_order_relaxed);
          if (pos2 == pos && enqueue_pos.load(std::memory_order_relaxed) == pos) {
            return false;
          }
          pos = pos2;
        }
      }
    }
    result = std::move(c->value);
    // (2) - this release-store synchronizes-with the acquire-load (3)
    c->sequence.store(pos + index_mask + 1, std::memory_order_release);
    return true;
  }
 
  void assign_value(T& v, const T& source) { v = source; }
  void assign_value(T& v, T&& source) { v = std::move(source); }
  template <class... Args>
  void assign_value(T& v, Args&&... args) {
    v = T{std::forward<Args>(args)...};
  }
 
  // TODO - add optional padding via policy
  struct cell {
    std::atomic<std::size_t> sequence;
    T value;
  };
 
  std::unique_ptr<cell[]> cells;
  const std::size_t index_mask;
  alignas(64) std::atomic<size_t> enqueue_pos;
  alignas(64) std::atomic<size_t> dequeue_pos;
};
} // namespace xenium
 
#ifdef _MSC_VER
  #pragma warning(pop)
#endif
 
#endif