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

folly::coro::makeUnorderedAsyncGeneratorFromAwaitable[Try]Range

Summary:
Allows waiting for a range of SemiAwaitables while receiving progress incrementally.
Often requested as a repeatedly-invokable collectAll.

Uses an external AsyncScope to avoid the following sharp edge with using an internal scope (demonstrated in V1 of this diff): if the AsyncGenerator is destroyed without being fully drained we'd have to block on the scope, which can deadlock.

Reviewed By: helfman

Differential Revision: D25924855

fbshipit-source-id: 46dd7c64d58560f57a7c18e35c159beb984286a7
parent e5ea61d1
...@@ -15,6 +15,8 @@ ...@@ -15,6 +15,8 @@
*/ */
#include <folly/ExceptionWrapper.h> #include <folly/ExceptionWrapper.h>
#include <folly/experimental/coro/AsyncPipe.h>
#include <folly/experimental/coro/AsyncScope.h>
#include <folly/experimental/coro/Mutex.h> #include <folly/experimental/coro/Mutex.h>
#include <folly/experimental/coro/detail/Barrier.h> #include <folly/experimental/coro/detail/Barrier.h>
#include <folly/experimental/coro/detail/BarrierTask.h> #include <folly/experimental/coro/detail/BarrierTask.h>
...@@ -214,6 +216,73 @@ auto collectAllImpl( ...@@ -214,6 +216,73 @@ auto collectAllImpl(
} }
} }
template <typename InputRange, typename IsTry>
auto makeUnorderedAsyncGeneratorFromAwaitableRangeImpl(
AsyncScope& scope, InputRange awaitables, IsTry) {
using Item =
async_generator_from_awaitable_range_item_t<InputRange, IsTry::value>;
return [](AsyncScope& scope,
InputRange awaitables) -> AsyncGenerator<Item&&> {
auto [results, pipe] = AsyncPipe<Item, false>::create();
const CancellationSource cancelSource;
auto guard = folly::makeGuard([&] { cancelSource.requestCancellation(); });
auto ex = co_await co_current_executor;
size_t expected = 0;
// Save the initial context and restore it after starting each task
// as the task may have modified the context before suspending and we
// want to make sure the next task is started with the same initial
// context.
const auto context = RequestContext::saveContext();
for (auto&& semiAwaitable : static_cast<InputRange&&>(awaitables)) {
scope.add(
[](auto semiAwaitable,
auto& cancelSource,
auto& pipe) -> Task<void> {
auto result = co_await co_withCancellation(
cancelSource.getToken(), co_awaitTry(std::move(semiAwaitable)));
if (!result.hasValue() && !IsTry::value) {
cancelSource.requestCancellation();
}
pipe.write(std::move(result));
}(static_cast<decltype(semiAwaitable)&&>(semiAwaitable),
cancelSource,
pipe)
.scheduleOn(ex));
++expected;
RequestContext::setContext(context);
}
while (true) {
CancellationCallback cancelCallback(
co_await co_current_cancellation_token,
[&]() noexcept { cancelSource.requestCancellation(); });
if constexpr (!IsTry::value) {
auto result = co_await co_awaitTry(results.next());
if (result.hasValue()) {
co_yield std::move(*result);
if (--expected) {
continue;
}
result = {}; // completion result
}
guard.dismiss();
co_yield co_result(std::move(result));
} else {
// Prevent AsyncPipe from receiving cancellation so we get the right
// number of OperationCancelleds.
auto result = co_await co_withCancellation({}, results.next());
co_yield std::move(*result);
if (--expected == 0) {
guard.dismiss();
co_return;
}
}
}
}(scope, std::move(awaitables));
}
} // namespace detail } // namespace detail
template <typename... SemiAwaitables> template <typename... SemiAwaitables>
...@@ -921,6 +990,26 @@ auto collectAllTryWindowed(InputRange awaitables, std::size_t maxConcurrency) ...@@ -921,6 +990,26 @@ auto collectAllTryWindowed(InputRange awaitables, std::size_t maxConcurrency)
co_return results; co_return results;
} }
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableRange(
AsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
false>&&> {
return detail::makeUnorderedAsyncGeneratorFromAwaitableRangeImpl(
scope, std::move(awaitables), bool_constant<false>{});
}
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
AsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
true>&&> {
return detail::makeUnorderedAsyncGeneratorFromAwaitableRangeImpl(
scope, std::move(awaitables), bool_constant<true>{});
}
} // namespace coro } // namespace coro
} // namespace folly } // namespace folly
......
...@@ -19,6 +19,8 @@ ...@@ -19,6 +19,8 @@
#include <folly/Try.h> #include <folly/Try.h>
#include <folly/Unit.h> #include <folly/Unit.h>
#include <folly/container/Access.h> #include <folly/container/Access.h>
#include <folly/experimental/coro/AsyncGenerator.h>
#include <folly/experimental/coro/AsyncScope.h>
#include <folly/experimental/coro/Coroutine.h> #include <folly/experimental/coro/Coroutine.h>
#include <folly/experimental/coro/Task.h> #include <folly/experimental/coro/Task.h>
#include <folly/experimental/coro/ViaIfAsync.h> #include <folly/experimental/coro/ViaIfAsync.h>
...@@ -230,6 +232,62 @@ auto collectAllTryRange(std::vector<SemiAwaitable> awaitables) ...@@ -230,6 +232,62 @@ auto collectAllTryRange(std::vector<SemiAwaitable> awaitables)
co_return co_await collectAllTryRange(awaitables | ranges::views::move); co_return co_await collectAllTryRange(awaitables | ranges::views::move);
} }
namespace detail {
template <typename InputRange, bool IsTry>
using async_generator_from_awaitable_range_item_t = conditional_t<
IsTry,
collect_all_try_range_component_t<range_reference_t<InputRange>>,
collect_all_range_component_t<range_reference_t<InputRange>>>;
}
////////////////////////////////////////////////////////////////////////////
// makeUnorderedAsyncGeneratorFromAwaitableRange(AsyncScope&,
// RangeOf<SemiAwaitable<T>>&&) -> AsyncGenerator<T&&>
// makeUnorderedAsyncGeneratorFromAwaitableTryRange(AsyncScope&,
// RangeOf<SemiAwaitable<T>>&&) -> AsyncGenerator<Try<T>&&>
// Returns an AsyncGenerator that yields results of passed-in awaitables in
// order of completion.
// Destroying or cancelling the AsyncGenerator cancels the remaining awaitables.
//
// makeUnorderedAsyncGeneratorFromAwaitableRange cancels all remaining
// awaitables when any of them fail with an exception. Any results obtained
// before the failure are still returned via the generator, then the first
// exception in time. makeUnorderedAsyncGeneratorFromAwaitableTryRange does not
// cancel awaitables when one fails, and yields all results even when cancelled.
//
// Awaitables are attached to the passed-in AsyncScope.
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableRange(
AsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
false>&&>;
template <typename InputRange>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
AsyncScope& scope, InputRange awaitables)
-> AsyncGenerator<detail::async_generator_from_awaitable_range_item_t<
InputRange,
true>&&>;
template <typename SemiAwaitable>
auto makeUnorderedAsyncGeneratorFromAwaitableRange(
AsyncScope& scope, std::vector<SemiAwaitable> awaitables)
-> decltype(makeUnorderedAsyncGeneratorFromAwaitableRange(
scope, awaitables | ranges::views::move)) {
co_return co_await makeUnorderedAsyncGeneratorFromAwaitableRange(
scope, awaitables | ranges::views::move);
}
template <typename SemiAwaitable>
auto makeUnorderedAsyncGeneratorFromAwaitableTryRange(
AsyncScope& 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,6 +819,113 @@ TEST_F(CollectAllRangeTest, VectorOfTaskWithExecutorUsage) { ...@@ -819,6 +819,113 @@ TEST_F(CollectAllRangeTest, VectorOfTaskWithExecutorUsage) {
}()); }());
} }
TEST_F(CollectAllRangeTest, GeneatorFromRange) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i;
};
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
for (int i = 5; i > 0; --i) {
co_yield makeTask(i);
}
};
auto results = folly::coro::makeUnorderedAsyncGeneratorFromAwaitableRange(
scope, generateTasks());
// co_await doesn't work inside EXPECT_EQ
EXPECT_TRUE(*(co_await results.next()) == 1);
EXPECT_TRUE(*(co_await results.next()) == 2);
EXPECT_TRUE(*(co_await results.next()) == 3);
EXPECT_TRUE(*(co_await results.next()) == 4);
EXPECT_TRUE(*(co_await results.next()) == 5);
EXPECT_FALSE(co_await results.next());
co_await scope.joinAsync();
}());
}
TEST_F(CollectAllRangeTest, GeneatorFromRangePartialConsume) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> { co_return i; };
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
for (int i = 5; i > 0; --i) {
co_yield makeTask(i);
}
};
auto results = folly::coro::makeUnorderedAsyncGeneratorFromAwaitableRange(
scope, generateTasks());
for (int i = 0; i < 3; ++i) {
co_await results.next();
}
co_await scope.joinAsync();
}());
}
TEST_F(CollectAllRangeTest, GeneatorFromRangeFailed) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i;
};
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
co_yield []() -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(350));
co_yield folly::coro::co_error(std::runtime_error("foo"));
}();
for (int i = 5; i > 0; --i) {
co_yield makeTask(i);
}
};
auto results = folly::coro::makeUnorderedAsyncGeneratorFromAwaitableRange(
scope, generateTasks());
// co_await doesn't work inside EXPECT_EQ
EXPECT_TRUE(*(co_await results.next()) == 1);
EXPECT_TRUE(*(co_await results.next()) == 2);
EXPECT_TRUE(*(co_await results.next()) == 3);
EXPECT_TRUE((co_await co_awaitTry(results.next()))
.hasException<std::runtime_error>());
co_await scope.joinAsync();
}());
}
TEST_F(CollectAllRangeTest, GeneatorFromRangeCancelled) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope 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) {
cancelSource.requestCancellation();
}
}
};
auto start = std::chrono::steady_clock::now();
auto results = folly::coro::makeUnorderedAsyncGeneratorFromAwaitableRange(
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) {
...@@ -1000,6 +1107,96 @@ TEST_F(CollectAllTryRangeTest, KeepsRequestContextOfChildTasksIndependent) { ...@@ -1000,6 +1107,96 @@ TEST_F(CollectAllTryRangeTest, KeepsRequestContextOfChildTasksIndependent) {
}()); }());
} }
TEST_F(CollectAllTryRangeTest, GeneatorFromRange) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(100 * i));
co_return i;
};
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
for (int i = 5; i > 0; --i) {
co_yield makeTask(i);
}
};
auto results =
folly::coro::makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, generateTasks());
// co_await doesn't work inside EXPECT_EQ
EXPECT_TRUE(**(co_await results.next()) == 1);
EXPECT_TRUE(**(co_await results.next()) == 2);
EXPECT_TRUE(**(co_await results.next()) == 3);
EXPECT_TRUE(**(co_await results.next()) == 4);
EXPECT_TRUE(**(co_await results.next()) == 5);
co_await scope.joinAsync();
}());
}
TEST_F(CollectAllTryRangeTest, GeneatorFromRangeFailed) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope scope;
auto makeTask = [](int i) -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(200 * i));
co_return i;
};
auto generateTasks =
[&]() -> folly::coro::Generator<folly::coro::Task<int>&&> {
co_yield []() -> folly::coro::Task<int> {
co_await folly::coro::sleep(std::chrono::milliseconds(700));
co_yield folly::coro::co_error(std::runtime_error("foo"));
}();
for (int i = 5; i > 0; --i) {
co_yield makeTask(i);
}
};
auto results =
folly::coro::makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, generateTasks());
// co_await doesn't work inside EXPECT_EQ
EXPECT_TRUE(**(co_await results.next()) == 1);
EXPECT_TRUE(**(co_await results.next()) == 2);
EXPECT_TRUE(**(co_await results.next()) == 3);
EXPECT_TRUE((co_await results.next())->hasException<std::runtime_error>());
EXPECT_TRUE(**(co_await results.next()) == 4);
EXPECT_TRUE(**(co_await results.next()) == 5);
co_await scope.joinAsync();
}());
}
TEST_F(CollectAllTryRangeTest, GeneatorFromRangeCancelled) {
folly::coro::blockingWait([]() -> folly::coro::Task<void> {
folly::coro::AsyncScope 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 = 1; i < 10; ++i) {
co_yield makeTask(i);
if (i == 4) {
cancelSource.requestCancellation();
}
}
};
auto start = std::chrono::steady_clock::now();
auto results =
folly::coro::makeUnorderedAsyncGeneratorFromAwaitableTryRange(
scope, generateTasks());
auto result = co_await folly::coro::co_withCancellation(
cancelSource.getToken(), 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 CollectAllWindowedTest : public testing::Test {}; class CollectAllWindowedTest : public testing::Test {};
TEST_F(CollectAllWindowedTest, ConcurrentTasks) { TEST_F(CollectAllWindowedTest, ConcurrentTasks) {
......
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