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(
}
}
template <typename InputRange, typename IsTry>
template <typename InputRange, typename IsTry, typename AsyncScope>
auto makeUnorderedAsyncGeneratorImpl(
AsyncScope& scope, InputRange awaitables, IsTry) {
using Item =
......@@ -230,21 +230,25 @@ auto makeUnorderedAsyncGeneratorImpl(
const auto context = RequestContext::saveContext();
for (auto&& semiAwaitable : static_cast<InputRange&&>(awaitablesParam)) {
scopeParam.add(
[](auto semiAwaitableParam,
auto& cancelSourceParam,
auto& p) -> Task<void> {
auto result = co_await co_withCancellation(
cancelSourceParam.getToken(),
co_awaitTry(std::move(semiAwaitableParam)));
if (!result.hasValue() && !IsTry::value) {
cancelSourceParam.requestCancellation();
}
p.write(std::move(result));
}(static_cast<decltype(semiAwaitable)&&>(semiAwaitable),
cancelSource,
pipe)
.scheduleOn(ex));
auto task = [](auto semiAwaitableParam,
auto& cancelSourceParam,
auto& p) -> Task<void> {
auto result = co_await co_awaitTry(std::move(semiAwaitableParam));
if (!result.hasValue() && !IsTry::value) {
cancelSourceParam.requestCancellation();
}
p.write(std::move(result));
}(static_cast<decltype(semiAwaitable)&&>(semiAwaitable),
cancelSource,
pipe);
if constexpr (std::is_same_v<AsyncScope, folly::coro::AsyncScope>) {
scopeParam.add(
co_withCancellation(cancelSource.getToken(), std::move(task))
.scheduleOn(ex));
} else {
static_assert(std::is_same_v<AsyncScope, CancellableAsyncScope>);
scopeParam.add(std::move(task).scheduleOn(ex), cancelSource.getToken());
}
++expected;
RequestContext::setContext(context);
}
......@@ -1021,6 +1025,26 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
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>
auto collectAny(SemiAwaitable&& awaitable, SemiAwaitables&&... awaitables)
-> folly::coro::Task<std::pair<
......
......@@ -288,6 +288,38 @@ auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
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)
// -> SemiAwaitable<std::vector<T>>
......
......@@ -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::AsyncScope scope;
folly::coro::CancellableAsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i;
......@@ -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::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { co_return i; };
......@@ -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::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
......@@ -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::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
......@@ -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 {};
TEST_F(CollectAllTryRangeTest, RangeOfVoidSomeFailing) {
......@@ -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::AsyncScope scope;
folly::coro::CancellableAsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i;
......@@ -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::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
......@@ -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::AsyncScope scope;
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