simple_channel.h

// Copyright 2019, Beeri 15.  All rights reserved.
// Author: Roman Gershman (romange@gmail.com)
//
#pragma once

#include "base/ProducerConsumerQueue.h"

#include <boost/fiber/context.hpp>

#include "util/fibers/event_count.h"
#include "util/fibers/fibers_ext.h"

namespace util {
namespace fibers_ext {

template <typename T> class SimpleChannel {
  typedef ::boost::fibers::context::wait_queue_t wait_queue_t;
  using spinlock_lock_t = ::boost::fibers::detail::spinlock_lock;

 public:
  SimpleChannel(size_t n) : q_(n) {}

  template <typename... Args> void Push(Args&&... recordArgs) noexcept;

  // Blocking call. Returns false if channel is closed, true otherwise with the popped value.
  bool Pop(T& dest);

  void StartClosing();

  template <typename... Args> bool TryPush(Args&&... args) noexcept {
    if (q_.write(std::forward<Args>(args)...)) {
      if (++throttled_pushes_ > q_.capacity() / 3) {
        pop_ec_.notify();
        throttled_pushes_ = 0;
      }
      return true;
    }
    return false;
  }

  bool TryPop(T& val) {
    if (q_.read(val)) {
      return true;
    }
    push_ec_.notify();
    return false;
  }

  bool IsClosing() const { return is_closing_.load(std::memory_order_relaxed); }

 private:
  unsigned throttled_pushes_ = 0;

  folly::ProducerConsumerQueue<T> q_;
  std::atomic_bool is_closing_{false};

  // Event counts provide almost negligible contention during fast-path (a single atomic add).
  EventCount push_ec_, pop_ec_;
};

template <typename T>
template <typename... Args>
void SimpleChannel<T>::Push(Args&&... args) noexcept {
  if (TryPush(std::forward<Args>(args)...))  // fast path.
    return;

  while (true) {
    EventCount::Key key = push_ec_.prepareWait();
    if (TryPush(std::forward<Args>(args)...)) {
      break;
    }
    push_ec_.wait(key.epoch());
  }
}

template <typename T> bool SimpleChannel<T>::Pop(T& dest) {
  if (TryPop(dest))  // fast path
    return true;

  while (true) {
    EventCount::Key key = pop_ec_.prepareWait();
    if (TryPop(dest)) {
      return true;
    }

    if (is_closing_.load(std::memory_order_acquire)) {
      return false;
    }

    pop_ec_.wait(key.epoch());
  }
}

template <typename T> void SimpleChannel<T>::StartClosing() {
  // Full barrier, StartClosing performance does not matter.
  is_closing_.store(true, std::memory_order_seq_cst);
  pop_ec_.notifyAll();
}

}  // namespace fibers_ext
}  // namespace util