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

diff --git a/folly/io/async/EventBase.cpp b/folly/io/async/EventBase.cpp
index db2efcef3209fae8ebf8489e0837e3b0d0ac7234..e250c3644bd74a015b09575d6b11a9c4c03eb2cf 100644 (file)
--- 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 ae26628d8ee52bc40401b08144a7a1b250bd3a8c..b450e97fee40762ed82c4ea83cd4797659fc1ee2 100644 (file)
--- 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 0c7b77f0b3497a3b553d481ca0d9ae1d9c82b638..29cf89fdd06e2175f3fb085bab57f60d87c32a81 100644 (file)
--- 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());
+}