Implement operator timed_wait

Summary: Currently the implementation just relies on futures::sleep. Ideally it should be able to use Executor's timer if available.

Reviewed By: yfeldblum

Differential Revision: D8404592

fbshipit-source-id: e952a0cc3d99bc528ea14ff5c68835d63a71e685

folly/experimental/coro/Utils.h

@@ -18,6 +18,10 @@
 #include <experimental/coroutine>
 #include <future>
 
+#include <folly/Optional.h>
+#include <folly/experimental/coro/Wait.h>
+#include <folly/futures/Future.h>
+
 namespace folly {
 namespace coro {
 
@@ -56,5 +60,82 @@ struct yield {
   void await_resume() {}
 };
 
+template <typename Awaitable>
+class TimedWaitAwaitable {
+ public:
+  static_assert(
+      std::is_same<Awaitable, std::decay_t<Awaitable>>::value,
+      "Awaitable should be decayed.");
+  using await_resume_return_type =
+      decltype((operator co_await(std::declval<Awaitable>())).await_resume());
+
+  TimedWaitAwaitable(Awaitable&& awaitable, std::chrono::milliseconds duration)
+      : awaitable_(std::move(awaitable)), duration_(duration) {}
+
+  bool await_ready() {
+    return false;
+  }
+
+  bool await_suspend(std::experimental::coroutine_handle<> ch) {
+    auto sharedState = std::make_shared<SharedState>(ch, storage_);
+    waitAndNotify(std::move(awaitable_), sharedState).detach();
+    futures::sleep(duration_).then(
+        [sharedState = std::move(sharedState)] { sharedState->setTimeout(); });
+    return true;
+  }
+
+  Optional<await_resume_return_type> await_resume() {
+    return std::move(storage_);
+  }
+
+ private:
+  class SharedState {
+   public:
+    SharedState(
+        std::experimental::coroutine_handle<> ch,
+        Optional<await_resume_return_type>& storage)
+        : ch_(std::move(ch)), storage_(storage) {}
+
+    void setValue(await_resume_return_type&& value) {
+      if (first_.exchange(true, std::memory_order_relaxed)) {
+        return;
+      }
+      assume(!storage_);
+      storage_ = std::move(value);
+      ch_();
+    }
+
+    void setTimeout() {
+      if (first_.exchange(true, std::memory_order_relaxed)) {
+        return;
+      }
+      ch_();
+    }
+
+   private:
+    std::atomic<bool> first_{false};
+    std::experimental::coroutine_handle<> ch_;
+    Optional<await_resume_return_type>& storage_;
+  };
+
+  static Wait waitAndNotify(
+      Awaitable awaitable,
+      std::shared_ptr<SharedState> sharedState) {
+    sharedState->setValue(co_await awaitable);
+  }
+
+  Awaitable awaitable_;
+  std::chrono::milliseconds duration_;
+  Optional<await_resume_return_type> storage_;
+};
+
+template <typename Awaitable>
+TimedWaitAwaitable<std::decay_t<Awaitable>> timed_wait(
+    Awaitable&& awaitable,
+    std::chrono::milliseconds duration) {
+  return TimedWaitAwaitable<std::decay_t<Awaitable>>(
+      std::forward<Awaitable>(awaitable), duration);
+}
+
 } // namespace coro
 } // namespace folly

folly/experimental/coro/Wait.h

@@ -53,6 +53,10 @@ class Wait {
 
   Wait(Wait&&) = default;
 
+  void detach() {
+    future_ = {};
+  }
+
   ~Wait() {
     if (future_.valid()) {
       future_.get();

folly/experimental/coro/test/CoroTest.cpp

@@ -177,4 +177,26 @@ TEST(Coro, CurrentExecutor) {
   EXPECT_EQ(42, future.get());
 }
 
+coro::Task<int> taskTimedWait() {
+  auto fastFuture =
+      futures::sleep(std::chrono::milliseconds{50}).then([] { return 42; });
+  auto fastResult = co_await coro::timed_wait(
+      std::move(fastFuture), std::chrono::milliseconds{100});
+  EXPECT_TRUE(fastResult);
+  EXPECT_EQ(42, *fastResult);
+
+  auto slowFuture =
+      futures::sleep(std::chrono::milliseconds{200}).then([] { return 42; });
+  auto slowResult = co_await coro::timed_wait(
+      std::move(slowFuture), std::chrono::milliseconds{100});
+  EXPECT_FALSE(slowResult);
+
+  co_return 42;
+}
+
+TEST(Coro, TimedWait) {
+  ManualExecutor executor;
+  EXPECT_EQ(42, via(&executor, taskTimedWait()).toFuture().getVia(&executor));
+}
+
 #endif

folly/futures/Future.h

@@ -1944,13 +1944,19 @@ class FutureRefAwaitable {
   folly::Future<T>& future_;
 };
 } // namespace detail
-} // namespace folly
 
 template <typename T>
-folly::detail::FutureAwaitable<T>
-/* implicit */ operator co_await(folly::Future<T>& future) {
-  return folly::detail::FutureRefAwaitable<T>(future);
+detail::FutureRefAwaitable<T>
+/* implicit */ operator co_await(Future<T>& future) {
+  return detail::FutureRefAwaitable<T>(future);
 }
+
+template <typename T>
+detail::FutureRefAwaitable<T>
+/* implicit */ operator co_await(Future<T>&& future) {
+  return detail::FutureRefAwaitable<T>(future);
+}
+} // namespace folly
 #endif
 
 #include <folly/futures/Future-inl.h>