scheduler.h

#pragma once
 
#include "function.h"
#include "future.h"
#include "thread_pool.h"
 
#include <mutex>
#include <chrono>
namespace coro {
 
 
 
template<typename CondVar>
class scheduler_t {
public:
 
    using notify_t = promise<void>::notify;
 
    using resume_cb = function<void(notify_t)>;
 
    using promise_t = promise<void>;
    using ident_t = const void *;
 
 
    scheduler_t() = default;
 
 
    template<std::convertible_to<CondVar> CondVarCfg>
    explicit scheduler_t(CondVarCfg &&cfg):_cond(std::forward<CondVarCfg>(cfg)) {}
 
    ~scheduler_t(){
        stop();
    }
 
 
 
    future<void> sleep_until(std::chrono::system_clock::time_point tp, ident_t ident = nullptr) {
        return [&](auto promise) {
            std::unique_lock lk(_mx);
            if (std::find(_blk.begin(), _blk.end(), ident) != _blk.end()) return ;
            _items.push_back({std::move(promise), tp, ident});
            std::push_heap(_items.begin(), _items.end(), compare_items);
            notify();
        };
    }
 
 
    template<typename A, typename B>
    future<void> sleep_for(std::chrono::duration<A,B> dur, ident_t ident = nullptr) {
        return sleep_until(std::chrono::system_clock::now()+dur, ident);
    }
 
 
 
 
 
    template<typename T, std::invocable<notify_t> ResumeCB>
    decltype(auto) run(future<T> &fut, ResumeCB &&cb) {
        bool stopflag = false;
        if (fut.await_ready() || !fut.set_callback([this,&stopflag]{
            std::lock_guard _(_mx);
            stopflag = true;
            notify();
        })) {
            return fut.await_resume();
        }
        auto cur = _current;
        trace::on_block([&] {
            worker(std::forward<ResumeCB>(cb), stopflag);
        }, fut);;
        _current = cur;
        return fut.await_resume();
    }
 
 
    template<typename T, std::invocable<notify_t> ResumeCB>
    decltype(auto) run(future<T> &&fut, ResumeCB &&cb) {
        return run(fut, std::forward<ResumeCB>(cb));
    }
 
 
    template<typename T>
    decltype(auto) run(future<T> &fut) {
        return run(fut,[](notify_t){});
    }
 
 
    template<typename T>
    decltype(auto) run(future<T> &&fut) {
        return run(fut);
    }
 
    using thread_pool_t = coro::thread_pool;
 
 
    static auto thread_pool(unsigned int threads) {
        return [thr = std::make_shared<thread_pool_t>(threads)](notify_t ntf){
            thr->enqueue(std::move(ntf));
        };
    }
 
    static auto thread_pool(std::shared_ptr<coro::thread_pool> pool) {
        return [pool](notify_t ntf){
            pool->enqueue(std::move(ntf));
        };
    }
 
 
    template<std::invocable<notify_t> ResumeCB>
    bool start(ResumeCB &&resumecb) {
        if (_thr.joinable()) return false;
        _stop = false;
        _thr = std::thread([this, rcb = std::move(resumecb)]() mutable {
            worker(std::move(rcb), _stop);
        });
        return true;
    }
 
 
    bool start() {
        return start([](notify_t){});
    }
 
    void stop() {
        std::unique_lock lk(_mx);
        _stop = true;
        lk.unlock();
        if (scheduler_t::_current != this) {
            notify();
            if (_thr.joinable()) _thr.join();
        } else {
            scheduler_t::_current = nullptr;
            if (_thr.joinable()) _thr.detach();
        }
    }
 
 
 
    static scheduler_t * current() {return _current;}
 
 
    class pending_cancel {
    private:
        struct deleter {
            ident_t _ident;
            void operator()(scheduler_t *sch) {
                sch->_blk.erase(std::remove(sch->_blk.begin(),sch->_blk.end(), _ident), sch->_blk.end());
            }
        };
 
        std::unique_ptr<scheduler_t, deleter> _ptr;
 
        pending_cancel(scheduler_t &sch, ident_t ident)
            :_ptr(&sch, {ident}) {};
 
        friend class scheduler_t;
    };
 
 
    pending_cancel cancel(ident_t ident) {
        promise<void>::notify d;
        std::lock_guard lk(_mx);
        d = cancel_lk(ident);
        _blk.push_back(ident);
        return pending_cancel(*this, ident);
    }
 
 
 
protected:
 
    struct item_t {
        promise_t _prom;
        std::chrono::system_clock::time_point _tp;
        ident_t _ident;
    };
 
    static bool compare_items(const item_t &a, const item_t &b) {
        return a._tp > b._tp;
    }
 
    std::mutex _mx;
    CondVar _cond;
    std::vector<item_t> _items;
    std::vector<ident_t> _blk;
    bool _stop = false;
    static thread_local scheduler_t *_current;
    std::thread _thr;
 
    void notify() {
        _cond.notify_all();
    }
 
    template<typename ResumeCB>
    void worker(ResumeCB &&rcb, bool &stop_flag) noexcept {
        std::unique_lock lk(_mx);
 
        _current = this;
 
        while (!stop_flag) {
            auto now = std::chrono::system_clock::now();
            if (!_items.empty() && _items.front()._tp < now) {
                auto d = _items.front()._prom();
                std::pop_heap(_items.begin(),_items.end(), compare_items);
                _items.pop_back();
                lk.unlock();
                rcb(std::move(d));
                if (_current != this) return;
                lk.lock();
            } else {
                wait_next(lk);
            }
        }
 
        _current = nullptr;
 
 
    }
 
    void wait_next(std::unique_lock<std::mutex> &lk) {
        if (!_items.empty()) {
            _cond.wait_until(lk, _items.front()._tp);
        } else {
            _cond.wait(lk);
        }
    }
 
    promise<void>::notify cancel_lk(ident_t ident) {
        auto iter = std::find_if(_items.begin(), _items.end(), [&](const item_t &item){
            return item._ident == ident;
        });
        if (iter == _items.end()) return {};
        if (iter == _items.begin()) {
            std::pop_heap(_items.begin(), _items.end(), compare_items);
            auto d = _items.back()._prom.cancel();
            _items.pop_back();
            return d;
        } else {
            auto d = iter->_prom.cancel();
            item_t &x = *iter;
            if (&x != &_items.back()) {
                std::swap(x, _items.back());
                std::make_heap(_items.begin(), _items.end(),  compare_items);
            }
            _items.pop_back();
            return d;
        }
    }
 
 
};
 
using scheduler = scheduler_t<std::condition_variable>;
 
template<typename CondVar>
inline thread_local scheduler_t<CondVar> *scheduler_t<CondVar>::_current = nullptr;
 
}