From 6645b95d6402b48e7b3abdc01146a7a61a9378c9 Mon Sep 17 00:00:00 2001
From: Melanie Subbiah <melaniesubbiah@fb.com>
Date: Fri, 22 Jul 2016 15:01:39 -0700
Subject: [PATCH] RequestContext changes to facilitate thread tracing

Summary: In order to trigger trace points when a new thread begins/finishes executing, I added functionality to call onSet and onUnset methods on all RequestData objects whenever setContext is called. The main question: will this approach cost us too much overhead?

Reviewed By: djwatson

Differential Revision: D3604948

fbshipit-source-id: 3b704ca0f2b713458427aa49be12f776939057f8
---
 folly/io/async/Request.h | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index 5c1b9699..7adb1a9c 100644
--- a/folly/io/async/Request.h
+++ b/folly/io/async/Request.h
@@ -34,6 +34,8 @@ namespace folly {
 class RequestData {
  public:
   virtual ~RequestData() = default;
+  virtual void onSet(){};
+  virtual void onUnset(){};
 };
 
 class RequestContext;
@@ -43,7 +45,7 @@ class RequestContext;
 // copied between threads.
 class RequestContext {
  public:
-  // Create a unique requext context for this request.
+  // Create a unique request context for this request.
   // It will be passed between queues / threads (where implemented),
   // so it should be valid for the lifetime of the request.
   static void create() {
@@ -105,6 +107,24 @@ class RequestContext {
     }
   }
 
+  void onSet() {
+    folly::RWSpinLock::ReadHolder guard(lock);
+    for (auto const& ent : data_) {
+      if (RequestData* data = ent.second.get()) {
+        data->onSet();
+      }
+    }
+  }
+
+  void onUnset() {
+    folly::RWSpinLock::ReadHolder guard(lock);
+    for (auto const& ent : data_) {
+      if (RequestData* data = ent.second.get()) {
+        data->onUnset();
+      }
+    }
+  }
+
   void clearContextData(const std::string& val) {
     folly::RWSpinLock::WriteHolder guard(lock);
     data_.erase(val);
@@ -122,8 +142,17 @@ class RequestContext {
   // multiple threads, or do post-send processing, etc.
   static std::shared_ptr<RequestContext>
   setContext(std::shared_ptr<RequestContext> ctx) {
-    using std::swap;
-    swap(ctx, getStaticContext());
+    auto& prev = getStaticContext();
+    if (ctx != prev) {
+      using std::swap;
+      if (ctx) {
+        ctx->onSet();
+      }
+      if (prev) {
+        prev->onUnset();
+      }
+      swap(ctx, prev);
+    }
     return ctx;
   }
 
-- 
2.34.1