work_stealing_pool.hpp

Go to the documentation of this file.

#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <deque>
#include <exception>
#include <functional>
#include <memory>
#include <mutex>
#include <thread>
#include <utility>
#include <vector>
 
namespace velographx {
 
struct WorkStealingStats {
  std::size_t submitted{0};
  std::size_t executed{0};
  std::size_t steal_attempts{0};
  std::size_t successful_steals{0};
  std::size_t local_steals{0};
  std::size_t remote_steals{0};
};
 
class WorkStealingPool {
 public:
  using Task = std::function<void()>;
 
  explicit WorkStealingPool(std::size_t threads = std::thread::hardware_concurrency(),
                            std::vector<std::size_t> queue_groups = {}) {
    if (threads == 0) threads = 1;
    if (queue_groups.size() != threads) queue_groups.assign(threads, 0);
    queue_groups_ = std::move(queue_groups);
    queues_.reserve(threads);
    for (std::size_t i = 0; i < threads; ++i) queues_.push_back(std::make_unique<Queue>());
    workers_.reserve(threads);
    for (std::size_t i = 0; i < threads; ++i) workers_.emplace_back([this, i] { worker(i); });
  }
 
  WorkStealingPool(const WorkStealingPool&) = delete;
  WorkStealingPool& operator=(const WorkStealingPool&) = delete;
 
  ~WorkStealingPool() {
    // Destructors must not propagate exceptions from user tasks. Public
    // wait_idle()/parallel_for() still surface the first task failure.
    wait_idle_no_throw();
    stop_.store(true, std::memory_order_release);
    cv_.notify_all();
    for (auto& worker : workers_) if (worker.joinable()) worker.join();
  }
 
  std::size_t size() const noexcept { return workers_.size(); }
  std::size_t queue_group(std::size_t index) const noexcept {
    return index < queue_groups_.size() ? queue_groups_[index] : 0;
  }
 
  void submit(Task task, std::size_t locality_hint = 0) {
    const auto queue = locality_hint % queues_.size();
    outstanding_.fetch_add(1, std::memory_order_acq_rel);
    submitted_.fetch_add(1, std::memory_order_relaxed);
    {
      std::lock_guard<std::mutex> lock(queues_[queue]->mutex);
      queues_[queue]->tasks.push_back(std::move(task));
    }
    cv_.notify_all();
  }
 
  template <class Fn>
  void parallel_for(std::size_t begin, std::size_t end, Fn&& fn,
                    std::size_t grain = 0) {
    if (end <= begin) return;
    if (grain == 0) grain = adaptive_grain(end - begin, size());
    std::size_t hint = 0;
    for (std::size_t first = begin; first < end; first += grain, ++hint) {
      const auto last = (first + grain < end) ? first + grain : end;
      submit([first, last, &fn] {
        for (std::size_t i = first; i < last; ++i) fn(i);
      }, hint);
    }
    wait_idle();
  }
 
  // Drain all submitted work. If one or more tasks fail, the pool records the
  // first exception, completes accounting for every task, and rethrows only
  // after the pool becomes idle. This keeps the pool reusable after failure.
  void wait_idle() {
    drain_until_idle();
    rethrow_pending_exception();
  }
 
  WorkStealingStats stats() const noexcept {
    return {submitted_.load(std::memory_order_relaxed), executed_.load(std::memory_order_relaxed),
            steal_attempts_.load(std::memory_order_relaxed), successful_steals_.load(std::memory_order_relaxed),
            local_steals_.load(std::memory_order_relaxed), remote_steals_.load(std::memory_order_relaxed)};
  }
 
  static std::size_t adaptive_grain(std::size_t work_items, std::size_t threads) noexcept {
    if (threads == 0) threads = 1;
    const std::size_t target_chunks = threads * 8;
    const std::size_t grain = (work_items + target_chunks - 1) / target_chunks;
    return grain == 0 ? 1 : grain;
  }
 
 private:
  void record_exception(std::exception_ptr error) noexcept {
    try {
      std::lock_guard<std::mutex> lock(exception_mutex_);
      if (!first_exception_) first_exception_ = std::move(error);
    } catch (...) {
      // A task exception must never escape a worker thread merely because the
      // exception-recording path itself encountered an exceptional runtime
      // condition.
    }
  }
 
  void complete_task(Task& task) noexcept {
    try {
      task();
    } catch (...) {
      record_exception(std::current_exception());
    }
    executed_.fetch_add(1, std::memory_order_relaxed);
    if (outstanding_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
      idle_cv_.notify_all();
    }
  }
 
  void drain_until_idle() {
    while (outstanding_.load(std::memory_order_acquire) != 0) {
      Task task;
      if (pop_any(task)) {
        complete_task(task);
        continue;
      }
      std::unique_lock<std::mutex> lock(wait_mutex_);
      idle_cv_.wait_for(lock, std::chrono::milliseconds(1), [this] {
        return outstanding_.load(std::memory_order_acquire) == 0;
      });
    }
  }
 
  void wait_idle_no_throw() noexcept {
    try {
      drain_until_idle();
    } catch (...) {
      // Destruction is a non-reporting boundary. Call wait_idle() explicitly
      // when task failures need to be observed by the caller.
    }
  }
 
  void rethrow_pending_exception() {
    std::exception_ptr error;
    {
      std::lock_guard<std::mutex> lock(exception_mutex_);
      error = std::exchange(first_exception_, {});
    }
    if (error) std::rethrow_exception(error);
  }
 
  struct Queue {
    std::mutex mutex;
    std::deque<Task> tasks;
  };
 
  bool pop_local(std::size_t index, Task& task) {
    std::lock_guard<std::mutex> lock(queues_[index]->mutex);
    if (queues_[index]->tasks.empty()) return false;
    task = std::move(queues_[index]->tasks.back());
    queues_[index]->tasks.pop_back();
    return true;
  }
 
  bool pop_any(Task& task) {
    for (auto& queue : queues_) {
      std::lock_guard<std::mutex> lock(queue->mutex);
      if (queue->tasks.empty()) continue;
      task = std::move(queue->tasks.front());
      queue->tasks.pop_front();
      return true;
    }
    return false;
  }
 
  bool try_steal_from(std::size_t thief, std::size_t victim, Task& task, bool local) {
    steal_attempts_.fetch_add(1, std::memory_order_relaxed);
    std::lock_guard<std::mutex> lock(queues_[victim]->mutex);
    if (queues_[victim]->tasks.empty()) return false;
    task = std::move(queues_[victim]->tasks.front());
    queues_[victim]->tasks.pop_front();
    successful_steals_.fetch_add(1, std::memory_order_relaxed);
    (local ? local_steals_ : remote_steals_).fetch_add(1, std::memory_order_relaxed);
    return true;
  }
 
  bool steal(std::size_t thief, Task& task) {
    const auto thief_group = queue_group(thief);
    for (std::size_t offset = 1; offset < queues_.size(); ++offset) {
      const std::size_t victim = (thief + offset) % queues_.size();
      if (queue_group(victim) != thief_group) continue;
      if (try_steal_from(thief, victim, task, true)) return true;
    }
    for (std::size_t offset = 1; offset < queues_.size(); ++offset) {
      const std::size_t victim = (thief + offset) % queues_.size();
      if (queue_group(victim) == thief_group) continue;
      if (try_steal_from(thief, victim, task, false)) return true;
    }
    return false;
  }
 
  void worker(std::size_t index) {
    while (!stop_.load(std::memory_order_acquire)) {
      Task task;
      if (pop_local(index, task) || steal(index, task)) {
        complete_task(task);
        continue;
      }
      std::unique_lock<std::mutex> lock(cv_mutex_);
      cv_.wait_for(lock, std::chrono::milliseconds(1), [this] {
        return stop_.load(std::memory_order_acquire) || outstanding_.load(std::memory_order_acquire) != 0;
      });
    }
  }
 
  std::vector<std::unique_ptr<Queue>> queues_;
  std::vector<std::thread> workers_;
  std::vector<std::size_t> queue_groups_;
  std::atomic<bool> stop_{false};
  std::atomic<std::size_t> outstanding_{0};
  std::atomic<std::size_t> submitted_{0};
  std::atomic<std::size_t> executed_{0};
  std::atomic<std::size_t> steal_attempts_{0};
  std::atomic<std::size_t> successful_steals_{0};
  std::atomic<std::size_t> local_steals_{0};
  std::atomic<std::size_t> remote_steals_{0};
  std::mutex cv_mutex_;
  std::condition_variable cv_;
  std::mutex wait_mutex_;
  std::condition_variable idle_cv_;
  std::mutex exception_mutex_;
  std::exception_ptr first_exception_;
};
 
}  // namespace velographx