Add SemiFuture support to folly::whileDo

Summary: whileDo is hard to use with SemiFutures, leading to potential use of inline executors just to make it work. This adds native support.

Reviewed By: yfeldblum

Differential Revision: D16988569

fbshipit-source-id: abaf36180d6ff8a5277eebb277f7f0a254315d36

folly/futures/Future-inl.h

@@ -2446,7 +2446,24 @@ Future<Unit> when(bool p, F&& thunk) {
 }
 
 template <class P, class F>
-Future<Unit> whileDo(P&& predicate, F&& thunk) {
+typename std::
+    enable_if<isSemiFuture<invoke_result_t<F>>::value, SemiFuture<Unit>>::type
+    whileDo(P&& predicate, F&& thunk) {
+  if (predicate()) {
+    auto future = thunk();
+    return std::move(future).deferExValue(
+        [predicate = std::forward<P>(predicate),
+         thunk = std::forward<F>(thunk)](auto&& ex, auto&&) mutable {
+          return whileDo(std::forward<P>(predicate), std::forward<F>(thunk))
+              .via(std::move(ex));
+        });
+  }
+  return makeSemiFuture();
+}
+
+template <class P, class F>
+typename std::enable_if<isFuture<invoke_result_t<F>>::value, Future<Unit>>::type
+whileDo(P&& predicate, F&& thunk) {
   if (predicate()) {
     auto future = thunk();
     return std::move(future).thenValue(

folly/futures/Future.h

@@ -1871,15 +1871,6 @@ class Future : private futures::detail::FutureBase<T> {
   template <class F>
   friend Future<Unit> when(bool p, F&& thunk);
 
-  /// Carry out the computation contained in the given future if
-  /// while the predicate continues to hold.
-  ///
-  /// thunk behaves like std::function<Future<T2>(void)>
-  ///
-  /// predicate behaves like std::function<bool(void)>
-  template <class P, class F>
-  friend Future<Unit> whileDo(P&& predicate, F&& thunk);
-
   template <class FT>
   friend Future<FT> futures::detail::convertFuture(
       SemiFuture<FT>&& sf,
@@ -2488,6 +2479,22 @@ auto unorderedReduce(Collection&& c, T&& initial, F&& func)
   return unorderedReduce(
       c.begin(), c.end(), std::forward<T>(initial), std::forward<F>(func));
 }
+
+/// Carry out the computation contained in the given future if
+/// while the predicate continues to hold.
+///
+/// if thunk behaves like std::function<Future<T2>(void)>
+///    returns Future<Unit>
+/// if thunk behaves like std::function<SemiFuture<T2>(void)>
+///    returns SemiFuture<Unit>
+/// predicate behaves like std::function<bool(void)>
+template <class P, class F>
+typename std::enable_if<isFuture<invoke_result_t<F>>::value, Future<Unit>>::type
+whileDo(P&& predicate, F&& thunk);
+template <class P, class F>
+typename std::
+    enable_if<isSemiFuture<invoke_result_t<F>>::value, SemiFuture<Unit>>::type
+    whileDo(P&& predicate, F&& thunk);
 } // namespace folly
 
 #if FOLLY_HAS_COROUTINES

folly/futures/test/WhileDoTest.cpp

@@ -18,6 +18,7 @@
 #include <mutex>
 #include <queue>
 
+#include <folly/executors/ManualExecutor.h>
 #include <folly/futures/Future.h>
 #include <folly/futures/Promise.h>
 #include <folly/portability/GTest.h>
@@ -58,53 +59,81 @@ inline std::function<bool(void)> makePred(int& i) {
   };
 }
 
-TEST(WhileDo, success) {
-  std::queue<std::shared_ptr<Promise<Unit>>> ps;
-  std::mutex ps_mutex;
+template <class F>
+inline void successTest(
+    std::queue<std::shared_ptr<Promise<Unit>>>& ps,
+    std::mutex& ps_mutex,
+    F& thunk) {
+  folly::ManualExecutor executor;
   int i = 0;
-  int interrupt = 0;
   bool complete = false;
   bool failure = false;
 
   auto pred = makePred(i);
-  auto thunk = makeThunk(ps, interrupt, ps_mutex);
   auto f = folly::whileDo(pred, thunk)
+               .via(&executor)
                .thenValue([&](auto&&) mutable { complete = true; })
                .thenError(folly::tag_t<FutureException>{}, [&](auto&& /* e */) {
                  failure = true;
                });
 
+  executor.drain();
   popAndFulfillPromise(ps, ps_mutex);
   EXPECT_FALSE(complete);
   EXPECT_FALSE(failure);
 
+  executor.drain();
   popAndFulfillPromise(ps, ps_mutex);
   EXPECT_FALSE(complete);
   EXPECT_FALSE(failure);
 
+  executor.drain();
   popAndFulfillPromise(ps, ps_mutex);
+  executor.drain();
   EXPECT_TRUE(f.isReady());
   EXPECT_TRUE(complete);
   EXPECT_FALSE(failure);
 }
 
-TEST(WhileDo, failure) {
+TEST(WhileDo, success) {
+  std::queue<std::shared_ptr<Promise<Unit>>> ps;
+  std::mutex ps_mutex;
+  int interrupt = 0;
+  auto thunk = makeThunk(ps, interrupt, ps_mutex);
+  successTest(ps, ps_mutex, thunk);
+}
+
+TEST(WhileDo, semiFutureSuccess) {
   std::queue<std::shared_ptr<Promise<Unit>>> ps;
   std::mutex ps_mutex;
-  int i = 0;
   int interrupt = 0;
+  auto thunk = [t = makeThunk(ps, interrupt, ps_mutex)]() {
+    return t().semi();
+  };
+  successTest(ps, ps_mutex, thunk);
+}
+
+template <class F>
+inline void failureTest(
+    std::queue<std::shared_ptr<Promise<Unit>>>& ps,
+    std::mutex& ps_mutex,
+    F& thunk) {
+  folly::ManualExecutor executor;
+  int i = 0;
   bool complete = false;
   bool failure = false;
 
   auto pred = makePred(i);
-  auto thunk = makeThunk(ps, interrupt, ps_mutex);
   auto f = folly::whileDo(pred, thunk)
+               .via(&executor)
                .thenValue([&](auto&&) mutable { complete = true; })
                .thenError(folly::tag_t<FutureException>{}, [&](auto&& /* e */) {
                  failure = true;
                });
 
+  executor.drain();
   popAndFulfillPromise(ps, ps_mutex);
+  executor.drain();
   EXPECT_FALSE(complete);
   EXPECT_FALSE(failure);
 
@@ -115,34 +144,77 @@ TEST(WhileDo, failure) {
   FutureException eggs("eggs");
   p2->setException(eggs);
 
+  executor.drain();
   EXPECT_TRUE(f.isReady());
   EXPECT_FALSE(complete);
   EXPECT_TRUE(failure);
 }
 
-TEST(WhileDo, interrupt) {
+TEST(WhileDo, failure) {
+  std::queue<std::shared_ptr<Promise<Unit>>> ps;
+  std::mutex ps_mutex;
+  int interrupt = 0;
+  auto thunk = makeThunk(ps, interrupt, ps_mutex);
+  failureTest(ps, ps_mutex, thunk);
+}
+
+TEST(WhileDo, semiFutureFailure) {
   std::queue<std::shared_ptr<Promise<Unit>>> ps;
   std::mutex ps_mutex;
   int interrupt = 0;
+  auto thunk = [t = makeThunk(ps, interrupt, ps_mutex)]() {
+    return t().semi();
+  };
+  failureTest(ps, ps_mutex, thunk);
+}
+
+template <class F>
+inline void interruptTest(
+    std::queue<std::shared_ptr<Promise<Unit>>>& ps,
+    std::mutex& ps_mutex,
+    int& interrupt,
+    F& thunk) {
+  folly::ManualExecutor executor;
   bool complete = false;
   bool failure = false;
 
   int i = 0;
   auto pred = makePred(i);
-  auto thunk = makeThunk(ps, interrupt, ps_mutex);
   auto f = folly::whileDo(pred, thunk)
+               .via(&executor)
                .thenValue([&](auto&&) mutable { complete = true; })
                .thenError(folly::tag_t<FutureException>{}, [&](auto&& /* e */) {
                  failure = true;
                });
 
+  executor.drain();
   EXPECT_EQ(0, interrupt);
 
   FutureException eggs("eggs");
   f.raise(eggs);
 
+  executor.drain();
   for (int j = 1; j <= 3; ++j) {
     EXPECT_EQ(1, interrupt);
     popAndFulfillPromise(ps, ps_mutex);
+    executor.drain();
   }
 }
+
+TEST(WhileDo, interrupt) {
+  std::queue<std::shared_ptr<Promise<Unit>>> ps;
+  std::mutex ps_mutex;
+  int interrupt = 0;
+  auto thunk = makeThunk(ps, interrupt, ps_mutex);
+  interruptTest(ps, ps_mutex, interrupt, thunk);
+}
+
+TEST(WhileDo, semiFutureInterrupt) {
+  std::queue<std::shared_ptr<Promise<Unit>>> ps;
+  std::mutex ps_mutex;
+  int interrupt = 0;
+  auto thunk = [t = makeThunk(ps, interrupt, ps_mutex)]() {
+    return t().semi();
+  };
+  interruptTest(ps, ps_mutex, interrupt, thunk);
+}