virtual_thread.hpp

浏览该文件的文档.

#ifndef NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__
#define NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__

 
#ifdef NEFORCE_STANDARD_20
#    include "NeForce/core/async/atomic.hpp"
#    include "NeForce/core/async/condition_variable.hpp"
#    include "NeForce/core/async/coroutine.hpp"
#    include "NeForce/core/container/queue.hpp"
#    include "NeForce/core/utility/optional.hpp"
NEFORCE_BEGIN_NAMESPACE__


struct virtual_thread_task {
    struct promise_type {
        exception_ptr exception_; 

        virtual_thread_task get_return_object() {
            return virtual_thread_task{coroutine_handle<promise_type>::from_promise(*this)};
        }

        suspend_never initial_suspend() { return suspend_never{}; }

        suspend_always final_suspend() noexcept { return suspend_always{}; }

        void return_void() {}

        void unhandled_exception() { exception_ = _NEFORCE current_exception(); }
    };
 
    coroutine_handle<promise_type> handle_; 

    virtual_thread_task(coroutine_handle<promise_type> h) :
    handle_(h) {}

    ~virtual_thread_task() {
        if (handle_ && !handle_.done()) {
            handle_.destroy();
        }
    }
 
    virtual_thread_task(const virtual_thread_task&) = delete;            
    virtual_thread_task& operator=(const virtual_thread_task&) = delete; 

    virtual_thread_task(virtual_thread_task&& other) noexcept :
    handle_(_NEFORCE exchange(other.handle_, nullptr)) {}

    virtual_thread_task& operator=(virtual_thread_task&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
        if (handle_ && !handle_.done()) {
            handle_.destroy();
        }
        handle_ = _NEFORCE exchange(other.handle_, nullptr);
        return *this;
    }
};

class virtual_thread_scheduler {
private:
    queue<coroutine_handle<>> task_queue_; 
    vector<thread> workers_;               
    mutex mutex_;                          
    condition_variable cv_;                
    atomic<bool> shutdown_;                

    virtual_thread_scheduler() :
    shutdown_(false) {}

    void worker_loop() {
        while (true) {
            coroutine_handle<> handle;
 
            {
                unique_lock<mutex> lock(mutex_);
                cv_.wait(lock, [this] { return shutdown_ || !task_queue_.empty(); });
 
                if (shutdown_ && task_queue_.empty()) {
                    return;
                }
 
                if (!task_queue_.empty()) {
                    handle = task_queue_.front();
                    task_queue_.pop();
                }
            }
 
            if (handle) {
                handle.resume();
            }
        }
    }
 
public:
    static virtual_thread_scheduler& get_instance() {
        static virtual_thread_scheduler instance;
        return instance;
    }

    void schedule(coroutine_handle<> handle) {
        {
            lock<mutex> lock(mutex_);
            task_queue_.push(handle);
        }
        cv_.notify_one();
    }

    void start_workers(size_t num_threads) {
        for (size_t i = 0; i < num_threads; ++i) {
            workers_.emplace_back([this] { worker_loop(); });
        }
    }

    void shutdown() {
        {
            lock<mutex> lock(mutex_);
            shutdown_ = true;
        }
        cv_.notify_all();
 
        for (auto& worker: workers_) {
            if (worker.joinable()) {
                worker.join();
            }
        }
    }

    ~virtual_thread_scheduler() { shutdown(); }
};

struct virtual_thread_awaiter {
    coroutine_handle<> handle_; 

    bool await_ready() const noexcept { return false; }

    void await_suspend(coroutine_handle<> handle) { virtual_thread_scheduler::get_instance().schedule(handle); }

    void await_resume() const noexcept {}
};
 

class virtual_thread {
private:
    optional<virtual_thread_task> task_; 

    template <typename Func>
    static virtual_thread_task create_task(Func func) {
        co_await suspend_never{};
        func();
    }
 
public:
    virtual_thread() = default;                                           
    virtual_thread(const virtual_thread&) = delete;                       
    virtual_thread& operator=(const virtual_thread&) = delete;            
    virtual_thread(virtual_thread&& other) noexcept = default;            
    virtual_thread& operator=(virtual_thread&& other) noexcept = default; 

    template <typename Func>
    static virtual_thread start(Func&& func) {
        virtual_thread vt;
        vt.task_ = virtual_thread::create_task(_NEFORCE forward<Func>(func));
        return vt;
    }

    static virtual_thread_awaiter yield() { return virtual_thread_awaiter{}; }

    static virtual_thread_task sleep(const int64_t ms) {
        co_await yield();
        this_thread::sleep_for(milliseconds(ms));
    }

    static void initialize(size_t num_threads) { virtual_thread_scheduler::get_instance().start_workers(num_threads); }

    static void shutdown() { virtual_thread_scheduler::get_instance().shutdown(); }
};
 // Coroutine
 
NEFORCE_END_NAMESPACE__
#endif
#endif // NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__