Support multiple inputs to ShallowCopyScopeGuard

Summary:
With the current constructor setup, overwriting multiple request datas involves chaining together multiple ShallowCopyScopeGuard. This is quite inefficient as it involves copying the requests-data/callbacks maps multiple times.

This diff introduces a new variadic constructor that accepts RequestDataItems (named pair of token and data)

Reviewed By: A5he

Differential Revision: D25317291

fbshipit-source-id: 8a9cb00a615ecf802f22727189e4dc5a0c056c73

folly/io/async/Request.h

@@ -24,6 +24,7 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <utility>
 
 namespace folly {
 
@@ -126,6 +127,8 @@ class RequestData {
   std::atomic<int> keepAliveCounter_{0};
 };
 
+using RequestDataItem = std::pair<RequestToken, std::unique_ptr<RequestData>>;
+
 // If you do not call create() to create a unique request context,
 // this default request context will always be returned, and is never
 // copied between threads.
@@ -417,6 +420,29 @@ struct ShallowCopyRequestContextScopeGuard {
     RequestContext::get()->overwriteContextData(val, std::move(data), true);
   }
 
+  /**
+   * Shallow copy then overwrite multiple RequestData instances
+   *
+   * Helper constructor which is more efficient than using multiple scope guards
+   * Accepts iterators to a container of <string/RequestToken, RequestData
+   * pointer> pairs
+   */
+  template <typename... TItems>
+  explicit ShallowCopyRequestContextScopeGuard(
+      RequestDataItem&& first,
+      TItems&&... rest)
+      : ShallowCopyRequestContextScopeGuard() {
+    auto rc = RequestContext::get();
+    auto overwriteContextData = [&rc](RequestDataItem&& item) {
+      rc->overwriteContextData(item.first, std::move(item.second), true);
+    };
+
+    overwriteContextData(std::move(first));
+
+    using _ = int[];
+    void(_{0, (void(overwriteContextData(std::forward<TItems>(rest))), 0)...});
+  }
+
   ~ShallowCopyRequestContextScopeGuard() {
     RequestContext::setContext(std::move(prev_));
   }

folly/io/async/test/RequestContextTest.cpp

@@ -360,6 +360,30 @@ TEST_F(RequestContextTest, ShallowCopyClear) {
   EXPECT_EQ(1, getData().unset_);
 }
 
+TEST_F(RequestContextTest, ShallowCopyMulti) {
+  RequestContextScopeGuard g0;
+  setData(1, "test1");
+  setData(2, "test2");
+  EXPECT_EQ(1, getData("test1").data_);
+  EXPECT_EQ(2, getData("test2").data_);
+  {
+    ShallowCopyRequestContextScopeGuard g1(
+        RequestDataItem{"test1", std::make_unique<TestData>(2)},
+        RequestDataItem{"test2", std::make_unique<TestData>(4)});
+
+    EXPECT_EQ(2, getData("test1").data_);
+    EXPECT_EQ(4, getData("test2").data_);
+    clearData("test1");
+    clearData("test2");
+    setData(4, "test1");
+    setData(8, "test2");
+    EXPECT_EQ(4, getData("test1").data_);
+    EXPECT_EQ(8, getData("test2").data_);
+  }
+  EXPECT_EQ(1, getData("test1").data_);
+  EXPECT_EQ(2, getData("test2").data_);
+}
+
 TEST_F(RequestContextTest, RootIdOnCopy) {
   auto ctxBase = std::make_shared<RequestContext>(0xab);
   EXPECT_EQ(0xab, ctxBase->getRootId());