Fix set/unset for default folly RequestContext

Summary:
My guess is that if we allow to set data in the default context, we also expect set/unset to work.
It currently doesn't. This fixes it (at least, the test repro passes), but not sure it's the right fix.

Reviewed By: djwatson

Differential Revision: D8910506

fbshipit-source-id: cedf1dd0ee91761d210137949da7477ed69fce31

folly/io/async/Request.cpp

@@ -144,14 +144,9 @@ std::shared_ptr<RequestContext> RequestContext::setContext(
   auto& curCtx = getStaticContext();
   if (ctx != curCtx) {
     FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
-    using std::swap;
-    if (curCtx) {
-      curCtx->onUnset();
-    }
-    swap(ctx, curCtx);
-    if (curCtx) {
-      curCtx->onSet();
-    }
+    get()->onUnset();
+    std::swap(ctx, curCtx);
+    get()->onSet();
   }
   return ctx;
 }

folly/io/async/test/RequestContextTest.cpp

@@ -44,6 +44,26 @@ class TestData : public RequestData {
   int data_;
 };
 
+RequestContext& getContext() {
+  auto* ctx = RequestContext::get();
+  EXPECT_TRUE(ctx != nullptr);
+  return *ctx;
+}
+
+void setData(int data = 0) {
+  getContext().setContextData("test", std::make_unique<TestData>(data));
+}
+
+bool hasData() {
+  return getContext().getContextData("test") != nullptr;
+}
+
+const TestData& getData() {
+  auto* ptr = dynamic_cast<TestData*>(getContext().getContextData("test"));
+  EXPECT_TRUE(ptr != nullptr);
+  return *ptr;
+}
+
 TEST(RequestContext, SimpleTest) {
   EventBase base;
 
@@ -81,6 +101,35 @@ TEST(RequestContext, SimpleTest) {
   EXPECT_TRUE(nullptr != RequestContext::get());
 }
 
+TEST(RequestContext, RequestContextScopeGuard) {
+  RequestContextScopeGuard g0;
+  setData(10);
+  {
+    RequestContextScopeGuard g1;
+    EXPECT_FALSE(hasData());
+    setData(20);
+    EXPECT_EQ(20, getData().data_);
+    EXPECT_EQ(1, getData().set_);
+    EXPECT_EQ(0, getData().unset_);
+  }
+  EXPECT_EQ(10, getData().data_);
+  EXPECT_EQ(2, getData().set_);
+  EXPECT_EQ(1, getData().unset_);
+}
+
+TEST(RequestContext, defaultContext) {
+  // Don't create a top level guard
+  // Regression test for set/onset used to not work with the default context
+  setData(10);
+  {
+    RequestContextScopeGuard g1;
+    EXPECT_FALSE(hasData());
+  }
+  EXPECT_EQ(10, getData().data_);
+  EXPECT_EQ(2, getData().set_);
+  EXPECT_EQ(1, getData().unset_);
+}
+
 TEST(RequestContext, setIfAbsentTest) {
   EXPECT_TRUE(RequestContext::get() != nullptr);