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

diff --git a/folly/futures/detail/Core.h b/folly/futures/detail/Core.h
index b2fb18d82baa399c07d76d753c5c439499f53da7..d3dba861613506f7727f627477037c9125cfa1f9 100644 (file)
--- 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 e4375f41ca130af2dc24f8f897ec77abf0e97613..aa2e2837e07c22516e91b4caa7698aa0c5080fab 100644 (file)
--- 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 f6c1ce8a0b25767323a16ae5c6c4268a1912ffc9..20b9da8ee672a3bb8a09636533789b1cc4738be9 100644 (file)
--- 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 acbd3c37c3ae118e12c1b843c731db40e6bbc71e..2a12980ead4f8c169b5b5bf5777964488c22ba7b 100644 (file)
--- 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;