fix race in Future::waitWithSemaphore

[folly.git] / folly / wangle / Future-inl.h

diff --git a/folly/wangle/Future-inl.h b/folly/wangle/Future-inl.h
index de8dbf4bdb6206ae7b133e8d01b75bec564c398e..f8751f5ef8ba83be8774f13aa5e00ce545732fcc 100644 (file)
--- a/folly/wangle/Future-inl.h
+++ b/folly/wangle/Future-inl.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <chrono>
+#include <thread>
 
 #include <folly/wangle/detail/State.h>
 #include <folly/Baton.h>
@@ -416,6 +417,12 @@ waitWithSemaphore(Future<T>&& f) {
     return std::move(t.value());
   });
   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.
+    std::this_thread::yield();
+  }
   return done;
 }
 
@@ -427,6 +434,12 @@ inline Future<void> waitWithSemaphore<void>(Future<void>&& f) {
     t.value();
   });
   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.
+    std::this_thread::yield();
+  }
   return done;
 }