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

folly/experimental/coro/AsyncScope.h

@@ -184,9 +184,19 @@ inline folly::SemiFuture<folly::Unit> AsyncScope::cleanup() noexcept {
 // A cancellable version of AsyncScope. Work added to this scope will be
 // provided a cancellation token for cancelling during join. See
 // 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 {
  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.
   std::size_t remaining() const noexcept { return scope_.remaining(); }
@@ -194,16 +204,22 @@ class CancellableAsyncScope {
   // Start the specified task/awaitable by co_awaiting it. The awaitable will be
   // provided a cancellation token to respond to cancelAndJoinAsync() in the
   // 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
   // in order to have any effect.
   //
   // See the documentation on AsyncScope::add.
   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(
         co_withCancellation(
-            cancellationSource_.getToken(),
+            token ? CancellationToken::merge(*token, cancellationToken_)
+                  : cancellationToken_,
             static_cast<Awaitable&&>(awaitable)),
         returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
   }
@@ -233,6 +249,7 @@ class CancellableAsyncScope {
 
  private:
   folly::CancellationSource cancellationSource_;
+  CancellationToken cancellationToken_;
   AsyncScope scope_;
 };
 

folly/experimental/coro/test/AsyncScopeTest.cpp

@@ -191,6 +191,34 @@ CO_TEST_F(CancellableAsyncScopeTest, CancelSuspendedWork) {
   // this correctly.
   co_await scope.cancelAndJoinAsync();
   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