wait_for.h

#ifndef ASYNCIO_WAIT_FOR_H
#define ASYNCIO_WAIT_FOR_H

#include <asyncio/concept/future.h>
#include <asyncio/concept/awaitable.h>
#include <asyncio/event_loop.h>
#include <asyncio/exception.h>
#include <asyncio/schedule_task.h>
#include <asyncio/result.h>
#include <chrono>

namespace ASYNCIO_NS {
namespace detail {

template<typename R, typename Duration>
struct WaitForAwaiter: NonCopyable {
    // 有结果就不用挂起
    constexpr bool await_ready() noexcept { return result_.has_value(); }
    
    // 返回结果值
    constexpr decltype(auto) await_resume() {
        return std::move(result_).result();
    }

    // 将调用co_await的协程状态设置为SUSPEND
    template<typename Promise>
    void await_suspend(std::coroutine_handle<Promise> continuation) noexcept {
        continuation_ = &continuation.promise();
        // set continuation_ to SUSPEND, don't schedule anymore, until it resume continuation_
        continuation_->set_state(Handle::SUSPEND);
    }

    // 构造函数：一个Future（Task），一个超时时间
    template<concepts::Awaitable Fut>
    WaitForAwaiter(Fut&& fut, Duration timeout)
            : timeout_handle_(*this, timeout)
            , wait_for_task_ {
                schedule_task(wait_for_task(no_wait_at_initial_suspend,/*将Task加入调度*/
                             std::forward<Fut>(fut)))
            } { }

private:
    // 加入调度的Task
    template<concepts::Awaitable Fut>
    Task<> wait_for_task(NoWaitAtInitialSuspend, Fut&& fut) {
        try {
            if constexpr (std::is_void_v<R>) { co_await std::forward<Fut>(fut); }
            else { result_.set_value(co_await std::forward<Fut>(fut));/*设置返回值*/ }
        } catch(...) {
            result_.unhandled_exception();
        }
        EventLoop& loop{get_event_loop()};
        loop.cancel_handle(timeout_handle_);    // Task已返回，取消超时句柄
        if (continuation_) {                    // 将等待Task的协程加入调度
            loop.call_soon(*continuation_);
        }
    }

private:
    struct TimeoutHandle: Handle {  // 等待超时句柄
        TimeoutHandle(WaitForAwaiter& awaiter, Duration timeout)
        : awaiter_(awaiter) {
            get_event_loop().call_later(timeout, *this);
        }
        void run() final { // timeout!
            awaiter_.wait_for_task_.cancel();   // 取消等待任务
            awaiter_.result_.set_exception(std::make_exception_ptr(TimeoutError{}));    // 设置超时异常
            get_event_loop().call_soon(*awaiter_.continuation_);    // 将等待Task的协程加入调度，因为已经超时了
        }

        WaitForAwaiter& awaiter_;
    };

    TimeoutHandle timeout_handle_;          // 超时句柄
    ScheduledTask<Task<>> wait_for_task_;   // 等待的Task
    Result<R> result_;              // 等待Task的返回值
    CoroHandle* continuation_{};    // 等待Task的协程的句柄
};

// 推导指引
// 调用WaitForAwaiter(Fut&&, Duration)这个构造函数时，推导出WaitForAwaiter的两个模板参数为AwaitResult<Fut>, Duration
template<concepts::Awaitable Fut, typename Duration>
WaitForAwaiter(Fut&&, Duration) -> WaitForAwaiter<AwaitResult<Fut>, Duration>;

template<concepts::Awaitable Fut, typename Duration>
struct WaitForAwaiterRegistry {
    WaitForAwaiterRegistry(Fut&& fut, Duration duration)
    : fut_(std::forward<Fut>(fut)), duration_(duration) { }

    auto operator co_await () && {
        return WaitForAwaiter{std::forward<Fut>(fut_), duration_};
    }
private:
    Fut fut_; // lift Awaitable's lifetime
    Duration duration_;
};

template<concepts::Awaitable Fut, typename Duration>
WaitForAwaiterRegistry(Fut&&, Duration) -> WaitForAwaiterRegistry<Fut, Duration>;

template<concepts::Awaitable Fut, typename Rep, typename Period>
auto wait_for(NoWaitAtInitialSuspend, Fut&& fut, std::chrono::duration<Rep, Period> timeout)
-> Task<AwaitResult<Fut>> { // lift awaitable type(WaitForAwaiterRegistry) to coroutine
    // 返回co_wait表达式的值
    co_return co_await WaitForAwaiterRegistry { std::forward<Fut>(fut), timeout };
}
}

template<concepts::Awaitable Fut, typename Rep, typename Period>
[[nodiscard("discard wait_for doesn't make sense")]]
Task<AwaitResult<Fut>> wait_for(Fut&& fut, std::chrono::duration<Rep, Period> timeout) {
    return detail::wait_for(no_wait_at_initial_suspend, std::forward<Fut>(fut), timeout);
}
}
#endif // ASYNCIO_WAIT_FOR_H