Use Synchronized in RequestContext

Summary:
[Folly] Use `Synchronized` in `RequestContext`.

Because we can. And it makes the code a tad simpler and also enforces access correctness a tad.

Also use `folly::make_unique` in `RequestContextTest` to keep balance between the explicit `new` and `delete` ops.

Reviewed By: markisaa

Differential Revision: D3781115

fbshipit-source-id: 63b41ddd8009e9546e3be5f89bdd23a4d791105c

diff --git a/folly/io/async/Request.cpp b/folly/io/async/Request.cpp
index dc46d9d87c841cccec2c9719fa19cab2388a3e6e..8d44a5568a64d551adac4a704ce180ff8465b982 100644 (file)
--- a/folly/io/async/Request.cpp
+++ b/folly/io/async/Request.cpp
@@ -23,6 +23,7 @@
 
 #include <glog/logging.h>
 
+#include <folly/MapUtil.h>
 #include <folly/SingletonThreadLocal.h>
 
 namespace folly {
@@ -30,77 +31,63 @@ namespace folly {
 void RequestContext::setContextData(
     const std::string& val,
     std::unique_ptr<RequestData> data) {
-  folly::RWSpinLock::WriteHolder guard(lock);
-  if (data_.find(val) != data_.end()) {
+  auto wlock = data_.wlock();
+  if (wlock->count(val)) {
     LOG_FIRST_N(WARNING, 1)
         << "Called RequestContext::setContextData with data already set";
 
-    data_[val] = nullptr;
+    (*wlock)[val] = nullptr;
   } else {
-    data_[val] = std::move(data);
+    (*wlock)[val] = std::move(data);
   }
 }
 
 bool RequestContext::setContextDataIfAbsent(
     const std::string& val,
     std::unique_ptr<RequestData> data) {
-  folly::RWSpinLock::UpgradedHolder guard(lock);
-  if (data_.find(val) != data_.end()) {
+  auto ulock = data_.ulock();
+  if (ulock->count(val)) {
     return false;
   }
 
-  folly::RWSpinLock::WriteHolder writeGuard(std::move(guard));
-  data_[val] = std::move(data);
+  auto wlock = ulock.moveFromUpgradeToWrite();
+  (*wlock)[val] = std::move(data);
   return true;
 }
 
 bool RequestContext::hasContextData(const std::string& val) const {
-  folly::RWSpinLock::ReadHolder guard(lock);
-  return data_.find(val) != data_.end();
+  return data_.rlock()->count(val);
 }
 
 RequestData* RequestContext::getContextData(const std::string& val) {
-  folly::RWSpinLock::ReadHolder guard(lock);
-  auto r = data_.find(val);
-  if (r == data_.end()) {
-    return nullptr;
-  } else {
-    return r->second.get();
-  }
+  return get_ref_default(*data_.rlock(), val, nullptr).get();
 }
 
 const RequestData* RequestContext::getContextData(
     const std::string& val) const {
-  folly::RWSpinLock::ReadHolder guard(lock);
-  auto r = data_.find(val);
-  if (r == data_.end()) {
-    return nullptr;
-  } else {
-    return r->second.get();
-  }
+  return get_ref_default(*data_.rlock(), val, nullptr).get();
 }
 
 void RequestContext::onSet() {
-  folly::RWSpinLock::ReadHolder guard(lock);
-  for (auto const& ent : data_) {
-    if (RequestData* data = ent.second.get()) {
+  auto rlock = data_.rlock();
+  for (auto const& ent : *rlock) {
+    if (auto& data = ent.second) {
       data->onSet();
     }
   }
 }
 
 void RequestContext::onUnset() {
-  folly::RWSpinLock::ReadHolder guard(lock);
-  for (auto const& ent : data_) {
-    if (RequestData* data = ent.second.get()) {
+  auto rlock = data_.rlock();
+  for (auto const& ent : *rlock) {
+    if (auto& data = ent.second) {
       data->onUnset();
     }
   }
 }
 
 void RequestContext::clearContextData(const std::string& val) {
-  folly::RWSpinLock::WriteHolder guard(lock);
-  data_.erase(val);
+  data_.wlock()->erase(val);
 }
 
 std::shared_ptr<RequestContext> RequestContext::setContext(

diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index 0f4d725633d2f87dfbbbbee088603092e095c2bf..48252a2fe9e24dbc0ee9abc5677f5074fc13bf8c 100644 (file)
--- a/folly/io/async/Request.h
+++ b/folly/io/async/Request.h
@@ -22,7 +22,9 @@
 
 #include <map>
 #include <memory>
+
 #include <folly/RWSpinLock.h>
+#include <folly/Synchronized.h>
 
 namespace folly {
 
@@ -103,8 +105,8 @@ class RequestContext {
  private:
   static std::shared_ptr<RequestContext>& getStaticContext();
 
-  mutable folly::RWSpinLock lock;
-  std::map<std::string, std::unique_ptr<RequestData>> data_;
+  using Data = std::map<std::string, std::unique_ptr<RequestData>>;
+  folly::Synchronized<Data, folly::RWSpinLock> data_;
 };
 
 class RequestContextScopeGuard {

diff --git a/folly/io/async/test/RequestContextTest.cpp b/folly/io/async/test/RequestContextTest.cpp
index 6b1cb6db159f33902e9ee4f4af53df53d3d92336..1f5c5e61f671c76035e34e2bfa5bbdbbb060464d 100644 (file)
--- a/folly/io/async/test/RequestContextTest.cpp
+++ b/folly/io/async/test/RequestContextTest.cpp
@@ -19,6 +19,7 @@
 #include <gtest/gtest.h>
 #include <thread>
 
+#include <folly/Memory.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/io/async/Request.h>
 
@@ -56,8 +57,7 @@ TEST(RequestContext, SimpleTest) {
   EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
 
   RequestContext::get()->setContextData(
-    "test",
-    std::unique_ptr<TestData>(new TestData(10)));
+      "test", folly::make_unique<TestData>(10));
   base.runInEventBaseThread([&](){
       EXPECT_TRUE(RequestContext::get() != nullptr);
       auto data = dynamic_cast<TestData*>(
@@ -84,15 +84,15 @@ TEST(RequestContext, setIfAbsentTest) {
   EXPECT_TRUE(RequestContext::get() != nullptr);
 
   RequestContext::get()->setContextData(
-      "test", std::unique_ptr<TestData>(new TestData(10)));
+      "test", folly::make_unique<TestData>(10));
   EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
-      "test", std::unique_ptr<TestData>(new TestData(20))));
+      "test", folly::make_unique<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))));
+      "test2", folly::make_unique<TestData>(20)));
   EXPECT_EQ(20,
             dynamic_cast<TestData*>(
                 RequestContext::get()->getContextData("test2"))->data_);
@@ -104,13 +104,13 @@ TEST(RequestContext, setIfAbsentTest) {
 TEST(RequestContext, testSetUnset) {
   RequestContext::create();
   auto ctx1 = RequestContext::saveContext();
-  ctx1->setContextData("test", std::unique_ptr<TestData>(new TestData(10)));
+  ctx1->setContextData("test", folly::make_unique<TestData>(10));
   auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
 
   // Override RequestContext
   RequestContext::create();
   auto ctx2 = RequestContext::saveContext();
-  ctx2->setContextData("test", std::unique_ptr<TestData>(new TestData(20)));
+  ctx2->setContextData("test", folly::make_unique<TestData>(20));
   auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
 
   // Check ctx1->onUnset was called
@@ -129,11 +129,3 @@ TEST(RequestContext, testSetUnset) {
   EXPECT_EQ(1, testData2->set_);
   EXPECT_EQ(1, testData2->unset_);
 }
-
-int main(int argc, char** argv) {
-  testing::InitGoogleTest(&argc, argv);
-  google::InitGoogleLogging(argv[0]);
-  google::ParseCommandLineFlags(&argc, &argv, true);
-
-  return RUN_ALL_TESTS();
-}