scheduler.h

#pragma once
 
#include "core.h"
 
#include <chrono>
#include <condition_variable>
#include <functional>
#include <loguru.hpp>
#include <mutex>
#include <optional>
#include <queue>
#include <string>
#include <thread>
#include <unordered_set>
 
namespace util {
 
namespace impl {
 
class Notifiable {
public:
    virtual void notify() = 0;
};
 
template <typename T>
void notifyScheduler(T& scheduler);
} // namespace impl

template <typename Clock = std::chrono::steady_clock>
class PeriodicScheduler : private impl::Notifiable {
public:
    using eventid_t = uint64_t;

    explicit PeriodicScheduler(const std::optional<std::string>& name = std::nullopt)
        : name(name), newEventAdded(false), quitting(false), nextID(0),
          thread(std::bind(&PeriodicScheduler::threadFn, this)) {}
 
    PeriodicScheduler(const PeriodicScheduler&) = delete;

    ~PeriodicScheduler() {
        {
            std::unique_lock lock(scheduleMutex);
            quitting = true;
            // wake up the thread with the cv so it can quit
            newEventAdded = true;
        }
        scheduleCV.notify_one();
        thread.join();
    }
 
    PeriodicScheduler& operator=(const PeriodicScheduler&) = delete;

    void clear() {
        std::unique_lock lock(scheduleMutex);
        schedule = decltype(schedule)();
        toRemove.clear();
        nextID = 0;
    }

    void removeEvent(eventid_t id) {
        std::unique_lock lock(scheduleMutex);
        toRemove.insert(id);
    }

    eventid_t scheduleEvent(std::chrono::milliseconds period,
                            const std::function<void()>& fn) {
        eventid_t id;
        {
            std::unique_lock lock(scheduleMutex);
            id = nextID++;
            auto now = Clock::now();
            schedule.push({id, now + period, period, fn});
            newEventAdded = true;
        }
        scheduleCV.notify_one();
        return id;
    }
 
private:
    friend void util::impl::notifyScheduler<>(PeriodicScheduler&);
 
    void notify() override {
        scheduleCV.notify_all();
    }

    void threadFn() {
        if (name.has_value()) {
            loguru::set_thread_name(name->c_str());
        }
        std::unique_lock lock(scheduleMutex);
        while (!quitting) {
            if (schedule.empty()) {
                scheduleCV.wait(lock, [&] { return newEventAdded; });
                newEventAdded = false;
            } else {
                auto now = Clock::now();
                schedule_t event = schedule.top();
                if (toRemove.find(event.id) != toRemove.end()) {
                    // if we should remove this id then remove the event and don't reschedule
                    schedule.pop();
                    toRemove.erase(event.id);
                } else if (event.nextSendTime <= now) {
                    // execute the event and reschedule it
                    schedule.pop();
                    event.fn();
                    event.nextSendTime += event.period;
                    schedule.push(event);
                } else {
                    // wait until the event should be executed
                    scheduleCV.wait_until(lock, event.nextSendTime,
                                          [&] { return newEventAdded; });
                    newEventAdded = false;
                }
            }
        }
    }

    struct schedule_t {
        eventid_t id;
        std::chrono::time_point<Clock> nextSendTime;
        std::chrono::milliseconds period;
        std::function<void()> fn;
        friend bool operator>(const PeriodicScheduler::schedule_t& t1,
                              const PeriodicScheduler::schedule_t& t2) {
            return t1.nextSendTime > t2.nextSendTime;
        }
    };
 
    std::optional<std::string> name;
    bool newEventAdded;
    bool quitting;
    eventid_t nextID = 0;
    std::thread thread;
    // we need a min priority queue
    std::priority_queue<schedule_t, std::vector<schedule_t>, std::greater<schedule_t>>
        schedule;
    std::mutex scheduleMutex;
    std::condition_variable scheduleCV;
    std::unordered_set<eventid_t> toRemove;
};

template <typename Clock = std::chrono::steady_clock>
class Watchdog : private impl::Notifiable {
public:
    Watchdog(const std::string& name, std::chrono::milliseconds duration,
             const std::function<void()>& callback, bool keepCallingOnDeath = false)
        : name(name), duration(duration), callback(callback),
          keepCallingOnDeath(keepCallingOnDeath), fed(false), quitting(false),
          thread(std::bind(&Watchdog::threadFn, this)) {}

    Watchdog(std::chrono::milliseconds duration, const std::function<void()>& callback,
             bool keepCallingOnDeath = false)
        : name(std::nullopt), duration(duration), callback(callback),
          keepCallingOnDeath(keepCallingOnDeath), fed(false), quitting(false),
          thread(std::bind(&Watchdog::threadFn, this)) {}
 
    Watchdog(const Watchdog&) = delete;
 
    ~Watchdog() {
        {
            std::lock_guard lock(mutex);
            quitting = true;
        }
        cv.notify_all();
        if (thread.joinable()) {
            thread.join();
        }
    }
 
    Watchdog& operator=(const Watchdog&) = delete;

    void feed() {
        {
            std::lock_guard lock(mutex);
            fed = true;
        }
        cv.notify_one();
    }
 
private:
    friend void util::impl::notifyScheduler<>(Watchdog&);
 
    std::optional<std::string> name;
    std::chrono::milliseconds duration;
    std::function<void()> callback;
    const bool keepCallingOnDeath;
    bool fed;
    bool quitting;
    std::mutex mutex;
    std::condition_variable cv;
    std::thread thread;
 
    void notify() override {
        cv.notify_all();
    }
 
    void threadFn() {
        if (name.has_value()) {
            loguru::set_thread_name(name->c_str());
        }
        std::unique_lock lock(mutex);
        while (!quitting) {
            std::chrono::time_point<Clock> wakeTime = Clock::now() + duration;
            if (cv.wait_until(lock, wakeTime, [&]() { return fed || quitting; })) {
                if (quitting) {
                    break;
                }
                fed = false;
            } else {
                callback();
                if (!keepCallingOnDeath) {
                    cv.wait(lock, [&]() { return fed || quitting; });
                }
            }
        }
    }
};

template <typename Clock = std::chrono::steady_clock>
class AsyncTask : private virtual impl::Notifiable {
public:
    AsyncTask(const std::optional<std::string>& name = std::nullopt)
        : name(name), running(false), quitting(false) {}
 
    AsyncTask(const AsyncTask&) = delete;
 
    virtual ~AsyncTask() {
        stop();
        if (thread.joinable()) {
            thread.join();
        }
    }
 
    AsyncTask& operator=(const AsyncTask&) = delete;

    virtual void start() {
        std::lock_guard lock(mutex);
        if (!running) {
            std::lock_guard threadLock(threadMutex);
            if (thread.joinable()) {
                thread.join();
            }
            running = true;
            quitting = false;
            thread = std::thread(&AsyncTask::run, this);
        }
    }

    virtual void stop() {
        bool isRunning = false;
        {
            std::lock_guard lock(mutex);
            if (isRunningInternal()) {
                quitting = true;
                isRunning = true;
            }
        }
        if (isRunning) {
            cv.notify_one();
            std::lock_guard threadLock(threadMutex);
            if (thread.joinable()) {
                thread.join();
            }
        }
    }

    bool isRunning() {
        std::lock_guard lock(mutex);
        return running;
    }
 
protected:
    virtual void task(std::unique_lock<std::mutex>& lock) = 0;

    bool isRunningInternal() {
        return running;
    }

    bool wait_until(std::unique_lock<std::mutex>& lock,
                    const std::chrono::time_point<Clock>& tp) {
        return cv.wait_until(lock, tp, [&]() { return quitting; });
    }

    template <typename Rep, typename Period>
    bool wait_for(std::unique_lock<std::mutex>& lock,
                  const std::chrono::duration<Rep, Period>& dur) {
        return wait_until(lock, Clock::now() + dur);
    }

    void wait_until_done(std::unique_lock<std::mutex>& lock) {
        return cv.wait(lock, [&]() { return quitting; });
    }

    void notify() override {
        cv.notify_all();
    }
 
private:
    std::optional<std::string> name;
    bool running;
    bool quitting;
    std::condition_variable cv;
    std::thread thread;
 
    // If acquiring both mutexes, acquire mutex first then threadMutex.
    std::mutex mutex;     // protects everything except thread
    std::mutex threadMutex; // protects only thread
 
    friend void util::impl::notifyScheduler<>(AsyncTask&);
 
    void run() {
        if (name.has_value()) {
            loguru::set_thread_name(name->c_str());
        }
        std::unique_lock<std::mutex> lock(mutex);
        // clear flags when exiting
        RAIIHelper r([&]() { running = quitting = false; });
        task(lock);
    }
};

template <typename Clock = std::chrono::steady_clock>
class PeriodicTask : public AsyncTask<Clock>, private virtual impl::Notifiable {
public:
    PeriodicTask(const std::chrono::milliseconds& period, const std::function<void()>& f)
        : AsyncTask<Clock>(), period(period), f(f) {}
 
protected:
    void task(std::unique_lock<std::mutex>& lock) override {
        auto event = scheduler.scheduleEvent(period, f);
        AsyncTask<Clock>::wait_until_done(lock);
        scheduler.removeEvent(event);
    }
 
    void notify() override {
        impl::notifyScheduler(scheduler);
        AsyncTask<Clock>::notify();
    }
 
private:
    std::chrono::milliseconds period;
    std::function<void()> f;
 
    inline static PeriodicScheduler<Clock> scheduler =
        PeriodicScheduler<Clock>("PeriodicTask_Scheduler");
 
    friend void util::impl::notifyScheduler<>(PeriodicTask&);
};
 
namespace impl {

template <typename T>
void notifyScheduler(T& scheduler) {
    Notifiable& n = scheduler;
    n.notify();
}
 
} // namespace impl
 
} // namespace util