Future version of fibers::Baton::wait

Summary: This makes it possible to add future APIs to primitives built using fibers::Baton.

Reviewed By: A5he

Differential Revision: D14396094

fbshipit-source-id: bf64a31c25ed4dd92b8c484cb4256078efdc6262

folly/futures/Future.cpp

@@ -57,5 +57,33 @@ Future<Unit> sleepUnsafe(Duration dur, Timekeeper* tk) {
   return sleep(dur, tk).toUnsafeFuture();
 }
 
+#if FOLLY_FUTURE_USING_FIBER
+namespace {
+class FutureWaiter : public fibers::Baton::Waiter {
+ public:
+  FutureWaiter(Promise<Unit> promise, std::unique_ptr<fibers::Baton> baton)
+      : promise_(std::move(promise)), baton_(std::move(baton)) {
+    baton_->setWaiter(*this);
+  }
+
+  void post() override {
+    promise_.setValue();
+    delete this;
+  }
+
+ private:
+  Promise<Unit> promise_;
+  std::unique_ptr<fibers::Baton> baton_;
+};
+} // namespace
+
+SemiFuture<Unit> wait(std::unique_ptr<fibers::Baton> baton) {
+  Promise<Unit> promise;
+  auto sf = promise.getSemiFuture();
+  new FutureWaiter(std::move(promise), std::move(baton));
+  return sf;
+}
+#endif
+
 } // namespace futures
 } // namespace folly

folly/futures/helpers.h

@@ -27,6 +27,10 @@
 
 namespace folly {
 
+namespace fibers {
+class Baton;
+}
+
 /// This namespace is for utility functions that would usually be static
 /// members of Future, except they don't make sense there because they don't
 /// depend on the template type (rather, on the type of their arguments in
@@ -141,6 +145,8 @@ auto mapTry(Executor& exec, Collection&& c, F&& func)
   return mapTry(exec, c.begin(), c.end(), std::forward<F>(func));
 }
 
+SemiFuture<Unit> wait(std::unique_ptr<fibers::Baton> baton);
+
 } // namespace futures
 
 /**

folly/futures/test/FutureTest.cpp

@@ -1625,3 +1625,21 @@ TEST(Future, ThenRecursion) {
 
   EXPECT_EQ(42, recursion(&executor, 100000).getVia(&executor));
 }
+
+TEST(Future, BatonWait) {
+  auto baton = std::make_unique<fibers::Baton>();
+  bool posted{false};
+
+  ManualExecutor executor;
+  auto postFuture = futures::sleep(std::chrono::milliseconds{100})
+                        .via(&executor)
+                        .thenValue([&posted, batonPtr = baton.get()](auto) {
+                          posted = true;
+                          batonPtr->post();
+                        });
+  futures::wait(std::move(baton))
+      .via(&executor)
+      .thenValue([&](auto) { EXPECT_TRUE(posted); })
+      .getVia(&executor);
+  EXPECT_TRUE(postFuture.isReady());
+}