fix accessing current rctx in `RequestData::onUnset` callbacks

Summary:
as titled

The issue here was that we were clearing out the value of current rctx from TLS, then calling `onUnset`. So calling `RequestContext::get` from within the `onUnset` callback resulted in returning the default global RCTX.

Reviewed By: yfeldblum, zhxchen17

Differential Revision: D22750092

fbshipit-source-id: 509676d047cf6aec05d08ff3bf656d6d1558fbf5

folly/io/async/Request.cpp

@@ -790,7 +790,8 @@ FOLLY_ALWAYS_INLINE
 /* static */ std::shared_ptr<RequestContext> RequestContext::setContextHazptr(
     std::shared_ptr<RequestContext>& newCtx,
     StaticContext& staticCtx) {
-  auto curCtx = std::move(staticCtx.first);
+  std::shared_ptr<RequestContext> prevCtx;
+  auto curCtx = staticCtx.first.get();
   bool checkCur = curCtx && curCtx->stateHazptr_.combined();
   bool checkNew = newCtx && newCtx->stateHazptr_.combined();
   if (checkCur && checkNew) {
@@ -806,6 +807,7 @@ FOLLY_ALWAYS_INLINE
         data->onUnset();
       }
     }
+    prevCtx = std::move(staticCtx.first);
     staticCtx.first = std::move(newCtx);
     staticCtx.second.store(staticCtx.first->rootId_, std::memory_order_relaxed);
     for (auto it = newcb.begin(); it != newcb.end(); ++it) {
@@ -819,16 +821,17 @@ FOLLY_ALWAYS_INLINE
     if (curCtx) {
       curCtx->stateHazptr_.onUnset();
     }
+    prevCtx = std::move(staticCtx.first);
     staticCtx.first = std::move(newCtx);
     if (staticCtx.first) {
-      staticCtx.first->stateHazptr_.onSet();
       staticCtx.second.store(
           staticCtx.first->rootId_, std::memory_order_relaxed);
+      staticCtx.first->stateHazptr_.onSet();
     } else {
       staticCtx.second.store(0, std::memory_order_relaxed);
     }
   }
-  return curCtx;
+  return prevCtx;
 }
 
 RequestContext::StaticContext& RequestContext::getStaticContext() {

folly/io/async/test/RequestContextHelper.h

@@ -31,14 +31,18 @@ class TestData : public RequestData {
 
   void onSet() override {
     set_++;
+    onSetRctx = RequestContext::get();
   }
 
   void onUnset() override {
     unset_++;
+    onUnSetRctx = RequestContext::get();
   }
 
   int set_ = 0, unset_ = 0;
   int data_;
+  RequestContext* onSetRctx = nullptr;
+  RequestContext* onUnSetRctx = nullptr;
 };
 
 } // namespace folly

folly/io/async/test/RequestContextTest.cpp

@@ -190,6 +190,7 @@ TEST_F(RequestContextTest, testSetUnset) {
 
   // onSet called in setContextData
   EXPECT_EQ(1, testData1->set_);
+  EXPECT_EQ(ctx1.get(), testData1->onSetRctx);
 
   // Override RequestContext
   RequestContext::create();
@@ -199,14 +200,19 @@ TEST_F(RequestContextTest, testSetUnset) {
 
   // onSet called in setContextData
   EXPECT_EQ(1, testData2->set_);
+  EXPECT_EQ(ctx2.get(), testData2->onSetRctx);
 
   // Check ctx1->onUnset was called
   EXPECT_EQ(1, testData1->unset_);
+  EXPECT_EQ(ctx1.get(), testData1->onUnSetRctx);
 
   RequestContext::setContext(ctx1);
   EXPECT_EQ(2, testData1->set_);
   EXPECT_EQ(1, testData1->unset_);
   EXPECT_EQ(1, testData2->unset_);
+  EXPECT_EQ(ctx1.get(), testData1->onSetRctx);
+  EXPECT_EQ(ctx1.get(), testData1->onUnSetRctx);
+  EXPECT_EQ(ctx2.get(), testData2->onUnSetRctx);
 
   RequestContext::setContext(ctx2);
   EXPECT_EQ(2, testData1->set_);