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

@@ -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();
     }
 

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();
 

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());
+}