Limit we should only co_await a Lazy in Lazy functions

There was an implicit requirement that we should only co_await a lazy in a Lazy function. Now we make it explicit to avoid make it confusing.
alibaba · Jun 6, 2023 · 5fc8fd0 · 5fc8fd0
1 parent b891e2d
commit 5fc8fd0
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 5 deletions.
diff --git a/async_simple/coro/Lazy.h b/async_simple/coro/Lazy.h
@@ -74,6 +74,10 @@ class LazyPromiseBase {
         bool await_ready() const noexcept { return false; }
         template <typename PromiseType>
         auto await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
+            static_assert(
+                std::is_base_of<LazyPromiseBase, PromiseType>::value,
+                "the final awaiter is only allowed to be called by Lazy");
+
             return h.promise()._continuation;
         }
         void await_resume() noexcept {}
@@ -82,7 +86,12 @@ class LazyPromiseBase {
     struct YieldAwaiter {
         YieldAwaiter(Executor* executor) : _executor(executor) {}
         bool await_ready() const noexcept { return false; }
-        void await_suspend(std::coroutine_handle<> handle) {
+        template <typename PromiseType>
+        void await_suspend(std::coroutine_handle<PromiseType> handle) {
+            static_assert(
+                std::is_base_of<LazyPromiseBase, PromiseType>::value,
+                "'co_await Yield' is only allowed to be called by Lazy");
+
             logicAssert(_executor,
                         "Yielding is only meaningful with an executor!");
             _executor->schedule(std::move(handle));
@@ -251,8 +260,15 @@ class LazyBase {
         using Base = detail::LazyAwaiterBase<T>;
         AwaiterBase(Handle coro) : Base(coro) {}
 
-        AS_INLINE auto await_suspend(
-            std::coroutine_handle<> continuation) noexcept(!reschedule) {
+        template <typename PromiseType>
+        AS_INLINE auto await_suspend(std::coroutine_handle<PromiseType>
+                                         continuation) noexcept(!reschedule) {
+            static_assert(
+                std::is_base_of<LazyPromiseBase, PromiseType>::value ||
+                    std::is_same_v<detail::DetachedCoroutine::promise_type,
+                                   PromiseType>,
+                "'co_await Yield' is only allowed to be called by Lazy");
+
             // current coro started, caller becomes my continuation
             this->_handle.promise()._continuation = continuation;
 

diff --git a/docs/docs.cn/Lazy.md b/docs/docs.cn/Lazy.md
@@ -52,7 +52,7 @@ Lazy<> task2() {
 - 调度器实现有 bug，被提交到调度器的任务最后不一定会被执行。
 - 被等待的任务执行过程中出现了异常，在这种情况下 `co_await` 会直接将该异常返回到当前正在等待的 Lazy，而不会执行之后的语句了。
 
-同时需要注意，使用 `co_await` 启动 Lazy 需要当前函数也为 C++20 无栈协程。
+同时需要注意，在用户代码中我们只应该在 `Lazy` 函数中 `co_await` 一个 `Lazy`。
 
 ### .start(callback) 启动
 

diff --git a/docs/docs.en/Lazy.md b/docs/docs.en/Lazy.md
@@ -52,7 +52,7 @@ Note that, we couldn't assume that the statements after a `co_await expression`
 - There is an bug in the scheduler. The task submitted to the scheduler wouldn't be promised to schedule.
 - There is an exception happened in the waited task. In this case, the current coroutine would return to its caller instead of executing the following statments.
 
-Another thing we need to note is that the function contains a `co_await` would be a C++20 stackless coroutine too.
+Note that we should only `co_await` a Lazy in a Lazy function in users code.
 
 ### .start(callback)