flow_work_stealing_queue.h

Go to the documentation of this file.

#pragma once
#include <condition_variable>
#include <deque>
#include <mutex>
#include <optional>
 
namespace flow {
 
  // A work-stealing queue. Similar to the concurrent queue, but it allows trySteal that pop the last element.
  template <typename T>
  class WorkStealingQueue {
    using allocator_type = typename std::deque<T>::allocator_type;
 
    mutable std::mutex mux_;
    std::condition_variable blocked_;
    std::deque<T> queue_;
 
  public:
    WorkStealingQueue() = default;

    explicit WorkStealingQueue(const allocator_type& allocator)
      : queue_(allocator) {
    }
 
    WorkStealingQueue(const WorkStealingQueue&) = delete;
    WorkStealingQueue(WorkStealingQueue&&) = delete;
    WorkStealingQueue& operator=(const WorkStealingQueue&) = delete;
    WorkStealingQueue& operator=(WorkStealingQueue&&) = delete;

    bool empty() const {
      std::lock_guard lock(mux_);
      return queue_.empty();
    }

    std::size_t size() const {
      std::lock_guard lock(mux_);
      return queue_.size();
    }

    void push(const T& value) {
      {
        std::lock_guard lock(mux_);
        queue_.push_back(value);
      }
      blocked_.notify_one();
    }

    void push(T&& value) {
      {
        std::lock_guard lock(mux_);
        queue_.push_back(std::move(value));
      }
      blocked_.notify_one();
    }

    template <typename... Args>
    void emplace(Args&&... args) {
      {
        std::lock_guard lock(mux_);
        queue_.emplace_back(std::forward<Args>(args)...);
      }
      blocked_.notify_one();
    }

    std::optional<T> tryFront() const {
      std::lock_guard lock(mux_);
      if (queue_.empty()) {
        return std::nullopt;
      }
      return queue_.front();
    }

    std::optional<T> tryPop() {
      std::lock_guard lock(mux_);
      if (queue_.empty()) {
        return std::nullopt;
      }
      std::optional<T> value = std::move(queue_.front());
 
      queue_.pop_front();
      return value;
    }

    T waitPop() {
      std::unique_lock lock(mux_);
      blocked_.wait(lock, [&]() { return !queue_.empty(); });
 
      T value = std::move(queue_.front());
      queue_.pop_front();
      return value;
    }

    std::optional<T> trySteal() {
      std::lock_guard lock(mux_);
      if (queue_.empty()) {
        return std::nullopt;
      }
      std::optional<T> value = std::move(queue_.back());
 
      queue_.pop_back();
      return value;
    }
  };
 
}