make wait() and friends return reference to self instead of new Future

Summary:
we still make the new Future, but assign it to ourselves. this avoids the following buggy pattern that people might expect to work

```
auto f = ...
f.wait();

// Careful. f.value() was moved out into the new Future, so you may have lost something
someOperationOn(f.value());

// Nope. We already set a callback internally in wait()
f.then(...);
```

Test Plan: unit

Reviewed By: davejwatson@fb.com

Subscribers: exa, yfeldblum, trunkagent, fbcode-common-diffs@, sammerat, cold-storage-diffs@, folly-diffs@, jsedgwick, aflock

FB internal diff: D1809040

Tasks: 6048284

Signature: t1:1809040:1422900812:1b416408eb5eaa71e88778c9c22ed8bfba087efe

folly/futures/Future-inl.h

@@ -688,27 +688,28 @@ Future<T> Future<T>::delayed(Duration dur, Timekeeper* tk) {
     });
 }
 
+namespace detail {
+
 template <class T>
-Future<T> Future<T>::wait() {
+void waitImpl(Future<T>& f) {
   Baton<> baton;
-  auto done = then([&](Try<T> t) {
+  f = f.then([&](Try<T> t) {
     baton.post();
     return makeFuture(std::move(t));
   });
   baton.wait();
-  while (!done.isReady()) {
-    // There's a race here between the return here and the actual finishing of
-    // the future. f is completed, but the setup may not have finished on done
-    // after the baton has posted.
+  // There's a race here between the return here and the actual finishing of
+  // the future. f is completed, but the setup may not have finished on done
+  // after the baton has posted.
+  while (!f.isReady()) {
     std::this_thread::yield();
   }
-  return done;
 }
 
 template <class T>
-Future<T> Future<T>::wait(Duration dur) {
+void waitImpl(Future<T>& f, Duration dur) {
   auto baton = std::make_shared<Baton<>>();
-  auto done = then([baton](Try<T> t) {
+  f = f.then([baton](Try<T> t) {
     baton->post();
     return makeFuture(std::move(t));
   });
@@ -716,26 +717,54 @@ Future<T> Future<T>::wait(Duration dur) {
   // true), then the returned Future is complete when it is returned to the
   // caller. We need to wait out the race for that Future to complete.
   if (baton->timed_wait(std::chrono::system_clock::now() + dur)) {
-    while (!done.isReady()) {
+    while (!f.isReady()) {
       std::this_thread::yield();
     }
   }
-  return done;
 }
 
 template <class T>
-Future<T>& Future<T>::waitVia(DrivableExecutor* e) & {
-  while (!isReady()) {
+void waitViaImpl(Future<T>& f, DrivableExecutor* e) {
+  while (!f.isReady()) {
     e->drive();
   }
+}
+
+} // detail
+
+template <class T>
+Future<T>& Future<T>::wait() & {
+  detail::waitImpl(*this);
   return *this;
 }
 
 template <class T>
-Future<T> Future<T>::waitVia(DrivableExecutor* e) && {
-  while (!isReady()) {
-    e->drive();
-  }
+Future<T>&& Future<T>::wait() && {
+  detail::waitImpl(*this);
+  return std::move(*this);
+}
+
+template <class T>
+Future<T>& Future<T>::wait(Duration dur) & {
+  detail::waitImpl(*this, dur);
+  return *this;
+}
+
+template <class T>
+Future<T>&& Future<T>::wait(Duration dur) && {
+  detail::waitImpl(*this, dur);
+  return std::move(*this);
+}
+
+template <class T>
+Future<T>& Future<T>::waitVia(DrivableExecutor* e) & {
+  detail::waitViaImpl(*this, e);
+  return *this;
+}
+
+template <class T>
+Future<T>&& Future<T>::waitVia(DrivableExecutor* e) && {
+  detail::waitViaImpl(*this, e);
   return std::move(*this);
 }
 

folly/futures/Future.h

@@ -410,14 +410,18 @@ class Future {
   /// now. The optional Timekeeper is as with futures::sleep().
   Future<T> delayed(Duration, Timekeeper* = nullptr);
 
-  /// Block until this Future is complete. Returns a new Future containing the
-  /// result.
-  Future<T> wait();
+  /// Block until this Future is complete. Returns a reference to this Future.
+  Future<T>& wait() &;
+
+  /// Overload of wait() for rvalue Futures
+  Future<T>&& wait() &&;
 
   /// Block until this Future is complete or until the given Duration passes.
-  /// Returns a new Future which either contains the result or is incomplete,
-  /// depending on whether the Duration passed.
-  Future<T> wait(Duration);
+  /// Returns a reference to this Future
+  Future<T>& wait(Duration) &;
+
+  /// Overload of wait(Duration) for rvalue Futures
+  Future<T>&& wait(Duration) &&;
 
   /// Call e->drive() repeatedly until the future is fulfilled. Examples
   /// of DrivableExecutor include EventBase and ManualExecutor. Returns a
@@ -426,7 +430,7 @@ class Future {
   Future<T>& waitVia(DrivableExecutor* e) &;
 
   /// Overload of waitVia() for rvalue Futures
-  Future<T> waitVia(DrivableExecutor* e) &&;
+  Future<T>&& waitVia(DrivableExecutor* e) &&;
 
  private:
   typedef detail::Core<T>* corePtr;

folly/futures/test/FutureTest.cpp

@@ -30,6 +30,7 @@
 #include <folly/futures/DrivableExecutor.h>
 #include <folly/MPMCQueue.h>
 
+#include <folly/io/async/EventBase.h>
 #include <folly/io/async/Request.h>
 
 using namespace folly;
@@ -37,6 +38,7 @@ using std::pair;
 using std::string;
 using std::unique_ptr;
 using std::vector;
+using std::chrono::milliseconds;
 
 #define EXPECT_TYPE(x, T) \
   EXPECT_TRUE((std::is_same<decltype(x), T>::value))
@@ -963,30 +965,78 @@ TEST(Future, wait) {
   EXPECT_EQ(result.load(), 42);
 }
 
+struct MoveFlag {
+  MoveFlag() = default;
+  MoveFlag(const MoveFlag&) = delete;
+  MoveFlag(MoveFlag&& other) noexcept {
+    other.moved = true;
+  }
+  bool moved{false};
+};
+
+TEST(Future, waitReplacesSelf) {
+  // wait
+  {
+    // lvalue
+    auto f1 = makeFuture(MoveFlag());
+    f1.wait();
+    EXPECT_FALSE(f1.value().moved);
+
+    // rvalue
+    auto f2 = makeFuture(MoveFlag()).wait();
+    EXPECT_FALSE(f2.value().moved);
+  }
+
+  // wait(Duration)
+  {
+    // lvalue
+    auto f1 = makeFuture(MoveFlag());
+    f1.wait(milliseconds(1));
+    EXPECT_FALSE(f1.value().moved);
+
+    // rvalue
+    auto f2 = makeFuture(MoveFlag()).wait(milliseconds(1));
+    EXPECT_FALSE(f2.value().moved);
+  }
+
+  // waitVia
+  {
+    folly::EventBase eb;
+    // lvalue
+    auto f1 = makeFuture(MoveFlag());
+    f1.waitVia(&eb);
+    EXPECT_FALSE(f1.value().moved);
+
+    // rvalue
+    auto f2 = makeFuture(MoveFlag()).waitVia(&eb);
+    EXPECT_FALSE(f2.value().moved);
+  }
+}
+
 TEST(Future, waitWithDuration) {
  {
   Promise<int> p;
   Future<int> f = p.getFuture();
-  auto t = f.wait(std::chrono::milliseconds(1));
-  EXPECT_FALSE(t.isReady());
+  f.wait(milliseconds(1));
+  EXPECT_FALSE(f.isReady());
   p.setValue(1);
-  EXPECT_TRUE(t.isReady());
+  EXPECT_TRUE(f.isReady());
  }
  {
   Promise<int> p;
   Future<int> f = p.getFuture();
   p.setValue(1);
-  auto t = f.wait(std::chrono::milliseconds(1));
-  EXPECT_TRUE(t.isReady());
+  f.wait(milliseconds(1));
+  EXPECT_TRUE(f.isReady());
  }
  {
   vector<Future<bool>> v_fb;
   v_fb.push_back(makeFuture(true));
   v_fb.push_back(makeFuture(false));
   auto f = whenAll(v_fb.begin(), v_fb.end());
-  auto t = f.wait(std::chrono::milliseconds(1));
-  EXPECT_TRUE(t.isReady());
-  EXPECT_EQ(2, t.value().size());
+  f.wait(milliseconds(1));
+  EXPECT_TRUE(f.isReady());
+  EXPECT_EQ(2, f.value().size());
  }
  {
   vector<Future<bool>> v_fb;
@@ -995,24 +1045,24 @@ TEST(Future, waitWithDuration) {
   v_fb.push_back(p1.getFuture());
   v_fb.push_back(p2.getFuture());
   auto f = whenAll(v_fb.begin(), v_fb.end());
-  auto t = f.wait(std::chrono::milliseconds(1));
-  EXPECT_FALSE(t.isReady());
+  f.wait(milliseconds(1));
+  EXPECT_FALSE(f.isReady());
   p1.setValue(true);
-  EXPECT_FALSE(t.isReady());
+  EXPECT_FALSE(f.isReady());
   p2.setValue(true);
-  EXPECT_TRUE(t.isReady());
+  EXPECT_TRUE(f.isReady());
  }
  {
-  auto t = makeFuture().wait(std::chrono::milliseconds(1));
-  EXPECT_TRUE(t.isReady());
+  auto f = makeFuture().wait(milliseconds(1));
+  EXPECT_TRUE(f.isReady());
  }
 
  {
    Promise<void> p;
    auto start = std::chrono::steady_clock::now();
-   auto f = p.getFuture().wait(std::chrono::milliseconds(100));
+   auto f = p.getFuture().wait(milliseconds(100));
    auto elapsed = std::chrono::steady_clock::now() - start;
-   EXPECT_GE(elapsed, std::chrono::milliseconds(100));
+   EXPECT_GE(elapsed, milliseconds(100));
    EXPECT_FALSE(f.isReady());
    p.setValue();
    EXPECT_TRUE(f.isReady());
@@ -1025,7 +1075,7 @@ TEST(Future, waitWithDuration) {
    folly::Baton<> b;
    auto t = std::thread([&]{
      b.post();
-     std::this_thread::sleep_for(std::chrono::milliseconds(100));
+     std::this_thread::sleep_for(milliseconds(100));
      p.setValue();
    });
    b.wait();