Add co_schedule() method to AsyncScope

Summary: Adds a new method `co_schedule()` that takes a task, schedules it on the current executor and adds it to the AsyncScope.

Reviewed By: andriigrynenko

Differential Revision: D30005811

fbshipit-source-id: 48c7aa240dc73c0b4dd54e539e620cf69c68c4fd

folly/experimental/coro/AsyncScope.h

@@ -18,6 +18,7 @@
 
 #include <folly/CancellationToken.h>
 #include <folly/experimental/coro/Coroutine.h>
+#include <folly/experimental/coro/CurrentExecutor.h>
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/detail/Barrier.h>
 #include <folly/experimental/coro/detail/BarrierTask.h>
@@ -224,6 +225,18 @@ class CancellableAsyncScope {
         returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
   }
 
+  // Schedules the given task on the current executor and adds it to the
+  // AsyncScope. The task will be provided a cancellation token to respond to
+  // cancelAndJoinAsync() in the future.
+  //
+  // Note that cancellation is cooperative, your task must handle cancellation
+  // in order to have any effect.
+  template <class T>
+  folly::coro::Task<void> co_schedule(folly::coro::Task<T>&& task) {
+    add(std::move(task).scheduleOn(co_await co_current_executor));
+    co_return;
+  }
+
   // Request cancellation for all started tasks that accepted a
   // CancellationToken in add().
   void requestCancellation() const noexcept {

folly/experimental/coro/test/AsyncScopeTest.cpp

@@ -221,4 +221,25 @@ CO_TEST_F(CancellableAsyncScopeTest, CancelSuspendedWork) {
   CO_ASSERT_EQ(0, scope3.remaining());
 }
 
+CO_TEST_F(CancellableAsyncScopeTest, CancelSuspendedWorkCoSchedule) {
+  using namespace std::chrono_literals;
+
+  auto makeTask = [&]() -> folly::coro::Task<> {
+    co_await folly::coro::sleep(300s);
+  };
+
+  folly::coro::CancellableAsyncScope scope;
+
+  CO_ASSERT_EQ(0, scope.remaining());
+  for (int i = 0; i < 10; ++i) {
+    co_await scope.co_schedule(makeTask());
+  }
+  CO_ASSERT_EQ(10, scope.remaining());
+
+  // Although we are suspended while sleeping, cancelAndJoinAsync will handle
+  // this correctly.
+  co_await scope.cancelAndJoinAsync();
+  CO_ASSERT_EQ(0, scope.remaining());
+}
+
 #endif // FOLLY_HAS_COROUTINES