Add support for co_awaitTry(task.scheduleOn())

Summary:
This adds support for applying the co_awaitTry() algorithm to a
TaskWithExecutor type.

This required a bit of a refactor of the `co_awaitTry()` implementtion
to not require that it's passed a SemiAwaitable but and to allow
Awaitable objects to also be passed to `co_awaitTry()`.

Renamed TaskWithExecutor::InlineAwaiter to InlineTryAwaitable and make
it private to better reflect its purpose.

No longer implement `Task::Awaiter::await_resume()` in terms of the
`await_resume_try()` method. This should avoid a call to the Try<T>
move-constructor by extracting the T result directly from the Try<T>
stored in the promise.

Reviewed By: andriigrynenko

Differential Revision: D24602729

fbshipit-source-id: d352e717877769502cc55a01731eb165ddad9ad8

folly/experimental/coro/Task.h

@@ -288,7 +288,7 @@ class FOLLY_NODISCARD TaskWithExecutor {
     [](TaskWithExecutor task,
        std::decay_t<F> cb) -> detail::InlineTaskDetached {
       try {
-        cb(co_await std::move(task).co_awaitTry());
+        cb(co_await folly::coro::co_awaitTry(std::move(task)));
       } catch (const std::exception& e) {
         cb(Try<StorageType>(exception_wrapper(std::current_exception(), e)));
       } catch (...) {
@@ -311,7 +311,7 @@ class FOLLY_NODISCARD TaskWithExecutor {
     [](TaskWithExecutor task,
        std::decay_t<F> cb) -> detail::InlineTaskDetached {
       try {
-        cb(co_await InlineAwaiter{std::exchange(task.coro_, {})});
+        cb(co_await InlineTryAwaitable{std::exchange(task.coro_, {})});
       } catch (const std::exception& e) {
         cb(Try<StorageType>(exception_wrapper(std::current_exception(), e)));
       } catch (...) {
@@ -336,7 +336,7 @@ class FOLLY_NODISCARD TaskWithExecutor {
     return sf;
   }
 
-  template <typename ResultCreator>
+ private:
   class Awaiter {
    public:
     explicit Awaiter(handle_t coro) noexcept : coro_(coro) {}
@@ -372,22 +372,26 @@ class FOLLY_NODISCARD TaskWithExecutor {
           });
     }
 
-    decltype(auto) await_resume() {
+    T await_resume() {
       // Eagerly destroy the coroutine-frame once we have retrieved the result.
       SCOPE_EXIT { std::exchange(coro_, {}).destroy(); };
-      ResultCreator resultCreator;
-      return resultCreator(std::move(coro_.promise().result()));
+      return std::move(coro_.promise().result()).value();
+    }
+
+    folly::Try<StorageType> await_resume_try() {
+      SCOPE_EXIT { std::exchange(coro_, {}).destroy(); };
+      return std::move(coro_.promise().result());
     }
 
    private:
     handle_t coro_;
   };
 
-  class InlineAwaiter {
+  class InlineTryAwaitable {
    public:
-    InlineAwaiter(handle_t coro) noexcept : coro_(coro) {}
+    InlineTryAwaitable(handle_t coro) noexcept : coro_(coro) {}
 
-    ~InlineAwaiter() {
+    ~InlineTryAwaitable() {
       if (coro_) {
         coro_.destroy();
       }
@@ -414,22 +418,9 @@ class FOLLY_NODISCARD TaskWithExecutor {
     handle_t coro_;
   };
 
-  struct ValueCreator {
-    T operator()(Try<StorageType>&& t) const { return std::move(t).value(); }
-  };
-
-  struct TryCreator {
-    Try<StorageType> operator()(Try<StorageType>&& t) const {
-      return std::move(t);
-    }
-  };
-
-  auto operator co_await() && noexcept {
-    return Awaiter<ValueCreator>{std::exchange(coro_, {})};
-  }
-
-  auto co_awaitTry() && noexcept {
-    return Awaiter<TryCreator>{std::exchange(coro_, {})};
+ public:
+  Awaiter operator co_await() && noexcept {
+    return Awaiter{std::exchange(coro_, {})};
   }
 
   friend TaskWithExecutor co_withCancellation(
@@ -566,9 +557,12 @@ class FOLLY_NODISCARD Task {
       return coro_;
     }
 
-    T await_resume() { return await_resume_try().value(); }
+    T await_resume() {
+      SCOPE_EXIT { std::exchange(coro_, {}).destroy(); };
+      return std::move(coro_.promise().result()).value();
+    }
 
-    auto await_resume_try() {
+    folly::Try<StorageType> await_resume_try() {
       SCOPE_EXIT { std::exchange(coro_, {}).destroy(); };
       return std::move(coro_.promise().result());
     }

folly/experimental/coro/ViaIfAsync.h

@@ -397,62 +397,106 @@ using semi_await_result_t = await_result_t<decltype(folly::coro::co_viaIfAsync(
 
 namespace detail {
 
-template <typename Awaiter>
+template <typename Awaitable>
 class TryAwaiter {
+  using Awaiter = awaiter_type_t<Awaitable>;
+
  public:
-  TryAwaiter(Awaiter&& awaiter) : awaiter_(std::move(awaiter)) {}
+  explicit TryAwaiter(Awaitable&& awaiter)
+      : awaiter_(folly::coro::get_awaiter(static_cast<Awaitable&&>(awaiter))) {}
 
-  bool await_ready() { return awaiter_.await_ready(); }
+  auto await_ready() noexcept(noexcept(std::declval<Awaiter&>().await_ready()))
+      -> decltype(std::declval<Awaiter&>().await_ready()) {
+    return awaiter_.await_ready();
+  }
 
   template <typename Promise>
-  auto await_suspend(std::experimental::coroutine_handle<Promise> coro) {
+  auto
+  await_suspend(std::experimental::coroutine_handle<Promise> coro) noexcept(
+      noexcept(std::declval<Awaiter&>().await_suspend(coro)))
+      -> decltype(std::declval<Awaiter&>().await_suspend(coro)) {
     return awaiter_.await_suspend(coro);
   }
 
-  auto await_resume() { return awaiter_.await_resume_try(); }
+  auto await_resume() noexcept(
+      noexcept(std::declval<Awaiter&>().await_resume_try()))
+      -> decltype(std::declval<Awaiter&>().await_resume_try()) {
+    return awaiter_.await_resume_try();
+  }
 
  private:
   Awaiter awaiter_;
 };
 
-template <typename Awaiter>
-auto makeTryAwaiter(Awaiter&& awaiter) {
-  return TryAwaiter<std::decay_t<Awaiter>>(std::move(awaiter));
-}
-
-template <typename SemiAwaitable>
-class TrySemiAwaitable {
+template <typename T>
+class TryAwaitable {
  public:
-  explicit TrySemiAwaitable(SemiAwaitable&& semiAwaitable)
-      : semiAwaitable_(std::move(semiAwaitable)) {}
-
-  friend auto co_viaIfAsync(
-      Executor::KeepAlive<> executor,
-      TrySemiAwaitable&& self) noexcept {
-    return makeTryAwaiter(get_awaiter(
-        co_viaIfAsync(std::move(executor), std::move(self.semiAwaitable_))));
+  template <typename T2>
+  explicit TryAwaitable(T2&& awaitable) noexcept(
+      std::is_nothrow_constructible_v<T, T2>)
+      : inner_(static_cast<T2&&>(awaitable)) {}
+
+  template <typename Factory>
+  explicit TryAwaitable(std::in_place_t, Factory&& factory)
+      : inner_(factory()) {}
+
+  template <
+      typename Self,
+      std::enable_if_t<
+          std::is_same_v<remove_cvref_t<Self>, TryAwaitable>,
+          int> = 0,
+      typename T2 = like_t<Self, T>,
+      std::enable_if_t<is_awaitable_v<T2>, int> = 0>
+  friend TryAwaiter<T2> operator co_await(Self&& self) {
+    return TryAwaiter<T2>{static_cast<Self&&>(self).inner_};
+  }
+
+  template <
+      typename T2 = T,
+      typename Result = decltype(folly::coro::co_withCancellation(
+          std::declval<const folly::CancellationToken&>(),
+          std::declval<T2>()))>
+  friend TryAwaitable<Result> co_withCancellation(
+      const folly::CancellationToken& cancelToken,
+      TryAwaitable&& awaitable) {
+    return TryAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                  return folly::coro::co_withCancellation(
+                                      cancelToken,
+                                      static_cast<T&&>(awaitable.inner_));
+                                }};
   }
 
-  friend auto co_withCancellation(
-      const CancellationToken& cancelToken,
-      TrySemiAwaitable&& awaitable) {
-    auto cancelAwaitable = folly::coro::co_withCancellation(
-        std::move(cancelToken),
-        static_cast<SemiAwaitable&&>(awaitable.semiAwaitable_));
-    return TrySemiAwaitable<decltype(cancelAwaitable)>(
-        std::move(cancelAwaitable));
+  template <
+      typename T2 = T,
+      typename Result = decltype(folly::coro::co_viaIfAsync(
+          std::declval<folly::Executor::KeepAlive<>>(),
+          std::declval<T2>()))>
+  friend TryAwaitable<Result> co_viaIfAsync(
+      folly::Executor::KeepAlive<> executor,
+      TryAwaitable&&
+          awaitable) noexcept(noexcept(folly::coro::
+                                           co_viaIfAsync(
+                                               std::declval<folly::Executor::
+                                                                KeepAlive<>>(),
+                                               std::declval<T2>()))) {
+    return TryAwaitable<Result>{std::in_place, [&]() -> decltype(auto) {
+                                  return folly::coro::co_viaIfAsync(
+                                      std::move(executor),
+                                      static_cast<T&&>(awaitable.inner_));
+                                }};
   }
 
  private:
-  SemiAwaitable semiAwaitable_;
+  T inner_;
 };
+
 } // namespace detail
 
-template <
-    typename SemiAwaitable,
-    typename = std::enable_if_t<is_semi_awaitable_v<SemiAwaitable>>>
-auto co_awaitTry(SemiAwaitable&& semiAwaitable) {
-  return detail::TrySemiAwaitable<SemiAwaitable>(std::move(semiAwaitable));
+template <typename Awaitable>
+detail::TryAwaitable<remove_cvref_t<Awaitable>> co_awaitTry(
+    Awaitable&& awaitable) {
+  return detail::TryAwaitable<remove_cvref_t<Awaitable>>{
+      static_cast<Awaitable&&>(awaitable)};
 }
 
 } // namespace coro