From badc3ebe7ef9d370e0d4b9d110f5dad6f6284d7e Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <alexshap@fb.com>
Date: Thu, 19 Nov 2015 20:12:41 -0800
Subject: [PATCH] Remove busy wait and support multiple wait

Summary: Remove busy wait from Future::wait.
If future.wait(timeout) has succeded we should
return a ready future, otherwise we should return a future
for the final result (not necessarily ready).

Reviewed By: yfeldblum

Differential Revision: D2646860

fb-gh-sync-id: 62671d09073ad86e84df8c9257e961d2a8c2a339
---
 folly/futures/Future-inl.h      | 30 +++++++++---------------------
 folly/futures/test/WaitTest.cpp | 13 +++++++++++++
 2 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index eeb0e2e5..6fb85a33 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -17,10 +17,10 @@
 #pragma once
 
 #include <algorithm>
+#include <cassert>
 #include <chrono>
 #include <random>
 #include <thread>
-
 #include <folly/experimental/fibers/Baton.h>
 #include <folly/Optional.h>
 #include <folly/Random.h>
@@ -936,18 +936,9 @@ void waitImpl(Future<T>& f) {
   if (f.isReady()) return;
 
   folly::fibers::Baton baton;
-  f = f.then([&](Try<T> t) {
-    baton.post();
-    return makeFuture(std::move(t));
-  });
+  f.setCallback_([&](const Try<T>& t) { baton.post(); });
   baton.wait();
-
-  // 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();
-  }
+  assert(f.isReady());
 }
 
 template <class T>
@@ -955,19 +946,16 @@ void waitImpl(Future<T>& f, Duration dur) {
   // short-circuit if there's nothing to do
   if (f.isReady()) return;
 
+  folly::MoveWrapper<Promise<T>> promise;
+  auto ret = promise->getFuture();
   auto baton = std::make_shared<folly::fibers::Baton>();
-  f = f.then([baton](Try<T> t) {
+  f.setCallback_([baton, promise](Try<T>&& t) mutable {
+    promise->setTry(std::move(t));
     baton->post();
-    return makeFuture(std::move(t));
   });
-
-  // Let's preserve the invariant that if we did not timeout (timed_wait returns
-  // 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.
+  f = std::move(ret);
   if (baton->timed_wait(dur)) {
-    while (!f.isReady()) {
-      std::this_thread::yield();
-    }
+    assert(f.isReady());
   }
 }
 
diff --git a/folly/futures/test/WaitTest.cpp b/folly/futures/test/WaitTest.cpp
index b2736f12..57f2f839 100644
--- a/folly/futures/test/WaitTest.cpp
+++ b/folly/futures/test/WaitTest.cpp
@@ -195,3 +195,16 @@ TEST(Wait, waitWithDuration) {
    t.join();
  }
 }
+
+TEST(Wait, multipleWait) {
+  auto f = futures::sleep(milliseconds(100));
+  for (size_t i = 0; i < 5; ++i) {
+    EXPECT_FALSE(f.isReady());
+    f.wait(milliseconds(3));
+  }
+  EXPECT_FALSE(f.isReady());
+  f.wait();
+  EXPECT_TRUE(f.isReady());
+  f.wait();
+  EXPECT_TRUE(f.isReady());
+}
-- 
2.34.1