Commit e2efc12f authored by Shai Szulanski's avatar Shai Szulanski Committed by Facebook GitHub Bot

Add error hint support to AsyncGenerator

Reviewed By: capickett

Differential Revision: D32375291

fbshipit-source-id: 4e88bd4b0caae41c84ef7c249b3108bed9afa80a
parent 160e1fa4
...@@ -262,7 +262,7 @@ class FOLLY_NODISCARD AsyncGenerator { ...@@ -262,7 +262,7 @@ class FOLLY_NODISCARD AsyncGenerator {
bool await_ready() noexcept { return !coro_; } bool await_ready() noexcept { return !coro_; }
template <typename Promise> template <typename Promise>
FOLLY_NOINLINE auto await_suspend( FOLLY_NOINLINE coroutine_handle<> await_suspend(
coroutine_handle<Promise> continuation) noexcept { coroutine_handle<Promise> continuation) noexcept {
auto& promise = coro_.promise(); auto& promise = coro_.promise();
...@@ -367,7 +367,9 @@ class FOLLY_NODISCARD AsyncGenerator { ...@@ -367,7 +367,9 @@ class FOLLY_NODISCARD AsyncGenerator {
namespace detail { namespace detail {
template <typename Reference, typename Value> template <typename Reference, typename Value>
class AsyncGeneratorPromise { class AsyncGeneratorPromise final
: public ExtendedCoroutinePromiseImpl<
AsyncGeneratorPromise<Reference, Value>> {
class YieldAwaiter { class YieldAwaiter {
public: public:
bool await_ready() noexcept { return false; } bool await_ready() noexcept { return false; }
...@@ -377,7 +379,12 @@ class AsyncGeneratorPromise { ...@@ -377,7 +379,12 @@ class AsyncGeneratorPromise {
// Pop AsyncStackFrame first as clearContext() clears the frame state. // Pop AsyncStackFrame first as clearContext() clears the frame state.
folly::popAsyncStackFrameCallee(promise.getAsyncFrame()); folly::popAsyncStackFrameCallee(promise.getAsyncFrame());
promise.clearContext(); promise.clearContext();
return promise.continuation_; if (promise.hasException()) {
auto [handle, frame] =
promise.continuation_.getErrorHandle(promise.getException());
return handle.getHandle();
}
return promise.continuation_.getHandle();
} }
void await_resume() noexcept {} void await_resume() noexcept {}
}; };
...@@ -504,12 +511,22 @@ class AsyncGeneratorPromise { ...@@ -504,12 +511,22 @@ class AsyncGeneratorPromise {
template <typename U> template <typename U>
auto await_transform(U&& value) { auto await_transform(U&& value) {
bypassExceptionThrowing_ =
bypassExceptionThrowing_ == BypassExceptionThrowing::REQUESTED
? BypassExceptionThrowing::ACTIVE
: BypassExceptionThrowing::INACTIVE;
return folly::coro::co_withAsyncStack(folly::coro::co_viaIfAsync( return folly::coro::co_withAsyncStack(folly::coro::co_viaIfAsync(
executor_.get_alias(), executor_.get_alias(),
folly::coro::co_withCancellation( folly::coro::co_withCancellation(
cancelToken_, static_cast<U&&>(value)))); cancelToken_, static_cast<U&&>(value))));
} }
template <typename Awaitable>
auto await_transform(NothrowAwaitable<Awaitable>&& awaitable) {
bypassExceptionThrowing_ = BypassExceptionThrowing::REQUESTED;
return await_transform(awaitable.unwrap());
}
auto await_transform(folly::coro::co_current_executor_t) noexcept { auto await_transform(folly::coro::co_current_executor_t) noexcept {
return ready_awaitable<folly::Executor*>{executor_.get()}; return ready_awaitable<folly::Executor*>{executor_.get()};
} }
...@@ -533,7 +550,7 @@ class AsyncGeneratorPromise { ...@@ -533,7 +550,7 @@ class AsyncGeneratorPromise {
executor_ = std::move(executor); executor_ = std::move(executor);
} }
void setContinuation(coroutine_handle<> continuation) noexcept { void setContinuation(ExtendedCoroutineHandle continuation) noexcept {
continuation_ = continuation; continuation_ = continuation;
} }
...@@ -568,6 +585,21 @@ class AsyncGeneratorPromise { ...@@ -568,6 +585,21 @@ class AsyncGeneratorPromise {
folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; } folly::AsyncStackFrame& getAsyncFrame() noexcept { return asyncFrame_; }
std::pair<ExtendedCoroutineHandle, AsyncStackFrame*> getErrorHandle(
exception_wrapper& ex) override {
if (bypassExceptionThrowing_ == BypassExceptionThrowing::ACTIVE) {
auto yieldAwaiter = yield_value(co_error(std::move(ex)));
DCHECK(!yieldAwaiter.await_ready());
return {
yieldAwaiter.await_suspend(
coroutine_handle<AsyncGeneratorPromise>::from_promise(*this)),
// yieldAwaiter.await_suspend pops a frame
getAsyncFrame().getParentFrame()};
}
return {
coroutine_handle<AsyncGeneratorPromise>::from_promise(*this), nullptr};
}
private: private:
void clearContext() noexcept { void clearContext() noexcept {
executor_ = {}; executor_ = {};
...@@ -583,7 +615,7 @@ class AsyncGeneratorPromise { ...@@ -583,7 +615,7 @@ class AsyncGeneratorPromise {
DONE, DONE,
}; };
coroutine_handle<> continuation_; ExtendedCoroutineHandle continuation_;
folly::AsyncStackFrame asyncFrame_; folly::AsyncStackFrame asyncFrame_;
folly::Executor::KeepAlive<> executor_; folly::Executor::KeepAlive<> executor_;
folly::CancellationToken cancelToken_; folly::CancellationToken cancelToken_;
...@@ -593,6 +625,12 @@ class AsyncGeneratorPromise { ...@@ -593,6 +625,12 @@ class AsyncGeneratorPromise {
}; };
State state_ = State::INVALID; State state_ = State::INVALID;
bool hasCancelTokenOverride_ = false; bool hasCancelTokenOverride_ = false;
enum class BypassExceptionThrowing : uint8_t {
INACTIVE,
ACTIVE,
REQUESTED,
} bypassExceptionThrowing_{BypassExceptionThrowing::INACTIVE};
}; };
} // namespace detail } // namespace detail
......
...@@ -629,4 +629,56 @@ TEST(AsyncGenerator, SafePoint) { ...@@ -629,4 +629,56 @@ TEST(AsyncGenerator, SafePoint) {
}()); }());
} }
TEST(AsyncGenerator, CoAwaitNothrow) {
auto res =
folly::coro::blockingWait(co_awaitTry([]() -> folly::coro::Task<void> {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
auto inner = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
throw std::runtime_error("");
}();
co_yield *co_await co_nothrow(inner.next());
try {
co_yield *co_await co_nothrow(inner.next());
} catch (...) {
ADD_FAILURE();
}
ADD_FAILURE();
}();
co_await co_nothrow(gen.next());
try {
co_await co_nothrow(gen.next());
} catch (...) {
ADD_FAILURE();
}
ADD_FAILURE();
}()));
EXPECT_TRUE(res.hasException<std::runtime_error>());
}
TEST(AsyncGenerator, CoAwaitNothrowMultiTask) {
auto res =
folly::coro::blockingWait(co_awaitTry([]() -> folly::coro::Task<void> {
auto gen = co_await
[]() -> folly::coro::Task<folly::coro::AsyncGenerator<int>> {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
throw std::runtime_error("");
}();
co_await co_nothrow(gen.next());
co_return gen;
}();
try {
co_await co_nothrow(gen.next());
} catch (...) {
ADD_FAILURE();
}
ADD_FAILURE();
}()));
EXPECT_TRUE(res.hasException<std::runtime_error>());
}
#endif #endif
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment