From 35a8cc207284214b981d289469b83c83780c9f54 Mon Sep 17 00:00:00 2001
From: Mirek Klimos <miro@fb.com>
Date: Mon, 18 Apr 2016 14:28:06 -0700
Subject: [PATCH] API to set folly::RequestContext for current scope

Summary:There're currently two ways to set RequestContext
- RequestContext::create() - creates new context and sets it
- RequestContext::setContext(context) - sets context previously captured by saveContext
In most cases, the RequestContext is set back after the request is processed but sometimes it's not (especially with RequestContext::create). We want to measure cpu time for a request by measuring the total cpu time when a RequestContext is set, so we need to make sure that it's properly reset after the thread is done with the request. Scope guards can help us with that.

Reviewed By: haijunz

Differential Revision: D3156698

fb-gh-sync-id: cfb678531810e8be5faaf02cb7803bd247557e42
fbshipit-source-id: cfb678531810e8be5faaf02cb7803bd247557e42
---
 folly/io/async/Request.h | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index 112f0fa8..5c1b9699 100644
--- a/folly/io/async/Request.h
+++ b/folly/io/async/Request.h
@@ -111,12 +111,15 @@ class RequestContext {
   }
 
   // The following API is used to pass the context through queues / threads.
-  // saveContext is called to geta shared_ptr to the context, and
+  // 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.
   //
+  // Whenever possible, use RequestContextScopeGuard instead of setContext
+  // to make sure that RequestContext is reset to the original value when
+  // we exit the scope.
+  //
   // A shared_ptr is used, because many request may fan out across
   // multiple threads, or do post-send processing, etc.
-
   static std::shared_ptr<RequestContext>
   setContext(std::shared_ptr<RequestContext> ctx) {
     using std::swap;
@@ -135,4 +138,24 @@ class RequestContext {
   std::map<std::string, std::unique_ptr<RequestData>> data_;
 };
 
+class RequestContextScopeGuard {
+ private:
+  std::shared_ptr<RequestContext> prev_;
+
+ public:
+  // Create a new RequestContext and reset to the original value when
+  // this goes out of scope.
+  RequestContextScopeGuard() : prev_(RequestContext::saveContext()) {
+    RequestContext::create();
+  }
+
+  // Set a RequestContext that was previously captured by saveContext(). It will
+  // be automatically reset to the original value when this goes out of scope.
+  explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
+      : prev_(RequestContext::setContext(std::move(ctx))) {}
+
+  ~RequestContextScopeGuard() {
+    RequestContext::setContext(std::move(prev_));
+  }
+};
 }
-- 
2.34.1