Add nesting to folly::RequestContext

Differential Revision: D7663240

fbshipit-source-id: d7025eb220ff3f87b05988090297573377ca009a

folly/io/async/Request.cpp

@@ -102,6 +102,19 @@ void RequestContext::onUnset() {
   }
 }
 
+std::shared_ptr<RequestContext> RequestContext::createChild() {
+  auto child = std::make_shared<RequestContext>();
+  auto rlock = state_.rlock();
+  for (const auto& entry : rlock->requestData_) {
+    auto& key = entry.first;
+    auto childData = entry.second->createChild();
+    if (childData) {
+      child->setContextData(key, std::move(childData));
+    }
+  }
+  return child;
+}
+
 void RequestContext::clearContextData(const std::string& val) {
   std::unique_ptr<RequestData> requestData;
   // Delete the RequestData after giving up the wlock just in case one of the

folly/io/async/Request.h

@@ -45,6 +45,12 @@ class RequestData {
   // instance overrides the hasCallback method to return true otherwise
   // the callback will not be executed
   virtual void onUnset() {}
+
+  // Create a child RequestData of this one, or return nullptr (in which case
+  // this RequestData will not exist in the child RequestContext)
+  virtual std::unique_ptr<RequestData> createChild() {
+    return nullptr;
+  }
 };
 
 // If you do not call create() to create a unique request context,
@@ -98,6 +104,8 @@ class RequestContext {
   void onSet();
   void onUnset();
 
+  std::shared_ptr<RequestContext> createChild();
+
   // The following API is used to pass the context through queues / threads.
   // saveContext is called to get a shared_ptr to the context, and
   // setContext is used to reset it on the other side of the queue.
@@ -155,4 +163,24 @@ class RequestContextScopeGuard {
     RequestContext::setContext(std::move(prev_));
   }
 };
+
+class RootRequestContextGuard : public RequestContextScopeGuard {
+ public:
+  RootRequestContextGuard() : RequestContextScopeGuard() {}
+};
+
+class NestedRequestContextGuard : public RequestContextScopeGuard {
+ public:
+  NestedRequestContextGuard() : RequestContextScopeGuard(createNested()) {}
+
+ private:
+  static std::shared_ptr<RequestContext> createNested() {
+    RequestContext* curr = RequestContext::get();
+    if (curr) {
+      return curr->createChild();
+    }
+    return std::make_shared<RequestContext>();
+  }
+};
+
 } // namespace folly

folly/io/async/test/RequestContextTest.cpp

@@ -47,11 +47,9 @@ class TestData : public RequestData {
 TEST(RequestContext, SimpleTest) {
   EventBase base;
 
-
   // There should always be a default context with get()
   EXPECT_TRUE(RequestContext::get() != nullptr);
 
-
   // but not with saveContext()
   EXPECT_EQ(RequestContext::saveContext(), nullptr);
   RequestContext::create();
@@ -62,20 +60,19 @@ TEST(RequestContext, SimpleTest) {
   EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
 
   RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
-  base.runInEventBaseThread([&](){
-      EXPECT_TRUE(RequestContext::get() != nullptr);
-      auto data = dynamic_cast<TestData*>(
-        RequestContext::get()->getContextData("test"))->data_;
-      EXPECT_EQ(10, data);
-      base.terminateLoopSoon();
-    });
-  auto th = std::thread([&](){
-      base.loopForever();
+  base.runInEventBaseThread([&]() {
+    EXPECT_TRUE(RequestContext::get() != nullptr);
+    auto data =
+        dynamic_cast<TestData*>(RequestContext::get()->getContextData("test"))
+            ->data_;
+    EXPECT_EQ(10, data);
+    base.terminateLoopSoon();
   });
+  auto th = std::thread([&]() { base.loopForever(); });
   th.join();
   EXPECT_TRUE(RequestContext::get() != nullptr);
-  auto a = dynamic_cast<TestData*>(
-    RequestContext::get()->getContextData("test"));
+  auto a =
+      dynamic_cast<TestData*>(RequestContext::get()->getContextData("test"));
   auto data = a->data_;
   EXPECT_EQ(10, data);
 
@@ -90,15 +87,17 @@ TEST(RequestContext, setIfAbsentTest) {
   RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
   EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
       "test", std::make_unique<TestData>(20)));
-  EXPECT_EQ(10,
-            dynamic_cast<TestData*>(
-                RequestContext::get()->getContextData("test"))->data_);
+  EXPECT_EQ(
+      10,
+      dynamic_cast<TestData*>(RequestContext::get()->getContextData("test"))
+          ->data_);
 
   EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
       "test2", std::make_unique<TestData>(20)));
-  EXPECT_EQ(20,
-            dynamic_cast<TestData*>(
-                RequestContext::get()->getContextData("test2"))->data_);
+  EXPECT_EQ(
+      20,
+      dynamic_cast<TestData*>(RequestContext::get()->getContextData("test2"))
+          ->data_);
 
   RequestContext::setContext(std::shared_ptr<RequestContext>());
   EXPECT_TRUE(nullptr != RequestContext::get());
@@ -154,3 +153,48 @@ TEST(RequestContext, deadlockTest) {
       "test", std::make_unique<DeadlockTestData>("test2"));
   RequestContext::get()->clearContextData("test");
 }
+
+TEST(RequestContext, Nested) {
+  class NestableRequestData : public RequestData {
+   public:
+    explicit NestableRequestData() {}
+    explicit NestableRequestData(int val) : val_(val) {}
+
+    bool hasCallback() override {
+      return false;
+    }
+
+    std::unique_ptr<RequestData> createChild() override {
+      return std::make_unique<NestableRequestData>(val_ + 1);
+    }
+
+    int val_{0};
+
+    static NestableRequestData* get() {
+      return static_cast<NestableRequestData*>(
+          RequestContext::get()->getContextData("nestable"));
+    }
+  };
+
+  RootRequestContextGuard root;
+  RequestContext::get()->setContextData(
+      "nestable", make_unique<NestableRequestData>());
+  RequestContext::get()->setContextData(
+      "unnestable", make_unique<TestData>(42));
+  EXPECT_EQ(NestableRequestData::get()->val_, 0);
+  EXPECT_NE(RequestContext::get()->getContextData("unnestable"), nullptr);
+  EXPECT_NE(NestableRequestData::get(), nullptr);
+  {
+    NestedRequestContextGuard nested1;
+    EXPECT_EQ(NestableRequestData::get()->val_, 1);
+    EXPECT_EQ(RequestContext::get()->getContextData("unnestable"), nullptr);
+    NestedRequestContextGuard nested2;
+    EXPECT_EQ(NestableRequestData::get()->val_, 2);
+    RootRequestContextGuard newRoot;
+    EXPECT_EQ(NestableRequestData::get(), nullptr);
+  }
+  {
+    NestedRequestContextGuard nested1;
+    EXPECT_EQ(NestableRequestData::get()->val_, 1);
+  }
+}