Lazy & RescheduleLazy 逻辑合并 #154

yanminhui · 2022-09-17T11:27:36Z

yanminhui
Sep 17, 2022
Collaborator

Lazy 与 RescheduleLazy 看起来很多代码是重叠的，除了决定是否要放到执行器执行外，很少有不同，似乎可以合并。

template <typename T, bool reschedule>
class LazyBase {
public:
    using promise_type = detail::LazyPromise<T>;
    using Handle = CoroHandle<promise_type>;
    using ValueType = T;

    struct AwaiterBase : public detail::LazyAwaiterBase<T> {
        using Base = detail::LazyAwaiterBase<T>;
        AwaiterBase(Handle coro) : Base(coro) {}

        AS_INLINE void await_suspend(
            std::coroutine_handle<> continuation) noexcept {
            auto h = Base::_handle;
            auto& pr = h.promise();
            pr._continuation = continuation;
            // current coro started, caller becomes my continuation
            if constexpr (!reschedule) {
                h.resume();
                return ;
            }
            // executor schedule performed if _executor exists
            pr._executor->schedule([h]() mutable {
                h.resume();
            });
        }
    };

    struct TryAwaiter : public AwaiterBase {
        TryAwaiter(Handle coro) : AwaiterBase(coro) {}
        AS_INLINE Try<T> await_resume() noexcept {
            return AwaiterBase::awaitResumeTry();
        };
    };

    struct ValueAwaiter : public AwaiterBase {
        ValueAwaiter(Handle coro) : AwaiterBase(coro) {}
        AS_INLINE T await_resume() { return AwaiterBase::awaitResume(); }
    };

    ~LazyBase() {
        if (_coro) {
            _coro.destroy();
            _coro = nullptr;
        }
    };
    explicit LazyBase(Handle coro) : _coro(coro) {}
    LazyBase(LazyBase&& other) : _coro(std::move(other._coro)) {
        other._coro = nullptr;
    }

    LazyBase(const LazyBase&) = delete;
    LazyBase& operator=(const LazyBase&) = delete;

    // Bind an executor to a Lazy, and convert it to RescheduleLazy.
    // You can only call via on rvalue, i.e. a Lazy is not accessible after
    // via() called.
    RescheduleLazy<T> via(Executor* ex) && {
        logicAssert(_coro.operator bool(),
                    "Lazy do not have a coroutine_handle");
        _coro.promise()._executor = ex;
        return RescheduleLazy<T>(std::exchange(_coro, nullptr));
    }

    Executor* getExecutor() { return _coro.promise()._executor; }

    // Bind an executor only. Don't re-schedule.
    //
    // Users shouldn't use `setEx` directly. `setEx` is designed
    // for internal purpose only. See uthread/Await.h/await for details.
    Lazy<T> setEx(Executor* ex) && {
        logicAssert(_coro.operator bool(),
                    "Lazy do not have a coroutine_handle");
        _coro.promise()._executor = ex;
        return Lazy<T>(std::exchange(_coro, nullptr));
    }

    template <typename F>
    void start(F&& callback) {
        // callback should take a single Try<T> as parameter, return value will
        // be ignored. a detached coroutine will not suspend at initial/final
        // suspend point.
        auto launchCoro = [](LazyBase lazy,
                             std::decay_t<F> cb) -> detail::DetachedCoroutine {
            cb(std::move(co_await lazy.coAwaitTry()));
        };
        [[maybe_unused]] auto detached =
            launchCoro(std::move(*this), std::forward<F>(callback));
    }

    bool isReady() const { return !_coro || _coro.done(); }

    auto operator co_await() {
        return ValueAwaiter(std::exchange(_coro, nullptr));
    }

    auto coAwaitTry() { return TryAwaiter(std::exchange(_coro, nullptr)); }

protected:
    Handle _coro;

    template <typename LazyType, typename IAlloc, typename OAlloc, bool Para>
    friend struct detail::CollectAllAwaiter;

    template <typename LazyType, typename IAlloc>
    friend struct detail::CollectAnyAwaiter;

    template <template <typename> typename LazyType, typename... Ts>
    friend struct detail::CollectAnyVariadicAwaiter;
};

template <typename T = void>
class [[nodiscard]] Lazy : public detail::LazyBase<T, /*reschedule=*/false> {
    using Base = detail::LazyBase<T, false>;

public:
    using Base::Base;

    auto coAwait(Executor* ex) {
        // derived lazy inherits executor
        Base::_coro.promise()._executor = ex;
        return typename Base::ValueAwaiter(std::exchange(Base::_coro, nullptr));
    }
};

template <typename T = void>
class [[nodiscard]] RescheduleLazy : public detail::LazyBase<T, /*reschedule=*/true> {
    using Base = detail::LazyBase<T, true>;

public:
    using Base::Base;

    void detach() {
        Base::start([](auto&& t) {
            if (t.hasError()) {
                std::rethrow_exception(t.getException());
            }
        });
    }
};

单元测试:

~/repos/async_simple/build$ ctest
Test project /Users/admin/repos/async_simple/build
    Start 1: run_async_simple_test
1/4 Test #1: run_async_simple_test ............   Passed    1.46 sec
    Start 2: run_async_simple_util_test
2/4 Test #2: run_async_simple_util_test .......   Passed    0.06 sec
    Start 3: run_async_simple_coro_test
3/4 Test #3: run_async_simple_coro_test .......   Passed   18.08 sec
    Start 4: run_async_simple_executor_test
4/4 Test #4: run_async_simple_executor_test ...   Passed    0.06 sec

100% tests passed, 0 tests failed out of 4

Total Test time (real) =  19.67 sec