Merge setContext and unsetShallowCopyContext

Summary:
There is no reason for setContext to call set/unset when both request context holds identical data.
So merge unsetShallowCopyContext implementation there.

Reviewed By: djwatson

Differential Revision: D9119718

fbshipit-source-id: 7b9c6badbfb29416b6b0032467984720b697a1e8

folly/io/async/Request.cpp

@@ -157,19 +157,47 @@ void RequestContext::clearContextData(const std::string& val) {
 }
 
 std::shared_ptr<RequestContext> RequestContext::setContext(
-    std::shared_ptr<RequestContext> ctx) {
+    std::shared_ptr<RequestContext> newCtx) {
   auto& curCtx = getStaticContext();
-  if (ctx != curCtx) {
-    FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
-    if (curCtx) {
+  if (newCtx != curCtx) {
+    FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), newCtx.get());
+
+    // Call set/unset for all request data that differs
+    if (newCtx && curCtx) {
+      auto newLock = newCtx->state_.rlock();
+      auto curLock = curCtx->state_.rlock();
+      auto niter = newLock->callbackData_.begin();
+      auto nend = newLock->callbackData_.end();
+      auto citer = curLock->callbackData_.begin();
+      auto cend = curLock->callbackData_.end();
+      while (true) {
+        if (niter == nend) {
+          if (citer == cend) {
+            break;
+          }
+          (*citer)->onUnset();
+          ++citer;
+        } else if (citer == cend || *niter < *citer) {
+          (*niter)->onSet();
+          ++niter;
+        } else if (*citer < *niter) {
+          (*citer)->onUnset();
+          ++citer;
+        } else {
+          DCHECK(*niter == *citer);
+          ++niter;
+          ++citer;
+        }
+      }
+    } else if (newCtx) {
+      newCtx->onSet();
+    } else if (curCtx) {
       curCtx->onUnset();
     }
-    std::swap(ctx, curCtx);
-    if (curCtx) {
-      curCtx->onSet();
-    }
+
+    std::swap(curCtx, newCtx);
   }
-  return ctx;
+  return newCtx;
 }
 
 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
@@ -197,47 +225,6 @@ RequestContext::setShallowCopyContext() {
   return child;
 }
 
-/* static */ void RequestContext::unsetShallowCopyContext(
-    std::shared_ptr<RequestContext> parent) {
-  auto& child = getStaticContext();
-
-  // Call set/unset for all request data that differs
-  if (parent && child) {
-    auto parentLock = parent->state_.rlock();
-    auto childLock = child->state_.rlock();
-    auto piter = parentLock->callbackData_.begin();
-    auto pend = parentLock->callbackData_.end();
-    auto citer = childLock->callbackData_.begin();
-    auto cend = childLock->callbackData_.end();
-    while (true) {
-      if (piter == pend) {
-        if (citer == cend) {
-          break;
-        }
-        (*citer)->onUnset();
-        ++citer;
-      } else if (citer == cend || *piter < *citer) {
-        (*piter)->onSet();
-        ++piter;
-      } else if (*citer < *piter) {
-        (*citer)->onUnset();
-        ++citer;
-      } else {
-        DCHECK(*piter == *citer);
-        ++piter;
-        ++citer;
-      }
-    }
-  } else if (parent) {
-    parent->onSet();
-  } else if (child) {
-    child->onUnset();
-  }
-
-  // Do not use setContext to avoid set/unset
-  child = parent;
-}
-
 RequestContext* RequestContext::get() {
   auto& context = getStaticContext();
   if (!context) {

folly/io/async/Request.h

@@ -148,9 +148,6 @@ class RequestContext {
   // then return the previous context (so it can be reset later).
   static std::shared_ptr<RequestContext> setShallowCopyContext();
 
-  // Reset the previously copied parent context
-  static void unsetShallowCopyContext(std::shared_ptr<RequestContext> ctx);
-
   // Similar to setContextData, except it overwrites the data
   // if already set (instead of warn + reset ptr).
   void overwriteContextData(
@@ -171,7 +168,7 @@ class RequestContext {
 
   struct State {
     std::map<std::string, RequestData::SharedPtr> requestData_;
-    // Note: shallow copy relies on this being ordered
+    // Note: setContext efficiency relies on this being ordered
     std::set<RequestData*> callbackData_;
   };
   folly::Synchronized<State> state_;
@@ -232,7 +229,7 @@ struct ShallowCopyRequestContextScopeGuard {
   }
 
   ~ShallowCopyRequestContextScopeGuard() {
-    RequestContext::unsetShallowCopyContext(std::move(prev_));
+    RequestContext::setContext(std::move(prev_));
   }
 
   ShallowCopyRequestContextScopeGuard(