virtual_thread.hpp

浏览该文件的文档.

#ifndef NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__
#define NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__

 
#include "NeForce/core/async/coroutine.hpp"
#ifdef NEFORCE_STANDARD_20
#    include "NeForce/core/async/atomic.hpp"
#    include "NeForce/core/async/condition_variable.hpp"
#    include "NeForce/core/async/mutex.hpp"
#    include "NeForce/core/async/thread.hpp"
#    include "NeForce/core/container/queue.hpp"
#    include "NeForce/core/container/vector.hpp"
#    include "NeForce/core/exception/exception_ptr.hpp"
#    include "NeForce/core/memory/aligned_buffer.hpp"
NEFORCE_BEGIN_NAMESPACE__


template <typename T = void>
struct virtual_thread_task;

template <typename T>
struct is_virtual_thread_task : false_type {};

template <typename T>
struct is_virtual_thread_task<virtual_thread_task<T>> : true_type {};

template <typename T>
constexpr bool is_virtual_thread_task_v = is_virtual_thread_task<T>::value;
 
 
NEFORCE_BEGIN_INNER__
 
struct task_shared_state_base {
    atomic<bool> scheduled_{false};
};
 
template <typename T>
struct task_shared_state : task_shared_state_base {
    aligned_buffer<T> result_buffer_;
    bool has_value_{false};
    exception_ptr exception_{nullptr};
 
    atomic<coroutine_handle<>> continuation_{nullptr};
    atomic<bool> completed_{false};
    mutex mtx_;
    condition_variable cv_;
    atomic<unsigned long> ref_count_{1};
 
    void add_ref() noexcept { ref_count_.fetch_add(1, memory_order_relaxed); }
 
    void release() noexcept {
        if (ref_count_.fetch_sub(1, memory_order_acq_rel) == 1) {
            if (has_value_) {
                result_buffer_.ptr()->~T();
            }
            delete this;
        }
    }
 
    static task_shared_state* create() { return new task_shared_state(); }
};
 
template <>
struct task_shared_state<void> : task_shared_state_base {
    exception_ptr exception_{nullptr};
 
    atomic<coroutine_handle<>> continuation_{nullptr};
    atomic<bool> completed_{false};
    mutex mtx_;
    condition_variable cv_;
    atomic<unsigned long> ref_count_{1};
 
    void add_ref() noexcept { ref_count_.fetch_add(1, memory_order_relaxed); }
 
    void release() noexcept {
        if (ref_count_.fetch_sub(1, memory_order_acq_rel) == 1) {
            delete this;
        }
    }
 
    static task_shared_state* create() { return new task_shared_state(); }
};
 
NEFORCE_END_INNER__
 

class virtual_thread_scheduler {
private:
    queue<coroutine_handle<>> task_queue_; 
    vector<thread> workers_;               
    mutex mutex_;                          
    condition_variable cv_;                
    atomic<bool> shutdown_{false};         
 
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(); }
 
private:
    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();
            }
        }
    }
};
 
 
NEFORCE_BEGIN_INNER__
 
struct yield_tag {};
 
struct sleep_tag {
    int64_t ms_;
};
 
void mark_continuation_scheduled(coroutine_handle<> cont);
 
NEFORCE_END_INNER__
 

template <>
struct virtual_thread_task<void> {
    struct promise_type {
        inner::task_shared_state<void>* shared_state_{inner::task_shared_state<void>::create()};

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

        suspend_never initial_suspend() noexcept { return {}; }

        auto final_suspend() noexcept {
            struct final_awaiter {
                bool await_ready() noexcept { return false; }
 
                bool await_suspend(coroutine_handle<promise_type> h) noexcept {
                    auto& p = h.promise();
                    auto* state = p.shared_state_;
 
                    state->completed_.store(true, memory_order_release);
                    state->cv_.notify_all();
 
                    auto cont = state->continuation_.exchange(nullptr, memory_order_acq_rel);
                    if (cont) {
                        try {
                            inner::mark_continuation_scheduled(cont);
                            virtual_thread_scheduler::get_instance().schedule(cont);
                            // NOLINTNEXTLINE(bugprone-empty-catch)
                        } catch (...) {
                            // ignore
                        }
                    }
                    return false;
                }
 
                void await_resume() noexcept {}
            };
            return final_awaiter{};
        }

        auto await_transform(inner::yield_tag /*unused*/) {
            struct yield_awaiter {
                inner::task_shared_state<void>* shared_state_;
 
                NEFORCE_NODISCARD bool await_ready() const noexcept { return false; }
 
                void await_suspend(coroutine_handle<> handle) const {
                    shared_state_->scheduled_.store(true, memory_order_release);
                    virtual_thread_scheduler::get_instance().schedule(handle);
                }
 
                void await_resume() const noexcept {}
            };
            return yield_awaiter{shared_state_};
        }

        auto await_transform(inner::sleep_tag tag) {
            struct sleep_awaiter_impl {
                inner::task_shared_state<void>* shared_state_;
                int64_t ms_;
 
                NEFORCE_NODISCARD bool await_ready() const noexcept { return ms_ <= 0; }
 
                void await_suspend(coroutine_handle<> handle) const {
                    shared_state_->scheduled_.store(true, memory_order_release);
                    thread([handle, ms = ms_] {
                        this_thread::sleep_for(milliseconds(ms));
                        virtual_thread_scheduler::get_instance().schedule(handle);
                    }).detach();
                }
 
                void await_resume() const noexcept {}
            };
            return sleep_awaiter_impl{shared_state_, tag.ms_};
        }

        template <typename Awaiter>
        decltype(auto) await_transform(Awaiter&& a) {
            return _NEFORCE forward<Awaiter>(a);
        }

        void return_void() {}

        void unhandled_exception() { shared_state_->exception_ = _NEFORCE current_exception(); }

        ~promise_type() {
            if (shared_state_ != nullptr) {
                shared_state_->release();
            }
        }
    };
 
    coroutine_handle<promise_type> handle_{nullptr};        
    inner::task_shared_state<void>* shared_state_{nullptr}; 

    virtual_thread_task() = default;

    explicit virtual_thread_task(coroutine_handle<promise_type> h) :
    handle_(h) {
        if (handle_) {
            shared_state_ = handle_.promise().shared_state_;
            if (shared_state_ != nullptr) {
                shared_state_->add_ref();
            }
        }
    }

    ~virtual_thread_task() {
        if (handle_) {
            bool completed = shared_state_ != nullptr && shared_state_->completed_.load(memory_order_acquire);
            bool was_scheduled = shared_state_ != nullptr && shared_state_->scheduled_.load(memory_order_acquire);
            if (!completed && !was_scheduled) {
                handle_.destroy();
            }
        }
        if (shared_state_ != nullptr) {
            shared_state_->release();
        }
    }
 
    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)),
    shared_state_(_NEFORCE exchange(other.shared_state_, nullptr)) {}

    virtual_thread_task& operator=(virtual_thread_task&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
        if (handle_) {
            bool completed = shared_state_ != nullptr && shared_state_->completed_.load(memory_order_acquire);
            bool was_scheduled = shared_state_ != nullptr && shared_state_->scheduled_.load(memory_order_acquire);
            if (!completed && !was_scheduled) {
                handle_.destroy();
            }
        }
        if (shared_state_ != nullptr) {
            shared_state_->release();
        }
        handle_ = _NEFORCE exchange(other.handle_, nullptr);
        shared_state_ = _NEFORCE exchange(other.shared_state_, nullptr);
        return *this;
    }

    bool await_ready() noexcept {
        return shared_state_ != nullptr && shared_state_->completed_.load(memory_order_acquire);
    }

    bool await_suspend(coroutine_handle<> caller) noexcept {
        shared_state_->continuation_.store(caller, memory_order_release);
 
        if (shared_state_->completed_.load(memory_order_acquire)) {
            if (shared_state_->continuation_.exchange(nullptr, memory_order_acq_rel)) {
                return false;
            }
        }
        return true;
    }

    void await_resume() {
        if (shared_state_->exception_) {
            rethrow_exception(shared_state_->exception_);
        }
    }

    void get_result() {
        if (!shared_state_->completed_.load(memory_order_acquire)) {
            unique_lock<mutex> lock(shared_state_->mtx_);
            shared_state_->cv_.wait(lock, [this] { return shared_state_->completed_.load(memory_order_acquire); });
        }
        if (shared_state_->exception_) {
            rethrow_exception(shared_state_->exception_);
        }
    }

    NEFORCE_NODISCARD bool is_done() const noexcept {
        return shared_state_ != nullptr && shared_state_->completed_.load(memory_order_acquire);
    }

    NEFORCE_NODISCARD bool valid() const noexcept { return shared_state_ != nullptr; }
};
 
 
NEFORCE_BEGIN_INNER__
 
inline void mark_continuation_scheduled(coroutine_handle<> cont) {
    auto vh = coroutine_handle<virtual_thread_task<void>::promise_type>::from_address(cont.address());
    vh.promise().shared_state_->scheduled_.store(true, memory_order_release);
}
 
NEFORCE_END_INNER__
 

template <typename T>
struct virtual_thread_task {
    struct promise_type {
        inner::task_shared_state<T>* shared_state_{inner::task_shared_state<T>::create()};

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

        suspend_never initial_suspend() noexcept { return {}; }

        auto final_suspend() noexcept {
            struct final_awaiter {
                bool await_ready() noexcept { return false; }
 
                bool await_suspend(coroutine_handle<promise_type> h) noexcept {
                    auto& p = h.promise();
                    auto* state = p.shared_state_;
 
                    state->completed_.store(true, memory_order_release);
                    state->cv_.notify_all();
 
                    auto cont = state->continuation_.exchange(nullptr, memory_order_acq_rel);
                    if (cont) {
                        try {
                            inner::mark_continuation_scheduled(cont);
                            virtual_thread_scheduler::get_instance().schedule(cont);
                            // NOLINTNEXTLINE(bugprone-empty-catch)
                        } catch (...) {
                            // ignore
                        }
                    }
                    return false;
                }
 
                void await_resume() noexcept {}
            };
            return final_awaiter{};
        }

        auto await_transform(inner::yield_tag /*unused*/) {
            struct yield_awaiter {
                inner::task_shared_state<T>* shared_state_;
 
                NEFORCE_NODISCARD bool await_ready() const noexcept { return false; }
 
                void await_suspend(coroutine_handle<> handle) const {
                    shared_state_->scheduled_.store(true, memory_order_release);
                    virtual_thread_scheduler::get_instance().schedule(handle);
                }
 
                void await_resume() const noexcept {}
            };
            return yield_awaiter{shared_state_};
        }

        auto await_transform(inner::sleep_tag tag) {
            struct sleep_awaiter_impl {
                inner::task_shared_state<T>* shared_state_;
                int64_t ms_;
 
                NEFORCE_NODISCARD bool await_ready() const noexcept { return ms_ <= 0; }
 
                void await_suspend(coroutine_handle<> handle) const {
                    shared_state_->scheduled_.store(true, memory_order_release);
                    thread([handle, ms = ms_] {
                        this_thread::sleep_for(milliseconds(ms));
                        virtual_thread_scheduler::get_instance().schedule(handle);
                    }).detach();
                }
 
                void await_resume() const noexcept {}
            };
            return sleep_awaiter_impl{shared_state_, tag.ms_};
        }

        void return_value(const T& value) {
            ::new (shared_state_->result_buffer_.addr()) T(value);
            shared_state_->has_value_ = true;
        }

        void return_value(T&& value) {
            ::new (shared_state_->result_buffer_.addr()) T(move(value));
            shared_state_->has_value_ = true;
        }

        template <typename Awaiter>
        decltype(auto) await_transform(Awaiter&& a) {
            return _NEFORCE forward<Awaiter>(a);
        }

        void unhandled_exception() { shared_state_->exception_ = _NEFORCE current_exception(); }

        ~promise_type() {
            if (shared_state_) {
                shared_state_->release();
            }
        }
    };
 
    coroutine_handle<promise_type> handle_{nullptr};     
    inner::task_shared_state<T>* shared_state_{nullptr}; 

    virtual_thread_task() = default;

    explicit virtual_thread_task(coroutine_handle<promise_type> h) :
    handle_(h) {
        if (handle_) {
            shared_state_ = handle_.promise().shared_state_;
            if (shared_state_) {
                shared_state_->add_ref();
            }
        }
    }

    ~virtual_thread_task() {
        if (handle_) {
            bool completed = shared_state_ && shared_state_->completed_.load(memory_order_acquire);
            bool was_scheduled = shared_state_ && shared_state_->scheduled_.load(memory_order_acquire);
            if (!completed && !was_scheduled) {
                handle_.destroy();
            }
        }
        if (shared_state_) {
            shared_state_->release();
        }
    }
 
    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)),
    shared_state_(_NEFORCE exchange(other.shared_state_, nullptr)) {}

    virtual_thread_task& operator=(virtual_thread_task&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
        if (handle_) {
            bool completed = shared_state_ && shared_state_->completed_.load(memory_order_acquire);
            bool was_scheduled = shared_state_ && shared_state_->scheduled_.load(memory_order_acquire);
            if (!completed && !was_scheduled) {
                handle_.destroy();
            }
        }
        if (shared_state_) {
            shared_state_->release();
        }
        handle_ = _NEFORCE exchange(other.handle_, nullptr);
        shared_state_ = _NEFORCE exchange(other.shared_state_, nullptr);
        return *this;
    }

    bool await_ready() noexcept { return shared_state_ && shared_state_->completed_.load(memory_order_acquire); }

    bool await_suspend(coroutine_handle<> caller) noexcept {
        shared_state_->continuation_.store(caller, memory_order_release);
 
        if (shared_state_->completed_.load(memory_order_acquire)) {
            if (shared_state_->continuation_.exchange(nullptr, memory_order_acq_rel)) {
                return false;
            }
        }
        return true;
    }

    T await_resume() {
        if (shared_state_->exception_) {
            rethrow_exception(shared_state_->exception_);
        }
        return move(*shared_state_->result_buffer_.ptr());
    }

    T get_result() {
        if (!shared_state_->completed_.load(memory_order_acquire)) {
            unique_lock<mutex> lock(shared_state_->mtx_);
            shared_state_->cv_.wait(lock, [this] { return shared_state_->completed_.load(memory_order_acquire); });
        }
        if (shared_state_->exception_) {
            rethrow_exception(shared_state_->exception_);
        }
        return move(*shared_state_->result_buffer_.ptr());
    }

    NEFORCE_NODISCARD bool is_done() const noexcept {
        return shared_state_ && shared_state_->completed_.load(memory_order_acquire);
    }

    NEFORCE_NODISCARD bool valid() const noexcept { return shared_state_ != nullptr; }
};

class virtual_thread {
private:
    template <typename Func>
    static virtual_thread_task<void> create_task(Func func) {
        func();
        co_return;
    }
 
public:
    template <typename Func>
    static auto start(Func&& func) {
        using result_type = invoke_result_t<decay_t<Func>>;
        if constexpr (is_virtual_thread_task_v<result_type>) {
            return _NEFORCE forward<Func>(func)();
        } else {
            return create_task(_NEFORCE forward<Func>(func));
        }
    }

    static inner::yield_tag yield() { return inner::yield_tag{}; }

    static inner::sleep_tag sleep(const int64_t ms) { return inner::sleep_tag{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(); }
};
 // VirtualThread
 
NEFORCE_END_NAMESPACE__
#endif
#endif // NEFORCE_CORE_ASYNC_VIRTUAL_THREAD_HPP__