From: Michael Bejda <mibpl@fb.com>
Date: Mon, 25 Jan 2016 17:45:49 +0000 (-0800)
Subject: Thread-safe RequestContext putIfAbsent operation
X-Git-Tag: deprecate-dynamic-initializer~143
X-Git-Url: http://demsky.eecs.uci.edu/git/?a=commitdiff_plain;h=1721cad02dfff0e9af2e3d79e370a84d5c1f7295;p=folly.git

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 a7c7bef9..24245e5c 100644
--- 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 529a15af..8e2a12ff 100644
--- 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]);