From 6c0b1ecfe84d791dec2156cbb7641afaed46b5b6 Mon Sep 17 00:00:00 2001
From: Alexey Spiridonov <lesha@meta.com>
Date: Wed, 16 Oct 2024 16:08:53 -0700
Subject: [PATCH] Add `co_awaitTry` support to `Baton`

Summary:
Since `Baton` doesn't currently support `co_awaitTry`, I found myself writing this sort of thing in my "coro/safe" implementation:

```
Try<ResultT> res;
if constexpr (std::is_same_v<
    decltype(inner), ::folly::coro::Baton::WaitOperation>) {
  co_await std::move(t);
} else {
  res = co_await co_awaitTry(std::move(t));
}
```

After this diff, the above could be abbreviated to:

```
auto res = co_await co_awaitTry(std::move(t));
```

Since awaiting `Baton` doesn't throw, I see no downside to supporting `co_awaitTry` here, besides the minor header bloat.

Reviewed By: andriigrynenko

Differential Revision: D64274928

fbshipit-source-id: 201f8b0765f19a769bd80b34ff4c86792c5b6df4
---
 folly/coro/BUCK               | 1 +
 folly/coro/Baton.h            | 5 +++++
 folly/coro/test/BatonTest.cpp | 3 ++-
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/folly/coro/BUCK b/folly/coro/BUCK
index 4b61e366abe..8944ca4f58e 100644
--- a/folly/coro/BUCK
+++ b/folly/coro/BUCK
@@ -96,6 +96,7 @@ cpp_library(
         "//folly/synchronization:atomic_util",
     ],
     exported_deps = [
+        "//folly:try",
         "//folly/experimental/coro:coroutine",
     ],
 )
diff --git a/folly/coro/Baton.h b/folly/coro/Baton.h
index 43b67e7a6f8..557e135c98f 100644
--- a/folly/coro/Baton.h
+++ b/folly/coro/Baton.h
@@ -18,6 +18,7 @@
 
 #include <atomic>
 
+#include <folly/Try.h>
 #include <folly/experimental/coro/Coroutine.h>
 
 #if FOLLY_HAS_COROUTINES
@@ -107,6 +108,10 @@ class Baton {
 
     void await_resume() noexcept {}
 
+    // Awaiting a baton doesn't throw, so supporting `co_awaitTry` here only
+    // serves to simplify generic code.
+    folly::Try<void> await_resume_try() noexcept { return {}; }
+
    protected:
     friend class Baton;
 
diff --git a/folly/coro/test/BatonTest.cpp b/folly/coro/test/BatonTest.cpp
index 0292357e8e4..3cc9995496d 100644
--- a/folly/coro/test/BatonTest.cpp
+++ b/folly/coro/test/BatonTest.cpp
@@ -90,7 +90,8 @@ TEST_F(BatonTest, MultiAwaitBaton) {
 
   auto makeTask2 = [&]() -> coro::Task<void> {
     reachedBeforeAwait2 = true;
-    co_await baton;
+    // Equivalent to `co_await baton`, we just want it to compile.
+    co_await co_awaitTry(baton.operator co_await());
     reachedAfterAwait2 = true;
   };