Add some missing customization points to NowTaskWithExecutor

Summary:
- Explicitly test that you can't await a moved `NowTaskWithExecutor`. The manual test showed the correct "deleted copy ctor" error for this.
- Plumb through `co_withCancellation` so that things like `timeout(nowTaskFn().scheduleOn(ex))` would work.
- Plumb through `unsafeMoveMustAwaitImmediately()` so that generic wrappers (like `Sequence.h` and `Noexcept.h`) can wrap `NowTask`s into other `MustAwaitImmediately` containers.

Reviewed By: ispeters

Differential Revision: D70355306

fbshipit-source-id: 7adae485aa377181947b7d38de1c4da67947cf29

Author: snarkmaster

folly/coro/safe/NowTask.h

@@ -48,14 +48,37 @@ template <typename T>
 class FOLLY_NODISCARD NowTaskWithExecutor final : private MustAwaitImmediately {
  protected:
   template <safe_alias, typename>
-  friend class BackgroundTask; // see `CanUnwrapNowTask` below
+  friend class BackgroundTask; // for `unwrapTaskWithExecutor`, remove later
   TaskWithExecutor<T> unwrapTaskWithExecutor() && { return std::move(inner_); }
 
   template <typename>
   friend class NowTask; // can construct
 
   explicit NowTaskWithExecutor(TaskWithExecutor<T> t) : inner_(std::move(t)) {}
 
+ public:
+  // Required for `await_result_t` to work.
+  // NB: Even though this is rvalue-qualified, this does not wrongly allow
+  // `co_await std::move(myNowTask().scheduleOn(ex));`, see `withExecutor` test
+  auto operator co_await() && noexcept {
+    return std::move(inner_).operator co_await();
+  }
+
+  NowTaskWithExecutor unsafeMoveMustAwaitImmediately() && {
+    return NowTaskWithExecutor{std::move(inner_)};
+  }
+
+  friend NowTaskWithExecutor co_withCancellation(
+      folly::CancellationToken cancelToken, NowTaskWithExecutor te) noexcept {
+    return NowTaskWithExecutor{
+        co_withCancellation(std::move(cancelToken), std::move(te.inner_))};
+  }
+
+  friend auto co_viaIfAsync(
+      Executor::KeepAlive<> executor, NowTaskWithExecutor te) noexcept {
+    return co_viaIfAsync(std::move(executor), std::move(te.inner_));
+  }
+
  private:
   TaskWithExecutor<T> inner_;
 };
@@ -91,7 +114,6 @@ class FOLLY_CORO_TASK_ATTRS NowTask final
   using promise_type = detail::NowTaskPromise<T>;
 
   // If `makeNowTask().scheduleOn()` is movable, it defeats our purpose.
-  FOLLY_NODISCARD
   NowTaskWithExecutor<T> scheduleOn(Executor::KeepAlive<> exec) && noexcept {
     return NowTaskWithExecutor<T>{
         std::move(*this).unwrap().scheduleOn(std::move(exec))};
@@ -121,13 +143,11 @@ class FOLLY_CORO_TASK_ATTRS NowTask final
   }
 
  protected:
-  // These 3 `friend`s (+ 1 above) are for `unwrap()`.  If this list grows,
-  // introduce a `CanUnwrapNowTask` passkey type.
   template <typename U>
-  friend auto toNowTask(NowTask<U> t);
+  friend auto toNowTask(NowTask<U> t); // for `unwrap`
   // `async_now_closure` wraps `NowTask`s into `NowTask`s
   template <auto>
-  friend auto detail::bind_captures_to_closure(auto&&, auto);
+  friend auto detail::bind_captures_to_closure(auto&&, auto); // for `unwrap`
 };
 
 // NB: `toNowTask(SafeTask)` is in `SafeTask.h` to avoid circular deps.

folly/coro/safe/test/NowTaskTest.cpp

@@ -22,6 +22,9 @@
 
 namespace folly::coro {
 
+static_assert(std::is_void_v<await_result_t<NowTaskWithExecutor<void>>>);
+static_assert(std::is_same_v<int, await_result_t<NowTaskWithExecutor<int>>>);
+
 Task<int> demoTask(int x) {
   co_return 1300 + x;
 }
@@ -83,6 +86,28 @@ CO_TEST(NowTaskTest, simple) {
   static_assert(!test_transform_moved_v<NowTask<int>&>);
 }
 
+template <typename T>
+inline constexpr bool test_transform_moved_with_executor_v = std::is_same_v<
+    detected_t<await_transform_result_t, T>,
+    StackAwareViaIfAsyncAwaitable<TaskWithExecutor<int>>>;
+
+CO_TEST(NowTaskTest, withExecutor) {
+  auto exec = co_await co_current_executor;
+  EXPECT_EQ(1337, co_await demoNowTask(37).scheduleOn(exec));
+
+  static_assert(test_transform_moved_with_executor_v<TaskWithExecutor<int>>);
+  static_assert(test_transform_moved_with_executor_v<TaskWithExecutor<int>&&>);
+  static_assert(test_transform_moved_with_executor_v<NowTaskWithExecutor<int>>);
+  static_assert(
+      !test_transform_moved_with_executor_v<NowTaskWithExecutor<int>&&>);
+#if 0 // The above asserts are a proxy for this manual test
+  auto te = demoNowTask(37).scheduleOn(exec);
+  co_await std::move(te);
+#endif
+  static_assert(
+      !test_transform_moved_with_executor_v<NowTaskWithExecutor<int>&>);
+}
+
 // `co_nothrow` isn't a function object, so we can't wrap it & pass prvalues
 template <typename T>
 using co_nothrow_result_t = decltype(co_nothrow(FOLLY_DECLVAL(T)));