From 6b6e6ae31c36422ef8bc6fabfb5394c16fc64ead Mon Sep 17 00:00:00 2001
From: Alex Landau <alandau@fb.com>
Date: Wed, 24 Jun 2015 16:20:08 -0700
Subject: [PATCH] Fix passing RequestContext to executor thread

Summary: If x->add() executes its lambda on a different thread and
doesn't pass the context on its own, the callback wouldn't
have the correct context set.

Reviewed By: @djwatson

Differential Revision: D2189318
---
 folly/futures/detail/Core.h       |  6 ++--
 folly/futures/test/FutureTest.cpp | 54 +++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/folly/futures/detail/Core.h b/folly/futures/detail/Core.h
index d21d74ce..f0018862 100644
--- a/folly/futures/detail/Core.h
+++ b/folly/futures/detail/Core.h
@@ -313,8 +313,6 @@ class Core {
   }
 
   void doCallback() {
-    RequestContext::setContext(context_);
-
     Executor* x = executor_;
     int8_t priority;
     if (x) {
@@ -333,20 +331,24 @@ class Core {
         if (LIKELY(x->getNumPriorities() == 1)) {
           x->add([this]() mutable {
             SCOPE_EXIT { detachOne(); };
+            RequestContext::setContext(context_);
             callback_(std::move(*result_));
           });
         } else {
           x->addWithPriority([this]() mutable {
             SCOPE_EXIT { detachOne(); };
+            RequestContext::setContext(context_);
             callback_(std::move(*result_));
           }, priority);
         }
       } catch (...) {
         --attached_; // Account for extra ++attached_ before try
+        RequestContext::setContext(context_);
         result_ = Try<T>(exception_wrapper(std::current_exception()));
         callback_(std::move(*result_));
       }
     } else {
+      RequestContext::setContext(context_);
       callback_(std::move(*result_));
     }
   }
diff --git a/folly/futures/test/FutureTest.cpp b/folly/futures/test/FutureTest.cpp
index 723b358e..d6e1990c 100644
--- a/folly/futures/test/FutureTest.cpp
+++ b/folly/futures/test/FutureTest.cpp
@@ -680,3 +680,57 @@ TEST(Future, thenDynamic) {
   p.setValue(2);
   EXPECT_EQ(f.get(), 5);
 }
+
+TEST(Future, RequestContext) {
+  class NewThreadExecutor : public Executor {
+   public:
+    ~NewThreadExecutor() override {
+      std::for_each(v_.begin(), v_.end(), [](std::thread& t){ t.join(); });
+    }
+    void add(Func f) override {
+      if (throwsOnAdd_) { throw std::exception(); }
+      v_.emplace_back(std::move(f));
+    }
+    void addWithPriority(Func f, int8_t prio) override { add(std::move(f)); }
+    uint8_t getNumPriorities() const override { return numPriorities_; }
+
+    void setHandlesPriorities() { numPriorities_ = 2; }
+    void setThrowsOnAdd() { throwsOnAdd_ = true; }
+   private:
+    std::vector<std::thread> v_;
+    uint8_t numPriorities_ = 1;
+    bool throwsOnAdd_ = false;
+  };
+
+  struct MyRequestData : RequestData {
+    MyRequestData(bool value = false) : value(value) {}
+    bool value;
+  };
+
+  NewThreadExecutor e;
+  RequestContext::create();
+  RequestContext::get()->setContextData("key",
+      folly::make_unique<MyRequestData>(true));
+  auto checker = [](int lineno) {
+    return [lineno](Try<int>&& t) {
+      auto d = static_cast<MyRequestData*>(
+        RequestContext::get()->getContextData("key"));
+      EXPECT_TRUE(d && d->value) << "on line " << lineno;
+    };
+  };
+
+  makeFuture(1).via(&e).then(checker(__LINE__));
+
+  e.setHandlesPriorities();
+  makeFuture(2).via(&e).then(checker(__LINE__));
+
+  Promise<int> p1, p2;
+  p1.getFuture().then(checker(__LINE__));
+
+  e.setThrowsOnAdd();
+  p2.getFuture().via(&e).then(checker(__LINE__));
+
+  RequestContext::create();
+  p1.setValue(3);
+  p2.setValue(4);
+}
-- 
2.34.1