thread_pool_backend_base.hpp

Go to the documentation of this file.

#pragma once
 
#include "../callable/bind.hpp"
#include "../callable/move_only_function.hpp"
#include "callbacks.hpp"
#include "deadline.hpp"
#include "indefinite_wait.hpp"
#include "polling_wait.hpp"
#include "shutdown_policy_backend.hpp"
#include "worker_context_guard.hpp"
#include "worker_count.hpp"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <optional>
#include <queue>
#include <string>
#include <system_error>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace threadschedule::detail
{
 
// ---------------------------------------------------------------------------
// thread_pool_backend_base
// ---------------------------------------------------------------------------
 
template <typename WaitPolicy>
class thread_pool_backend_base
{
public:
  using task_type = std::function<void()>;
  using queued_task = detail::move_only_function<void()>;
 
  struct statistics
  {
    size_t total_threads;
    size_t active_threads;
    size_t pending_tasks;
    size_t completed_tasks;
    double tasks_per_second;
    std::chrono::microseconds avg_task_time;
  };
 
  explicit thread_pool_backend_base(size_t num_threads = default_worker_count(), bool register_workers = false)
      : num_threads_(checked_worker_count(num_threads)), register_workers_(register_workers), stop_(false),
        start_time_(std::chrono::steady_clock::now())
  {
    workers_.reserve(num_threads_);
 
    try
      {
        for (size_t i = 0; i < num_threads_; ++i)
          workers_.emplace_back(&thread_pool_backend_base::worker_function, this, i);
        startup_.wait(num_threads_);
      }
    catch (...)
      {
        stop_.store(true, std::memory_order_release);
        condition_.notify_all();
        for (auto& worker : workers_)
          if (worker.joinable())
            worker.join();
        throw;
      }
  }
 
  thread_pool_backend_base(thread_pool_backend_base const&) = delete;
  auto operator=(thread_pool_backend_base const&) -> thread_pool_backend_base& = delete;
 
  ~thread_pool_backend_base()
  {
    shutdown(shutdown_policy_backend::drain);
  }
 
 
  template <typename F, typename... Args>
  auto
  try_submit(F&& f, Args&&... args) -> expected<std::future<bind_result_t<F, Args...>>, std::error_code>
  {
    using return_type = bind_result_t<F, Args...>;
 
    auto task = std::make_shared<std::packaged_task<return_type()>>(
        detail::bind_args(std::forward<F>(f), std::forward<Args>(args)...));
 
    std::future<return_type> result = task->get_future();
 
    {
      std::lock_guard<std::mutex> lock(queue_mutex_);
      if (stop_)
        return unexpected(std::make_error_code(std::errc::operation_canceled));
      tasks_.emplace([task]() { (*task)(); });
    }
 
    condition_.notify_one();
    return result;
  }
 
  template <typename F, typename... Args>
  auto
  submit(F&& f, Args&&... args) -> std::future<bind_result_t<F, Args...>>
  {
    auto result = try_submit(std::forward<F>(f), std::forward<Args>(args)...);
    if (!result.has_value())
      throw std::runtime_error("Pool is shutting down");
    return std::move(result.value());
  }
 
  template <typename F, typename... Args>
  void
  post(F&& f, Args&&... args)
  {
    auto r = try_post(std::forward<F>(f), std::forward<Args>(args)...);
    if (!r.has_value())
      throw std::runtime_error("Pool is shutting down");
  }
 
  template <typename F, typename... Args>
  auto
  try_post(F&& f, Args&&... args) -> expected<void, std::error_code>
  {
    queued_task task(detail::bind_args(std::forward<F>(f), std::forward<Args>(args)...));
    {
      std::lock_guard<std::mutex> lock(queue_mutex_);
      if (stop_)
        return unexpected(std::make_error_code(std::errc::operation_canceled));
      tasks_.push(std::move(task));
    }
    condition_.notify_one();
    return {};
  }
 
  template <typename Iterator>
  auto
  try_submit_batch(Iterator begin, Iterator end) -> expected<std::vector<std::future<void>>, std::error_code>
  {
    std::vector<std::future<void>> futures;
    size_t const batch_size_hint = detail::multipass_range_size(begin, end);
    futures.reserve(batch_size_hint);
    std::vector<std::shared_ptr<std::packaged_task<void()>>> prepared;
    prepared.reserve(batch_size_hint);
 
    for (auto it = begin; it != end; ++it)
      {
        auto task = std::make_shared<std::packaged_task<void()>>(*it);
        futures.push_back(task->get_future());
        prepared.push_back(std::move(task));
      }
 
    bool enqueued = false;
    try
      {
        std::lock_guard<std::mutex> lock(queue_mutex_);
        if (stop_)
          return unexpected(std::make_error_code(std::errc::operation_canceled));
 
        for (auto const& task : prepared)
          {
            tasks_.emplace([task]() { (*task)(); });
            enqueued = true;
          }
      }
    catch (...)
      {
        if (enqueued)
          condition_.notify_all();
        throw;
      }
 
    condition_.notify_all();
    return futures;
  }
 
  template <typename Iterator>
  auto
  submit_batch(Iterator begin, Iterator end) -> std::vector<std::future<void>>
  {
    auto result = try_submit_batch(begin, end);
    if (!result.has_value())
      throw std::runtime_error("Pool is shutting down");
    return std::move(result.value());
  }
 
  template <typename Iterator, typename F>
  void
  parallel_for_each(Iterator begin, Iterator end, F&& func)
  {
    if (is_current_worker())
      detail::throw_worker_deadlock();
    detail::parallel_for_each_chunked(*this, begin, end, std::forward<F>(func), num_threads_);
  }
 
 
 
  [[nodiscard]] auto
  size() const noexcept -> size_t
  {
    return num_threads_;
  }
 
  [[nodiscard]] auto
  pending_tasks() const -> size_t
  {
    std::lock_guard<std::mutex> lock(queue_mutex_);
    return tasks_.size();
  }
 
 
 
  auto
  configure_threads(std::string const& name_prefix, native_scheduling_policy policy = native_scheduling_policy::other,
                    native_thread_priority priority = native_thread_priority::normal())
      -> expected<void, std::error_code>
  {
    return detail::configure_worker_threads(workers_, name_prefix, policy, priority,
                                            register_workers_ ? &runtime_registry() : nullptr);
  }
 
  auto
  configure_threads(native_thread_config const& config) -> expected<void, std::error_code>
  {
    return detail::configure_worker_threads(workers_, config, register_workers_ ? &runtime_registry() : nullptr);
  }
 
  auto
  set_affinity(native_thread_affinity const& affinity) -> expected<void, std::error_code>
  {
    return detail::set_worker_affinity(workers_, affinity);
  }
 
  auto
  distribute_across_cpus() -> expected<void, std::error_code>
  {
    return detail::distribute_workers_across_cpus(workers_);
  }
 
 
 
  void
  wait_for_tasks()
  {
    if (is_current_worker())
      detail::throw_worker_deadlock();
    std::unique_lock<std::mutex> lock(queue_mutex_);
    task_finished_condition_.wait(lock, [this]
                                    { return tasks_.empty() && active_tasks_.load(std::memory_order_acquire) == 0; });
  }
 
  [[nodiscard]] auto
  is_current_worker() const noexcept -> bool
  {
    return current_pool == this;
  }
 
  void
  shutdown(shutdown_policy_backend policy = shutdown_policy_backend::drain)
  {
    if (is_current_worker())
      detail::throw_worker_deadlock();
    std::lock_guard<std::recursive_timed_mutex> shutdown_lock(shutdown_mutex_);
    if (workers_.empty())
      return;
    std::queue<queued_task> discarded;
    {
      std::lock_guard<std::mutex> lock(queue_mutex_);
      if (!stop_)
        stop_ = true;
      if (policy == shutdown_policy_backend::drop_pending && !tasks_.empty())
        tasks_.swap(discarded);
    }
    shutdown_completed_all_ = shutdown_completed_all_ && discarded.empty();
 
    condition_.notify_all();
    task_finished_condition_.notify_all();
    {
      std::queue<queued_task> empty;
      discarded.swap(empty);
    }
 
    for (auto& worker : workers_)
      {
        if (worker.joinable())
          worker.join();
      }
 
    workers_.clear();
    shutdown_completed_at_ = std::chrono::steady_clock::now();
  }
 
  auto
  shutdown_for(std::chrono::milliseconds timeout) -> bool
  {
    if (is_current_worker())
      detail::throw_worker_deadlock();
    auto const deadline = shutdown_deadline_after(timeout);
    std::unique_lock<std::recursive_timed_mutex> shutdown_lock(shutdown_mutex_, std::defer_lock);
    if (deadline == std::chrono::steady_clock::time_point::max())
      shutdown_lock.lock();
    else if (!shutdown_lock.try_lock_until(deadline))
      return false;
 
    std::unique_lock<std::mutex> lock(queue_mutex_);
    if (workers_.empty())
      return shutdown_completed_all_ && shutdown_completed_at_ <= deadline;
    if (!stop_)
      stop_ = true;
    condition_.notify_all();
    bool const drained = task_finished_condition_.wait_until(
        lock, deadline, [this] { return tasks_.empty() && active_tasks_.load(std::memory_order_acquire) == 0; });
    std::queue<queued_task> discarded;
    if (!drained)
      tasks_.swap(discarded);
    shutdown_completed_all_ = shutdown_completed_all_ && discarded.empty();
    lock.unlock();
 
    condition_.notify_all();
    task_finished_condition_.notify_all();
    {
      std::queue<queued_task> empty;
      discarded.swap(empty);
    }
    if (deadline == std::chrono::steady_clock::time_point::max())
      {
        for (auto& worker : workers_)
          if (worker.joinable())
            worker.join();
        workers_.clear();
      }
    if (drained && shutdown_completed_at_ == std::chrono::steady_clock::time_point{})
      shutdown_completed_at_ = std::chrono::steady_clock::now();
    return drained && shutdown_completed_all_;
  }
 
 
 
  [[nodiscard]] auto
  get_statistics() const -> statistics
  {
    auto const now = std::chrono::steady_clock::now();
    auto const elapsed = std::chrono::duration_cast<std::chrono::seconds>(now - start_time_);
 
    std::lock_guard<std::mutex> lock(queue_mutex_);
    statistics stats;
    stats.total_threads = num_threads_;
    stats.active_threads = active_tasks_.load(std::memory_order_acquire);
    stats.pending_tasks = tasks_.size();
    stats.completed_tasks = completed_tasks_.load(std::memory_order_acquire);
 
    if (elapsed.count() > 0)
      {
        stats.tasks_per_second = static_cast<double>(stats.completed_tasks) / elapsed.count();
      }
    else
      {
        stats.tasks_per_second = 0.0;
      }
 
    auto const total_task_time = total_task_time_.load(std::memory_order_acquire);
    if (stats.completed_tasks > 0)
      {
        stats.avg_task_time = std::chrono::microseconds(total_task_time / stats.completed_tasks);
      }
    else
      {
        stats.avg_task_time = std::chrono::microseconds(0);
      }
 
    return stats;
  }
 
 
 
  void
  set_on_task_start(task_start_callback cb)
  {
    std::lock_guard<std::mutex> lock(trace_mutex_);
    on_task_start_ = task_start_callback_storage(std::move(cb));
  }
 
  template <typename Callback,
            std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<Callback>, task_start_callback>, int> = 0>
  void
  set_on_task_start(Callback&& cb)
  {
    static_assert(std::is_invocable_r_v<void, Callback&, std::chrono::steady_clock::time_point, std::thread::id>,
                  "Task start callback must accept (time_point, std::thread::id)");
    std::lock_guard<std::mutex> lock(trace_mutex_);
    on_task_start_ = detail::make_copyable_function<void(std::chrono::steady_clock::time_point, std::thread::id)>(
        std::forward<Callback>(cb));
  }
 
  void
  set_on_task_end(task_end_callback cb)
  {
    std::lock_guard<std::mutex> lock(trace_mutex_);
    on_task_end_ = task_end_callback_storage(std::move(cb));
  }
 
  template <typename Callback,
            std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<Callback>, task_end_callback>, int> = 0>
  void
  set_on_task_end(Callback&& cb)
  {
    static_assert(std::is_invocable_r_v<void, Callback&, std::chrono::steady_clock::time_point, std::thread::id,
                                        std::chrono::microseconds>,
                  "Task end callback must accept (time_point, std::thread::id, "
                  "std::chrono::microseconds)");
    std::lock_guard<std::mutex> lock(trace_mutex_);
    on_task_end_ = detail::make_copyable_function<void(std::chrono::steady_clock::time_point, std::thread::id,
                                                       std::chrono::microseconds)>(std::forward<Callback>(cb));
  }
 
 
private:
  size_t num_threads_;
  bool register_workers_;
  detail::worker_startup_latch startup_;
  std::vector<detail::thread_backend> workers_;
  std::queue<queued_task> tasks_;
 
  mutable std::mutex queue_mutex_;
  std::condition_variable condition_;
  std::condition_variable task_finished_condition_;
  std::recursive_timed_mutex shutdown_mutex_;
  std::atomic<bool> stop_;
  bool shutdown_completed_all_{ true };
  std::chrono::steady_clock::time_point shutdown_completed_at_{};
  std::atomic<size_t> active_tasks_{ 0 };
  std::atomic<size_t> completed_tasks_{ 0 };
  std::atomic<uint64_t> total_task_time_{ 0 };
 
  std::mutex trace_mutex_;
  task_start_callback_storage on_task_start_;
  task_end_callback_storage on_task_end_;
 
  std::chrono::steady_clock::time_point start_time_;
  inline static thread_local thread_pool_backend_base* current_pool = nullptr;
 
  void
  worker_function(size_t worker_id)
  {
    detail::worker_context_guard<thread_pool_backend_base> worker_context(current_pool, this);
    std::optional<registration_guard_backend> reg_guard;
    try
      {
        if (register_workers_)
          reg_guard.emplace("pool_worker_" + std::to_string(worker_id), "threadschedule.pool");
        startup_.arrive();
      }
    catch (...)
      {
        startup_.arrive(std::current_exception());
        return;
      }
 
    while (true)
      {
        queued_task task;
        bool found_task = false;
 
        {
          std::unique_lock<std::mutex> lock(queue_mutex_);
 
          if (WaitPolicy::wait(condition_, lock, [this] { return stop_ || !tasks_.empty(); }))
            {
              if (stop_ && tasks_.empty())
                {
                  return;
                }
 
              if (!tasks_.empty())
                {
                  task = std::move(tasks_.front());
                  tasks_.pop();
                  found_task = true;
                  active_tasks_.fetch_add(1, std::memory_order_relaxed);
                }
            }
          else if (stop_)
            {
              return;
            }
        }
 
        if (found_task)
          {
            auto const start_time = std::chrono::steady_clock::now();
            auto const tid = std::this_thread::get_id();
 
            try
              {
                task_start_callback_storage on_task_start;
                {
                  std::lock_guard<std::mutex> tl(trace_mutex_);
                  on_task_start = on_task_start_;
                }
                if (on_task_start)
                  on_task_start(start_time, tid);
              }
            catch (...)
              {
              }
 
            // See work_stealing_pool_backend::worker_function for rationale.
            try
              {
                task();
              }
            catch (...)
              {
              }
            auto const end_time = std::chrono::steady_clock::now();
 
            auto const task_duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
            total_task_time_.fetch_add(task_duration.count(), std::memory_order_relaxed);
 
            try
              {
                task_end_callback_storage on_task_end;
                {
                  std::lock_guard<std::mutex> tl(trace_mutex_);
                  on_task_end = on_task_end_;
                }
                if (on_task_end)
                  on_task_end(end_time, tid, task_duration);
              }
            catch (...)
              {
              }
 
            {
              std::lock_guard<std::mutex> lock(queue_mutex_);
              active_tasks_.fetch_sub(1, std::memory_order_relaxed);
            }
            completed_tasks_.fetch_add(1, std::memory_order_relaxed);
 
            task_finished_condition_.notify_all();
          }
      }
  }
};
 
using thread_pool_backend = thread_pool_backend_base<indefinite_wait>;
 
using polling_pool_backend = thread_pool_backend_base<polling_wait<>>;
 
} // namespace threadschedule::detail