scheduled_pool.hpp

Go to the documentation of this file.

#pragma once

 
#include "expected.hpp"
#include "thread_pool.hpp"
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <future>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <thread>
 
namespace threadschedule
{

class ScheduledTaskHandle
{
  public:
    explicit ScheduledTaskHandle(uint64_t id) : id_(id), cancelled_(std::make_shared<std::atomic<bool>>(false))
    {
    }
 
    void cancel()
    {
        cancelled_->store(true, std::memory_order_release);
    }
 
    [[nodiscard]] auto is_cancelled() const noexcept -> bool
    {
        return cancelled_->load(std::memory_order_acquire);
    }
 
    [[nodiscard]] auto id() const noexcept -> uint64_t
    {
        return id_;
    }
 
  private:
    uint64_t id_;
    std::shared_ptr<std::atomic<bool>> cancelled_;
 
    template <typename>
    friend class ScheduledThreadPoolT;
    [[nodiscard]] auto get_cancel_flag() const -> std::shared_ptr<std::atomic<bool>>
    {
        return cancelled_;
    }
};

template <typename PoolType = ThreadPool>
class ScheduledThreadPoolT
{
  public:
    using Task = std::function<void()>;
    using OneShotTask = detail::move_callable<void()>;
    using PeriodicTask = detail::copyable_callable<void()>;
    using TimePoint = std::chrono::steady_clock::time_point;
    using Duration = std::chrono::steady_clock::duration;
 
    struct ScheduledTaskInfo
    {
        uint64_t id;
        TimePoint next_run;
        Duration interval; // Zero for one-time tasks
        OneShotTask one_shot_task;
        std::shared_ptr<PeriodicTask> periodic_task;
        std::shared_ptr<std::atomic<bool>> cancelled;
        bool periodic;
    };

    explicit ScheduledThreadPoolT(size_t worker_threads = std::thread::hardware_concurrency())
        : pool_(worker_threads), stop_(false), next_task_id_(1)
    {
        std::promise<Tid> scheduler_started;
        auto scheduler_ready = scheduler_started.get_future();
 
        scheduler_thread_ = ThreadWrapper([this, started = std::move(scheduler_started)]() mutable {
            started.set_value(ThreadInfo::get_thread_id());
            scheduler_loop();
        });
 
        scheduler_tid_ = scheduler_ready.get();
        (void)ThreadInfo(scheduler_tid_).set_name("ts_sched_pool");
    }
 
    ScheduledThreadPoolT(ScheduledThreadPoolT const&) = delete;
    auto operator=(ScheduledThreadPoolT const&) -> ScheduledThreadPoolT& = delete;
 
    ~ScheduledThreadPoolT()
    {
        shutdown();
    }

    auto schedule_after(Duration delay, Task task) -> ScheduledTaskHandle
    {
        auto run_time = std::chrono::steady_clock::now() + delay;
        return insert_one_shot_task(run_time, detail::make_move_callable<void()>(std::move(task)));
    }
 
    template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, Task>, int> = 0>
    auto schedule_after(Duration delay, F&& task) -> ScheduledTaskHandle
    {
        auto run_time = std::chrono::steady_clock::now() + delay;
        return insert_one_shot_task(run_time, detail::make_move_callable<void()>(std::forward<F>(task)));
    }

    auto schedule_at(TimePoint time_point, Task task) -> ScheduledTaskHandle
    {
        return insert_one_shot_task(time_point, detail::make_move_callable<void()>(std::move(task)));
    }
 
    template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, Task>, int> = 0>
    auto schedule_at(TimePoint time_point, F&& task) -> ScheduledTaskHandle
    {
        return insert_one_shot_task(time_point, detail::make_move_callable<void()>(std::forward<F>(task)));
    }

    auto schedule_periodic(Duration interval, Task task) -> ScheduledTaskHandle
    {
        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>, int> = 0>
    auto schedule_periodic(Duration interval, F&& task) -> ScheduledTaskHandle
    {
        return schedule_periodic_after(Duration::zero(), interval, std::forward<F>(task));
    }

    auto schedule_periodic_after(Duration initial_delay, Duration interval, Task task) -> ScheduledTaskHandle
    {
        auto const run_time = std::chrono::steady_clock::now() + initial_delay;
        return insert_periodic_task(run_time, interval, detail::make_copyable_callable<void()>(std::move(task)));
    }
 
    template <typename F, std::enable_if_t<!std::is_same_v<detail::remove_cvref_t<F>, Task>, int> = 0>
    auto schedule_periodic_after(Duration initial_delay, Duration interval, F&& task) -> ScheduledTaskHandle
    {
        auto const run_time = std::chrono::steady_clock::now() + initial_delay;
        return insert_periodic_task(run_time, interval,
                                    detail::make_copyable_callable<void()>(std::forward<F>(task)));
    }

    static void cancel(ScheduledTaskHandle& handle)
    {
        handle.cancel();
    }

    [[nodiscard]] auto scheduled_count() const -> size_t
    {
        std::lock_guard<std::mutex> lock(mutex_);
        return scheduled_tasks_.size();
    }

    [[nodiscard]] auto thread_pool() -> PoolType&
    {
        return pool_;
    }

    void shutdown()
    {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            if (stop_)
                return;
            stop_ = true;
        }
 
        condition_.notify_one();
 
        if (scheduler_thread_.joinable())
        {
            scheduler_thread_.join();
        }
 
        pool_.shutdown();
    }

    auto configure_threads(std::string const& name_prefix, SchedulingPolicy policy = SchedulingPolicy::OTHER,
                           ThreadPriority priority = ThreadPriority::normal())
    {
        return pool_.configure_threads(name_prefix, policy, priority);
    }
 
    [[nodiscard]] auto scheduler_thread_info() const -> std::optional<ThreadInfo>
    {
        if (!scheduler_thread_.joinable() || scheduler_tid_ == Tid{})
            return std::nullopt;
        return ThreadInfo(scheduler_tid_);
    }
 
    auto configure_scheduler_thread(std::string const& name, SchedulingPolicy policy = SchedulingPolicy::OTHER,
                                    ThreadPriority priority = ThreadPriority::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);
    }
 
  private:
    PoolType pool_;
    ThreadWrapper scheduler_thread_;
    Tid scheduler_tid_{};
 
    mutable std::mutex mutex_;
    std::condition_variable condition_;
    std::atomic<bool> stop_;
 
    std::multimap<TimePoint, ScheduledTaskInfo> scheduled_tasks_;
    std::atomic<uint64_t> next_task_id_;
 
    auto insert_one_shot_task(TimePoint run_time, OneShotTask task) -> ScheduledTaskHandle
    {
        std::lock_guard<std::mutex> lock(mutex_);
 
        uint64_t const task_id = next_task_id_++;
        ScheduledTaskHandle handle(task_id);
 
        if (stop_)
        {
            handle.cancel();
            return handle;
        }
 
        ScheduledTaskInfo info;
        info.id = task_id;
        info.next_run = run_time;
        info.interval = Duration::zero();
        info.one_shot_task = std::move(task);
        info.cancelled = handle.get_cancel_flag();
        info.periodic = false;
 
        scheduled_tasks_.insert({run_time, std::move(info)});
        condition_.notify_one();
 
        return handle;
    }
 
    auto insert_periodic_task(TimePoint run_time, Duration interval, PeriodicTask task) -> ScheduledTaskHandle
    {
        std::lock_guard<std::mutex> lock(mutex_);
 
        uint64_t const task_id = next_task_id_++;
        ScheduledTaskHandle handle(task_id);
 
        if (stop_)
        {
            handle.cancel();
            return handle;
        }
 
        ScheduledTaskInfo info;
        info.id = task_id;
        info.next_run = run_time;
        info.interval = interval;
        info.periodic_task = std::make_shared<PeriodicTask>(std::move(task));
        info.cancelled = handle.get_cancel_flag();
        info.periodic = true;
 
        scheduled_tasks_.insert({run_time, std::move(info)});
        condition_.notify_one();
 
        return handle;
    }
 
    void scheduler_loop()
    {
        while (true)
        {
            std::unique_lock<std::mutex> lock(mutex_);
 
            // 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)
            {
                condition_.wait_until(lock, it->first, [this] { return stop_.load(); });
 
                if (stop_)
                    return;
 
                continue;
            }
 
            // Extract the task info
            ScheduledTaskInfo info = std::move(it->second);
            scheduled_tasks_.erase(it);
 
            // Check if cancelled
            if (info.cancelled->load(std::memory_order_acquire))
            {
                continue;
            }
 
            // Schedule for execution in the thread pool
            try
            {
                auto cancelled_flag = info.cancelled;
                if (info.periodic)
                {
                    auto periodic_task = info.periodic_task;
                    pool_.post([periodic_task, cancelled_flag]() {
                        if (!cancelled_flag->load(std::memory_order_acquire))
                        {
                            (*periodic_task)();
                        }
                    });
 
                    if (!info.cancelled->load(std::memory_order_acquire))
                    {
                        info.next_run += info.interval;
                        scheduled_tasks_.insert({info.next_run, std::move(info)});
                    }
                }
                else
                {
                    auto one_shot_task = std::move(info.one_shot_task);
                    pool_.post([task = std::move(one_shot_task), cancelled_flag]() mutable {
                        if (!cancelled_flag->load(std::memory_order_acquire))
                        {
                            task();
                        }
                    });
                }
            }
            catch (...)
            {
                // Thread pool might be shutting down
            }
        }
    }
};

using ScheduledThreadPool = ScheduledThreadPoolT<ThreadPool>;
using ScheduledHighPerformancePool = ScheduledThreadPoolT<HighPerformancePool>;
using ScheduledFastThreadPool = ScheduledThreadPoolT<FastThreadPool>;
using ScheduledLightweightPool = ScheduledThreadPoolT<LightweightPool>;
 
} // namespace threadschedule