From 6e8fcd168c0edf827e6f828bf81f52b661e8d17c Mon Sep 17 00:00:00 2001
From: Mirek Klimos <miro@fb.com>
Date: Wed, 3 Aug 2016 14:40:28 -0700
Subject: [PATCH] Unset RequestContext properly in EventBase::runLoopCallbacks

Summary: We need to make sure RequestContext is unset properly for correct attribution of events to requests in BPF. djwatson, is there a reason not to set RequestContext when running runLoopCallbacks() from EventBase destructor?

Reviewed By: yfeldblum

Differential Revision: D3640289

fbshipit-source-id: bc48e936618adb1a1619de004b8479f58d3b683d
---
 folly/io/async/EventBase.cpp          |  8 +++-----
 folly/io/async/EventBase.h            |  2 +-
 folly/io/async/test/EventBaseTest.cpp | 16 ++++++++++++++++
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/folly/io/async/EventBase.cpp b/folly/io/async/EventBase.cpp
index db2efcef..e250c364 100644
--- a/folly/io/async/EventBase.cpp
+++ b/folly/io/async/EventBase.cpp
@@ -223,7 +223,7 @@ EventBase::~EventBase() {
     delete &runBeforeLoopCallbacks_.front();
   }
 
-  (void) runLoopCallbacks(false);
+  (void)runLoopCallbacks();
 
   if (!fnRunner_->consumeUntilDrained()) {
     LOG(ERROR) << "~EventBase(): Unable to drain notification queue";
@@ -656,7 +656,7 @@ bool EventBase::tryRunAfterDelay(
   return true;
 }
 
-bool EventBase::runLoopCallbacks(bool setContext) {
+bool EventBase::runLoopCallbacks() {
   if (!loopCallbacks_.empty()) {
     bumpHandlingTime();
     // Swap the loopCallbacks_ list with a temporary list on our stack.
@@ -674,9 +674,7 @@ bool EventBase::runLoopCallbacks(bool setContext) {
     while (!currentCallbacks.empty()) {
       LoopCallback* callback = &currentCallbacks.front();
       currentCallbacks.pop_front();
-      if (setContext) {
-        RequestContext::setContext(callback->context_);
-      }
+      folly::RequestContextScopeGuard rctx(callback->context_);
       callback->runLoopCallback();
     }
 
diff --git a/folly/io/async/EventBase.h b/folly/io/async/EventBase.h
index ae26628d..b450e97f 100644
--- a/folly/io/async/EventBase.h
+++ b/folly/io/async/EventBase.h
@@ -670,7 +670,7 @@ class EventBase : private boost::noncopyable,
   bool loopBody(int flags = 0);
 
   // executes any callbacks queued by runInLoop(); returns false if none found
-  bool runLoopCallbacks(bool setContext = true);
+  bool runLoopCallbacks();
 
   void initNotificationQueue();
 
diff --git a/folly/io/async/test/EventBaseTest.cpp b/folly/io/async/test/EventBaseTest.cpp
index 0c7b77f0..29cf89fd 100644
--- a/folly/io/async/test/EventBaseTest.cpp
+++ b/folly/io/async/test/EventBaseTest.cpp
@@ -1853,3 +1853,19 @@ TEST(EventBaseTest, DrivableExecutorTest) {
 
   t.join();
 }
+
+TEST(EventBaseTest, RequestContextTest) {
+  EventBase evb;
+  auto defaultCtx = RequestContext::get();
+
+  {
+    RequestContextScopeGuard rctx;
+    auto context = RequestContext::get();
+    EXPECT_NE(defaultCtx, context);
+    evb.runInLoop([context] { EXPECT_EQ(context, RequestContext::get()); });
+  }
+
+  EXPECT_EQ(defaultCtx, RequestContext::get());
+  evb.loop();
+  EXPECT_EQ(defaultCtx, RequestContext::get());
+}
-- 
2.34.1