access pointers folly::RequestContext currently running on all threads in the process

Summary:
What:
exposed an API `folly::RequestContext::getRootsFromAllThreads` that reports all "root" RequestContexts for all threads in the process.
(where "root" is defined as the original RequestContext created for the request but which might have been overridden with `ShallowCopyRequestContextScopeGuard)`

Reviewed By: andriigrynenko

Differential Revision: D18849187

fbshipit-source-id: f3c67becf1e38fb8d2198c81ed0ce77bac860d63

folly/io/async/Request.cpp

@@ -32,6 +32,11 @@ DEFINE_bool(
 
 namespace folly {
 
+namespace {
+using SingletonT =
+    SingletonThreadLocal<RequestContext::StaticContext, RequestContext>;
+}
+
 RequestToken::RequestToken(const std::string& str) {
   auto& cache = getCache();
   {
@@ -455,8 +460,26 @@ RequestContext::StateHazptr::Combined* RequestContext::StateHazptr::expand(
   return new Combined(dataCapacity, callbackCapacity, *c);
 }
 
-RequestContext::RequestContext() {
-  useHazptr_ = FLAGS_reqctx_use_hazptr;
+RequestContext::RequestContext()
+    : useHazptr_(FLAGS_reqctx_use_hazptr),
+      rootId_(reinterpret_cast<intptr_t>(this)) {}
+
+RequestContext::RequestContext(const RequestContext& ctx, RootTag)
+    : RequestContext(ctx) {
+  rootId_ = reinterpret_cast<intptr_t>(this);
+}
+
+RequestContext::RequestContext(const RequestContext& ctx, ChildTag)
+    : RequestContext(ctx) {}
+
+/* static */ std::shared_ptr<RequestContext> RequestContext::copyAsRoot(
+    const RequestContext& ctx) {
+  return std::make_shared<RequestContext>(ctx, RootTag{});
+}
+
+/* static */ std::shared_ptr<RequestContext> RequestContext::copyAsChild(
+    const RequestContext& ctx) {
+  return std::make_shared<RequestContext>(ctx, ChildTag{});
 }
 
 bool RequestContext::doSetContextDataLock(
@@ -647,17 +670,20 @@ void exec_set_difference(const TData& data, const TData& other, TExec&& exec) {
   auto newCtx = std::move(newCtx_); // enforce that it is really moved-from
 
   auto& staticCtx = getStaticContext();
-  if (newCtx == staticCtx) {
+  if (newCtx == staticCtx.first) {
     return newCtx;
   }
 
   FOLLY_SDT(
-      folly, request_context_switch_before, staticCtx.get(), newCtx.get());
+      folly,
+      request_context_switch_before,
+      staticCtx.first.get(),
+      newCtx.get());
 
   if ((newCtx.get() && newCtx->useHazptr()) ||
-      (staticCtx.get() && staticCtx->useHazptr())) {
+      (staticCtx.first.get() && staticCtx.first->useHazptr())) {
     DCHECK(!newCtx.get() || newCtx->useHazptr());
-    DCHECK(!staticCtx.get() || staticCtx->useHazptr());
+    DCHECK(!staticCtx.first.get() || staticCtx.first->useHazptr());
     return RequestContext::setContextHazptr(newCtx, staticCtx);
   } else {
     return RequestContext::setContextLock(newCtx, staticCtx);
@@ -667,43 +693,47 @@ void exec_set_difference(const TData& data, const TData& other, TExec&& exec) {
 FOLLY_ALWAYS_INLINE
 /* static */ std::shared_ptr<RequestContext> RequestContext::setContextLock(
     std::shared_ptr<RequestContext>& newCtx,
-    std::shared_ptr<RequestContext>& staticCtx) {
+    StaticContext& staticCtx) {
   auto curCtx = staticCtx;
-  if (newCtx && curCtx) {
+  if (newCtx && curCtx.first) {
     // Only call set/unset for all request data that differs
     auto ret = folly::acquireLocked(
-        as_const(newCtx->state_), as_const(curCtx->state_));
+        as_const(newCtx->state_), as_const(curCtx.first->state_));
     auto& newLock = std::get<0>(ret);
     auto& curLock = std::get<1>(ret);
     auto& newData = newLock->callbackData_;
     auto& curData = curLock->callbackData_;
     exec_set_difference(
         curData, newData, [](RequestData* data) { data->onUnset(); });
-    staticCtx = newCtx;
+    staticCtx.first = newCtx;
+    staticCtx.second = newCtx->rootId_;
     exec_set_difference(
         newData, curData, [](RequestData* data) { data->onSet(); });
   } else {
-    if (curCtx) {
-      curCtx->onUnset();
+    if (curCtx.first) {
+      curCtx.first->onUnset();
     }
-    staticCtx = newCtx;
+    staticCtx.first = newCtx;
     if (newCtx) {
+      staticCtx.second = newCtx->rootId_;
       newCtx->onSet();
+    } else {
+      staticCtx.second = 0;
     }
   }
-  return curCtx;
+  return curCtx.first;
 }
 
 FOLLY_ALWAYS_INLINE
 /* static */ std::shared_ptr<RequestContext> RequestContext::setContextHazptr(
     std::shared_ptr<RequestContext>& newCtx,
-    std::shared_ptr<RequestContext>& staticCtx) {
+    StaticContext& staticCtx) {
   auto curCtx = std::move(staticCtx);
-  bool checkCur = curCtx && curCtx->stateHazptr_.combined();
+  bool checkCur = curCtx.first && curCtx.first->stateHazptr_.combined();
   bool checkNew = newCtx && newCtx->stateHazptr_.combined();
   if (checkCur && checkNew) {
     hazptr_array<2> h;
-    auto curc = h[0].get_protected(curCtx->stateHazptr_.combined_);
+    auto curc = h[0].get_protected(curCtx.first->stateHazptr_.combined_);
     auto newc = h[1].get_protected(newCtx->stateHazptr_.combined_);
     auto& curcb = curc->callbackData_;
     auto& newcb = newc->callbackData_;
@@ -714,7 +744,8 @@ FOLLY_ALWAYS_INLINE
         data->onUnset();
       }
     }
-    staticCtx = std::move(newCtx);
+    staticCtx.first = std::move(newCtx);
+    staticCtx.second = staticCtx.first->rootId_;
     for (auto it = newcb.begin(); it != newcb.end(); ++it) {
       DCHECK(it.key());
       auto data = it.key();
@@ -723,34 +754,48 @@ FOLLY_ALWAYS_INLINE
       }
     }
   } else {
-    if (curCtx) {
-      curCtx->stateHazptr_.onUnset();
+    if (curCtx.first) {
+      curCtx.first->stateHazptr_.onUnset();
     }
-    staticCtx = std::move(newCtx);
-    if (staticCtx) {
-      staticCtx->stateHazptr_.onSet();
+    staticCtx.first = std::move(newCtx);
+    if (staticCtx.first) {
+      staticCtx.first->stateHazptr_.onSet();
+      staticCtx.second = staticCtx.first->rootId_;
+    } else {
+      staticCtx.second = 0;
     }
   }
-  return curCtx;
+  return curCtx.first;
 }
 
-std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
-  using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
+RequestContext::StaticContext& RequestContext::getStaticContext() {
   return SingletonT::get();
 }
 
+/* static */ std::vector<intptr_t> RequestContext::getRootIdsFromAllThreads() {
+  std::vector<intptr_t> result;
+  for (const auto& e : SingletonT::accessAllThreads()) {
+    result.push_back(e.second);
+  }
+  return result;
+}
+
 /* static */ std::shared_ptr<RequestContext>
 RequestContext::setShallowCopyContext() {
-  auto& parent = getStaticContext();
-  auto child = parent ? std::make_shared<RequestContext>(*parent)
+  auto& parent = getStaticContext().first;
+  auto child = parent ? RequestContext::copyAsChild(*parent)
                       : std::make_shared<RequestContext>();
+  if (!parent) {
+    child->rootId_ = 0;
+  }
   // Do not use setContext to avoid global set/unset
+  // Also rootId does not change so do not bother setting it.
   std::swap(child, parent);
   return child;
 }
 
 RequestContext* RequestContext::get() {
-  auto& context = getStaticContext();
+  auto& context = getStaticContext().first;
   if (!context) {
     static RequestContext defaultContext;
     return std::addressof(defaultContext);

folly/io/async/Request.h

@@ -148,6 +148,13 @@ class RequestData {
 class RequestContext {
  public:
   RequestContext();
+  RequestContext(RequestContext&& ctx) = delete;
+  RequestContext& operator=(const RequestContext&) = delete;
+  RequestContext& operator=(RequestContext&&) = delete;
+
+  // copy ctor is disabled, use copyAsRoot/copyAsChild instead.
+  static std::shared_ptr<RequestContext> copyAsRoot(const RequestContext& ctx);
+  static std::shared_ptr<RequestContext> copyAsChild(const RequestContext& ctx);
 
   // Create a unique request context for this request.
   // It will be passed between queues / threads (where implemented),
@@ -159,6 +166,12 @@ class RequestContext {
   // Get the current context.
   static RequestContext* get();
 
+  intptr_t getRootId() const {
+    return rootId_;
+  }
+
+  static std::vector<intptr_t> getRootIdsFromAllThreads();
+
   // The following APIs are used to add, remove and access RequestData instance
   // in the RequestContext instance, normally used for per-RequestContext
   // tracking or callback on set and unset. These APIs are Thread-safe.
@@ -238,18 +251,28 @@ class RequestContext {
       std::shared_ptr<RequestContext>&& ctx);
 
   static std::shared_ptr<RequestContext> saveContext() {
-    return getStaticContext();
+    return getStaticContext().first;
   }
 
  private:
-  static std::shared_ptr<RequestContext>& getStaticContext();
+  struct ChildTag {};
+  struct RootTag {};
+  RequestContext(const RequestContext& ctx) = default;
+
+ public:
+  RequestContext(const RequestContext& ctx, RootTag tag);
+  RequestContext(const RequestContext& ctx, ChildTag tag);
+  using StaticContext = std::pair<std::shared_ptr<RequestContext>, intptr_t>;
+
+ private:
+  static StaticContext& getStaticContext();
 
   static std::shared_ptr<RequestContext> setContextLock(
       std::shared_ptr<RequestContext>& newCtx,
-      std::shared_ptr<RequestContext>& staticCtx);
+      StaticContext& staticCtx);
   static std::shared_ptr<RequestContext> setContextHazptr(
       std::shared_ptr<RequestContext>& newCtx,
-      std::shared_ptr<RequestContext>& staticCtx);
+      StaticContext& staticCtx);
 
   // Start shallow copy guard implementation details:
   // All methods are private to encourage proper use
@@ -388,6 +411,8 @@ class RequestContext {
   }; // StateHazptr
   StateHazptr stateHazptr_;
   bool useHazptr_;
+  // Shallow copies keep a note of the root context
+  intptr_t rootId_;
 };
 
 /**

folly/io/async/test/RequestContextTest.cpp

@@ -82,8 +82,13 @@ TEST_F(RequestContextTest, SimpleTest) {
   EXPECT_EQ(RequestContext::saveContext(), nullptr);
   RequestContext::create();
   EXPECT_NE(RequestContext::saveContext(), nullptr);
+  auto rootids = RequestContext::getRootIdsFromAllThreads();
+  EXPECT_EQ(1, rootids.size());
+  EXPECT_EQ(RequestContext::get()->getRootId(), rootids[0]);
+  EXPECT_EQ(reinterpret_cast<intptr_t>(RequestContext::get()), rootids[0]);
   RequestContext::create();
   EXPECT_NE(RequestContext::saveContext(), nullptr);
+  EXPECT_NE(RequestContext::get()->getRootId(), rootids[0]);
 
   EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
 
@@ -94,6 +99,10 @@ TEST_F(RequestContextTest, SimpleTest) {
                     RequestContext::get()->getContextData(testtoken))
                     ->data_;
     EXPECT_EQ(10, data);
+    rootids = RequestContext::getRootIdsFromAllThreads();
+    EXPECT_EQ(2, rootids.size());
+    EXPECT_EQ(RequestContext::get()->getRootId(), rootids[0]);
+    EXPECT_EQ(RequestContext::get()->getRootId(), rootids[1]);
     base.terminateLoopSoon();
   });
   auto th = std::thread([&]() { base.loopForever(); });
@@ -236,14 +245,23 @@ TEST_F(RequestContextTest, sharedGlobalTest) {
     }
   };
 
-  RequestContextScopeGuard g0;
-  RequestContext::get()->setContextData(
-      "test", std::make_unique<GlobalTestData>());
+  intptr_t root = 0;
   {
-    RequestContextScopeGuard g1;
+    RequestContextScopeGuard g0;
     RequestContext::get()->setContextData(
         "test", std::make_unique<GlobalTestData>());
+    auto root0 = RequestContext::saveContext().get()->getRootId();
+    EXPECT_EQ(RequestContext::getRootIdsFromAllThreads()[0], root0);
+    {
+      RequestContextScopeGuard g1;
+      RequestContext::get()->setContextData(
+          "test", std::make_unique<GlobalTestData>());
+      auto root1 = RequestContext::saveContext().get()->getRootId();
+      EXPECT_EQ(RequestContext::getRootIdsFromAllThreads()[0], root1);
+    }
+    EXPECT_EQ(RequestContext::getRootIdsFromAllThreads()[0], root0);
   }
+  EXPECT_EQ(RequestContext::getRootIdsFromAllThreads().front(), root);
 }
 
 TEST_F(RequestContextTest, ShallowCopyBasic) {
@@ -251,34 +269,49 @@ TEST_F(RequestContextTest, ShallowCopyBasic) {
   setData(123, "immutable");
   EXPECT_EQ(123, getData("immutable").data_);
   EXPECT_FALSE(hasData());
+  EXPECT_EQ(0, RequestContext::getRootIdsFromAllThreads()[0]);
 
   {
     ShallowCopyRequestContextScopeGuard g1;
     EXPECT_EQ(123, getData("immutable").data_);
     setData(789);
     EXPECT_EQ(789, getData().data_);
+    EXPECT_EQ(0, RequestContext::getRootIdsFromAllThreads()[0]);
   }
 
   EXPECT_FALSE(hasData());
   EXPECT_EQ(123, getData("immutable").data_);
   EXPECT_EQ(1, getData("immutable").set_);
   EXPECT_EQ(0, getData("immutable").unset_);
+  EXPECT_EQ(0, RequestContext::getRootIdsFromAllThreads()[0]);
 }
 
 TEST_F(RequestContextTest, ShallowCopyOverwrite) {
   RequestContextScopeGuard g0;
   setData(123);
   EXPECT_EQ(123, getData().data_);
+  auto rootid = RequestContext::get()->getRootId();
+  EXPECT_EQ(rootid, RequestContext::getRootIdsFromAllThreads()[0]);
   {
     ShallowCopyRequestContextScopeGuard g1(
         "test", std::make_unique<TestData>(789));
     EXPECT_EQ(789, getData().data_);
     EXPECT_EQ(1, getData().set_);
     EXPECT_EQ(0, getData().unset_);
+    // should have inherited parent's rootid
+    EXPECT_EQ(rootid, RequestContext::getRootIdsFromAllThreads()[0]);
+
+    {
+      // rootId is preserved for shallow copies of shallow copies
+      ShallowCopyRequestContextScopeGuard g2;
+      EXPECT_EQ(rootid, RequestContext::getRootIdsFromAllThreads()[0]);
+    }
+    EXPECT_EQ(rootid, RequestContext::getRootIdsFromAllThreads()[0]);
   }
   EXPECT_EQ(123, getData().data_);
   EXPECT_EQ(2, getData().set_);
   EXPECT_EQ(1, getData().unset_);
+  EXPECT_EQ(rootid, RequestContext::getRootIdsFromAllThreads()[0]);
 }
 
 TEST_F(RequestContextTest, ShallowCopyDefaultContext) {
@@ -310,3 +343,16 @@ TEST_F(RequestContextTest, ShallowCopyClear) {
   EXPECT_EQ(2, getData().set_);
   EXPECT_EQ(1, getData().unset_);
 }
+
+TEST_F(RequestContextTest, RootIdOnCopy) {
+  auto ctxBase = std::make_shared<RequestContext>();
+  EXPECT_EQ(reinterpret_cast<intptr_t>(ctxBase.get()), ctxBase->getRootId());
+  {
+    auto ctx = RequestContext::copyAsRoot(*ctxBase);
+    EXPECT_EQ(reinterpret_cast<intptr_t>(ctx.get()), ctx->getRootId());
+  }
+  {
+    auto ctx = RequestContext::copyAsChild(*ctxBase);
+    EXPECT_EQ(reinterpret_cast<intptr_t>(ctxBase.get()), ctx->getRootId());
+  }
+}