virtual_thread.hpp

浏览该文件的文档.

#ifndef MSTL_CORE_ASYNC_VIRTUAL_THREAD_HPP__
#define MSTL_CORE_ASYNC_VIRTUAL_THREAD_HPP__

 
#include "MSTL/core/container/queue.hpp"
#include "MSTL/core/utility/optional.hpp"
#include "MSTL/core/async/condition_variable.hpp"
#include "MSTL/core/async/atomic.hpp"
#ifdef MSTL_STANDARD_20__
#include <coroutine>
MSTL_BEGIN_NAMESPACE__


struct virtual_thread_task {
    struct promise_type {
        _MSTL exception_ptr exception_;  

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

        std::suspend_never initial_suspend() {
            return std::suspend_never{};
        }

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

        void return_void() {}

        void unhandled_exception() {
            exception_ = _MSTL current_exception();
        }
    };
 
    std::coroutine_handle<promise_type> handle_;  

    virtual_thread_task(std::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_(_MSTL exchange(other.handle_, nullptr)) {}

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

class virtual_thread_scheduler {
private:
    queue<std::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) {
            std::coroutine_handle<> handle;
 
            {
                smart_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(std::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 {
    std::coroutine_handle<> handle_;  

    bool await_ready() const noexcept {
        return false;
    }

    void await_suspend(std::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 std::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(_MSTL 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
 
MSTL_END_NAMESPACE__
#endif
#endif // MSTL_CORE_ASYNC_VIRTUAL_THREAD_HPP__