From 809da17c47ed3a16a086b6cbface7499e9886b35 Mon Sep 17 00:00:00 2001
From: Dave Watson <davejwatson@fb.com>
Date: Thu, 30 Oct 2014 16:16:56 -0700
Subject: [PATCH] save/restore request context in future

Summary:
Generic 'threadlocal' like object that follows async calls around.

Note: Finagle's futures do this also.  It is super useful.

Test Plan: I should write a unittest?

Reviewed By: haijunz@fb.com, hans@fb.com

Subscribers: trunkagent, dawidp, doug, fugalh, njormrod, folly-diffs@

FB internal diff: D1650843

Tasks: 4698780

Signature: t1:1650843:1414711295:c7439733680ab4903eb9c05dcf8bf52100bf3f07
---
 folly/wangle/detail/Core.h       |  6 +++++
 folly/wangle/test/FutureTest.cpp | 42 ++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/folly/wangle/detail/Core.h b/folly/wangle/detail/Core.h
index f3aa6d24..f83e4f71 100644
--- a/folly/wangle/detail/Core.h
+++ b/folly/wangle/detail/Core.h
@@ -30,6 +30,8 @@
 #include <folly/wangle/Executor.h>
 #include <folly/wangle/detail/FSM.h>
 
+#include <folly/io/async/Request.h>
+
 namespace folly { namespace wangle { namespace detail {
 
 // As of GCC 4.8.1, the std::function in libstdc++ optimizes only for pointers
@@ -80,6 +82,7 @@ class Core : protected FSM<State> {
         throw std::logic_error("setCallback called twice");
       }
 
+      context_ = RequestContext::saveContext();
       callback_ = std::move(func);
     };
 
@@ -194,6 +197,8 @@ class Core : protected FSM<State> {
       // TODO(5306911) we should probably try/catch
       calledBack_ = true;
       Executor* x = executor_;
+
+      RequestContext::setContext(context_);
       if (x) {
         MoveWrapper<std::function<void(Try<T>&&)>> cb(std::move(callback_));
         MoveWrapper<folly::Optional<Try<T>>> val(std::move(result_));
@@ -217,6 +222,7 @@ class Core : protected FSM<State> {
 
   folly::Optional<Try<T>> result_;
   std::function<void(Try<T>&&)> callback_;
+  std::shared_ptr<RequestContext> context_{nullptr};
   std::atomic<bool> calledBack_ {false};
   std::atomic<unsigned char> detached_ {0};
   std::atomic<bool> active_ {true};
diff --git a/folly/wangle/test/FutureTest.cpp b/folly/wangle/test/FutureTest.cpp
index 8d1f95d3..1ecca30c 100644
--- a/folly/wangle/test/FutureTest.cpp
+++ b/folly/wangle/test/FutureTest.cpp
@@ -28,6 +28,9 @@
 #include <folly/wangle/Future.h>
 #include <folly/wangle/ManualExecutor.h>
 
+#include <folly/io/async/Request.h>
+
+using namespace folly;
 using namespace folly::wangle;
 using std::pair;
 using std::string;
@@ -908,3 +911,42 @@ TEST(Future, detachRace) {
   f.reset();
   t1.join();
 }
+
+class TestData : public RequestData {
+ public:
+  explicit TestData(int data) : data_(data) {}
+  virtual ~TestData() {}
+  int data_;
+};
+
+TEST(Future, context) {
+
+  // Start a new context
+  RequestContext::create();
+
+  EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
+
+  // Set some test data
+  RequestContext::get()->setContextData(
+    "test",
+    std::unique_ptr<TestData>(new TestData(10)));
+
+  // Start a future
+  Promise<void> p;
+  auto future = p.getFuture().then([&]{
+    // Check that the context followed the future
+    EXPECT_TRUE(RequestContext::get() != nullptr);
+    auto a = dynamic_cast<TestData*>(
+      RequestContext::get()->getContextData("test"));
+    auto data = a->data_;
+    EXPECT_EQ(10, data);
+  });
+
+  // Clear the context
+  RequestContext::setContext(nullptr);
+
+  EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
+
+  // Fulfil the promise
+  p.setValue();
+}
-- 
2.34.1