Make timed_wait work for semi awaitables

Reviewed By: lewissbaker

Differential Revision: D15363390

fbshipit-source-id: 98c30928773992c5bb7f040e6f5afb8ee6fef713

folly/experimental/coro/Collect-inl.h

@@ -18,6 +18,7 @@
 #include <folly/experimental/coro/Mutex.h>
 #include <folly/experimental/coro/detail/Barrier.h>
 #include <folly/experimental/coro/detail/BarrierTask.h>
+#include <folly/experimental/coro/detail/Helpers.h>
 
 namespace folly {
 namespace coro {
@@ -33,25 +34,6 @@ inline Unit getValueOrUnit(Try<void>&& value) {
   return Unit{};
 }
 
-// Helper class that can be used to annotate Awaitable objects that will
-// guarantee that they will be resumed on the correct executor so that
-// when the object is awaited within a Task<T> it doesn't automatically
-// wrap the Awaitable in something that forces a reschedule onto the
-// executor.
-template <typename Awaitable>
-class UnsafeResumeInlineSemiAwaitable {
- public:
-  explicit UnsafeResumeInlineSemiAwaitable(Awaitable&& awaitable) noexcept
-      : awaitable_(awaitable) {}
-
-  Awaitable&& viaIfAsync(folly::Executor::KeepAlive<>) && noexcept {
-    return static_cast<Awaitable&&>(awaitable_);
-  }
-
- private:
-  Awaitable awaitable_;
-};
-
 template <typename SemiAwaitable, typename Result>
 detail::BarrierTask makeCollectAllTask(
     folly::Executor* executor,

folly/experimental/coro/TimedWait.h

+/*
+ * Copyright 2017-present Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <experimental/coroutine>
+#include <type_traits>
+
+#include <folly/Optional.h>
+#include <folly/experimental/coro/Baton.h>
+#include <folly/experimental/coro/Task.h>
+#include <folly/experimental/coro/Traits.h>
+#include <folly/experimental/coro/detail/Helpers.h>
+#include <folly/futures/Future.h>
+
+namespace folly {
+namespace coro {
+template <typename Awaitable>
+Task<Optional<lift_unit_t<detail::decay_rvalue_reference_t<
+    detail::lift_lvalue_reference_t<semi_await_result_t<Awaitable>>>>>>
+timed_wait(Awaitable awaitable, Duration duration) {
+  auto posted = std::make_shared<std::atomic<bool>>(false);
+  Baton baton;
+  Try<lift_unit_t<detail::decay_rvalue_reference_t<
+      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&&) {
+        if (!posted->exchange(true, std::memory_order_relaxed)) {
+          baton.post();
+        }
+      });
+
+  co_invoke(
+      [awaitable = std::move(
+           awaitable)]() mutable -> Task<semi_await_result_t<Awaitable>> {
+        co_return co_await std::move(awaitable);
+      })
+      .scheduleOn(co_await co_current_executor)
+      .start([posted, &baton, &result](auto&& r) {
+        if (!posted->exchange(true, std::memory_order_relaxed)) {
+          result = std::move(r);
+          baton.post();
+        }
+      });
+
+  co_await detail::UnsafeResumeInlineSemiAwaitable{get_awaiter(baton)};
+
+  if (!result.hasValue() && !result.hasException()) {
+    co_return folly::none;
+  }
+  co_return *result;
+}
+
+} // namespace coro
+} // namespace folly

folly/experimental/coro/Utils.h

@@ -16,14 +16,8 @@
 #pragma once
 
 #include <experimental/coroutine>
-#include <future>
 #include <type_traits>
 
-#include <folly/Optional.h>
-#include <folly/experimental/coro/Traits.h>
-#include <folly/experimental/coro/Wait.h>
-#include <folly/futures/Future.h>
-
 namespace folly {
 namespace coro {
 
@@ -58,103 +52,5 @@ class AwaitableReady<void> {
   void await_suspend(std::experimental::coroutine_handle<>) noexcept {}
   void await_resume() noexcept {}
 };
-
-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 = await_result_t<Awaitable>;
-
-  TimedWaitAwaitable(Awaitable&& awaitable, std::chrono::milliseconds duration)
-      : awaitable_(std::move(awaitable)), duration_(duration) {}
-
-  bool await_ready() {
-    return false;
-  }
-
-  FOLLY_CORO_AWAIT_SUSPEND_NONTRIVIAL_ATTRIBUTES bool await_suspend(
-      std::experimental::coroutine_handle<> ch) {
-    auto sharedState = std::make_shared<SharedState>(ch, storage_);
-    waitAndNotify(std::move(awaitable_), sharedState).detach();
-    futures::sleepUnsafe(duration_).thenValue(
-        [sharedState = std::move(sharedState)](Unit) {
-          sharedState->setTimeout();
-        });
-    return true;
-  }
-
-  Optional<await_resume_return_type> await_resume() {
-    if (!storage_.hasValue() && !storage_.hasException()) {
-      return folly::none;
-    }
-    return std::move(storage_).value();
-  }
-
- private:
-  class SharedState {
-   public:
-    SharedState(
-        std::experimental::coroutine_handle<> ch,
-        Try<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_.hasValue() && !storage_.hasException());
-      tryEmplace(storage_, static_cast<await_resume_return_type&&>(value));
-      ch_();
-    }
-
-    void setException(exception_wrapper e) {
-      if (first_.exchange(true, std::memory_order_relaxed)) {
-        return;
-      }
-      assume(!storage_.hasValue() && !storage_.hasException());
-      storage_.emplaceException(std::move(e));
-      ch_();
-    }
-
-    void setTimeout() {
-      if (first_.exchange(true, std::memory_order_relaxed)) {
-        return;
-      }
-      ch_();
-    }
-
-   private:
-    std::atomic<bool> first_{false};
-    std::experimental::coroutine_handle<> ch_;
-    Try<await_resume_return_type>& storage_;
-  };
-
-  static Wait waitAndNotify(
-      Awaitable awaitable,
-      std::shared_ptr<SharedState> sharedState) {
-    try {
-      sharedState->setValue(co_await std::forward<Awaitable>(awaitable));
-    } catch (const std::exception& e) {
-      sharedState->setException(exception_wrapper(std::current_exception(), e));
-    } catch (...) {
-      sharedState->setException(exception_wrapper(std::current_exception()));
-    }
-  }
-
-  Awaitable awaitable_;
-  std::chrono::milliseconds duration_;
-  Try<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/detail/Helpers.h

+/*
+ * Copyright 2019-present Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/Executor.h>
+
+namespace folly {
+namespace coro {
+namespace detail {
+// Helper class that can be used to annotate Awaitable objects that will
+// guarantee that they will be resumed on the correct executor so that
+// when the object is awaited within a Task<T> it doesn't automatically
+// wrap the Awaitable in something that forces a reschedule onto the
+// executor.
+template <typename Awaitable>
+class UnsafeResumeInlineSemiAwaitable {
+ public:
+  explicit UnsafeResumeInlineSemiAwaitable(Awaitable&& awaitable) noexcept
+      : awaitable_(awaitable) {}
+
+  Awaitable&& viaIfAsync(folly::Executor::KeepAlive<>) && noexcept {
+    return static_cast<Awaitable&&>(awaitable_);
+  }
+
+ private:
+  Awaitable awaitable_;
+};
+} // namespace detail
+} // namespace coro
+} // namespace folly

folly/experimental/coro/test/CoroTest.cpp

@@ -22,6 +22,7 @@
 #include <folly/executors/ManualExecutor.h>
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Task.h>
+#include <folly/experimental/coro/TimedWait.h>
 #include <folly/experimental/coro/Utils.h>
 #include <folly/fibers/Semaphore.h>
 #include <folly/io/async/ScopedEventBaseThread.h>
@@ -252,7 +253,7 @@ TEST(Coro, CurrentExecutor) {
   EXPECT_EQ(42, coro::blockingWait(std::move(task)));
 }
 
-coro::Task<void> taskTimedWait() {
+coro::Task<void> taskTimedWaitFuture() {
   auto ex = co_await coro::co_current_executor;
   auto fastFuture =
       futures::sleep(std::chrono::milliseconds{50}).via(ex).thenValue([](Unit) {
@@ -296,8 +297,46 @@ coro::Task<void> taskTimedWait() {
   co_return;
 }
 
-TEST(Coro, TimedWait) {
-  coro::blockingWait(taskTimedWait());
+TEST(Coro, TimedWaitFuture) {
+  coro::blockingWait(taskTimedWaitFuture());
+}
+
+coro::Task<void> taskTimedWaitTask() {
+  auto fastTask = []() -> coro::Task<int> {
+    co_await futures::sleep(std::chrono::milliseconds{50});
+    co_return 42;
+  }();
+  auto fastResult = co_await coro::timed_wait(
+      std::move(fastTask), std::chrono::milliseconds{100});
+  EXPECT_TRUE(fastResult);
+  EXPECT_EQ(42, *fastResult);
+
+  struct ExpectedException : public std::runtime_error {
+    ExpectedException() : std::runtime_error("ExpectedException") {}
+  };
+
+  auto throwingTask = []() -> coro::Task<void> {
+    co_await futures::sleep(std::chrono::milliseconds{50});
+    throw ExpectedException();
+  }();
+  EXPECT_THROW(
+      (void)co_await coro::timed_wait(
+          std::move(throwingTask), std::chrono::milliseconds{100}),
+      ExpectedException);
+
+  auto slowTask = []() -> coro::Task<int> {
+    co_await futures::sleep(std::chrono::milliseconds{200});
+    co_return 42;
+  }();
+  auto slowResult = co_await coro::timed_wait(
+      std::move(slowTask), std::chrono::milliseconds{100});
+  EXPECT_FALSE(slowResult);
+
+  co_return;
+}
+
+TEST(Coro, TimedWaitTask) {
+  coro::blockingWait(taskTimedWaitTask());
 }
 
 template <int value>