Thread-safe RequestContext putIfAbsent operation

Summary:
Adds a thread-safe putIfAbsent operation to the RequestContext. The current setContextData() is not sufficent to do it safely.
Just like setContextData, this method is unfair, as a high volume of reads will block the spinlock.

Reviewed By: fugalh

Differential Revision: D2850752

fb-gh-sync-id: 2ff22ea9e9bd8f27f6ae7a57214a6dbc4fdcd4c5

folly/io/async/Request.h

@@ -76,6 +76,20 @@ class RequestContext {
     }
   }
 
+  // Unlike setContextData, this method does not panic if the key is already
+  // present. Returns true iff the new value has been inserted.
+  bool setContextDataIfAbsent(const std::string& val,
+                              std::unique_ptr<RequestData> data) {
+    folly::RWSpinLock::UpgradedHolder guard(lock);
+    if (data_.find(val) != data_.end()) {
+      return false;
+    }
+
+    folly::RWSpinLock::WriteHolder writeGuard(std::move(guard));
+    data_[val] = std::move(data);
+    return true;
+  }
+
   bool hasContextData(const std::string& val) {
     folly::RWSpinLock::ReadHolder guard(lock);
     return data_.find(val) != data_.end();

folly/io/async/test/RequestContextTest.cpp

@@ -73,6 +73,27 @@ TEST(RequestContext, SimpleTest) {
   EXPECT_TRUE(nullptr != RequestContext::get());
 }
 
+TEST(RequestContext, setIfAbsentTest) {
+  EXPECT_TRUE(RequestContext::get() != nullptr);
+
+  RequestContext::get()->setContextData(
+      "test", std::unique_ptr<TestData>(new TestData(10)));
+  EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
+      "test", std::unique_ptr<TestData>(new TestData(20))));
+  EXPECT_EQ(10,
+            dynamic_cast<TestData*>(
+                RequestContext::get()->getContextData("test"))->data_);
+
+  EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
+      "test2", std::unique_ptr<TestData>(new TestData(20))));
+  EXPECT_EQ(20,
+            dynamic_cast<TestData*>(
+                RequestContext::get()->getContextData("test2"))->data_);
+
+  RequestContext::setContext(std::shared_ptr<RequestContext>());
+  EXPECT_TRUE(nullptr != RequestContext::get());
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::InitGoogleLogging(argv[0]);