Fix co_awaitTry() with Task::scheduleOn()

Summary:
Fixes compilation of `co_await co_awaitTry(someTask.scheduleOn(ex))` when
awaited within other `Task` coroutines by adding the `await_resume_try()`
method to ViaIfAsyncAwaiter.

Also added some unit-tests for this case.

Reviewed By: andriigrynenko

Differential Revision: D24650827

fbshipit-source-id: 687717d9f73f1fdcfb1f2ff17e4ddca46c16b189

folly/experimental/coro/ViaIfAsync.h

@@ -238,12 +238,22 @@ class ViaIfAsyncAwaiter {
         viaCoroutine_.getWrappedCoroutineWithSavedContext(continuation));
   }
 
-  decltype(auto) await_resume() noexcept(
-      noexcept(std::declval<Awaiter&>().await_resume())) {
+  auto await_resume() noexcept(
+      noexcept(std::declval<Awaiter&>().await_resume()))
+      -> decltype(std::declval<Awaiter&>().await_resume()) {
     viaCoroutine_.destroy();
     return awaiter_.await_resume();
   }
 
+  template <
+      typename Awaiter2 = Awaiter,
+      typename Result = decltype(std::declval<Awaiter2&>().await_resume_try())>
+  Result await_resume_try() noexcept(
+      noexcept(std::declval<Awaiter&>().await_resume_try())) {
+    viaCoroutine_.destroy();
+    return awaiter_.await_resume_try();
+  }
+
   detail::ViaCoroutine viaCoroutine_;
   Awaiter awaiter_;
 };

folly/experimental/coro/test/TaskTest.cpp

@@ -531,4 +531,47 @@ TEST_F(TaskTest, MakeTask) {
     EXPECT_EQ(s2.i, 0);
   }());
 }
+
+TEST_F(TaskTest, CoAwaitTryWithScheduleOn) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    auto t = []() -> folly::coro::Task<int> { co_return 42; }();
+
+    folly::Try<int> result = co_await folly::coro::co_awaitTry(
+        std::move(t).scheduleOn(folly::getGlobalCPUExecutor()));
+    EXPECT_EQ(42, result.value());
+  }());
+}
+
+TEST_F(TaskTest, CoAwaitTryWithScheduleOnAndCancellation) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    folly::CancellationSource cancelSrc;
+
+    auto makeTask = [&]() -> folly::coro::Task<int> {
+      auto ct = co_await folly::coro::co_current_cancellation_token;
+      EXPECT_FALSE(ct.isCancellationRequested());
+      cancelSrc.requestCancellation();
+      EXPECT_TRUE(ct.isCancellationRequested());
+      co_return 42;
+    };
+
+    {
+      folly::Try<int> result = co_await folly::coro::co_withCancellation(
+          cancelSrc.getToken(),
+          folly::coro::co_awaitTry(
+              makeTask().scheduleOn(folly::getGlobalCPUExecutor())));
+      EXPECT_EQ(42, result.value());
+    }
+
+    cancelSrc = {};
+
+    {
+      folly::Try<int> result =
+          co_await folly::coro::co_awaitTry(folly::coro::co_withCancellation(
+              cancelSrc.getToken(),
+              makeTask().scheduleOn(folly::getGlobalCPUExecutor())));
+      EXPECT_EQ(42, result.value());
+    }
+  }());
+}
+
 #endif