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

diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index a7c7bef9a3365710bc07c81d7e6dd51ed8b33f7a..24245e5ca98938b156f5daae5aba9bcd2a347b9c 100644 (file)
--- a/folly/io/async/Request.h
+++ b/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();

diff --git a/folly/io/async/test/RequestContextTest.cpp b/folly/io/async/test/RequestContextTest.cpp
index 529a15aff41be90ff18adf12c13b517c7f87b92f..8e2a12ff6c78a82df1acf827cce8f9719622d5c5 100644 (file)
--- a/folly/io/async/test/RequestContextTest.cpp
+++ b/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]);