From d89b4d34c1bb7d3dd7c1a8d6f2a58c148f0b122a Mon Sep 17 00:00:00 2001
From: Mirek Klimos <miro@fb.com>
Date: Thu, 4 Aug 2016 12:09:50 -0700
Subject: [PATCH] RequestContext::create should call onUnset callback

Summary: melaniesubbiah introduced onSet / onUnset callbacks on RequestData in D3604948; we need unset() to be called when an RC is overriden with RequestContext::create() so that things work as expected. Also, change the order of calling onSet / onUnset - from RequestData perspective, it shouldn't look like there are two contexts set at the same time

Reviewed By: palmtenor

Differential Revision: D3667017

fbshipit-source-id: b9bfb858fe65ffb11de8e6d6f13b8f4cf6266bc9
---
 folly/io/async/Request.cpp                 |  6 ++--
 folly/io/async/Request.h                   |  2 +-
 folly/io/async/test/RequestContextTest.cpp | 36 ++++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/folly/io/async/Request.cpp b/folly/io/async/Request.cpp
index 939ac2e0..8ea7f3e3 100644
--- a/folly/io/async/Request.cpp
+++ b/folly/io/async/Request.cpp
@@ -108,12 +108,12 @@ std::shared_ptr<RequestContext> RequestContext::setContext(
   auto& prev = getStaticContext();
   if (ctx != prev) {
     using std::swap;
-    if (ctx) {
-      ctx->onSet();
-    }
     if (prev) {
       prev->onUnset();
     }
+    if (ctx) {
+      ctx->onSet();
+    }
     swap(ctx, prev);
   }
   return ctx;
diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index 194305fc..0f4d7256 100644
--- a/folly/io/async/Request.h
+++ b/folly/io/async/Request.h
@@ -47,7 +47,7 @@ class RequestContext {
   // It will be passed between queues / threads (where implemented),
   // so it should be valid for the lifetime of the request.
   static void create() {
-    getStaticContext() = std::make_shared<RequestContext>();
+    setContext(std::make_shared<RequestContext>());
   }
 
   // Get the current context.
diff --git a/folly/io/async/test/RequestContextTest.cpp b/folly/io/async/test/RequestContextTest.cpp
index 8e2a12ff..6b1cb6db 100644
--- a/folly/io/async/test/RequestContextTest.cpp
+++ b/folly/io/async/test/RequestContextTest.cpp
@@ -28,6 +28,13 @@ class TestData : public RequestData {
  public:
   explicit TestData(int data) : data_(data) {}
   ~TestData() override {}
+  void onSet() override {
+    set_++;
+  }
+  void onUnset() override {
+    unset_++;
+  }
+  int set_ = 0, unset_ = 0;
   int data_;
 };
 
@@ -94,6 +101,35 @@ TEST(RequestContext, setIfAbsentTest) {
   EXPECT_TRUE(nullptr != RequestContext::get());
 }
 
+TEST(RequestContext, testSetUnset) {
+  RequestContext::create();
+  auto ctx1 = RequestContext::saveContext();
+  ctx1->setContextData("test", std::unique_ptr<TestData>(new TestData(10)));
+  auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
+
+  // Override RequestContext
+  RequestContext::create();
+  auto ctx2 = RequestContext::saveContext();
+  ctx2->setContextData("test", std::unique_ptr<TestData>(new TestData(20)));
+  auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
+
+  // Check ctx1->onUnset was called
+  EXPECT_EQ(0, testData1->set_);
+  EXPECT_EQ(1, testData1->unset_);
+
+  RequestContext::setContext(ctx1);
+  EXPECT_EQ(1, testData1->set_);
+  EXPECT_EQ(1, testData1->unset_);
+  EXPECT_EQ(0, testData2->set_);
+  EXPECT_EQ(1, testData2->unset_);
+
+  RequestContext::setContext(ctx2);
+  EXPECT_EQ(1, testData1->set_);
+  EXPECT_EQ(2, testData1->unset_);
+  EXPECT_EQ(1, testData2->set_);
+  EXPECT_EQ(1, testData2->unset_);
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::InitGoogleLogging(argv[0]);
-- 
2.34.1