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

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

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

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) {

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;