Avoid set/onset calls while shallow copying RequestContext

Summary:
Per title
I found that it was easier to bypass setContext (and not extend RequestContextScopeGuard) to achieve this result.

setShallowCopyContext now directly set the copy as current
the new unsetShallowCopyContext exclusively calls set/onset for the context data that was overriden. This assumes that cost will be fine because the callbackData_ sets are small or empty.
Similar reason they were split from requestData_ in the first place, for RequestContextScopeGuard efficiency.

(Note: this ignores all push blocking failures!)

Reviewed By: andriigrynenko

Differential Revision: D8943668

fbshipit-source-id: ad99177429de6b7b65bf95fa0e94334d73c3750d

folly/io/async/Request.cpp

@@ -177,21 +177,59 @@ std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
   return SingletonT::get();
 }
 
-/* static */ std::shared_ptr<RequestContext> RequestContext::shallowCopy() {
+/* static */ std::shared_ptr<RequestContext>
+RequestContext::setShallowCopyContext() {
   auto* parent = get();
   auto child = std::make_shared<RequestContext>();
-  if (parent) {
-    auto rlock = parent->state_.rlock();
-    auto wlock = child->state_.wlock();
-    wlock->callbackData_ = rlock->callbackData_;
-    for (const auto& entry : rlock->requestData_) {
-      wlock->requestData_[entry.first] =
+
+  {
+    auto parentLock = parent->state_.rlock();
+    auto childLock = child->state_.wlock();
+    childLock->callbackData_ = parentLock->callbackData_;
+    for (const auto& entry : parentLock->requestData_) {
+      childLock->requestData_[entry.first] =
           RequestData::constructPtr(entry.second.get());
     }
   }
+
+  // Do not use setContext to avoid global set/unset
+  std::swap(child, getStaticContext());
   return child;
 }
 
+/* static */ void RequestContext::unsetShallowCopyContext(
+    std::shared_ptr<RequestContext> ctx) {
+  // Do not use setContext to avoid set/unset
+  std::swap(ctx, getStaticContext());
+  // For readability
+  auto child = std::move(ctx);
+  // Handle the case where parent is actually default context
+  auto* parent = get();
+
+  // Call set/unset for all request data that differs
+  {
+    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 (piter != pend || citer != cend) {
+      if (piter == pend || *citer < *piter) {
+        (*citer)->onUnset();
+        ++citer;
+      } else if (citer == cend || *piter < *citer) {
+        (*piter)->onSet();
+        ++piter;
+      } else {
+        DCHECK(*piter == *citer);
+        ++piter;
+        ++citer;
+      }
+    }
+  }
+}
+
 RequestContext* RequestContext::get() {
   auto& context = getStaticContext();
   if (!context) {

folly/io/async/Request.h

@@ -140,17 +140,23 @@ class RequestContext {
  private:
   static std::shared_ptr<RequestContext>& getStaticContext();
 
+  // Start shallow copy guard implementation details:
+  // All methods are private to encourage proper use
   friend struct ShallowCopyRequestContextScopeGuard;
 
-  // Private to encourage using shallow copy guard
-  static std::shared_ptr<RequestContext> shallowCopy();
+  // This sets a shallow copy of the current context as current,
+  // 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).
-  // Private to encourage using shallow copy guard
   void overwriteContextData(
       const std::string& val,
       std::unique_ptr<RequestData> data);
+  // End shallow copy guard
 
   enum class DoSetBehaviour {
     SET,
@@ -165,6 +171,7 @@ class RequestContext {
 
   struct State {
     std::map<std::string, RequestData::SharedPtr> requestData_;
+    // Note: shallow copy relies on this being ordered
     std::set<RequestData*> callbackData_;
   };
   folly::Synchronized<State> state_;
@@ -205,11 +212,11 @@ class RequestContextScopeGuard {
  * This allows to overwrite a specific RequestData pointer for the
  * scope's duration, without breaking others.
  *
- * TODO: currently calls onSet and onUnset incorrectly, will fix on top
+ * Only modified pointers will have their set/onset methods called
  */
-struct ShallowCopyRequestContextScopeGuard : public RequestContextScopeGuard {
+struct ShallowCopyRequestContextScopeGuard {
   ShallowCopyRequestContextScopeGuard()
-      : RequestContextScopeGuard(RequestContext::shallowCopy()) {}
+      : prev_(RequestContext::setShallowCopyContext()) {}
 
   /**
    * Shallow copy then overwrite one specific RequestData
@@ -223,6 +230,22 @@ struct ShallowCopyRequestContextScopeGuard : public RequestContextScopeGuard {
       : ShallowCopyRequestContextScopeGuard() {
     RequestContext::get()->overwriteContextData(val, std::move(data));
   }
+
+  ~ShallowCopyRequestContextScopeGuard() {
+    RequestContext::unsetShallowCopyContext(std::move(prev_));
+  }
+
+  ShallowCopyRequestContextScopeGuard(
+      const ShallowCopyRequestContextScopeGuard&) = delete;
+  ShallowCopyRequestContextScopeGuard& operator=(
+      const ShallowCopyRequestContextScopeGuard&) = delete;
+  ShallowCopyRequestContextScopeGuard(ShallowCopyRequestContextScopeGuard&&) =
+      delete;
+  ShallowCopyRequestContextScopeGuard& operator=(
+      ShallowCopyRequestContextScopeGuard&&) = delete;
+
+ private:
+  std::shared_ptr<RequestContext> prev_;
 };
 
 } // namespace folly

folly/io/async/test/RequestContextTest.cpp

@@ -227,9 +227,8 @@ TEST(RequestContext, ShallowCopyBasic) {
 
   EXPECT_FALSE(hasData());
   EXPECT_EQ(123, getData("immutable").data_);
-  // TODO: Should be 1/0
-  EXPECT_EQ(3, getData("immutable").set_);
-  EXPECT_EQ(2, getData("immutable").unset_);
+  EXPECT_EQ(1, getData("immutable").set_);
+  EXPECT_EQ(0, getData("immutable").unset_);
 }
 
 TEST(RequestContext, ShallowCopyOverwrite) {
@@ -244,9 +243,8 @@ TEST(RequestContext, ShallowCopyOverwrite) {
     EXPECT_EQ(0, getData().unset_);
   }
   EXPECT_EQ(123, getData().data_);
-  // TODO: Should be 2/1
-  EXPECT_EQ(3, getData().set_);
-  EXPECT_EQ(2, getData().unset_);
+  EXPECT_EQ(2, getData().set_);
+  EXPECT_EQ(1, getData().unset_);
 }
 
 TEST(RequestContext, ShallowCopyDefaultContext) {
@@ -275,7 +273,6 @@ TEST(RequestContext, ShallowCopyClear) {
     EXPECT_EQ(789, getData().data_);
   }
   EXPECT_EQ(123, getData().data_);
-  // TODO: Should be 2/1
-  EXPECT_EQ(3, getData().set_);
-  EXPECT_EQ(2, getData().unset_);
+  EXPECT_EQ(2, getData().set_);
+  EXPECT_EQ(1, getData().unset_);
 }