Commit 2f5a71dd authored by Shai Szulanski's avatar Shai Szulanski Committed by Facebook GitHub Bot

Allow using CancellableAsyncScope with external cancellation token

Summary: Right now the internal cancellation signal will be silently ignored, which hurts usability. Add mechanism for merging with external source and comment describing the pitfall.

Reviewed By: capickett

Differential Revision: D29935016

fbshipit-source-id: 9311930ff9fbfc0470fdcfb5d425f36e7f0aff06
parent c17ed205
...@@ -184,9 +184,19 @@ inline folly::SemiFuture<folly::Unit> AsyncScope::cleanup() noexcept { ...@@ -184,9 +184,19 @@ inline folly::SemiFuture<folly::Unit> AsyncScope::cleanup() noexcept {
// A cancellable version of AsyncScope. Work added to this scope will be // A cancellable version of AsyncScope. Work added to this scope will be
// provided a cancellation token for cancelling during join. See // provided a cancellation token for cancelling during join. See
// add() and cancelAndJoinAsync() for more information. // add() and cancelAndJoinAsync() for more information.
//
// Note that Task and AsyncGenerator will ignore the internal cancellation
// signal if they already have a cancellation token (i.e. if someone has already
// called co_withCancellation on them.)
// If you need an external cancellation signal as well, pass that token to this
// constructor or to add() instead of attaching it to the Awaitable.
class CancellableAsyncScope { class CancellableAsyncScope {
public: public:
CancellableAsyncScope() noexcept = default; CancellableAsyncScope() noexcept
: cancellationToken_(cancellationSource_.getToken()) {}
explicit CancellableAsyncScope(CancellationToken&& token)
: cancellationToken_(CancellationToken::merge(
cancellationSource_.getToken(), std::move(token))) {}
// Query the number of tasks added to the scope that have not yet completed. // Query the number of tasks added to the scope that have not yet completed.
std::size_t remaining() const noexcept { return scope_.remaining(); } std::size_t remaining() const noexcept { return scope_.remaining(); }
...@@ -194,16 +204,22 @@ class CancellableAsyncScope { ...@@ -194,16 +204,22 @@ class CancellableAsyncScope {
// Start the specified task/awaitable by co_awaiting it. The awaitable will be // Start the specified task/awaitable by co_awaiting it. The awaitable will be
// provided a cancellation token to respond to cancelAndJoinAsync() in the // provided a cancellation token to respond to cancelAndJoinAsync() in the
// future. // future.
// An additional cancellation token may be passed in to apply to the
// awaitable; it will be merged with the internal token.
// //
// Note that cancellation is cooperative, your task must handle cancellation // Note that cancellation is cooperative, your task must handle cancellation
// in order to have any effect. // in order to have any effect.
// //
// See the documentation on AsyncScope::add. // See the documentation on AsyncScope::add.
template <typename Awaitable> template <typename Awaitable>
void add(Awaitable&& awaitable, void* returnAddress = nullptr) { void add(
Awaitable&& awaitable,
std::optional<CancellationToken> token = std::nullopt,
void* returnAddress = nullptr) {
scope_.add( scope_.add(
co_withCancellation( co_withCancellation(
cancellationSource_.getToken(), token ? CancellationToken::merge(*token, cancellationToken_)
: cancellationToken_,
static_cast<Awaitable&&>(awaitable)), static_cast<Awaitable&&>(awaitable)),
returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS()); returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
} }
...@@ -233,6 +249,7 @@ class CancellableAsyncScope { ...@@ -233,6 +249,7 @@ class CancellableAsyncScope {
private: private:
folly::CancellationSource cancellationSource_; folly::CancellationSource cancellationSource_;
CancellationToken cancellationToken_;
AsyncScope scope_; AsyncScope scope_;
}; };
......
...@@ -191,6 +191,34 @@ CO_TEST_F(CancellableAsyncScopeTest, CancelSuspendedWork) { ...@@ -191,6 +191,34 @@ CO_TEST_F(CancellableAsyncScopeTest, CancelSuspendedWork) {
// this correctly. // this correctly.
co_await scope.cancelAndJoinAsync(); co_await scope.cancelAndJoinAsync();
CO_ASSERT_EQ(0, scope.remaining()); CO_ASSERT_EQ(0, scope.remaining());
folly::CancellationSource source;
folly::coro::CancellableAsyncScope scope2(source.getToken());
CO_ASSERT_EQ(0, scope2.remaining());
for (int i = 0; i < 10; ++i) {
scope2.add(makeTask().scheduleOn(folly::getGlobalCPUExecutor()));
}
CO_ASSERT_EQ(10, scope2.remaining());
source.requestCancellation();
co_await scope2.joinAsync();
CO_ASSERT_EQ(0, scope2.remaining());
source = {};
folly::coro::CancellableAsyncScope scope3;
CO_ASSERT_EQ(0, scope3.remaining());
for (int i = 0; i < 10; ++i) {
scope3.add(
makeTask().scheduleOn(folly::getGlobalCPUExecutor()),
source.getToken());
}
CO_ASSERT_EQ(10, scope3.remaining());
source.requestCancellation();
co_await scope3.joinAsync();
CO_ASSERT_EQ(0, scope3.remaining());
} }
#endif // FOLLY_HAS_COROUTINES #endif // FOLLY_HAS_COROUTINES
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