scheduled_pool_backend_base.hpp

Go to the documentation of this file.

#pragma once
 
#include "../pool/backend.hpp"
#include "../pool/worker_count.hpp"
#include "../time.hpp"
#include "scheduled_task_backend.hpp"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <future>
#include <limits>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <thread>
#include <type_traits>
#include <utility>
 
namespace threadschedule::detail
{
 
template <typename PoolType = thread_pool_backend>
class scheduled_pool_backend_base
{
public:
  using task_type = std::function<void()>;
  using one_shot_task_type = detail::move_only_function<void()>;
  using periodic_task_type = detail::move_only_function<void()>;
  using time_point = std::chrono::steady_clock::time_point;
  using duration = std::chrono::steady_clock::duration;
 
  struct periodic_task_state
  {
    explicit periodic_task_state(periodic_task_type task_value) : task(std::move(task_value)) {}
 
    periodic_task_type task;
    std::atomic<bool> running{ false };
  };
 
  struct scheduled_dispatch_state
  {
    explicit scheduled_dispatch_state(std::shared_ptr<detail::scheduled_cancellation_state> cancellation_value,
                                      std::shared_ptr<periodic_task_state> periodic_value = {})
        : cancellation(std::move(cancellation_value)), periodic(std::move(periodic_value))
    {
    }
 
    ~scheduled_dispatch_state()
    {
      if (started.load(std::memory_order_acquire))
        return;
      cancellation->pool_stopped.store(true, std::memory_order_release);
      if (periodic)
        periodic->running.store(false, std::memory_order_release);
    }
 
    std::shared_ptr<detail::scheduled_cancellation_state> cancellation;
    std::shared_ptr<periodic_task_state> periodic;
    std::atomic<bool> started{ false };
  };
 
  class scheduled_task_info
  {
  public:
    enum class kind : std::uint8_t
    {
      one_shot,
      periodic
    };
 
    [[nodiscard]] static auto
    one_shot(std::uint64_t id, time_point run_time, one_shot_task_type task,
             std::shared_ptr<detail::scheduled_cancellation_state> cancellation) -> scheduled_task_info
    {
      return scheduled_task_info(id, run_time, std::move(task), std::move(cancellation));
    }
 
    [[nodiscard]] static auto
    periodic(std::uint64_t id, time_point run_time, duration interval, periodic_task_type task,
             std::shared_ptr<detail::scheduled_cancellation_state> cancellation) -> scheduled_task_info
    {
      if (interval <= duration::zero())
        throw std::invalid_argument("scheduled_pool periodic interval must be positive");
      return scheduled_task_info(id, run_time, interval, std::move(task), std::move(cancellation));
    }
 
    [[nodiscard]] auto
    type() const noexcept -> kind
    {
      return kind_;
    }
    [[nodiscard]] auto
    cancellation() const noexcept -> std::shared_ptr<detail::scheduled_cancellation_state> const&
    {
      return cancellation_;
    }
    [[nodiscard]] auto
    periodic_state() const noexcept -> std::shared_ptr<periodic_task_state> const&
    {
      return periodic_task_;
    }
    [[nodiscard]] auto
    take_one_shot() noexcept -> one_shot_task_type
    {
      return std::move(one_shot_task_);
    }
    [[nodiscard]] auto
    next_run() const noexcept -> time_point
    {
      return next_run_;
    }
 
    void
    advance(time_point now)
    {
      auto const interval = *interval_;
      auto advanced = detail::checked_deadline_after(next_run_, interval);
      if (!advanced)
        throw std::system_error(advanced.error(), "scheduled_pool periodic deadline");
      next_run_ = advanced.value();
      if (next_run_ <= now)
        {
          auto const quotient = (now - next_run_) / interval;
          auto const maximum = (std::numeric_limits<decltype(quotient)>::max)();
          if (quotient == maximum)
            throw std::system_error(std::make_error_code(std::errc::value_too_large),
                                    "scheduled_pool periodic deadline");
          auto const missed = quotient + 1;
          if (missed > maximum / interval.count())
            throw std::system_error(std::make_error_code(std::errc::value_too_large),
                                    "scheduled_pool periodic deadline");
          advanced = detail::checked_deadline_after(next_run_, interval * missed);
          if (!advanced)
            throw std::system_error(advanced.error(), "scheduled_pool periodic deadline");
          next_run_ = advanced.value();
        }
    }
 
  private:
    scheduled_task_info(std::uint64_t id, time_point run_time, one_shot_task_type task,
                        std::shared_ptr<detail::scheduled_cancellation_state> cancellation)
        : id_(id), next_run_(run_time), one_shot_task_(std::move(task)), cancellation_(std::move(cancellation)),
          kind_(kind::one_shot)
    {
    }
 
    scheduled_task_info(std::uint64_t id, time_point run_time, duration interval, periodic_task_type task,
                        std::shared_ptr<detail::scheduled_cancellation_state> cancellation)
        : id_(id), next_run_(run_time), interval_(interval),
          periodic_task_(std::make_shared<periodic_task_state>(std::move(task))),
          cancellation_(std::move(cancellation)), kind_(kind::periodic)
    {
    }
 
    std::uint64_t id_;
    time_point next_run_;
    std::optional<duration> interval_;
    one_shot_task_type one_shot_task_;
    std::shared_ptr<periodic_task_state> periodic_task_;
    std::shared_ptr<detail::scheduled_cancellation_state> cancellation_;
    kind kind_;
  };
 
  explicit scheduled_pool_backend_base(size_t worker_threads = default_worker_count())
      : pool_(worker_threads), stop_(false), next_task_id_(1)
  {
    start_scheduler();
  }
 
  template <typename T = PoolType, std::enable_if_t<std::is_constructible_v<T, size_t, bool>, int> = 0>
  scheduled_pool_backend_base(size_t worker_threads, bool register_workers)
      : pool_(worker_threads, register_workers), stop_(false), next_task_id_(1)
  {
    start_scheduler();
  }
 
  scheduled_pool_backend_base(scheduled_pool_backend_base const&) = delete;
  auto operator=(scheduled_pool_backend_base const&) -> scheduled_pool_backend_base& = delete;
 
  ~scheduled_pool_backend_base()
  {
    shutdown();
  }
 
  auto
  schedule_after(duration delay, task_type task) -> scheduled_task_backend
  {
    auto run_time = detail::checked_deadline_after(std::chrono::steady_clock::now(), delay);
    if (!run_time)
      throw std::system_error(run_time.error(), "scheduled_pool delay");
    return insert_one_shot_task(run_time.value(), detail::make_move_only_function<void()>(std::move(task)));
  }
 
  template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, task_type>, int> = 0>
  auto
  schedule_after(duration delay, F&& task) -> scheduled_task_backend
  {
    auto run_time = detail::checked_deadline_after(std::chrono::steady_clock::now(), delay);
    if (!run_time)
      throw std::system_error(run_time.error(), "scheduled_pool delay");
    return insert_one_shot_task(run_time.value(), detail::make_move_only_function<void()>(std::forward<F>(task)));
  }
 
  auto
  schedule_at(time_point time_point, task_type task) -> scheduled_task_backend
  {
    return insert_one_shot_task(time_point, detail::make_move_only_function<void()>(std::move(task)));
  }
 
  template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, task_type>, int> = 0>
  auto
  schedule_at(time_point time_point, F&& task) -> scheduled_task_backend
  {
    return insert_one_shot_task(time_point, detail::make_move_only_function<void()>(std::forward<F>(task)));
  }
 
  auto
  schedule_periodic(duration interval, task_type task) -> scheduled_task_backend
  {
    return schedule_periodic_after(duration::zero(), interval, std::move(task));
  }
 
  template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, task_type>, int> = 0>
  auto
  schedule_periodic(duration interval, F&& task) -> scheduled_task_backend
  {
    return schedule_periodic_after(duration::zero(), interval, std::forward<F>(task));
  }
 
  auto
  schedule_periodic_after(duration initial_delay, duration interval, task_type task) -> scheduled_task_backend
  {
    if (interval <= duration::zero())
      throw std::invalid_argument("scheduled_pool periodic interval must be positive");
    auto const run_time = detail::checked_deadline_after(std::chrono::steady_clock::now(), initial_delay);
    if (!run_time)
      throw std::system_error(run_time.error(), "scheduled_pool initial delay");
    auto const next_run = detail::checked_deadline_after(run_time.value(), interval);
    if (!next_run)
      throw std::system_error(next_run.error(), "scheduled_pool periodic deadline");
    return insert_periodic_task(run_time.value(), interval, detail::make_move_only_function<void()>(std::move(task)));
  }
 
  template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, task_type>, int> = 0>
  auto
  schedule_periodic_after(duration initial_delay, duration interval, F&& task) -> scheduled_task_backend
  {
    if (interval <= duration::zero())
      throw std::invalid_argument("scheduled_pool periodic interval must be positive");
    auto const run_time = detail::checked_deadline_after(std::chrono::steady_clock::now(), initial_delay);
    if (!run_time)
      throw std::system_error(run_time.error(), "scheduled_pool initial delay");
    auto const next_run = detail::checked_deadline_after(run_time.value(), interval);
    if (!next_run)
      throw std::system_error(next_run.error(), "scheduled_pool periodic deadline");
    return insert_periodic_task(run_time.value(), interval,
                                detail::make_move_only_function<void()>(std::forward<F>(task)));
  }
 
  static void
  cancel(scheduled_task_backend& handle)
  {
    handle.cancel();
  }
 
  [[nodiscard]] auto
  scheduled_count() const -> size_t
  {
    std::lock_guard<std::mutex> lock(mutex_);
    return scheduled_tasks_.size();
  }
 
  [[nodiscard]] auto
  is_current_context() const noexcept -> bool
  {
    return pool_.is_current_worker() || current_scheduler == this;
  }
 
  [[nodiscard]] auto
  thread_pool() -> PoolType&
  {
    return pool_;
  }
 
  void
  shutdown(shutdown_policy_backend policy = shutdown_policy_backend::drain)
  {
    if (is_current_context())
      detail::throw_worker_deadlock();
    std::lock_guard<std::recursive_mutex> shutdown_lock(shutdown_mutex_);
    std::multimap<time_point, scheduled_task_info> discarded;
    {
      std::lock_guard<std::mutex> lock(mutex_);
      if (stop_)
        return;
      stop_ = true;
      for (auto const& task : scheduled_tasks_)
        task.second.cancellation()->pool_stopped.store(true, std::memory_order_release);
      scheduled_tasks_.swap(discarded);
    }
 
    discarded.clear();
 
    condition_.notify_one();
 
    {
      std::lock_guard<std::mutex> scheduler_lock(scheduler_mutex_);
      if (scheduler_thread_.joinable())
        {
          scheduler_thread_.join();
          scheduler_tid_ = native_thread_id{};
        }
    }
 
    pool_.shutdown(policy);
  }
 
  auto
  configure_threads(std::string const& name_prefix, native_scheduling_policy policy = native_scheduling_policy::other,
                    native_thread_priority priority = native_thread_priority::normal())
  {
    return pool_.configure_threads(name_prefix, policy, priority);
  }
 
  auto
  configure_threads(native_thread_config const& config) -> expected<void, std::error_code>
  {
    return pool_.configure_threads(config);
  }
 
  [[nodiscard]] auto
  scheduler_thread_info() const -> std::optional<thread_info>
  {
    if (current_scheduler == this)
      return thread_info(thread_info::get_thread_id());
    std::lock_guard<std::mutex> scheduler_lock(scheduler_mutex_);
    if (!scheduler_thread_.joinable() || scheduler_tid_ == native_thread_id{})
      return std::nullopt;
    return thread_info(scheduler_tid_);
  }
 
  auto
  configure_scheduler_thread(std::string const& name, native_scheduling_policy policy = native_scheduling_policy::other,
                             native_thread_priority priority = native_thread_priority::normal())
      -> expected<void, std::error_code>
  {
    auto info = scheduler_thread_info();
    if (!info.has_value())
      return unexpected(std::make_error_code(std::errc::no_such_process));
    return detail::configure_thread(info.value(), name, policy, priority);
  }
 
  auto
  configure_scheduler_thread(native_thread_config const& config) -> expected<void, std::error_code>
  {
    auto info = scheduler_thread_info();
    if (!info.has_value())
      return unexpected(std::make_error_code(std::errc::no_such_process));
    return info->configure(config);
  }
 
private:
  PoolType pool_;
  detail::thread_backend scheduler_thread_;
  native_thread_id scheduler_tid_{};
 
  mutable std::mutex mutex_;
  std::condition_variable condition_;
  std::atomic<bool> stop_;
  std::recursive_mutex shutdown_mutex_;
  mutable std::mutex scheduler_mutex_;
 
  std::multimap<time_point, scheduled_task_info> scheduled_tasks_;
  std::atomic<uint64_t> next_task_id_;
 
  inline static thread_local scheduled_pool_backend_base* current_scheduler = nullptr;
 
  void
  start_scheduler()
  {
    std::promise<native_thread_id> scheduler_started;
    auto scheduler_ready = scheduler_started.get_future();
 
    scheduler_thread_ = detail::thread_backend(
        [this, started = std::move(scheduler_started)]() mutable
          {
            current_scheduler = this;
            started.set_value(thread_info::get_thread_id());
            scheduler_loop();
            current_scheduler = nullptr;
          });
 
    scheduler_tid_ = scheduler_ready.get();
    (void)thread_info(scheduler_tid_).set_name("ts_sched_pool");
  }
 
  auto
  insert_one_shot_task(time_point run_time, one_shot_task_type task) -> scheduled_task_backend
  {
    std::lock_guard<std::mutex> lock(mutex_);
 
    uint64_t const task_id = next_task_id_++;
    scheduled_task_backend handle(task_id);
 
    if (stop_)
      {
        handle.cancel();
        return handle;
      }
 
    scheduled_tasks_.insert(
        { run_time, scheduled_task_info::one_shot(task_id, run_time, std::move(task), handle.get_cancellation()) });
    condition_.notify_one();
 
    return handle;
  }
 
  auto
  insert_periodic_task(time_point run_time, duration interval, periodic_task_type task) -> scheduled_task_backend
  {
    std::lock_guard<std::mutex> lock(mutex_);
 
    uint64_t const task_id = next_task_id_++;
    scheduled_task_backend handle(task_id);
 
    if (stop_)
      {
        handle.cancel();
        return handle;
      }
 
    scheduled_tasks_.insert({ run_time, scheduled_task_info::periodic(task_id, run_time, interval, std::move(task),
                                                                      handle.get_cancellation()) });
    condition_.notify_one();
 
    return handle;
  }
 
  void
  scheduler_loop()
  {
    while (true)
      {
        std::unique_lock<std::mutex> lock(mutex_);
 
        if (stop_)
          return;
 
        // Wait until we have tasks or need to stop
        if (scheduled_tasks_.empty())
          {
            condition_.wait(lock, [this] { return stop_ || !scheduled_tasks_.empty(); });
 
            if (stop_)
              return;
          }
 
        // Get the next task to execute
        auto const now = std::chrono::steady_clock::now();
        auto it = scheduled_tasks_.begin();
 
        if (it == scheduled_tasks_.end())
          {
            continue;
          }
 
        // Wait until it's time to execute
        if (it->first > now)
          {
            // wait_until releases the mutex while retaining a reference to
            // its deadline. Shutdown may clear scheduled_tasks_ during that
            // wait, so never pass a map key by reference here.
            auto const next_run = it->first;
            condition_.wait_until(lock, next_run);
 
            if (stop_)
              return;
 
            continue;
          }
 
        // Extract the task info
        scheduled_task_info info = std::move(it->second);
        scheduled_tasks_.erase(it);
        lock.unlock();
 
        // Check if cancelled
        if (info.cancellation()->user_cancelled.load(std::memory_order_acquire)
            || info.cancellation()->pool_stopped.load(std::memory_order_acquire))
          {
            continue;
          }
 
        // Schedule for execution in the thread pool
        try
          {
            auto cancellation = info.cancellation();
            if (info.type() == scheduled_task_info::kind::periodic)
              {
                auto periodic_task = info.periodic_state();
                bool const already_running = periodic_task->running.exchange(true, std::memory_order_acq_rel);
                if (!already_running)
                  {
                    auto dispatch = std::make_shared<scheduled_dispatch_state>(cancellation, periodic_task);
                    try
                      {
                        pool_.post(
                            [periodic_task, cancellation, dispatch]()
                              {
                                dispatch->started.store(true, std::memory_order_release);
                                try
                                  {
                                    if (!cancellation->user_cancelled.load(std::memory_order_acquire))
                                      periodic_task->task();
                                  }
                                catch (...)
                                  {
                                    periodic_task->running.store(false, std::memory_order_release);
                                    throw;
                                  }
                                periodic_task->running.store(false, std::memory_order_release);
                              });
                      }
                    catch (...)
                      {
                        periodic_task->running.store(false, std::memory_order_release);
                        throw;
                      }
                  }
 
                info.advance(std::chrono::steady_clock::now());
                std::lock_guard<std::mutex> schedule_lock(mutex_);
                if (!stop_ && !cancellation->user_cancelled.load(std::memory_order_acquire))
                  scheduled_tasks_.insert({ info.next_run(), std::move(info) });
              }
            else
              {
                auto one_shot_task = info.take_one_shot();
                auto dispatch = std::make_shared<scheduled_dispatch_state>(cancellation);
                pool_.post(
                    [task = std::move(one_shot_task), cancellation, dispatch]() mutable
                      {
                        dispatch->started.store(true, std::memory_order_release);
                        if (!cancellation->user_cancelled.load(std::memory_order_acquire))
                          {
                            task();
                          }
                      });
              }
          }
        catch (...)
          {
            info.cancellation()->pool_stopped.store(true, std::memory_order_release);
            // Thread pool might be shutting down
          }
      }
  }
};
 
using scheduled_pool_backend = scheduled_pool_backend_base<thread_pool_backend>;
using scheduled_work_stealing_pool_backend = scheduled_pool_backend_base<work_stealing_pool_backend>;
using scheduled_polling_pool_backend = scheduled_pool_backend_base<polling_pool_backend>;
using scheduled_lightweight_pool_backend = scheduled_pool_backend_base<lightweight_pool_backend>;
 
} // namespace threadschedule::detail