co_await_result -- a result<T> analog of co_awaitTry

Summary: See docblock. The goal is to lower the barrier for adopting `result` in new code.

Reviewed By: ispeters

Differential Revision: D71752310

fbshipit-source-id: bf588066387760eecb65f075b2fec3bcc5864ed8

Author: snarkmaster

folly/coro/AwaitResult.h

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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 <folly/coro/ViaIfAsync.h>
+#include <folly/result/try.h>
+
+#if FOLLY_HAS_RESULT
+
+namespace folly::coro {
+
+/// `co_await_result` is the `result<T>` analog of the older `co_awaitTry`.
+///
+/// In a `folly::coro` async coroutine, use `co_await_result` like so:
+///
+///   result<int> res = co_await co_await_result(taskReturningInt());
+///   if (auto* ex = get_exception<MyError>(res)) {
+///     /* handle ex */
+///   } else {
+///     sum += co_await co_ready(res); // efficiently propagate unhandled error
+///   }
+///
+/// Contrast that with related async coro vocabulary:
+///  - `co_yield co_result(r)` from `Result.h` -- propagate `result<T>` or
+///    `Try<T>` to the awaiter of the current coro.
+///  - `auto& v = co_await co_ready(r)` from `Ready.h` -- given a `result<T>`,
+///    unpack the value, or propagate any error to our awaiter.
+///
+/// The purpose of `co_await_result` is to handle errors from a child task via
+/// `result<T>`, rather than through `try {} catch {}`.  Some reasons to do so:
+///   - Your error-handling APIs (logging, retry, etc) use `result<T>`.
+///   - You wish to avoid the ~microsecond cost of thrown exceptions,
+///     applicable only when your error path is hot, and the child uses
+///     `co_yield` instead of `throw` to propagate exceptions.
+
+namespace detail {
+
+// This reuses the `await_resume_try` machinery in the hope that the compiler
+// will be able to optimize away the `Try` -> `result` conversion.  In some
+// cases, this may incur an extra move-copy, so if this ends up being hot,
+// adding a specialized `await_resume_result` interface may be justified.
+template <typename Awaitable>
+class ResultAwaiter {
+ private:
+  static_assert(is_awaitable_try<Awaitable&&>);
+
+  using Awaiter = awaiter_type_t<Awaitable>;
+  Awaiter awaiter_;
+
+ public:
+  explicit ResultAwaiter(Awaitable&& awaiter)
+      : awaiter_(get_awaiter(static_cast<Awaitable&&>(awaiter))) {}
+
+  // clang-format off
+  auto await_ready() FOLLY_DETAIL_FORWARD_BODY(awaiter_.await_ready())
+
+  template <typename Promise>
+  auto await_suspend(coroutine_handle<Promise> coro)
+      FOLLY_DETAIL_FORWARD_BODY(awaiter_.await_suspend(coro))
+
+  auto await_resume()
+      FOLLY_DETAIL_FORWARD_BODY(try_to_result(awaiter_.await_resume_try()))
+  // clang-format on
+};
+
+template <typename T>
+class [[FOLLY_ATTR_CLANG_CORO_AWAIT_ELIDABLE]] ResultAwaitable
+    : public CommutativeWrapperAwaitable<ResultAwaitable, T> {
+ public:
+  using CommutativeWrapperAwaitable<ResultAwaitable, T>::
+      CommutativeWrapperAwaitable;
+
+  template <
+      typename Self,
+      std::enable_if_t<
+          std::is_same_v<remove_cvref_t<Self>, ResultAwaitable>,
+          int> = 0,
+      typename T2 = like_t<Self, T>,
+      std::enable_if_t<is_awaitable_v<T2>, int> = 0>
+  friend ResultAwaiter<T2> operator co_await(Self && self) {
+    return ResultAwaiter<T2>{static_cast<Self&&>(self).inner_};
+  }
+
+  using folly_private_noexcept_awaitable_t = std::true_type;
+};
+
+} // namespace detail
+
+// IMPORTANT: If you need an `Awaitable&&` overload, you must bifurcate this
+// API on `must_await_immediately_v`, see `co_awaitTry` for an example.
+template <typename Awaitable>
+detail::ResultAwaitable<Awaitable> co_await_result(
+    [[FOLLY_ATTR_CLANG_CORO_AWAIT_ELIDABLE_ARGUMENT]] Awaitable awaitable) {
+  return detail::ResultAwaitable<Awaitable>{
+      mustAwaitImmediatelyUnsafeMover(std::move(awaitable))()};
+}
+
+} // namespace folly::coro
+
+#endif

folly/coro/Result.h

@@ -22,6 +22,7 @@
 #include <folly/ExceptionWrapper.h>
 #include <folly/OperationCancelled.h>
 #include <folly/Try.h>
+#include <folly/result/result.h>
 
 namespace folly {
 namespace coro {
@@ -56,6 +57,14 @@ class co_result final {
     assert(!result_.hasException() || result_.exception());
   }
 
+#if FOLLY_HAS_RESULT
+  // Covered in `AwaitResultTest.cpp`, unlike the rest of this file, which is
+  // covered in `TaskTest.cpp`.
+  explicit co_result(folly::result<T>&& result) noexcept(
+      std::is_nothrow_move_constructible<T>::value)
+      : co_result(result_to_try(std::move(result))) {}
+#endif
+
   const Try<T>& result() const { return result_; }
 
   Try<T>& result() { return result_; }

folly/coro/test/AwaitResultTest.cpp

@@ -0,0 +1,123 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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/coro/AwaitResult.h>
+#include <folly/coro/GtestHelpers.h>
+#include <folly/coro/Result.h>
+#include <folly/coro/safe/NowTask.h>
+
+/// Besides `co_await_result`, this test also covers the `folly::result<T>`
+/// integration with `coro::co_result` from `coro/Result.h`.
+
+#if FOLLY_HAS_RESULT
+
+namespace folly::coro {
+
+CO_TEST(AwaitResultTest, co_await_result_of_error) {
+  auto voidErrorTask = []() -> NowTask<void> {
+    co_yield co_error(std::runtime_error("foo"));
+  };
+  { // Capture the error
+    auto res = co_await co_await_result(voidErrorTask());
+    EXPECT_STREQ("foo", get_exception<std::runtime_error>(res)->what());
+  }
+  { // Also test `coro::co_result` integration
+    auto res = co_await co_await_result([&]() -> NowTask<void> {
+      co_yield co_result(co_await co_await_result(voidErrorTask()));
+    }());
+    EXPECT_STREQ("foo", get_exception<std::runtime_error>(res)->what());
+  }
+}
+
+CO_TEST(AwaitResultTest, co_await_result_of_value) {
+  // Return a move-only thing to make sure we don't copy
+  auto valueTask = []() -> NowTask<std::unique_ptr<int>> {
+    co_return std::make_unique<int>(1337);
+  };
+  { // Capture the value
+    auto res = co_await co_await_result(valueTask());
+    // Real code should use `co_await co_ready(res)` to unpack!
+    EXPECT_EQ(1337, *res.value_or_throw());
+  }
+  { // Also test `coro::co_result` integration
+    auto res = co_await co_await_result([&]() -> NowTask<std::unique_ptr<int>> {
+      co_yield co_result(co_await co_await_result(valueTask()));
+    }());
+    EXPECT_EQ(1337, *res.value_or_throw());
+  }
+}
+
+CO_TEST(AwaitResultTest, co_await_result_of_void) {
+  int numAwaited = 0;
+  auto voidTask = [&]() -> NowTask<void> {
+    ++numAwaited;
+    co_return;
+  };
+  { // Capturing a "value" completion
+    auto res = co_await co_await_result(voidTask());
+    static_assert(std::is_same_v<decltype(res), result<>>);
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(1, numAwaited);
+  }
+  { // Also test `coro::co_result` integration
+    auto res = co_await co_await_result([&]() -> NowTask<void> {
+      co_yield co_result(co_await co_await_result(voidTask()));
+    }());
+    static_assert(std::is_same_v<decltype(res), result<>>);
+    EXPECT_TRUE(res.has_value());
+    EXPECT_EQ(2, numAwaited);
+  }
+}
+
+CO_TEST(AwaitResultTest, co_await_result_stopped) {
+  auto stoppedTask = [&]() -> NowTask<void> {
+    co_yield co_cancelled;
+    LOG(FATAL) << "not reached";
+  };
+  { // Capturing a "stopped" completion
+    auto res = co_await co_await_result(stoppedTask());
+    EXPECT_TRUE(res.has_stopped());
+  }
+  { // Also test `coro::co_result` integration
+    auto res = co_await co_await_result([&]() -> NowTask<void> {
+      co_yield co_result(co_await co_await_result(stoppedTask()));
+      LOG(FATAL) << "not reached";
+    }());
+    EXPECT_TRUE(res.has_stopped());
+  }
+}
+
+struct ThrowingMove {
+  ThrowingMove(const ThrowingMove&) = default;
+  ThrowingMove& operator=(const ThrowingMove&) = default;
+  ThrowingMove(ThrowingMove&&) {}
+  ThrowingMove& operator=(ThrowingMove&&) { return *this; }
+};
+struct NothrowMove {};
+
+template <typename T>
+using TestAwaiter = detail::ResultAwaiter<TaskWithExecutor<T>>;
+
+// While it's unclear if anything cares about `await_resume()` `noexcept`
+// discipline, it's easy enough to check.
+static_assert(noexcept(FOLLY_DECLVAL(TestAwaiter<int>).await_resume()));
+static_assert(noexcept(FOLLY_DECLVAL(TestAwaiter<NothrowMove>).await_resume()));
+static_assert(
+    !noexcept(FOLLY_DECLVAL(TestAwaiter<ThrowingMove>).await_resume()));
+
+} // namespace folly::coro
+
+#endif

folly/coro/test/NoexceptTest.cpp

@@ -260,11 +260,6 @@ CO_TEST(NoexceptTest, AsNoexceptNowTask) {
   co_await checkAsNoexcept<NowTask<void>>();
 }
 
-// XXX test AsNoexcept more -- generally need to cover most of the compositional
-// properties above :(
-//   - cover WITH and WITHOUT executor
-//   - crib more from `checkFatalOnThrow` and `NowTaskIsImmediate` above
-
 } // namespace folly::coro
 
 #endif