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

Make co_awaitTry(AsyncGenerator) return Try<NextResult<T>>

Summary:
There are two problems with the current approach of returning Try<T>:
- It is impossible to write generic algorithms like coro::timeout that convert any awaitable into a Task of its await result without throwing exceptions because there's no way to reconstruct the expected return type. More generally, we want the property that the await_try_result_t::element_type matches the await_result_t so we can make drop-in replacements by wrapping in functions like timeout.
- There's no way to both avoid moving yielded values and avoid throwing exceptions because Try doesn't support references (and an earlier diff adding this support was rejected), which means the two performance optimizations avaioable to users of AsyncGenerator are mutually exclusive

We fix this to restore the aforementioned invariant by wrapping the existing result type. This is a marginal inefficiency, so if we notice regressions as a result we can specialize these Try instantiations to consolidate the storage. For now we do not except this to matter.

Reviewed By: andriigrynenko

Differential Revision: D29680441

fbshipit-source-id: 4ef74f4645d990b623bb95a297718fb576a9b977
parent 7fc541e8
......@@ -201,14 +201,14 @@ class TransformProcessorBase : public IChannelCallback {
while (true) {
auto outputResult =
co_await folly::coro::co_awaitTry(outputGen->next());
if (!outputResult.hasValue() && !outputResult.hasException()) {
if (!outputResult.hasException() && !outputResult->has_value()) {
break;
}
if (cancelToken.isCancellationRequested()) {
co_return CloseResult();
}
if (outputResult.hasValue()) {
sender_->senderPush(std::move(outputResult.value()));
if (!outputResult.hasException()) {
sender_->senderPush(std::move(outputResult->value()));
} else {
// The transform coroutine threw an exception. We will close the
// output receiver.
......
......@@ -292,16 +292,15 @@ class FOLLY_NODISCARD AsyncGenerator {
}
}
folly::Try<Value> await_resume_try() {
folly::Try<Value> result;
folly::Try<NextResult> await_resume_try() {
if (coro_) {
if (coro_.promise().hasValue()) {
result.emplace(coro_.promise().getRvalue());
return folly::Try<NextResult>(NextResult{coro_});
} else if (coro_.promise().hasException()) {
result.emplaceException(coro_.promise().getException());
return folly::Try<NextResult>(coro_.promise().getException());
}
}
return result;
return folly::Try<NextResult>(NextResult{});
}
private:
......@@ -469,6 +468,23 @@ class AsyncGeneratorPromise {
}
}
YieldAwaiter yield_value(
co_result<typename AsyncGenerator<Reference, Value>::NextResult>&&
res) noexcept {
DCHECK(res.result().hasValue() || res.result().hasException());
if (res.result().hasException()) {
return yield_value(co_error(res.result().exception()));
} else if (res.result().hasValue()) {
if (res.result()->has_value()) {
return yield_value(std::move(res.result()->value()));
} else {
return_void();
return {};
}
}
return yield_value(co_error(UsingUninitializedTry{}));
}
variant_awaitable<YieldAwaiter, ready_awaitable<>> await_transform(
co_safe_point_t) noexcept {
if (cancelToken_.isCancellationRequested()) {
......@@ -506,7 +522,8 @@ class AsyncGeneratorPromise {
void setCancellationToken(folly::CancellationToken cancelToken) noexcept {
// Only keep the first cancellation token.
// ie. the inner-most cancellation scope of the consumer's calling context.
// ie. the inner-most cancellation scope of the consumer's calling
// context.
if (!hasCancelTokenOverride_) {
cancelToken_ = std::move(cancelToken);
hasCancelTokenOverride_ = true;
......
......@@ -256,12 +256,12 @@ auto makeUnorderedAsyncGeneratorImpl(
if constexpr (!IsTry::value) {
auto result = co_await co_awaitTry(results.next());
if (result.hasValue()) {
co_yield std::move(*result);
if (result.hasValue() && result->has_value()) {
co_yield std::move(**result);
if (--expected) {
continue;
}
result = {}; // completion result
result.emplace(); // completion result
}
guard.dismiss();
co_yield co_result(std::move(result));
......
......@@ -585,13 +585,12 @@ TEST(AsyncGeneraor, CoAwaitTry) {
CHECK(false);
}();
folly::Try<std::string> item1 =
co_await folly::coro::co_awaitTry(gen.next());
auto item1 = co_await folly::coro::co_awaitTry(gen.next());
CHECK(item1.hasValue());
CHECK(item1.value() == "foo");
CHECK(*item1.value() == "foo");
auto item2 = co_await folly::coro::co_awaitTry(gen.next());
CHECK(item2.hasValue());
CHECK(item2.value() == "bar");
CHECK(*item2.value() == "bar");
auto item3 = co_await folly::coro::co_awaitTry(gen.next());
CHECK(item3.hasException());
CHECK(item3.exception().is_compatible_with<SomeError>());
......
......@@ -169,7 +169,7 @@ TEST(Timeout, AsyncGenerator) {
// Completing synchronously with a value.
auto result = co_await coro::timeout(
[]() -> coro::AsyncGenerator<int> { co_yield 42; }().next(), 1s);
EXPECT_EQ(42, result);
EXPECT_EQ(42, *result);
// Test that it handles failing synchronously
auto tryResult = co_await coro::co_awaitTry(coro::timeout(
......
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