Commit bdd4cbe1 authored by Lewis Baker's avatar Lewis Baker Committed by Facebook GitHub Bot

Add co_awaitTry() support to AsyncGenerator

Summary:
Add the ability for an AsyncGenerator to yield an error without needing to throw an exception and for consumers to retrieve the error without rethrowing the exception.

Added a new `folly::coro::co_error` type that acts as a simple wrapper around a `folly::exception_wrapper` that can be used to indicate that this value is an error.

This allows `AsyncGenerator` bodies to now `co_yield` a `co_error` as a way of completing with an error without needing to throw the exception.

For example:
```
class SomeError : public std::exception { ... };
folly::coro::AsyncGenerator<int> example() {
  co_yield 42;
  co_yield folly::coro::co_error(SomeError{});
}
```

This diff also adds support for `folly::coro::co_awaitTry()` to the `AsyncGenerator::next()` method.

For example:
```
folly::coro::Task<void> consumer() {
  folly::coro::AsyncGenerator<int> gen = example();
  folly::Try<int> result =
    co_await folly::coro::co_awaitTry(gen.next());
  // use result
}
```

Added a benchmark to compare performance of the new interfaces vs the existing throwing-based error handling.

Reviewed By: yfeldblum

Differential Revision: D18015926

fbshipit-source-id: 27d2e16e7dde9b4c871846f5d17b40ec3737330a
parent 3153b1b0
...@@ -17,8 +17,11 @@ ...@@ -17,8 +17,11 @@
#pragma once #pragma once
#include <folly/CancellationToken.h> #include <folly/CancellationToken.h>
#include <folly/ExceptionWrapper.h>
#include <folly/Traits.h> #include <folly/Traits.h>
#include <folly/Try.h>
#include <folly/experimental/coro/CurrentExecutor.h> #include <folly/experimental/coro/CurrentExecutor.h>
#include <folly/experimental/coro/Error.h>
#include <folly/experimental/coro/Invoke.h> #include <folly/experimental/coro/Invoke.h>
#include <folly/experimental/coro/Utils.h> #include <folly/experimental/coro/Utils.h>
#include <folly/experimental/coro/ViaIfAsync.h> #include <folly/experimental/coro/ViaIfAsync.h>
...@@ -57,9 +60,22 @@ class AsyncGeneratorPromise { ...@@ -57,9 +60,22 @@ class AsyncGeneratorPromise {
}; };
public: public:
AsyncGeneratorPromise() noexcept {}
~AsyncGeneratorPromise() { ~AsyncGeneratorPromise() {
if (hasValue_) { switch (state_) {
value_.destruct(); case State::VALUE:
folly::coro::detail::deactivate(value_);
break;
case State::EXCEPTION_WRAPPER:
folly::coro::detail::deactivate(exceptionWrapper_);
break;
case State::EXCEPTION_PTR:
folly::coro::detail::deactivate(exceptionPtr_);
break;
case State::DONE:
case State::INVALID:
break;
} }
} }
...@@ -78,16 +94,16 @@ class AsyncGeneratorPromise { ...@@ -78,16 +94,16 @@ class AsyncGeneratorPromise {
} }
YieldAwaiter final_suspend() noexcept { YieldAwaiter final_suspend() noexcept {
DCHECK(!hasValue_); DCHECK(!hasValue());
return {}; return {};
} }
YieldAwaiter yield_value(Reference&& value) noexcept( YieldAwaiter yield_value(Reference&& value) noexcept(
std::is_nothrow_move_constructible<Reference>::value) { std::is_nothrow_move_constructible<Reference>::value) {
DCHECK(!hasValue_); DCHECK(state_ == State::INVALID);
value_.construct(static_cast<Reference&&>(value)); folly::coro::detail::activate(value_, static_cast<Reference&&>(value));
hasValue_ = true; state_ = State::VALUE;
return {}; return YieldAwaiter{};
} }
// In the case where 'Reference' is not actually a reference-type we // In the case where 'Reference' is not actually a reference-type we
...@@ -103,19 +119,29 @@ class AsyncGeneratorPromise { ...@@ -103,19 +119,29 @@ class AsyncGeneratorPromise {
int> = 0> int> = 0>
YieldAwaiter yield_value(U&& value) noexcept( YieldAwaiter yield_value(U&& value) noexcept(
std::is_nothrow_constructible_v<Reference, U>) { std::is_nothrow_constructible_v<Reference, U>) {
DCHECK(!hasValue_); DCHECK(state_ == State::INVALID);
value_.construct(static_cast<U&&>(value)); folly::coro::detail::activate(value_, static_cast<U&&>(value));
hasValue_ = true; state_ = State::VALUE;
return {};
}
YieldAwaiter yield_value(co_error&& error) noexcept {
DCHECK(state_ == State::INVALID);
folly::coro::detail::activate(
exceptionWrapper_, std::move(error.exception()));
state_ = State::EXCEPTION_WRAPPER;
return {}; return {};
} }
void unhandled_exception() noexcept { void unhandled_exception() noexcept {
DCHECK(!hasValue_); DCHECK(state_ == State::INVALID);
exception_ = std::current_exception(); folly::coro::detail::activate(exceptionPtr_, std::current_exception());
state_ = State::EXCEPTION_PTR;
} }
void return_void() noexcept { void return_void() noexcept {
DCHECK(!hasValue_); DCHECK(state_ == State::INVALID);
state_ = State::DONE;
} }
template <typename U> template <typename U>
...@@ -131,7 +157,7 @@ class AsyncGeneratorPromise { ...@@ -131,7 +157,7 @@ class AsyncGeneratorPromise {
} }
auto await_transform(folly::coro::co_current_cancellation_token_t) noexcept { auto await_transform(folly::coro::co_current_cancellation_token_t) noexcept {
return AwaitableReady<folly::CancellationToken>{cancelToken_}; return AwaitableReady<const folly::CancellationToken&>{cancelToken_};
} }
void setCancellationToken(folly::CancellationToken cancelToken) noexcept { void setCancellationToken(folly::CancellationToken cancelToken) noexcept {
...@@ -152,27 +178,42 @@ class AsyncGeneratorPromise { ...@@ -152,27 +178,42 @@ class AsyncGeneratorPromise {
continuation_ = continuation; continuation_ = continuation;
} }
bool hasException() const noexcept {
return state_ == State::EXCEPTION_WRAPPER || state_ == State::EXCEPTION_PTR;
}
folly::exception_wrapper getException() noexcept {
DCHECK(hasException());
if (state_ == State::EXCEPTION_WRAPPER) {
return std::move(exceptionWrapper_.get());
} else {
return exception_wrapper::from_exception_ptr(
std::move(exceptionPtr_.get()));
}
}
void throwIfException() { void throwIfException() {
DCHECK(!hasValue_); if (state_ == State::EXCEPTION_WRAPPER) {
if (exception_) { exceptionWrapper_.get().throw_exception();
std::rethrow_exception(exception_); } else if (state_ == State::EXCEPTION_PTR) {
std::rethrow_exception(std::move(exceptionPtr_.get()));
} }
} }
decltype(auto) getRvalue() noexcept { decltype(auto) getRvalue() noexcept {
DCHECK(hasValue_); DCHECK(hasValue());
return std::move(value_).get(); return std::move(value_).get();
} }
void clearValue() noexcept { void clearValue() noexcept {
if (hasValue_) { if (hasValue()) {
hasValue_ = false; state_ = State::INVALID;
value_.destruct(); folly::coro::detail::deactivate(value_);
} }
} }
bool hasValue() const noexcept { bool hasValue() const noexcept {
return hasValue_; return state_ == State::VALUE;
} }
private: private:
...@@ -184,12 +225,28 @@ class AsyncGeneratorPromise { ...@@ -184,12 +225,28 @@ class AsyncGeneratorPromise {
return executor; return executor;
} }
enum class State : std::uint8_t {
INVALID,
VALUE,
EXCEPTION_PTR,
EXCEPTION_WRAPPER,
DONE,
};
std::experimental::coroutine_handle<> continuation_; std::experimental::coroutine_handle<> continuation_;
folly::Executor::KeepAlive<> executor_; folly::Executor::KeepAlive<> executor_;
folly::CancellationToken cancelToken_; folly::CancellationToken cancelToken_;
std::exception_ptr exception_; union {
ManualLifetime<Reference> value_; // Store both an exception_wrapper and an exception_ptr to
bool hasValue_ = false; // avoid needing to rethrow the exception in unhandled_exception()
// to ensure we capture the type-information into the exception_wrapper
// just in case we are only going to rethrow it back into the awaiting
// coroutine.
ManualLifetime<folly::exception_wrapper> exceptionWrapper_;
ManualLifetime<std::exception_ptr> exceptionPtr_;
ManualLifetime<Reference> value_;
};
State state_ = State::INVALID;
bool hasCancelTokenOverride_ = false; bool hasCancelTokenOverride_ = false;
}; };
...@@ -436,7 +493,7 @@ class FOLLY_NODISCARD AsyncGenerator { ...@@ -436,7 +493,7 @@ class FOLLY_NODISCARD AsyncGenerator {
NextResult await_resume() { NextResult await_resume() {
if (!coro_) { if (!coro_) {
return NextResult{}; return NextResult{};
} else if (coro_.done()) { } else if (!coro_.promise().hasValue()) {
coro_.promise().throwIfException(); coro_.promise().throwIfException();
return NextResult{}; return NextResult{};
} else { } else {
...@@ -444,6 +501,18 @@ class FOLLY_NODISCARD AsyncGenerator { ...@@ -444,6 +501,18 @@ class FOLLY_NODISCARD AsyncGenerator {
} }
} }
folly::Try<Value> await_resume_try() {
folly::Try<Value> result;
if (coro_) {
if (coro_.promise().hasValue()) {
result.emplace(coro_.promise().getRvalue());
} else if (coro_.promise().hasException()) {
result.emplaceException(coro_.promise().getException());
}
}
return result;
}
private: private:
friend NextSemiAwaitable; friend NextSemiAwaitable;
explicit NextAwaitable(handle_t coro) noexcept : coro_(coro) {} explicit NextAwaitable(handle_t coro) noexcept : coro_(coro) {}
......
/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <folly/Benchmark.h>
#include <folly/Portability.h>
#if FOLLY_HAS_COROUTINES
#include <folly/experimental/coro/AsyncGenerator.h>
#include <folly/experimental/coro/BlockingWait.h>
#include <folly/experimental/coro/Task.h>
#include <folly/experimental/coro/ViaIfAsync.h>
#include <folly/ExceptionWrapper.h>
#include <exception>
struct SomeError : std::exception {};
BENCHMARK(asyncGeneratorThrowError, iters) {
folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
for (size_t iter = 0; iter < iters; ++iter) {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
throw SomeError{};
}();
auto item1 = co_await gen.next();
try {
auto item2 = co_await gen.next();
std::terminate();
} catch (const SomeError&) {
}
}
}());
}
BENCHMARK(asyncGeneratorThrowErrorAwaitTry, iters) {
folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
for (size_t iter = 0; iter < iters; ++iter) {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
throw SomeError{};
}();
auto try1 = co_await folly::coro::co_awaitTry(gen.next());
auto try2 = co_await folly::coro::co_awaitTry(gen.next());
if (!try2.hasException() ||
!try2.exception().is_compatible_with<SomeError>()) {
std::terminate();
}
}
}());
}
BENCHMARK(asyncGeneratorYieldError, iters) {
folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
for (size_t iter = 0; iter < iters; ++iter) {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
co_yield folly::coro::co_error(SomeError{});
}();
auto item1 = co_await gen.next();
try {
auto item2 = co_await gen.next();
std::terminate();
} catch (const SomeError&) {
}
}
}());
}
BENCHMARK(asyncGeneratorYieldErrorAwaitTry, iters) {
folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
for (size_t iter = 0; iter < iters; ++iter) {
auto gen = []() -> folly::coro::AsyncGenerator<int> {
co_yield 42;
co_yield folly::coro::co_error(SomeError{});
}();
auto try1 = co_await folly::coro::co_awaitTry(gen.next());
auto try2 = co_await folly::coro::co_awaitTry(gen.next());
if (!try2.hasException() ||
!try2.exception().is_compatible_with<SomeError>()) {
std::terminate();
}
}
}());
}
#endif
int main(int argc, char** argv) {
gflags::ParseCommandLineFlags(&argc, &argv, true);
folly::runBenchmarks();
return 0;
}
...@@ -491,4 +491,48 @@ TEST_F(AsyncGeneratorTest, SymmetricTransfer) { ...@@ -491,4 +491,48 @@ TEST_F(AsyncGeneratorTest, SymmetricTransfer) {
}()); }());
} }
TEST(AsyncGenerator, YieldCoError) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto gen = []() -> folly::coro::AsyncGenerator<std::string> {
co_yield "foo";
co_yield "bar";
co_yield folly::coro::co_error(SomeError{});
CHECK(false);
}();
auto item1 = co_await gen.next();
CHECK(item1.value() == "foo");
auto item2 = co_await gen.next();
CHECK(item2.value() == "bar");
try {
(void)co_await gen.next();
CHECK(false);
} catch (SomeError) {
}
}());
}
TEST(AsyncGeneraor, CoAwaitTry) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
auto gen = []() -> folly::coro::AsyncGenerator<std::string> {
co_yield "foo";
co_yield "bar";
co_yield folly::coro::co_error(SomeError{});
CHECK(false);
}();
folly::Try<std::string> item1 =
co_await folly::coro::co_awaitTry(gen.next());
CHECK(item1.hasValue());
CHECK(item1.value() == "foo");
auto item2 = co_await folly::coro::co_awaitTry(gen.next());
CHECK(item2.hasValue());
CHECK(item2.value() == "bar");
auto item3 = co_await folly::coro::co_awaitTry(gen.next());
CHECK(item3.hasException());
CHECK(item3.exception().is_compatible_with<SomeError>());
}());
}
#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