dispatch_queue.hpp

#pragma once
 
#include <functional>
#include <utility>
 
#include "function_result.hpp"
#include "task.hpp"
#include "promise.hpp"
#include "worker_pool.hpp"
 
namespace dispatch_queue {
 
class dispatch_queue {
public:
    dispatch_queue();

    dispatch_queue(int thread_count);

    template<typename Fn>
    dispatch_queue(int thread_count, Fn&& worker_init) {
        if (thread_count < 0) {
            thread_count = std::thread::hardware_concurrency();
        }
        if (thread_count > 0) {
            worker_pool = std::make_unique<detail::worker_pool>(task_queue, thread_count, worker_init);
        }
    }
 
    dispatch_queue(const dispatch_queue&) = delete;
    dispatch_queue& operator=(const dispatch_queue&) = delete;

    ~dispatch_queue();

    template<typename F, typename... Args, typename Ret = detail::function_result<F, Args...>>
    task<Ret> dispatch(F&& f, Args&&... args) {
        return dispatch_internal(false, std::forward<F>(f), std::forward<Args>(args)...);
    }

    template<typename F, typename... Args, typename Ret = detail::function_result<F, Args...>>
    task<Ret> dispatch_main(F&& f, Args&&... args) {
        return dispatch_internal(true, std::forward<F>(f), std::forward<Args>(args)...);
    }

    bool is_threaded() const;

    int thread_count() const;

    size_t size() const;

    bool empty() const;

    void clear();

    void main_loop();

    void wait();

    template<class Rep, class Period>
    bool wait_for(const std::chrono::duration<Rep, Period>& timeout_duration) {
        if (worker_pool) {
            return worker_pool->wait_for(timeout_duration);
        }
        else {
            return true;
        }
    }

    template<class Clock, class Duration>
    bool wait_until(const std::chrono::time_point<Clock, Duration>& timeout_time) {
        if (worker_pool) {
            return worker_pool->wait_until(timeout_time);
        }
        else {
            return true;
        }
    }

    void shutdown();
 
#ifdef __cpp_lib_coroutine
private:
    struct dispatch_awaiter {
        dispatch_queue& dispatch_queue;
 
        bool await_ready() const noexcept { return false; }
        void await_suspend(std::coroutine_handle<> cont) const {
            dispatch_queue.dispatch([cont]{
                cont();
                if (cont.done()) {
                    cont.destroy();
                }
            });
        }
        void await_resume() {}
    };
 
    struct dispatch_main_awaiter {
        dispatch_queue& dispatch_queue;
 
        bool await_ready() const noexcept { return false; }
        void await_suspend(std::coroutine_handle<> cont) const {
            dispatch_queue.dispatch_main([cont]{
                cont();
                if (cont.done()) {
                    cont.destroy();
                }
            });
        }
        void await_resume() {}
    };
public:
    dispatch_awaiter dispatch() {
        return dispatch_awaiter(*this);
    }

    dispatch_main_awaiter dispatch_main() {
        return dispatch_main_awaiter(*this);
    }
#endif
 
private:
    std::unique_ptr<detail::worker_pool> worker_pool;
    detail::pending_task_queue task_queue;
 
    template<typename F, typename... Args, typename Ret = detail::function_result<F, Args...>>
    task<Ret> dispatch_internal(bool run_on_main_loop, F&& f, Args&&... args) {
        auto work = std::bind(std::move(f), std::forward<Args>(args)...);
        if (worker_pool) {
            auto future = detail::task_future<Ret>::create_pending();
            worker_pool->enqueue_task({ future->wrap(work) }, run_on_main_loop);
            return task<Ret>(future);
        }
        else if (run_on_main_loop) {
            auto future = detail::task_future<Ret>::create_pending();
            task_queue.push({ future->wrap(work) }, run_on_main_loop);
            return task<Ret>(future);
        }
        else {
            auto future = detail::task_future<Ret>::create(work);
            return task<Ret>(future);
        }
    }
};
 
} // end namespace dispatch_queue