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

Add CancellableAsyncScope overloads for makeUnorderedAsyncGenerator

Reviewed By: capickett

Differential Revision: D29942172

fbshipit-source-id: 98649f6bf911d2a7a2f0ac687fa51fe76e0b777c
parent dc491606
...@@ -211,7 +211,7 @@ auto collectAllImpl( ...@@ -211,7 +211,7 @@ auto collectAllImpl(
} }
} }
template <typename InputRange, typename IsTry> template <typename InputRange, typename IsTry, typename AsyncScope>
auto makeUnorderedAsyncGeneratorImpl( auto makeUnorderedAsyncGeneratorImpl(
AsyncScope& scope, InputRange awaitables, IsTry) { AsyncScope& scope, InputRange awaitables, IsTry) {
using Item = using Item =
...@@ -230,21 +230,25 @@ auto makeUnorderedAsyncGeneratorImpl( ...@@ -230,21 +230,25 @@ auto makeUnorderedAsyncGeneratorImpl(
const auto context = RequestContext::saveContext(); const auto context = RequestContext::saveContext();
for (auto&& semiAwaitable : static_cast<InputRange&&>(awaitablesParam)) { for (auto&& semiAwaitable : static_cast<InputRange&&>(awaitablesParam)) {
scopeParam.add( auto task = [](auto semiAwaitableParam,
[](auto semiAwaitableParam, auto& cancelSourceParam,
auto& cancelSourceParam, auto& p) -> Task<void> {
auto& p) -> Task<void> { auto result = co_await co_awaitTry(std::move(semiAwaitableParam));
auto result = co_await co_withCancellation( if (!result.hasValue() && !IsTry::value) {
cancelSourceParam.getToken(), cancelSourceParam.requestCancellation();
co_awaitTry(std::move(semiAwaitableParam))); }
if (!result.hasValue() && !IsTry::value) { p.write(std::move(result));
cancelSourceParam.requestCancellation(); }(static_cast<decltype(semiAwaitable)&&>(semiAwaitable),
} cancelSource,
p.write(std::move(result)); pipe);
}(static_cast<decltype(semiAwaitable)&&>(semiAwaitable), if constexpr (std::is_same_v<AsyncScope, folly::coro::AsyncScope>) {
cancelSource, scopeParam.add(
pipe) co_withCancellation(cancelSource.getToken(), std::move(task))
.scheduleOn(ex)); .scheduleOn(ex));
} else {
static_assert(std::is_same_v<AsyncScope, CancellableAsyncScope>);
scopeParam.add(std::move(task).scheduleOn(ex), cancelSource.getToken());
}
++expected; ++expected;
RequestContext::setContext(context); RequestContext::setContext(context);
} }
...@@ -1021,6 +1025,26 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange( ...@@ -1021,6 +1025,26 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, std::move(awaitables), bool_constant<true>{}); scope, std::move(awaitables), bool_constant<true>{});
} }
template <typename InputRange>
auto makeUnorderedAsyncGenerator(
CancellableAsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
false>&&> {
return detail::makeUnorderedAsyncGeneratorImpl(
scope, std::move(awaitables), bool_constant<false>{});
}
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
CancellableAsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
true>&&> {
return detail::makeUnorderedAsyncGeneratorImpl(
scope, std::move(awaitables), bool_constant<true>{});
}
template <typename SemiAwaitable, typename... SemiAwaitables> template <typename SemiAwaitable, typename... SemiAwaitables>
auto collectAny(SemiAwaitable&& awaitable, SemiAwaitables&&... awaitables) auto collectAny(SemiAwaitable&& awaitable, SemiAwaitables&&... awaitables)
-> folly::coro::Task<std::pair< -> folly::coro::Task<std::pair<
......
...@@ -288,6 +288,38 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange( ...@@ -288,6 +288,38 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, awaitables | ranges::views::move); scope, awaitables | ranges::views::move);
} }
// Can also be used with CancellableAsyncScope
template <typename InputRange>
auto makeUnorderedAsyncGenerator(
CancellableAsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
false>&&>;
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
CancellableAsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
true>&&>;
template <typename SemiAwaitable>
auto makeUnorderedAsyncGenerator(
CancellableAsyncScope& scope, std::vector<SemiAwaitable> awaitables)
-> decltype(makeUnorderedAsyncGenerator(
scope, awaitables | ranges::views::move)) {
co_return co_await makeUnorderedAsyncGenerator(
scope, awaitables | ranges::views::move);
}
template <typename SemiAwaitable>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
CancellableAsyncScope& scope, std::vector<SemiAwaitable> awaitables)
-> decltype(makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, awaitables | ranges::views::move)) {
co_return co_await makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, awaitables | ranges::views::move);
}
/////////////////////////////////////////////////////////////////////////////// ///////////////////////////////////////////////////////////////////////////////
// collectAllWindowed(RangeOf<SemiAwaitable<T>>&&, size_t maxConcurrency) // collectAllWindowed(RangeOf<SemiAwaitable<T>>&&, size_t maxConcurrency)
// -> SemiAwaitable<std::vector<T>> // -> SemiAwaitable<std::vector<T>>
......
...@@ -819,9 +819,9 @@ TEST_F(CollectAllRangeTest, VectorOfTaskWithExecutorUsage) { ...@@ -819,9 +819,9 @@ TEST_F(CollectAllRangeTest, VectorOfTaskWithExecutorUsage) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneatorFromRange) { TEST_F(CollectAllRangeTest, GeneratorFromRange) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::CancellableAsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i)); co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i; co_return i;
...@@ -846,7 +846,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRange) { ...@@ -846,7 +846,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRange) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneatorFromRangePartialConsume) { TEST_F(CollectAllRangeTest, GeneratorFromRangePartialConsume) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { co_return i; }; auto makeTask = [](int i) -> folly::coro::Task<int> { co_return i; };
...@@ -866,7 +866,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangePartialConsume) { ...@@ -866,7 +866,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangePartialConsume) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneatorFromRangeFailed) { TEST_F(CollectAllRangeTest, GeneratorFromRangeFailed) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
...@@ -896,7 +896,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangeFailed) { ...@@ -896,7 +896,7 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangeFailed) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneatorFromRangeCancelled) { TEST_F(CollectAllRangeTest, GeneratorFromRangeCancelled) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
...@@ -926,6 +926,36 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangeCancelled) { ...@@ -926,6 +926,36 @@ TEST_F(CollectAllRangeTest, GeneatorFromRangeCancelled) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneratorFromRangeCancelledFromScope) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::CancellableAsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(1000 * i));
co_return i;
};
folly::CancellationSource cancelSource;
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
for (int i = 0; i < 10; ++i) {
co_yield makeTask(i);
if (i == 4) {
scope.requestCancellation();
}
}
};
auto start = std::chrono::steady_clock::now();
auto results =
folly::coro::makeUnorderedAsyncGenerator(scope, generateTasks());
auto result = co_await folly::coro::co_withCancellation(
cancelSource.getToken(), co_awaitTry(results.next()));
auto end = std::chrono::steady_clock::now();
EXPECT_LT(end - start, std::chrono::milliseconds(1000));
EXPECT_TRUE(result.hasException<folly::OperationCancelled>());
co_await scope.joinAsync();
}());
}
class CollectAllTryRangeTest : public testing::Test {}; class CollectAllTryRangeTest : public testing::Test {};
TEST_F(CollectAllTryRangeTest, RangeOfVoidSomeFailing) { TEST_F(CollectAllTryRangeTest, RangeOfVoidSomeFailing) {
...@@ -1107,9 +1137,9 @@ TEST_F(CollectAllTryRangeTest, KeepsRequestContextOfChildTasksIndependent) { ...@@ -1107,9 +1137,9 @@ TEST_F(CollectAllTryRangeTest, KeepsRequestContextOfChildTasksIndependent) {
}()); }());
} }
TEST_F(CollectAllTryRangeTest, GeneatorFromRange) { TEST_F(CollectAllTryRangeTest, GeneratorFromRange) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::CancellableAsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i)); co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i; co_return i;
...@@ -1134,7 +1164,7 @@ TEST_F(CollectAllTryRangeTest, GeneatorFromRange) { ...@@ -1134,7 +1164,7 @@ TEST_F(CollectAllTryRangeTest, GeneatorFromRange) {
}()); }());
} }
TEST_F(CollectAllTryRangeTest, GeneatorFromRangeFailed) { TEST_F(CollectAllTryRangeTest, GeneratorFromRangeFailed) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
...@@ -1166,7 +1196,7 @@ TEST_F(CollectAllTryRangeTest, GeneatorFromRangeFailed) { ...@@ -1166,7 +1196,7 @@ TEST_F(CollectAllTryRangeTest, GeneatorFromRangeFailed) {
}()); }());
} }
TEST_F(CollectAllTryRangeTest, GeneatorFromRangeCancelled) { TEST_F(CollectAllTryRangeTest, GeneratorFromRangeCancelled) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> { folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope; folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { auto makeTask = [](int i) -> folly::coro::Task<int> {
......
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