From b6f62e958838d8855980c6713db536a125a0a5fb Mon Sep 17 00:00:00 2001
From: Hans Fugal <fugalh@fb.com>
Date: Wed, 8 Oct 2014 12:06:38 -0700
Subject: [PATCH] Fix via

Summary: Sometimes you just have to take a step back. :-P

Test Plan: All the unit tests including the one that had been disabled, now pass.

Reviewed By: hannesr@fb.com

Subscribers: meisner, trunkagent, net-systems@, fugalh, exa, njormrod, davejwatson

FB internal diff: D1596368

Tasks: 4920689, 4480567, 5306911
---
 folly/wangle/Future-inl.h        | 39 ++++----------------------------
 folly/wangle/detail/State.h      | 20 +++++++++++++---
 folly/wangle/test/FutureTest.cpp |  2 +-
 3 files changed, 22 insertions(+), 39 deletions(-)

diff --git a/folly/wangle/Future-inl.h b/folly/wangle/Future-inl.h
index f404067b..e66e504c 100644
--- a/folly/wangle/Future-inl.h
+++ b/folly/wangle/Future-inl.h
@@ -193,42 +193,11 @@ template <class T>
 template <typename Executor>
 inline Future<T> Future<T>::via(Executor* executor) {
   throwIfInvalid();
-  MoveWrapper<Promise<T>> promise;
-
-  auto f = promise->getFuture();
-  // We are obligated to return a cold future.
-  f.deactivate();
-  // But we also need to make this one cold for via to at least work some of
-  // the time. (see below)
-  deactivate();
-
-  then([=](Try<T>&& t) mutable {
-    MoveWrapper<Try<T>> tw(std::move(t));
-    // There is a race here.
-    // When the promise is fulfilled, and the future is still inactive, when
-    // the future is activated (e.g. by destruction) the callback will happen
-    // in that context, not in the intended context (which has already left
-    // the building).
-    //
-    // Currently, this will work fine because all the temporaries are
-    // destructed in an order that is compatible with this implementation:
-    //
-    //   makeFuture().via(x).then(a).then(b);
-    //
-    // However, this will not work reliably:
-    //
-    //   auto f2 = makeFuture().via(x);
-    //   f2.then(a).then(b);
-    //
-    // Because the first temporary is destructed on the first line, and the
-    // executor is fed. But by the time f2 is destructed, the executor
-    // may have already fulfilled the promise on the other thread.
-    //
-    // TODO(#4920689) fix it
-    executor->add([=]() mutable { promise->fulfilTry(std::move(*tw)); });
-  });
 
-  return f;
+  this->deactivate();
+  state_->setExecutor(executor);
+
+  return std::move(*this);
 }
 
 template <class T>
diff --git a/folly/wangle/detail/State.h b/folly/wangle/detail/State.h
index 0f7f4bf1..c0d921f6 100644
--- a/folly/wangle/detail/State.h
+++ b/folly/wangle/detail/State.h
@@ -26,6 +26,7 @@
 #include <folly/wangle/Try.h>
 #include <folly/wangle/Promise.h>
 #include <folly/wangle/Future.h>
+#include <folly/wangle/Executor.h>
 
 namespace folly { namespace wangle { namespace detail {
 
@@ -141,14 +142,26 @@ class State {
 
   bool isActive() { return active_; }
 
+  void setExecutor(Executor* x) {
+    std::lock_guard<decltype(mutex_)> lock(mutex_);
+    executor_ = x;
+  }
+
  private:
   void maybeCallback() {
     std::lock_guard<decltype(mutex_)> lock(mutex_);
     if (!calledBack_ &&
         value_ && callback_ && isActive()) {
-      // TODO we should probably try/catch here
-      callback_(std::move(*value_));
-      calledBack_ = true;
+      // TODO(5306911) we should probably try/catch here
+      if (executor_) {
+        MoveWrapper<folly::Optional<Try<T>>> val(std::move(value_));
+        MoveWrapper<std::function<void(Try<T>&&)>> cb(std::move(callback_));
+        executor_->add([cb, val]() mutable { (*cb)(std::move(**val)); });
+        calledBack_ = true;
+      } else {
+        callback_(std::move(*value_));
+        calledBack_ = true;
+      }
     }
   }
 
@@ -173,6 +186,7 @@ class State {
   bool calledBack_ = false;
   unsigned char detached_ = 0;
   bool active_ = true;
+  Executor* executor_ = nullptr;
 
   // this lock isn't meant to protect all accesses to members, only the ones
   // that need to be threadsafe: the act of setting value_ and callback_, and
diff --git a/folly/wangle/test/FutureTest.cpp b/folly/wangle/test/FutureTest.cpp
index 0f045af3..877efcfb 100644
--- a/folly/wangle/test/FutureTest.cpp
+++ b/folly/wangle/test/FutureTest.cpp
@@ -808,7 +808,7 @@ TEST(Future, viaRaces) {
 }
 
 // TODO(#4920689)
-TEST(Future, DISABLED_viaRaces_2stage) {
+TEST(Future, viaRaces_2stage) {
   ManualExecutor x;
   Promise<void> p;
   auto tid = std::this_thread::get_id();
-- 
2.34.1