Fix folly::coro::timeout to work with AsyncGenerator

Summary: The current signature results in an attempt to construct a `Try<semi_await_result_t<SemiAwaitable>>` from a `semi_await_try_result_t<SemiAwaitable>::element_type`. While the inner types match for Task, they don't for AsyncGenerator - the former is `AsyncGenerator<T>::NextResult`, which is not constructible from a T. (We could alternatively add that constructor, but that would result in timeout returning a `Task<AsyncGenerator<T>::NextResult>`, which is undesirably indirect).

Reviewed By: yfeldblum

Differential Revision: D28475605

fbshipit-source-id: c30a60e3ec942f801653a08b020796d47d4ef763

folly/experimental/coro/Timeout-inl.h

@@ -23,7 +23,7 @@
 namespace folly::coro {
 
 template <typename SemiAwaitable, typename Duration>
-Task<semi_await_result_t<SemiAwaitable>> timeout(
+Task<typename semi_await_try_result_t<SemiAwaitable>::element_type> timeout(
     SemiAwaitable semiAwaitable, Duration timeoutDuration, Timekeeper* tk) {
   CancellationSource cancelSource;
   folly::coro::Baton baton;

folly/experimental/coro/Timeout.h

@@ -43,7 +43,7 @@ namespace folly::coro {
 // If a timekeeper is provided then uses that timekeeper to start the timer,
 // otherwise uses the process' default TimeKeeper if 'tk' is null.
 template <typename SemiAwaitable, typename Duration>
-Task<semi_await_result_t<SemiAwaitable>> timeout(
+Task<typename semi_await_try_result_t<SemiAwaitable>::element_type> timeout(
     SemiAwaitable semiAwaitable,
     Duration timeoutDuration,
     Timekeeper* tk = nullptr);

folly/experimental/coro/ViaIfAsync.h

@@ -602,6 +602,12 @@ detail::TryAwaitable<remove_cvref_t<Awaitable>> co_awaitTry(
       static_cast<Awaitable&&>(awaitable)};
 }
 
+template <typename T>
+using semi_await_try_result_t =
+    await_result_t<decltype(folly::coro::co_viaIfAsync(
+        std::declval<folly::Executor::KeepAlive<>>(),
+        folly::coro::co_awaitTry(std::declval<T>())))>;
+
 } // namespace coro
 } // namespace folly
 

folly/experimental/coro/test/TimeoutTest.cpp

@@ -16,6 +16,7 @@
 
 #include <folly/Portability.h>
 
+#include <folly/experimental/coro/AsyncGenerator.h>
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
 #include <folly/experimental/coro/Sleep.h>
@@ -163,4 +164,22 @@ TEST(Timeout, CancelParent) {
   }());
 }
 
+TEST(Timeout, AsyncGenerator) {
+  coro::blockingWait([]() -> coro::Task<> {
+    // Completing synchronously with a value.
+    auto result = co_await coro::timeout(
+        []() -> coro::AsyncGenerator<int> { co_yield 42; }().next(), 1s);
+    EXPECT_EQ(42, result);
+
+    // Test that it handles failing synchronously
+    auto tryResult = co_await coro::co_awaitTry(coro::timeout(
+        []() -> coro::AsyncGenerator<int> {
+          co_yield coro::co_error(std::runtime_error{"bad value"});
+        }()
+                    .next(),
+        1s));
+    EXPECT_TRUE(tryResult.hasException<std::runtime_error>());
+  }());
+}
+
 #endif // FOLLY_HAS_COROUTINES