From db37af8437399c0578a30e5bf553f9fe231b5c82 Mon Sep 17 00:00:00 2001
From: Matt Dordal <mnd@fb.com>
Date: Tue, 13 May 2014 10:37:45 -0700
Subject: [PATCH] add waitWithSemaphore to folly::wangle

Summary:
It may be useful to wait for a future to finish. This adds a utility function
to do so, returning a completed future.

NB: While it doesn't matter which thread executes the `then`, there does need
to be two threads. If not, this will deadlock forever. I'm not sure if there's
a way to detect/prevent that.

Test Plan: added some unit tests.

Reviewed By: hans@fb.com

FB internal diff: D1319330
---
 folly/wangle/Future-inl.h        | 13 ++++++++++++
 folly/wangle/Future.h            |  9 ++++++++
 folly/wangle/test/FutureTest.cpp | 36 ++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/folly/wangle/Future-inl.h b/folly/wangle/Future-inl.h
index 8d8af267..10d844c2 100644
--- a/folly/wangle/Future-inl.h
+++ b/folly/wangle/Future-inl.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include "detail.h"
+#include <folly/LifoSem.h>
 
 namespace folly { namespace wangle {
 
@@ -418,6 +419,18 @@ whenN(InputIterator first, InputIterator last, size_t n) {
   return ctx->p.getFuture();
 }
 
+template <typename F>
+typename F::value_type waitWithSemaphore(F&& f) {
+  LifoSem sem;
+  Try<typename F::value_type> done;
+  f.then([&](Try<typename F::value_type> &&t) {
+    done = std::move(t);
+    sem.post();
+  });
+  sem.wait();
+  return done.value();
+}
+
 }}
 
 // I haven't included a Future<T&> specialization because I don't forsee us
diff --git a/folly/wangle/Future.h b/folly/wangle/Future.h
index f20ef409..a0030f84 100644
--- a/folly/wangle/Future.h
+++ b/folly/wangle/Future.h
@@ -313,6 +313,15 @@ Future<std::vector<std::pair<
   Try<typename std::iterator_traits<InputIterator>::value_type::value_type>>>>
 whenN(InputIterator first, InputIterator last, size_t n);
 
+/** Wait for the given future to complete on a semaphore. Returns the result of
+ * the given future.
+ *
+ * NB if the promise for the future would be fulfilled in the same thread that
+ * you call this, it will deadlock.
+ */
+template <class F>
+typename F::value_type waitWithSemaphore(F&& f);
+
 }} // folly::wangle
 
 #include "Future-inl.h"
diff --git a/folly/wangle/test/FutureTest.cpp b/folly/wangle/test/FutureTest.cpp
index a362e9bf..d9256445 100644
--- a/folly/wangle/test/FutureTest.cpp
+++ b/folly/wangle/test/FutureTest.cpp
@@ -15,10 +15,12 @@
  */
 
 #include <algorithm>
+#include <atomic>
 #include <folly/small_vector.h>
 #include <gtest/gtest.h>
 #include <memory>
 #include <string>
+#include <thread>
 #include <type_traits>
 #include <unistd.h>
 #include "folly/wangle/Executor.h"
@@ -621,3 +623,37 @@ TEST(Future, throwIfFailed) {
       EXPECT_NO_THROW(t.throwIfFailed());
     });
 }
+
+TEST(Future, waitWithSemaphoreImmediate) {
+  waitWithSemaphore(makeFuture());
+  auto done = waitWithSemaphore(makeFuture(42));
+  EXPECT_EQ(done, 42);
+}
+
+TEST(Future, waitWithSemaphore) {
+  Promise<int> p;
+  Future<int> f = p.getFuture();
+  std::atomic<bool> flag{false};
+  std::atomic<int> result{1};
+  std::atomic<std::thread::id> id;
+
+  std::thread t([&](Future<int>&& tf){
+      auto n = tf.then([&](Try<int> && t) {
+          id = std::this_thread::get_id();
+          return t.value();
+        });
+      flag = true;
+      result.store(waitWithSemaphore(std::move(n)));
+      LOG(INFO) << result;
+    },
+    std::move(f)
+    );
+  while(!flag){}
+  EXPECT_EQ(result.load(), 1);
+  p.setValue(42);
+  t.join();
+  // validate that the continuation ended up executing in this thread, which
+  // is more to ensure that this test actually tests what it should
+  EXPECT_EQ(id, std::this_thread::get_id());
+  EXPECT_EQ(result.load(), 42);
+}
-- 
2.34.1