coro safe_point

Summary:
Add a new concept of a safe-point to folly::Task<...> with implementation-defined semantics.

Example usage:
```lang=c++
Task<> co_something() {
  while (is_work_available()) {
    co_await co_safe_point;
    do_work();
  }
}
```

Previously, a coroutine function which would cancel itself if cancellation is requested would need to spell cancellation pass-through a bit less elegantly:
```lang=c++
Task<> co_something() {
  while (is_work_available()) {
    auto token = co_await co_current_cancellation_token;
    if (token.isCancellationRequested()) {
      co_yield co_cancel;
    }
    do_work();
  }
}
```

Initial semantics include:
* If cancellation has been requested, finish the coroutine with cancellation.
* Otherwise, continue the coroutine.

In the future, it may be wise to add additional behavior:
* If the coroutine has been running for too long since its previous resumption, reschedule it.

Reviewed By: aary

Differential Revision: D25861265

fbshipit-source-id: 6da896f1529d652bfe222cf3a70e9dbe19778510

folly/experimental/coro/AsyncGenerator.h

@@ -144,6 +144,14 @@ class AsyncGeneratorPromise {
     }
   }
 
+  AwaitableVariant<YieldAwaiter, AwaitableReady<void>> await_transform(
+      co_safe_point_t) noexcept {
+    if (cancelToken_.isCancellationRequested()) {
+      return yield_value(co_cancelled);
+    }
+    return AwaitableReady<void>{};
+  }
+
   void unhandled_exception() noexcept {
     DCHECK(state_ == State::INVALID);
     folly::coro::detail::activate(exceptionPtr_, std::current_exception());

folly/experimental/coro/CurrentExecutor.h

@@ -159,5 +159,37 @@ using co_current_cancellation_token_t = detail::co_current_cancellation_token_;
 inline constexpr co_current_cancellation_token_t co_current_cancellation_token{
     co_current_cancellation_token_t::secret_::token_};
 
+//  co_safe_point_t
+//  co_safe_point
+//
+//  A semi-awaitable type and value which, when awaited in an async coroutine
+//  supporting safe-points, causes a safe-point to be reached.
+//
+//  Example:
+//
+//    co_await co_safe_point; // a safe-point is reached
+//
+//  At this safe-point:
+//  - If cancellation has been requested then the coroutine is terminated with
+//    cancellation.
+//  - To aid overall system concurrency, the coroutine may be rescheduled onto
+//    the current executor.
+//  - Otherwise, the coroutine is resumed.
+//
+//  Recommended for use wherever cancellation is checked and handled via early
+//  termination.
+//
+//  Technical note: behavior is typically implemented in some overload
+//  of await_transform in the coroutine's promise type, or in the awaitable
+//  or awaiter it returns. Example:
+//
+//      struct /* some coroutine type */ {
+//        struct promise_type {
+//          /* some awaiter */ await_transform(co_safe_point_t) noexcept;
+//        };
+//      };
+class co_safe_point_t final {};
+FOLLY_INLINE_VARIABLE constexpr co_safe_point_t co_safe_point{};
+
 } // namespace coro
 } // namespace folly

folly/experimental/coro/Task.h

@@ -78,6 +78,15 @@ class TaskPromiseBase {
  protected:
   TaskPromiseBase() noexcept {}
 
+  template <typename Promise>
+  AwaitableVariant<FinalAwaiter, AwaitableReady<void>> do_safe_point(
+      Promise& promise) noexcept {
+    if (cancelToken_.isCancellationRequested()) {
+      return promise.yield_value(co_cancelled);
+    }
+    return AwaitableReady<void>{};
+  }
+
  public:
   static void* operator new(std::size_t size) {
     return ::folly_coro_async_malloc(size);
@@ -137,6 +146,7 @@ class TaskPromise : public TaskPromiseBase {
       !std::is_rvalue_reference_v<T>,
       "Task<T&&> is not supported. "
       "Consider using Task<T> or Task<std::unique_ptr<T>> instead.");
+  friend class TaskPromiseBase;
 
   using StorageType = detail::lift_lvalue_reference_t<T>;
 
@@ -182,6 +192,12 @@ class TaskPromise : public TaskPromiseBase {
     return final_suspend();
   }
 
+  using TaskPromiseBase::await_transform;
+
+  auto await_transform(co_safe_point_t) noexcept {
+    return do_safe_point(*this);
+  }
+
  private:
   Try<StorageType> result_;
 };
@@ -189,6 +205,8 @@ class TaskPromise : public TaskPromiseBase {
 template <>
 class TaskPromise<void> : public TaskPromiseBase {
  public:
+  friend class TaskPromiseBase;
+
   using StorageType = void;
 
   TaskPromise() noexcept = default;
@@ -218,6 +236,12 @@ class TaskPromise<void> : public TaskPromiseBase {
     return final_suspend();
   }
 
+  using TaskPromiseBase::await_transform;
+
+  auto await_transform(co_safe_point_t) noexcept {
+    return do_safe_point(*this);
+  }
+
  private:
   Try<void> result_;
 };

folly/experimental/coro/Utils.h

@@ -18,6 +18,9 @@
 
 #include <experimental/coroutine>
 #include <type_traits>
+#include <variant>
+
+#include <folly/Utility.h>
 
 namespace folly {
 namespace coro {
@@ -49,5 +52,63 @@ class AwaitableReady<void> {
   void await_suspend(std::experimental::coroutine_handle<>) noexcept {}
   void await_resume() noexcept {}
 };
+
+namespace detail {
+
+struct await_suspend_return_coroutine_fn {
+  template <typename A, typename P>
+  std::experimental::coroutine_handle<> operator()(
+      A& a,
+      std::experimental::coroutine_handle<P> coro) const
+      noexcept(noexcept(a.await_suspend(coro))) {
+    using result = decltype(a.await_suspend(coro));
+    auto noop = std::experimental::noop_coroutine();
+    if constexpr (std::is_same_v<void, result>) {
+      a.await_suspend(coro);
+      return noop;
+    } else if constexpr (std::is_same_v<bool, result>) {
+      return a.await_suspend(coro) ? noop : coro;
+    } else {
+      return a.await_suspend(coro);
+    }
+  }
+};
+inline constexpr await_suspend_return_coroutine_fn
+    await_suspend_return_coroutine{};
+
+template <typename... A>
+class AwaitableVariant : private std::variant<A...> {
+ private:
+  using base = std::variant<A...>;
+
+  template <typename P = void>
+  using handle = std::experimental::coroutine_handle<P>;
+
+  template <typename Visitor>
+  auto visit(Visitor v) {
+    return std::visit(v, static_cast<base&>(*this));
+  }
+
+ public:
+  using base::base; // assume there are no valueless-by-exception instances
+
+  auto await_ready() noexcept(
+      (noexcept(FOLLY_DECLVAL(A&).await_ready()) && ...)) {
+    return visit([&](auto& a) { return a.await_ready(); });
+  }
+  template <typename P>
+  auto await_suspend(handle<P> coro) noexcept(
+      (noexcept(FOLLY_DECLVAL(A&).await_suspend(coro)) && ...)) {
+    auto impl = await_suspend_return_coroutine;
+    return visit([&](auto& a) { return impl(a, coro); });
+  }
+  auto await_resume() noexcept(
+      (noexcept(FOLLY_DECLVAL(A&).await_resume()) && ...)) {
+    return visit([&](auto& a) { return a.await_resume(); });
+  }
+};
+
+} // namespace detail
+
 } // namespace coro
 } // namespace folly

folly/experimental/coro/test/AsyncGeneratorTest.cpp

@@ -599,4 +599,35 @@ TEST(AsyncGeneraor, CoAwaitTry) {
   }());
 }
 
+TEST(AsyncGenerator, SafePoint) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    enum class step_type {
+      init,
+      before_continue_sp,
+      after_continue_sp,
+      before_cancel_sp,
+      after_cancel_sp,
+    };
+    step_type step = step_type::init;
+
+    folly::CancellationSource cancelSrc;
+    auto gen = [&]() -> folly::coro::AsyncGenerator<int> {
+      step = step_type::before_continue_sp;
+      co_await folly::coro::co_safe_point;
+      step = step_type::after_continue_sp;
+
+      cancelSrc.requestCancellation();
+
+      step = step_type::before_cancel_sp;
+      co_await folly::coro::co_safe_point;
+      step = step_type::after_cancel_sp;
+    }();
+
+    auto result = co_await folly::coro::co_awaitTry(
+        folly::coro::co_withCancellation(cancelSrc.getToken(), gen.next()));
+    EXPECT_THROW(result.value(), folly::OperationCancelled);
+    EXPECT_EQ(step_type::before_cancel_sp, step);
+  }());
+}
+
 #endif

folly/experimental/coro/test/TaskTest.cpp

@@ -591,4 +591,35 @@ TEST_F(TaskTest, Moved) {
   }
 }
 
+TEST_F(TaskTest, SafePoint) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    enum class step_type {
+      init,
+      before_continue_sp,
+      after_continue_sp,
+      before_cancel_sp,
+      after_cancel_sp,
+    };
+    step_type step = step_type::init;
+
+    folly::CancellationSource cancelSrc;
+    auto makeTask = [&]() -> folly::coro::Task<void> {
+      step = step_type::before_continue_sp;
+      co_await folly::coro::co_safe_point;
+      step = step_type::after_continue_sp;
+
+      cancelSrc.requestCancellation();
+
+      step = step_type::before_cancel_sp;
+      co_await folly::coro::co_safe_point;
+      step = step_type::after_cancel_sp;
+    };
+
+    auto result = co_await folly::coro::co_awaitTry( //
+        folly::coro::co_withCancellation(cancelSrc.getToken(), makeTask()));
+    EXPECT_THROW(result.value(), folly::OperationCancelled);
+    EXPECT_EQ(step_type::before_cancel_sp, step);
+  }());
+}
+
 #endif