From 67fcea204a38d79224d19aeab682420f2397510e Mon Sep 17 00:00:00 2001
From: Alex Yarmula <ayarmula@fb.com>
Date: Thu, 6 Oct 2016 07:21:22 -0700
Subject: [PATCH] use previous Executor after delayed

Summary:
When `delayed` is called on the Future, the underlying `futures::sleep` call runs on a timer thread, and the resulting callback is called on the same thread. Therefore, in the following sequence:

  f.via(&someExecutor).within(one_ms).then([&]() { /* [1] */ })

The code in [1] is not running in someExecutor. This can cause confusion by users of the library who expect the initial `via` to be sticky.

This change returns to the prior `Executor` after `delayed` is finished.

Reviewed By: yfeldblum

Differential Revision: D3979179

fbshipit-source-id: 936ff5626e8ac377ffb15babf573349466984e3a
---
 folly/futures/Future-inl.h            |  9 +++---
 folly/futures/test/TimekeeperTest.cpp | 42 ++++++++++++++++++---------
 2 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index 1cd4035e..14ae5fc4 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -972,10 +972,11 @@ Future<T> Future<T>::within(Duration dur, E e, Timekeeper* tk) {
 template <class T>
 Future<T> Future<T>::delayed(Duration dur, Timekeeper* tk) {
   return collectAll(*this, futures::sleep(dur, tk))
-    .then([](std::tuple<Try<T>, Try<Unit>> tup) {
-      Try<T>& t = std::get<0>(tup);
-      return makeFuture<T>(std::move(t));
-    });
+      .then([](std::tuple<Try<T>, Try<Unit>> tup) {
+        Try<T>& t = std::get<0>(tup);
+        return makeFuture<T>(std::move(t));
+      })
+      .via(getExecutor());
 }
 
 namespace detail {
diff --git a/folly/futures/test/TimekeeperTest.cpp b/folly/futures/test/TimekeeperTest.cpp
index 664dd737..e4aeba4d 100644
--- a/folly/futures/test/TimekeeperTest.cpp
+++ b/folly/futures/test/TimekeeperTest.cpp
@@ -168,20 +168,36 @@ TEST(Timekeeper, chainedInterruptTest) {
   EXPECT_FALSE(test);
 }
 
-TEST(Timekeeper, executor) {
-  class ExecutorTester : public Executor {
-   public:
-    void add(Func f) override {
-      count++;
-      f();
-    }
-    std::atomic<int> count{0};
+namespace {
+class ExecutorTester : public Executor {
+ public:
+  void add(Func f) override {
+    count++;
+    f();
+  }
+  std::atomic<int> count{0};
   };
-
-  auto f = makeFuture();
-  ExecutorTester tester;
-  f.via(&tester).within(one_ms).then([&](){}).wait();
-  EXPECT_EQ(2, tester.count);
+  }
+
+  TEST(Timekeeper, executor) {
+    auto f = makeFuture();
+    ExecutorTester tester;
+    f.via(&tester).within(one_ms).then([&]() {}).wait();
+    EXPECT_EQ(2, tester.count);
+  }
+
+  TEST(Timekeeper, executorSameAfterDelayed) {
+    ExecutorTester tester;
+    // make sure we're using the same Executor (tester) after delayed returns
+    // by comparing the executor count between invocations
+    makeFuture()
+        .via(&tester)
+        .delayed(one_ms)
+        .then([&]() { return tester.count.load(); })
+        .then([&](int countBefore) {
+          EXPECT_EQ(1, tester.count.load() - countBefore);
+        })
+        .wait();
 }
 
 // TODO(5921764)
-- 
2.34.1