From cadfe2cdced5c15f1a3e836dba01e931d09d2763 Mon Sep 17 00:00:00 2001
From: Lee Howes <lwh@fb.com>
Date: Wed, 27 Dec 2017 13:10:51 -0800
Subject: [PATCH] Add getVia and getTryVia to SemiFuture.

Summary: Add getVia and getTryVia to SemiFuture to allow driving chains of work conveniently in the current thread.

Reviewed By: yfeldblum

Differential Revision: D6631898

fbshipit-source-id: 324ef342a44d4ef502188b3cffde17103f0e6cb2
---
 folly/futures/Future-inl.h            | 48 ++++++++++++++++++++++++---
 folly/futures/Future.h                | 29 +++++++++++++---
 folly/futures/test/SemiFutureTest.cpp | 11 ++++++
 3 files changed, 78 insertions(+), 10 deletions(-)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index c7c66b58..25f8f464 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -788,11 +788,6 @@ Future<T>::onError(F&& func) {
   return f;
 }
 
-template <class T>
-Try<T>& Future<T>::getTryVia(DrivableExecutor* e) {
-  return waitVia(e).getTry();
-}
-
 template <class Func>
 auto via(Executor* x, Func&& func)
     -> Future<typename isFuture<decltype(std::declval<Func>()())>::Inner> {
@@ -1392,6 +1387,21 @@ void waitViaImpl(Future<T>& f, DrivableExecutor* e) {
   assert(f.isReady());
 }
 
+template <class T>
+void waitViaImpl(SemiFuture<T>& f, DrivableExecutor* e) {
+  // Set callback so to ensure that the via executor has something on it
+  // so that once the preceding future triggers this callback, drive will
+  // always have a callback to satisfy it
+  if (f.isReady()) {
+    return;
+  }
+  f = std::move(f).via(e).then([](T&& t) { return std::move(t); });
+  while (!f.isReady()) {
+    e->drive();
+  }
+  assert(f.isReady());
+}
+
 } // namespace detail
 } // namespace futures
 
@@ -1419,6 +1429,18 @@ SemiFuture<T>&& SemiFuture<T>::wait(Duration dur) && {
   return std::move(*this);
 }
 
+template <class T>
+SemiFuture<T>& SemiFuture<T>::waitVia(DrivableExecutor* e) & {
+  futures::detail::waitViaImpl(*this, e);
+  return *this;
+}
+
+template <class T>
+SemiFuture<T>&& SemiFuture<T>::waitVia(DrivableExecutor* e) && {
+  futures::detail::waitViaImpl(*this, e);
+  return std::move(*this);
+}
+
 template <class T>
 T SemiFuture<T>::get() && {
   return std::move(wait().value());
@@ -1440,6 +1462,17 @@ Try<T> SemiFuture<T>::getTry() && {
   return std::move(this->core_->getTry());
 }
 
+template <class T>
+T SemiFuture<T>::getVia(DrivableExecutor* e) && {
+  return std::move(waitVia(e).value());
+}
+
+template <class T>
+Try<T> SemiFuture<T>::getTryVia(DrivableExecutor* e) && {
+  waitVia(e);
+  return std::move(this->core_->getTry());
+}
+
 template <class T>
 Future<T>& Future<T>::wait() & {
   futures::detail::waitImpl(*this);
@@ -1503,6 +1536,11 @@ T Future<T>::getVia(DrivableExecutor* e) {
   return std::move(waitVia(e).value());
 }
 
+template <class T>
+Try<T>& Future<T>::getTryVia(DrivableExecutor* e) {
+  return waitVia(e).getTry();
+}
+
 namespace futures {
 namespace detail {
 template <class T>
diff --git a/folly/futures/Future.h b/folly/futures/Future.h
index 20a56b7e..4818e1ed 100644
--- a/folly/futures/Future.h
+++ b/folly/futures/Future.h
@@ -256,6 +256,16 @@ class SemiFuture : private futures::detail::FutureBase<T> {
   /// Try of the value (moved out).
   Try<T> getTry() &&;
 
+  /// Call e->drive() repeatedly until the future is fulfilled. Examples
+  /// of DrivableExecutor include EventBase and ManualExecutor. Returns the
+  /// value (moved out), or throws the exception.
+  T getVia(DrivableExecutor* e) &&;
+
+  /// Call e->drive() repeatedly until the future is fulfilled. Examples
+  /// of DrivableExecutor include EventBase and ManualExecutor. Returns the
+  /// Try of the value (moved out).
+  Try<T> getTryVia(DrivableExecutor* e) &&;
+
   /// Block until this Future is complete. Returns a reference to this Future.
   SemiFuture<T>& wait() &;
 
@@ -269,6 +279,15 @@ class SemiFuture : private futures::detail::FutureBase<T> {
   /// Overload of wait(Duration) for rvalue Futures
   SemiFuture<T>&& wait(Duration) &&;
 
+  /// Call e->drive() repeatedly until the future is fulfilled. Examples
+  /// of DrivableExecutor include EventBase and ManualExecutor. Returns a
+  /// reference to this SemiFuture so that you can chain calls if desired.
+  /// value (moved out), or throws the exception.
+  SemiFuture<T>& waitVia(DrivableExecutor* e) &;
+
+  /// Overload of waitVia() for rvalue Futures
+  SemiFuture<T>&& waitVia(DrivableExecutor* e) &&;
+
   /// Returns an inactive Future which will call back on the other side of
   /// executor (when it is activated).
   ///
@@ -409,16 +428,16 @@ class Future : private futures::detail::FutureBase<T> {
   // movable
   Future& operator=(Future&&) noexcept;
 
-  /// Call e->drive() repeatedly until the future is fulfilled. Examples
-  /// of DrivableExecutor include EventBase and ManualExecutor. Returns a
-  /// reference to the Try of the value.
-  Try<T>& getTryVia(DrivableExecutor* e);
-
   /// Call e->drive() repeatedly until the future is fulfilled. Examples
   /// of DrivableExecutor include EventBase and ManualExecutor. Returns the
   /// value (moved out), or throws the exception.
   T getVia(DrivableExecutor* e);
 
+  /// Call e->drive() repeatedly until the future is fulfilled. Examples
+  /// of DrivableExecutor include EventBase and ManualExecutor. Returns a
+  /// reference to the Try of the value.
+  Try<T>& getTryVia(DrivableExecutor* e);
+
   /// Unwraps the case of a Future<Future<T>> instance, and returns a simple
   /// Future<T> instance.
   template <class F = T>
diff --git a/folly/futures/test/SemiFutureTest.cpp b/folly/futures/test/SemiFutureTest.cpp
index ce24b4e7..1c44fe8d 100644
--- a/folly/futures/test/SemiFutureTest.cpp
+++ b/folly/futures/test/SemiFutureTest.cpp
@@ -275,6 +275,17 @@ TEST(SemiFuture, SimpleDefer) {
   ASSERT_EQ(innerResult, 17);
 }
 
+TEST(SemiFuture, DeferWithGetVia) {
+  std::atomic<int> innerResult{0};
+  EventBase e2;
+  Promise<folly::Unit> p;
+  auto f = p.getFuture();
+  auto sf = std::move(f).semi().defer([&]() { innerResult = 17; });
+  p.setValue();
+  std::move(sf).getVia(&e2);
+  ASSERT_EQ(innerResult, 17);
+}
+
 TEST(SemiFuture, DeferWithVia) {
   std::atomic<int> innerResult{0};
   EventBase e2;
-- 
2.34.1