Add error hint to CoroutineHandle (resubmit)

Summary: Provies single-await co_nothrow wrapper

Reviewed By: capickett

Differential Revision: D32405920

fbshipit-source-id: b0b5c55149aebe4246422fabeffb7ed433b4e20d

folly/experimental/coro/Coroutine.h

@@ -45,6 +45,11 @@
 
 #if FOLLY_HAS_COROUTINES
 
+namespace folly {
+class exception_wrapper;
+struct AsyncStackFrame;
+} // namespace folly
+
 namespace folly::coro {
 
 #if __has_include(<coroutine>) && !defined(LLVM_COROUTINES)
@@ -166,10 +171,17 @@ class variant_awaitable : private std::variant<A...> {
 
 #endif // __has_include(<variant>)
 
+class ExtendedCoroutineHandle;
+
 // Extended promise interface folly::coro types are expected to implement
 class ExtendedCoroutinePromise {
  public:
   virtual coroutine_handle<> getHandle() = 0;
+  // Types may provide a more efficient resumption path when they know they will
+  // be receiving an error result from the awaitee.
+  // If they do, they might also update the active stack frame.
+  virtual std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
+      exception_wrapper&);
 
  protected:
   ~ExtendedCoroutinePromise() = default;
@@ -212,6 +224,14 @@ class ExtendedCoroutineHandle {
 
   ExtendedCoroutinePromise* getPromise() const noexcept { return extended_; }
 
+  std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
+      exception_wrapper& ex) {
+    if (extended_) {
+      return extended_->getErrorHandle(ex);
+    }
+    return {basic_, nullptr};
+  }
+
   explicit operator bool() const noexcept { return !!extended_; }
 
  private:
@@ -228,6 +248,11 @@ class ExtendedCoroutineHandle {
   ExtendedCoroutinePromise* extended_{nullptr};
 };
 
+inline std::pair<ExtendedCoroutineHandle, AsyncStackFrame*>
+ExtendedCoroutinePromise::getErrorHandle(exception_wrapper&) {
+  return {getHandle(), nullptr};
+}
+
 } // namespace folly::coro
 
 #endif

folly/experimental/coro/Task.h

@@ -66,7 +66,7 @@ class TaskPromiseBase {
     template <typename Promise>
     FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES coroutine_handle<>
     await_suspend(coroutine_handle<Promise> coro) noexcept {
-      TaskPromiseBase& promise = coro.promise();
+      auto& promise = coro.promise();
       // If the continuation has been exchanged, then we expect that the
       // exchanger will handle the lifetime of the async stack. See
       // ScopeExitTaskPromise's FinalAwaiter for more details.
@@ -76,6 +76,11 @@ class TaskPromiseBase {
       if (promise.ownsAsyncFrame_) {
         folly::popAsyncStackFrameCallee(promise.asyncFrame_);
       }
+      if (promise.result_.hasException()) {
+        auto [handle, frame] =
+            promise.continuation_.getErrorHandle(promise.result_.exception());
+        return handle.getHandle();
+      }
       return promise.continuation_.getHandle();
     }
 
@@ -85,7 +90,8 @@ class TaskPromiseBase {
   friend class FinalAwaiter;
 
  protected:
-  TaskPromiseBase() noexcept {}
+  TaskPromiseBase() noexcept = default;
+  ~TaskPromiseBase() = default;
 
   template <typename Promise>
   variant_awaitable<FinalAwaiter, ready_awaitable<>> do_safe_point(
@@ -111,12 +117,23 @@ class TaskPromiseBase {
 
   template <typename Awaitable>
   auto await_transform(Awaitable&& awaitable) {
+    bypassExceptionThrowing_ =
+        bypassExceptionThrowing_ == BypassExceptionThrowing::REQUESTED
+        ? BypassExceptionThrowing::ACTIVE
+        : BypassExceptionThrowing::INACTIVE;
+
     return folly::coro::co_withAsyncStack(folly::coro::co_viaIfAsync(
         executor_.get_alias(),
         folly::coro::co_withCancellation(
             cancelToken_, static_cast<Awaitable&&>(awaitable))));
   }
 
+  template <typename Awaitable>
+  auto await_transform(NothrowAwaitable<Awaitable>&& awaitable) {
+    bypassExceptionThrowing_ = BypassExceptionThrowing::REQUESTED;
+    return await_transform(awaitable.unwrap());
+  }
+
   auto await_transform(co_current_executor_t) noexcept {
     return ready_awaitable<folly::Executor*>{executor_.get()};
   }
@@ -161,6 +178,13 @@ class TaskPromiseBase {
   folly::CancellationToken cancelToken_;
   bool hasCancelTokenOverride_ = false;
   bool ownsAsyncFrame_ = true;
+
+ protected:
+  enum class BypassExceptionThrowing : uint8_t {
+    INACTIVE,
+    ACTIVE,
+    REQUESTED,
+  } bypassExceptionThrowing_{BypassExceptionThrowing::INACTIVE};
 };
 
 template <typename T>
@@ -220,6 +244,20 @@ class TaskPromise final : public TaskPromiseBase,
     return do_safe_point(*this);
   }
 
+  std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
+      exception_wrapper& ex) override {
+    if (bypassExceptionThrowing_ == BypassExceptionThrowing::ACTIVE) {
+      auto finalAwaiter = yield_value(co_error(std::move(ex)));
+      DCHECK(!finalAwaiter.await_ready());
+      return {
+          finalAwaiter.await_suspend(
+              coroutine_handle<TaskPromise>::from_promise(*this)),
+          // finalAwaiter.await_suspend pops a frame
+          getAsyncFrame().getParentFrame()};
+    }
+    return {coroutine_handle<TaskPromise>::from_promise(*this), nullptr};
+  }
+
  private:
   Try<StorageType> result_;
 };
@@ -265,6 +303,20 @@ class TaskPromise<void> final
     return do_safe_point(*this);
   }
 
+  std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
+      exception_wrapper& ex) override {
+    if (bypassExceptionThrowing_ == BypassExceptionThrowing::ACTIVE) {
+      auto finalAwaiter = yield_value(co_error(std::move(ex)));
+      DCHECK(!finalAwaiter.await_ready());
+      return {
+          finalAwaiter.await_suspend(
+              coroutine_handle<TaskPromise>::from_promise(*this)),
+          // finalAwaiter.await_suspend pops a frame
+          getAsyncFrame().getParentFrame()};
+    }
+    return {coroutine_handle<TaskPromise>::from_promise(*this), nullptr};
+  }
+
  private:
   Try<void> result_;
 };

folly/experimental/coro/ViaIfAsync.h

@@ -61,7 +61,7 @@ class ViaCoroutinePromiseBase {
     executor_ = std::move(executor);
   }
 
-  void setContinuation(coroutine_handle<> continuation) noexcept {
+  void setContinuation(ExtendedCoroutineHandle continuation) noexcept {
     continuation_ = continuation;
   }
 
@@ -83,15 +83,18 @@ class ViaCoroutinePromiseBase {
   void executeContinuation() noexcept {
     RequestContextScopeGuard contextScope{std::move(context_)};
     if (asyncFrame_ != nullptr) {
-      folly::resumeCoroutineWithNewAsyncStackRoot(continuation_, *asyncFrame_);
+      folly::resumeCoroutineWithNewAsyncStackRoot(
+          continuation_.getHandle(), *asyncFrame_);
     } else {
       continuation_.resume();
     }
   }
 
  protected:
+  virtual ~ViaCoroutinePromiseBase() = default;
+
   folly::Executor::KeepAlive<> executor_;
-  coroutine_handle<> continuation_;
+  ExtendedCoroutineHandle continuation_;
   folly::AsyncStackFrame* asyncFrame_ = nullptr;
   std::shared_ptr<RequestContext> context_;
 };
@@ -99,10 +102,14 @@ class ViaCoroutinePromiseBase {
 template <bool IsStackAware>
 class ViaCoroutine {
  public:
-  class promise_type : public ViaCoroutinePromiseBase {
+  class promise_type final : public ViaCoroutinePromiseBase,
+                             public ExtendedCoroutinePromiseImpl<promise_type> {
     struct FinalAwaiter {
       bool await_ready() noexcept { return false; }
 
+      // This code runs immediately after the inner awaitable resumes its fake
+      // continuation, and it schedules the real continuation on the awaiter's
+      // executor
       FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES void await_suspend(
           coroutine_handle<promise_type> h) noexcept {
         auto& promise = h.promise();
@@ -134,6 +141,16 @@ class ViaCoroutine {
       DCHECK(this->asyncFrame_ != nullptr);
       return *this->asyncFrame_;
     }
+
+    std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
+        exception_wrapper& ex) override {
+      auto [handle, frame] = continuation_.getErrorHandle(ex);
+      setContinuation(handle);
+      if (frame && IsStackAware) {
+        asyncFrame_ = frame;
+      }
+      return {promise_type::getHandle(), nullptr};
+    }
   };
 
   ViaCoroutine(ViaCoroutine&& other) noexcept
@@ -155,7 +172,7 @@ class ViaCoroutine {
     coro_.promise().setExecutor(std::move(executor));
   }
 
-  void setContinuation(coroutine_handle<> continuation) noexcept {
+  void setContinuation(ExtendedCoroutineHandle continuation) noexcept {
     coro_.promise().setContinuation(continuation);
   }
 
@@ -610,6 +627,81 @@ using semi_await_try_result_t =
         std::declval<folly::Executor::KeepAlive<>>(),
         folly::coro::co_awaitTry(std::declval<T>())))>;
 
+namespace detail {
+
+template <typename T>
+class NothrowAwaitable {
+ public:
+  template <typename T2>
+  explicit NothrowAwaitable(T2&& awaitable) noexcept(
+      std::is_nothrow_constructible_v<T, T2>)
+      : inner_(static_cast<T2&&>(awaitable)) {}
+
+  template <typename Factory>
+  explicit NothrowAwaitable(std::in_place_t, Factory&& factory)
+      : inner_(factory()) {}
+
+  T&& unwrap() { return std::move(inner_); }
+
+  template <
+      typename T2 = T,
+      typename Result = decltype(folly::coro::co_withCancellation(
+          std::declval<const folly::CancellationToken&>(), std::declval<T2>()))>
+  friend NothrowAwaitable<Result> co_withCancellation(
+      const folly::CancellationToken& cancelToken,
+      NothrowAwaitable&& awaitable) {
+    return NothrowAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                      return folly::coro::co_withCancellation(
+                                          cancelToken,
+                                          static_cast<T&&>(awaitable.inner_));
+                                    }};
+  }
+
+  template <
+      typename T2 = T,
+      typename Result =
+          decltype(folly::coro::co_withAsyncStack(std::declval<T2>()))>
+  friend NothrowAwaitable<Result>
+  tag_invoke(cpo_t<co_withAsyncStack>, NothrowAwaitable&& awaitable) noexcept(
+      noexcept(folly::coro::co_withAsyncStack(std::declval<T2>()))) {
+    return NothrowAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                      return folly::coro::co_withAsyncStack(
+                                          static_cast<T&&>(awaitable.inner_));
+                                    }};
+  }
+
+  template <
+      typename T2 = T,
+      typename Result = decltype(folly::coro::co_viaIfAsync(
+          std::declval<folly::Executor::KeepAlive<>>(), std::declval<T2>()))>
+  friend NothrowAwaitable<Result> co_viaIfAsync(
+      folly::Executor::KeepAlive<> executor,
+      NothrowAwaitable&&
+          awaitable) noexcept(noexcept(folly::coro::
+                                           co_viaIfAsync(
+                                               std::declval<folly::Executor::
+                                                                KeepAlive<>>(),
+                                               std::declval<T2>()))) {
+    return NothrowAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                      return folly::coro::co_viaIfAsync(
+                                          std::move(executor),
+                                          static_cast<T&&>(awaitable.inner_));
+                                    }};
+  }
+
+ private:
+  T inner_;
+};
+
+} // namespace detail
+
+template <typename Awaitable>
+detail::NothrowAwaitable<remove_cvref_t<Awaitable>> co_nothrow(
+    Awaitable&& awaitable) {
+  return detail::NothrowAwaitable<remove_cvref_t<Awaitable>>{
+      static_cast<Awaitable&&>(awaitable)};
+}
+
 } // namespace coro
 } // namespace folly
 

folly/experimental/coro/test/ScopeExitTest.cpp

@@ -129,4 +129,20 @@ TEST_F(ScopeExitTest, NonMoveableState) {
   EXPECT_EQ(count, 2);
 }
 
+TEST_F(ScopeExitTest, OneExitActionThroughNoThrow) {
+  EXPECT_THROW(
+      folly::coro::blockingWait([this]() -> Task<> {
+        co_await co_scope_exit([this]() -> Task<> {
+          ++count;
+          co_return;
+        });
+        co_await co_nothrow([]() -> Task<> {
+          throw std::runtime_error("Something bad happened!");
+          co_return;
+        }());
+      }()),
+      std::runtime_error);
+  EXPECT_EQ(count, 1);
+}
+
 #endif // FOLLY_HAS_COROUTINES

folly/experimental/coro/test/TaskTest.cpp

@@ -653,4 +653,72 @@ TEST_F(TaskTest, SafePoint) {
   }());
 }
 
+TEST_F(TaskTest, CoAwaitNothrow) {
+  auto res =
+      folly::coro::blockingWait(co_awaitTry([]() -> folly::coro::Task<void> {
+        auto t = []() -> folly::coro::Task<int> { co_return 42; }();
+
+        int result = co_await folly::coro::co_nothrow(std::move(t));
+        EXPECT_EQ(42, result);
+
+        t = []() -> folly::coro::Task<int> {
+          co_yield folly::coro::co_error(std::runtime_error(""));
+        }();
+
+        try {
+          result = co_await folly::coro::co_nothrow(std::move(t));
+        } catch (...) {
+          ADD_FAILURE();
+        }
+        ADD_FAILURE();
+      }()));
+  EXPECT_TRUE(res.hasException<std::runtime_error>());
+}
+
+TEST_F(TaskTest, CoAwaitNothrowWithScheduleOn) {
+  auto res =
+      folly::coro::blockingWait(co_awaitTry([]() -> folly::coro::Task<void> {
+        auto t = []() -> folly::coro::Task<int> { co_return 42; }();
+
+        int result = co_await folly::coro::co_nothrow(
+            std::move(t).scheduleOn(folly::getGlobalCPUExecutor()));
+        EXPECT_EQ(42, result);
+
+        t = []() -> folly::coro::Task<int> {
+          co_yield folly::coro::co_error(std::runtime_error(""));
+        }();
+
+        try {
+          result = co_await folly::coro::co_nothrow(
+              std::move(t).scheduleOn(folly::getGlobalCPUExecutor()));
+        } catch (...) {
+          ADD_FAILURE();
+        }
+        ADD_FAILURE();
+      }()));
+  EXPECT_TRUE(res.hasException<std::runtime_error>());
+}
+
+TEST_F(TaskTest, CoAwaitThrowAfterNothrow) {
+  auto res =
+      folly::coro::blockingWait(co_awaitTry([]() -> folly::coro::Task<void> {
+        auto t = []() -> folly::coro::Task<int> { co_return 42; }();
+
+        int result = co_await folly::coro::co_nothrow(std::move(t));
+        EXPECT_EQ(42, result);
+
+        t = []() -> folly::coro::Task<int> {
+          co_yield folly::coro::co_error(std::runtime_error(""));
+        }();
+
+        try {
+          result = co_await std::move(t);
+          ADD_FAILURE();
+        } catch (...) {
+          throw std::logic_error("translated");
+        }
+      }()));
+  EXPECT_TRUE(res.hasException<std::logic_error>());
+}
+
 #endif