Fold most NowTask logic into TaskWrapper + fixes

Summary:
This improves `TaskWrapper` and makes `NowTask` nearly trivial. Up-stack, this also removes a bunch of duplicative protocol-forwarding code from `SafeTask`.

Main refinements:
  - Much better doc coverage for `TaskWrapper`.
  - Add object slicing detection on methods that explicitly unwrap.
  - Add `promise_type` to `OpaqueTaskWrapperCrtp` since it seems desirable for it to always be a coroutine type.
  - Use configuration structs to make the setup of the CRTP bases cleaner and more flexible.
  - Provide `TaskWithExecutorWrapperCrtp` and a `TaskWithExecutorT` config so that `TaskWrapperCrtp` can reference it.
  - Stop forwarding unsafe APIs like `scheduleOn` and `semi`.
  - Rename some identifiers for brevity / readability. In particular, what used to be `protected` `unwrap()` became `unwrapTask()` and `unwrapTaskWithExecutor()` to avoid conflicting with the awaitable wrapper from `ViaIfAsync.h` and with the other `unwrap()` from `Task.h`.
  - Expand test coverage.

Reviewed By: ispeters

Differential Revision: D70805317

fbshipit-source-id: 19c7ee55ca1302b4094fe6ce1eff15cc2b56dc3e

folly/coro/TaskWrapper.h

@@ -19,32 +19,77 @@
 #include <folly/coro/AwaitImmediately.h>
 #include <folly/coro/Task.h>
 
-/// The header provides base classes for wrapping `folly::coro::Task` with
-/// custom functionality.  These work by composition, which avoids the
-/// pitfalls of inheritance -- your custom wrapper will not be "is-a-Task",
-/// and will not implicitly "object slice" to a `Task`.
+/// `TaskWrapper.h` provides base classes for wrapping `folly::coro::Task` with
+/// custom functionality.  These work by composition, which avoids the pitfalls
+/// of inheritance -- your custom wrapper will not be "is-a-Task", and will not
+/// implicitly "object slice" to a `Task`.
 ///
-/// Keep in mind that some destructive APIs, like `.semi()`, effectively
-/// unwrap the `Task`.  If this is important for your use-case, you may need
-/// to add features (e.g.  `TaskWithExecutorWrapper`, on-unwrap callbacks).
+/// The point of this header is to uniformly forward the large API surface of
+/// `Task`, `TaskWithExecutor`, and `TaskPromise`, leaving just the "new logic"
+/// in each wrapper's implementation.
 ///
-/// The point of this header is to uniformly forward the large API surface
-/// of `TaskPromise` & `Task`, leaving just the "new logic" in each wrapper.
-/// As `Task.h` evolves, a central `TaskWrapper.h` is easier to maintain.
+///   - `OpaqueTaskWrapperCrtp` makes your type a coroutine (`promise_type`)
+///     that can `co_await` other `folly::coro` objects.  However, your type
+///     itself is NOT made (semi)awaitable.  This can be useful if you want to
+///     restrict how your type can be awaited.
 ///
-/// You'll derive from `TaskWrapperPromise` -- which must reference a
-/// derived class of `TaskWrapperCrtp` that is your new user-facing coro.
+///   - `TaskWrapperCrtp` is a coroutine, and is semi-awaitable by other
+///     `folly::coro` coroutines. It has the following features:
+///        * `co_await`ability (using `co_viaIfAsync`)
+///        * Interoperates with `folly::coro` awaitable wrappers like
+///          `co_awaitTry` and `co_nothrow`.
+///        * `co_withCancellation` to add a cancellation token
+///        * `co_withExecutor` to add a cancellation token
+///        * Basic reflection via `folly/coro/Traits.h`
+///        * Empty base optimization for zero runtime overhead
 ///
-/// Please use `FOLLY_CORO_TASK_ATTRS` on your `Task` wrapper class.
-/// Caveat: this assumes that the coro's caller will outlive it.  That is
-/// true for `Task`, and almost certainly true of all sensible wrapper
-/// types.  If you have an analog of `TaskWithExecutor`, also mark that
-/// `FOLLY_NODISCARD`.
+///   - `TaskWithExecutorWrapperCrtp` is awaitable, but not a coroutine.  It
+///     has the same features, except for `co_withExecutor`.
 ///
-/// To discourage inheritance and object-slicing bugs, mark your derived
-/// wrappers `final` -- they can be wrapped recursively.
+/// ### WARNING: Do not blindly forward more APIs in `TaskWrapper.h`!
 ///
-/// Read `TaskWrapperTest.cpp` for examples of a minimal & recursive wrapper.
+/// Several existing wrappers are immediately-awaitable (`AwaitImmediately.h`).
+/// For those tasks (e.g. `NowTask`), API forwarding is risky:
+///   - Do NOT forward `semi()`, `scheduleOn()`, `start*()`, `unwrap()`, or
+///     other destructive methods, or CPOs that take the awaitable
+///     by-reference.  All of those make it trivial to accidentally break the
+///     immediately-awaitable invariant, and cause lifetime bugs.
+///   - When forwarding an API, use either a static method or CPO.  Then,
+///     either ONLY take the awaitable by-value, or bifurcate the API on
+///     `must_await_immediately_v<Awaitable>`, grep for examples.
+///
+/// If you **have** to forward an unsafe API, here are some suggestions:
+///   - Only add them in your wrapper.
+///   - Add them via `UnsafeTaskWrapperCrtp` deriving from `TaskWrapperCrtp`.
+///   - Add boolean flags to the configuration struct, and gate the methods via
+///     `enable_if`.  NB: You probably cannot gate these on `Derived` **not**
+///     being `MustAwaitImmediately`, since CRTP bases see an incomplete type.
+///
+/// ### WARNING: Beware of object slicing in "unwrapping" APIs
+///
+/// Start by reading "A note on object slicing" in `AwaitImmediately.h`.
+///
+/// If your wrapper is adding new members, or customizing object lifecycle
+/// (dtor / copy / move / assignment), then you must:
+///   - Write a custom `getUnsafeMover()`.
+///   - Overload the protected `unsafeTask()` and `unsafeTaskWithExecutor()` to
+///     reduce slicing risk.
+///   - Take care not to slice down to the `Crtp` bases.
+///
+/// ### How to implement a wrapper
+///
+/// First, read the WARNINGs above. Then, follow one of the "Tiny" examples
+/// in `TaskWrapperTest.cpp`. The important things are:
+///   - Actually read the "object slicing" warning above!
+///   - In most cases, you'll need to both implement a task, and customize its
+///     `TaskWithExecutorT`.  If you leave that as `coro::TaskWithExecutor`,
+///     some users will accidentally avoid your wrapper's effects.
+///   - Tag `YourTaskWithExecutor` with `FOLLY_NODISCARD`.
+///   - Tag `YourTask` with the `FOLLY_CORO_TASK_ATTRS` attribute.  Caveat:
+///     This assumes that the coro's caller will outlive it.  That is true for
+///     `Task`, and almost certainly true of all sensible wrapper types.
+///   - Mark your wrappers `final` to discourage inheritance and object-slicing
+///     bugs.  They can still be wrapped recursively.
 ///
 /// Future: Once this has a benchmark, see if `FOLLY_ALWAYS_INLINE` makes
 /// any difference on the wrapped functions (it shouldn't).
@@ -53,45 +98,36 @@
 
 namespace folly::coro {
 
-namespace detail {
-template <typename, typename, typename>
-class TaskPromiseWrapperBase;
-} // namespace detail
-
-template <typename, typename, typename>
-class TaskWrapperCrtp;
-
 namespace detail {
 
 template <typename Wrapper>
-using task_wrapper_underlying_semiawaitable_t =
-    typename Wrapper::TaskWrapperUnderlyingSemiAwaitable;
+using task_wrapper_inner_semiawaitable_t =
+    typename Wrapper::folly_private_task_wrapper_inner_t;
 
 template <typename SemiAwaitable, typename T>
 inline constexpr bool is_task_or_wrapper_v =
     (!std::is_same_v<nonesuch, SemiAwaitable> && // Does not wrap Task
      (std::is_same_v<SemiAwaitable, Task<T>> || // Wraps Task
       is_task_or_wrapper_v<
-          detected_t<task_wrapper_underlying_semiawaitable_t, SemiAwaitable>,
+          detected_t<task_wrapper_inner_semiawaitable_t, SemiAwaitable>,
           T>));
 
 template <typename Wrapper>
-using task_wrapper_underlying_promise_t =
-    typename Wrapper::TaskWrapperUnderlyingPromise;
+using task_wrapper_inner_promise_t = typename Wrapper::TaskWrapperInnerPromise;
 
 template <typename Promise, typename T>
 inline constexpr bool is_task_promise_or_wrapper_v =
     (!std::is_same_v<nonesuch, Promise> && // Does not wrap TaskPromise
      (std::is_same_v<Promise, TaskPromise<T>> || // Wraps TaskPromise
       is_task_promise_or_wrapper_v<
-          detected_t<task_wrapper_underlying_promise_t, Promise>,
+          detected_t<task_wrapper_inner_promise_t, Promise>,
           T>));
 
-template <typename T, typename WrappedSemiAwaitable, typename Promise>
+template <typename T, typename WrapperTask, typename Promise>
 class TaskPromiseWrapperBase {
  protected:
   static_assert(
-      is_task_or_wrapper_v<WrappedSemiAwaitable, T>,
+      is_task_or_wrapper_v<WrapperTask, T>,
       "SemiAwaitable must be a sequence of wrappers ending in Task<T>");
   static_assert(
       is_task_promise_or_wrapper_v<Promise, T>,
@@ -103,10 +139,10 @@ class TaskPromiseWrapperBase {
   ~TaskPromiseWrapperBase() = default;
 
  public:
-  using TaskWrapperUnderlyingPromise = Promise;
+  using TaskWrapperInnerPromise = Promise;
 
-  WrappedSemiAwaitable get_return_object() noexcept {
-    return WrappedSemiAwaitable{promise_.get_return_object()};
+  WrapperTask get_return_object() noexcept {
+    return WrapperTask{promise_.get_return_object()};
   }
 
   static void* operator new(std::size_t size) {
@@ -155,23 +191,23 @@ class TaskPromiseWrapperBase {
   }
 };
 
-template <typename T, typename SemiAwaitable, typename Promise>
+template <typename T, typename WrapperTask, typename Promise>
 class TaskPromiseWrapper
-    : public TaskPromiseWrapperBase<T, SemiAwaitable, Promise> {
+    : public TaskPromiseWrapperBase<T, WrapperTask, Promise> {
  protected:
   TaskPromiseWrapper() noexcept = default;
   ~TaskPromiseWrapper() = default;
 
  public:
-  template <typename U = T> // see `returnImplicitCtor` test
+  template <typename U = T> // see "`co_return` with implicit ctor" test
   auto return_value(U&& value) {
     return this->promise_.return_value(std::forward<U>(value));
   }
 };
 
-template <typename SemiAwaitable, typename Promise>
-class TaskPromiseWrapper<void, SemiAwaitable, Promise>
-    : public TaskPromiseWrapperBase<void, SemiAwaitable, Promise> {
+template <typename WrapperTask, typename Promise>
+class TaskPromiseWrapper<void, WrapperTask, Promise>
+    : public TaskPromiseWrapperBase<void, WrapperTask, Promise> {
  protected:
   TaskPromiseWrapper() noexcept = default;
   ~TaskPromiseWrapper() = default;
@@ -182,57 +218,155 @@ class TaskPromiseWrapper<void, SemiAwaitable, Promise>
 
 } // namespace detail
 
-// Inherit from this instead of `TaskWrapperCrtp` if you don't want your
-// wrapped task to be awaitable without being `unwrap()`ped first.
-template <typename Derived, typename T, typename SemiAwaitable>
+// Inherit from `OpaqueTaskWrapperCrtp` instead of `TaskWrapperCrtp` if you
+// don't want your wrapped task to become a regular semi-awaitable.
+template <typename Derived, typename Cfg>
 class OpaqueTaskWrapperCrtp {
+ public:
+  using promise_type = typename Cfg::PromiseT;
+
+  using folly_private_must_await_immediately_t =
+      must_await_immediately_t<typename Cfg::InnerTaskT>;
+  using folly_private_task_wrapper_inner_t = typename Cfg::InnerTaskT;
+
+  // Do NOT add any protocols here, see `TaskWrapperCrtp` instead.
+
  private:
+  using Inner = folly_private_task_wrapper_inner_t;
   static_assert(
-      detail::is_task_or_wrapper_v<SemiAwaitable, T>,
-      "*TaskWrapperCrtp must wrap a sequence of wrappers ending in Task<T>");
+      detail::is_task_or_wrapper_v<Inner, typename Cfg::ValueT>,
+      "*TaskWrapper must wrap a sequence of wrappers ending in Task<T>");
 
-  SemiAwaitable task_;
+  Inner task_;
 
  protected:
-  template <typename, typename, typename>
+  template <typename, typename, typename> // can construct
   friend class ::folly::coro::detail::TaskPromiseWrapperBase;
+  friend class MustAwaitImmediatelyUnsafeMover< // can construct
+      Derived,
+      detail::unsafe_mover_for_must_await_immediately_t<Inner>>;
+
+  explicit OpaqueTaskWrapperCrtp(Inner t) : task_(std::move(t)) {
+    static_assert(
+        must_await_immediately_v<Derived> ||
+            !must_await_immediately_v<typename Cfg::TaskWithExecutorT>,
+        "`TaskWithExecutorT` must `AddMustAwaitImmediately` because the inner "
+        "task did");
+  }
 
-  explicit OpaqueTaskWrapperCrtp(SemiAwaitable t) : task_(std::move(t)) {}
+  // See "A note on object slicing" above `mustAwaitImmediatelyUnsafeMover`
+  Inner unwrapTask() && {
+    static_assert(sizeof(Inner) == sizeof(Derived));
+    return std::move(task_);
+  }
+};
 
-  SemiAwaitable unwrap() && { return std::move(task_); }
+// IMPORTANT: Read "Do not blindly forward more APIs" in the file docblock.  In
+// a nutshell, adding methods, or by-ref CPOs, can compromise the safety of
+// immediately-awaitable wrappers, so DON'T DO THAT.
+template <typename Derived, typename Cfg>
+class TaskWrapperCrtp : public OpaqueTaskWrapperCrtp<Derived, Cfg> {
+  using OpaqueTaskWrapperCrtp<Derived, Cfg>::OpaqueTaskWrapperCrtp;
 
  public:
-  using TaskWrapperUnderlyingSemiAwaitable = SemiAwaitable;
+  // For `NowTask` & `SafeTask` API-compatibility, DO NOT add `scheduleOn()`.
+  // Use `co_withExecutor(ex, task())` instead of `task().scheduleOn(ex)`,
+  //
+  // Pass `tw` by-value, since `&&` would break immediately-awaitable types
+  friend typename Cfg::TaskWithExecutorT co_withExecutor(
+      Executor::KeepAlive<> executor, Derived tw) noexcept {
+    return typename Cfg::TaskWithExecutorT{
+        co_withExecutor(std::move(executor), std::move(tw).unwrapTask())};
+  }
+
+  // Pass `tw` by-value, since `&&` would break immediately-awaitable types
+  friend Derived co_withCancellation(
+      CancellationToken cancelToken, Derived tw) noexcept {
+    return Derived{co_withCancellation(
+        std::move(cancelToken), std::move(tw).unwrapTask())};
+  }
+
+  // Pass `tw` by-value, since `&&` would break immediately-awaitable types
+  friend auto co_viaIfAsync(
+      Executor::KeepAlive<> executor, Derived tw) noexcept {
+    return co_viaIfAsync(std::move(executor), std::move(tw).unwrapTask());
+  }
+
+  auto getUnsafeMover(ForMustAwaitImmediately p) && {
+    // See "A note on object slicing" above `mustAwaitImmediatelyUnsafeMover`
+    static_assert(sizeof(Derived) == sizeof(typename Cfg::InnerTaskT));
+    return MustAwaitImmediatelyUnsafeMover{
+        (Derived*)nullptr, std::move(*this).unwrapTask().getUnsafeMover(p)};
+  }
+
+  using folly_private_task_wrapper_crtp_base = TaskWrapperCrtp;
 };
 
-template <typename Derived, typename T, typename SemiAwaitable>
-class TaskWrapperCrtp
-    : public OpaqueTaskWrapperCrtp<Derived, T, SemiAwaitable> {
+// IMPORTANT: Read "Do not blindly forward more APIs" in the file docblock.  In
+// a nutshell, adding methods, or by-ref CPOs, can compromise the safety of
+// immediately-awaitable wrappers, so DON'T DO THAT.
+template <typename Derived, typename Cfg>
+class TaskWithExecutorWrapperCrtp {
+ private:
+  using Inner = typename Cfg::InnerTaskWithExecutorT;
+  Inner inner_;
+
  protected:
-  template <typename, typename, typename>
-  friend class ::folly::coro::detail::TaskPromiseWrapperBase;
+  friend class MustAwaitImmediatelyUnsafeMover< // can construct
+      Derived,
+      detail::unsafe_mover_for_must_await_immediately_t<Inner>>;
+
+  // See "A note on object slicing" above `mustAwaitImmediatelyUnsafeMover`
+  Inner unwrapTaskWithExecutor() && {
+    static_assert(sizeof(Inner) == sizeof(Derived));
+    return std::move(inner_);
+  }
+
+  // Our task can construct us, and that logic lives in the CRTP base
+  friend typename Cfg::WrapperTaskT::folly_private_task_wrapper_crtp_base;
 
-  using OpaqueTaskWrapperCrtp<Derived, T, SemiAwaitable>::OpaqueTaskWrapperCrtp;
+  explicit TaskWithExecutorWrapperCrtp(Inner t) : inner_(std::move(t)) {}
 
  public:
-  // NB: In the future, this might ALSO produce a wrapped object.
-  FOLLY_NODISCARD
-  TaskWithExecutor<T> scheduleOn(Executor::KeepAlive<> executor) && noexcept {
-    return std::move(*this).unwrap().scheduleOn(std::move(executor));
+  // Required for `await_result_t` and `get_awaiter` to work.  That, of course,
+  // lets the power-user violate the "immediately awaitable" invariant, but
+  // they are essential for metaprogramming.  Let's hope that users won't
+  // carelessly invoke `operator co_await()` or `get_awaiter`!
+  //
+  // An alternate solution for **just** `await_result_t` would have been to
+  // specialize `struct await_result` to look for a special member type on
+  // these.  We can revisit if this proves a frequent footgun.
+  //
+  // NB: This does not let this **naively** wrong code compile -- that goes
+  // through `await_transform()`.  See `NowTaskTest.cpp` for proof.
+  //   auto t = co_withExecutor(ex, someNowTask());
+  //   co_await std::move(t);
+  auto operator co_await() && noexcept {
+    return std::move(inner_).operator co_await();
   }
 
-  FOLLY_NOINLINE auto semi() && { return std::move(*this).unwrap().semi(); }
-
+  // Pass `te` by-value, since `&&` would break immediately-awaitable types
   friend Derived co_withCancellation(
-      folly::CancellationToken cancelToken, Derived&& tw) noexcept {
+      CancellationToken cancelToken, Derived te) noexcept {
     return Derived{
-        co_withCancellation(std::move(cancelToken), std::move(tw).unwrap())};
+        co_withCancellation(std::move(cancelToken), std::move(te.inner_))};
   }
 
+  // Pass `te` by-value, since `&&` would break immediately-awaitable types
   friend auto co_viaIfAsync(
-      folly::Executor::KeepAlive<> executor, Derived&& tw) noexcept {
-    return co_viaIfAsync(std::move(executor), std::move(tw).unwrap());
+      Executor::KeepAlive<> executor, Derived te) noexcept {
+    return co_viaIfAsync(std::move(executor), std::move(te.inner_));
+  }
+
+  auto getUnsafeMover(ForMustAwaitImmediately p) && {
+    // See "A note on object slicing" above `mustAwaitImmediatelyUnsafeMover`
+    static_assert(sizeof(Derived) == sizeof(Inner));
+    return MustAwaitImmediatelyUnsafeMover{
+        (Derived*)nullptr, std::move(inner_).getUnsafeMover(p)};
   }
+
+  using folly_private_must_await_immediately_t =
+      must_await_immediately_t<Inner>;
 };
 
 } // namespace folly::coro