Add folly::coro::AsyncScope::remaining()

Summary:
Add a method for querying the number of remaining tasks
that have been added to an AsyncScope but that have not
yet completed.

This can be useful for implementing some load-shedding
strategies based on the current number of active tasks
added to a scope.

Reviewed By: andriigrynenko

Differential Revision: D21461637

fbshipit-source-id: 628c236e0c99fa952e63f84366186246e5286bc8

folly/experimental/coro/AsyncScope.h

@@ -63,6 +63,9 @@ class AsyncScope {
   // completes before calling the destructor.
   ~AsyncScope();
 
+  // Query the number of tasks added to the scope that have not yet completed.
+  std::size_t remaining() const noexcept;
+
   // Start the specified task/awaitable by co_awaiting it.
   //
   // Exceptions
@@ -141,6 +144,11 @@ inline AsyncScope::~AsyncScope() {
       "cleanup() not yet complete");
 }
 
+inline std::size_t AsyncScope::remaining() const noexcept {
+  const std::size_t count = barrier_.remaining();
+  return count > 1 ? (count - 1) : 0;
+}
+
 template <typename Awaitable>
 inline void AsyncScope::add(Awaitable&& awaitable) {
   assert(

folly/experimental/coro/detail/Barrier.h

@@ -46,6 +46,17 @@ class Barrier {
     assert(SIZE_MAX - oldCount >= count);
   }
 
+  // Query the number of remaining tasks that the barrier is waiting
+  // for. This indicates the number of arrive() calls that must be
+  // made before the Barrier will be released.
+  //
+  // Note that this should just be used as an approximate guide
+  // for the number of outstanding tasks. This value may be out
+  // of date immediately upon being returned.
+  std::size_t remaining() const noexcept {
+    return count_.load(std::memory_order_acquire);
+  }
+
   [[nodiscard]] std::experimental::coroutine_handle<> arrive() noexcept {
     const std::size_t oldCount = count_.fetch_sub(1, std::memory_order_acq_rel);
 

folly/experimental/coro/test/AsyncScopeTest.cpp

@@ -83,4 +83,30 @@ TEST_F(AsyncScopeTest, StartChildTasksAfterCleanupStarted) {
   }());
 }
 
+TEST_F(AsyncScopeTest, QueryRemainingCount) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<> {
+    folly::coro::Baton baton;
+
+    auto makeTask = [&]() -> folly::coro::Task<> { co_await baton; };
+
+    auto executor = co_await folly::coro::co_current_executor;
+
+    folly::coro::AsyncScope scope;
+
+    CHECK_EQ(0, scope.remaining());
+
+    for (int i = 0; i < 10; ++i) {
+      scope.add(makeTask().scheduleOn(executor));
+    }
+
+    CHECK_EQ(10, scope.remaining());
+
+    baton.post();
+
+    co_await scope.joinAsync();
+
+    CHECK_EQ(0, scope.remaining());
+  }());
+}
+
 #endif // FOLLY_HAS_COROUTINES