From 4d499e087e0445ff71a12eb1d5071cf33529d046 Mon Sep 17 00:00:00 2001
From: Mirek Klimos <miro@fb.com>
Date: Fri, 13 May 2016 12:02:35 -0700
Subject: [PATCH] Replace RequestContext::setContext with
 RequestContextScopeGuard in folly

Summary: To make sure RequestContext is properly unset when we stop processing request on a thread. This changes the API in Fibers, NotificationQueue, HHWheelTimer and AsyncTimeout, and fixes RequestContext handling in Futures (reset RC after the callback is done)

Reviewed By: andriigrynenko

Differential Revision: D3279644

fbshipit-source-id: a6a1c2840cdce179334aa1a3b28fa514cd5358c1
---
 folly/futures/detail/Core.h        | 8 ++++----
 folly/io/async/AsyncTimeout.cpp    | 5 +----
 folly/io/async/HHWheelTimer.cpp    | 4 +---
 folly/io/async/NotificationQueue.h | 5 +----
 4 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/folly/futures/detail/Core.h b/folly/futures/detail/Core.h
index b2fb18d8..d3dba861 100644
--- a/folly/futures/detail/Core.h
+++ b/folly/futures/detail/Core.h
@@ -334,28 +334,28 @@ class Core {
         if (LIKELY(x->getNumPriorities() == 1)) {
           x->add([this]() mutable {
             SCOPE_EXIT { detachOne(); };
-            RequestContext::setContext(context_);
+            RequestContextScopeGuard rctx(context_);
             SCOPE_EXIT { callback_ = {}; };
             callback_(std::move(*result_));
           });
         } else {
           x->addWithPriority([this]() mutable {
             SCOPE_EXIT { detachOne(); };
-            RequestContext::setContext(context_);
+            RequestContextScopeGuard rctx(context_);
             SCOPE_EXIT { callback_ = {}; };
             callback_(std::move(*result_));
           }, priority);
         }
       } catch (...) {
         --attached_; // Account for extra ++attached_ before try
-        RequestContext::setContext(context_);
+        RequestContextScopeGuard rctx(context_);
         result_ = Try<T>(exception_wrapper(std::current_exception()));
         SCOPE_EXIT { callback_ = {}; };
         callback_(std::move(*result_));
       }
     } else {
       SCOPE_EXIT { detachOne(); };
-      RequestContext::setContext(context_);
+      RequestContextScopeGuard rctx(context_);
       SCOPE_EXIT { callback_ = {}; };
       callback_(std::move(*result_));
     }
diff --git a/folly/io/async/AsyncTimeout.cpp b/folly/io/async/AsyncTimeout.cpp
index e4375f41..aa2e2837 100644
--- a/folly/io/async/AsyncTimeout.cpp
+++ b/folly/io/async/AsyncTimeout.cpp
@@ -157,12 +157,9 @@ void AsyncTimeout::libeventCallback(libevent_fd_t fd, short events, void* arg) {
   // this can't possibly fire if timeout->eventBase_ is nullptr
   timeout->timeoutManager_->bumpHandlingTime();
 
-  auto old_ctx =
-    RequestContext::setContext(timeout->context_);
+  RequestContextScopeGuard rctx(timeout->context_);
 
   timeout->timeoutExpired();
-
-  RequestContext::setContext(old_ctx);
 }
 
 } // folly
diff --git a/folly/io/async/HHWheelTimer.cpp b/folly/io/async/HHWheelTimer.cpp
index f6c1ce8a..20b9da8e 100644
--- a/folly/io/async/HHWheelTimer.cpp
+++ b/folly/io/async/HHWheelTimer.cpp
@@ -214,10 +214,8 @@ void HHWheelTimer::timeoutExpired() noexcept {
       count_--;
       cb->wheel_ = nullptr;
       cb->expiration_ = milliseconds(0);
-      auto old_ctx =
-        RequestContext::setContext(cb->context_);
+      RequestContextScopeGuard rctx(cb->context_);
       cb->timeoutExpired();
-      RequestContext::setContext(old_ctx);
     }
   }
   if (count_ > 0) {
diff --git a/folly/io/async/NotificationQueue.h b/folly/io/async/NotificationQueue.h
index acbd3c37..2a12980e 100644
--- a/folly/io/async/NotificationQueue.h
+++ b/folly/io/async/NotificationQueue.h
@@ -705,8 +705,7 @@ void NotificationQueue<MessageT>::Consumer::consumeMessages(
       auto& data = queue_->queue_.front();
 
       MessageT msg(std::move(data.first));
-      auto old_ctx =
-        RequestContext::setContext(data.second);
+      RequestContextScopeGuard rctx(std::move(data.second));
       queue_->queue_.pop_front();
 
       // Check to see if the queue is empty now.
@@ -728,8 +727,6 @@ void NotificationQueue<MessageT>::Consumer::consumeMessages(
       messageAvailable(std::move(msg));
       destroyedFlagPtr_ = nullptr;
 
-      RequestContext::setContext(old_ctx);
-
       // If the callback was destroyed before it returned, we are done
       if (callbackDestroyed) {
         return;
-- 
2.34.1