Fix timed_wait to not hold executor keep-alive for too long

Reviewed By: yfeldblum

Differential Revision: D16426555

fbshipit-source-id: 0211c58b6919dc412b9cd7a029f0543d0060b40b

folly/experimental/coro/TimedWait.h

@@ -37,11 +37,11 @@ timed_wait(Awaitable awaitable, Duration duration) {
       detail::lift_lvalue_reference_t<semi_await_result_t<Awaitable>>>>>
       result;
 
-  futures::sleep(duration)
-      .via(co_await co_current_executor)
-      .setCallback_([posted, &baton](auto&&, auto&&) {
+  futures::sleepUnsafe(duration).setCallback_(
+      [posted, &baton, executor = co_await co_current_executor](
+          auto&&, auto&&) {
         if (!posted->exchange(true, std::memory_order_relaxed)) {
-          baton.post();
+          executor->add([&baton] { baton.post(); });
         }
       });
 

folly/experimental/coro/test/CoroTest.cpp

@@ -339,6 +339,18 @@ TEST(Coro, TimedWaitTask) {
   coro::blockingWait(taskTimedWaitTask());
 }
 
+TEST(Coro, TimedWaitKeepAlive) {
+  auto start = std::chrono::steady_clock::now();
+  coro::blockingWait([]() -> coro::Task<void> {
+    co_await coro::timed_wait(
+        futures::sleep(std::chrono::milliseconds{100}),
+        std::chrono::seconds{60});
+    co_return;
+  }());
+  auto duration = std::chrono::steady_clock::now() - start;
+  EXPECT_LE(duration, std::chrono::seconds{30});
+}
+
 template <int value>
 struct AwaitableInt {
   bool await_ready() const {