Commit 214d5f7c authored by Yedidya Feldblum's avatar Yedidya Feldblum Committed by Facebook Github Bot 1

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
parent b90a2ba5
...@@ -23,6 +23,7 @@ ...@@ -23,6 +23,7 @@
#include <glog/logging.h> #include <glog/logging.h>
#include <folly/MapUtil.h>
#include <folly/SingletonThreadLocal.h> #include <folly/SingletonThreadLocal.h>
namespace folly { namespace folly {
...@@ -30,77 +31,63 @@ namespace folly { ...@@ -30,77 +31,63 @@ namespace folly {
void RequestContext::setContextData( void RequestContext::setContextData(
const std::string& val, const std::string& val,
std::unique_ptr<RequestData> data) { std::unique_ptr<RequestData> data) {
folly::RWSpinLock::WriteHolder guard(lock); auto wlock = data_.wlock();
if (data_.find(val) != data_.end()) { if (wlock->count(val)) {
LOG_FIRST_N(WARNING, 1) LOG_FIRST_N(WARNING, 1)
<< "Called RequestContext::setContextData with data already set"; << "Called RequestContext::setContextData with data already set";
data_[val] = nullptr; (*wlock)[val] = nullptr;
} else { } else {
data_[val] = std::move(data); (*wlock)[val] = std::move(data);
} }
} }
bool RequestContext::setContextDataIfAbsent( bool RequestContext::setContextDataIfAbsent(
const std::string& val, const std::string& val,
std::unique_ptr<RequestData> data) { std::unique_ptr<RequestData> data) {
folly::RWSpinLock::UpgradedHolder guard(lock); auto ulock = data_.ulock();
if (data_.find(val) != data_.end()) { if (ulock->count(val)) {
return false; return false;
} }
folly::RWSpinLock::WriteHolder writeGuard(std::move(guard)); auto wlock = ulock.moveFromUpgradeToWrite();
data_[val] = std::move(data); (*wlock)[val] = std::move(data);
return true; return true;
} }
bool RequestContext::hasContextData(const std::string& val) const { bool RequestContext::hasContextData(const std::string& val) const {
folly::RWSpinLock::ReadHolder guard(lock); return data_.rlock()->count(val);
return data_.find(val) != data_.end();
} }
RequestData* RequestContext::getContextData(const std::string& val) { RequestData* RequestContext::getContextData(const std::string& val) {
folly::RWSpinLock::ReadHolder guard(lock); return get_ref_default(*data_.rlock(), val, nullptr).get();
auto r = data_.find(val);
if (r == data_.end()) {
return nullptr;
} else {
return r->second.get();
}
} }
const RequestData* RequestContext::getContextData( const RequestData* RequestContext::getContextData(
const std::string& val) const { const std::string& val) const {
folly::RWSpinLock::ReadHolder guard(lock); return get_ref_default(*data_.rlock(), val, nullptr).get();
auto r = data_.find(val);
if (r == data_.end()) {
return nullptr;
} else {
return r->second.get();
}
} }
void RequestContext::onSet() { void RequestContext::onSet() {
folly::RWSpinLock::ReadHolder guard(lock); auto rlock = data_.rlock();
for (auto const& ent : data_) { for (auto const& ent : *rlock) {
if (RequestData* data = ent.second.get()) { if (auto& data = ent.second) {
data->onSet(); data->onSet();
} }
} }
} }
void RequestContext::onUnset() { void RequestContext::onUnset() {
folly::RWSpinLock::ReadHolder guard(lock); auto rlock = data_.rlock();
for (auto const& ent : data_) { for (auto const& ent : *rlock) {
if (RequestData* data = ent.second.get()) { if (auto& data = ent.second) {
data->onUnset(); data->onUnset();
} }
} }
} }
void RequestContext::clearContextData(const std::string& val) { void RequestContext::clearContextData(const std::string& val) {
folly::RWSpinLock::WriteHolder guard(lock); data_.wlock()->erase(val);
data_.erase(val);
} }
std::shared_ptr<RequestContext> RequestContext::setContext( std::shared_ptr<RequestContext> RequestContext::setContext(
......
...@@ -22,7 +22,9 @@ ...@@ -22,7 +22,9 @@
#include <map> #include <map>
#include <memory> #include <memory>
#include <folly/RWSpinLock.h> #include <folly/RWSpinLock.h>
#include <folly/Synchronized.h>
namespace folly { namespace folly {
...@@ -103,8 +105,8 @@ class RequestContext { ...@@ -103,8 +105,8 @@ class RequestContext {
private: private:
static std::shared_ptr<RequestContext>& getStaticContext(); static std::shared_ptr<RequestContext>& getStaticContext();
mutable folly::RWSpinLock lock; using Data = std::map<std::string, std::unique_ptr<RequestData>>;
std::map<std::string, std::unique_ptr<RequestData>> data_; folly::Synchronized<Data, folly::RWSpinLock> data_;
}; };
class RequestContextScopeGuard { class RequestContextScopeGuard {
......
...@@ -19,6 +19,7 @@ ...@@ -19,6 +19,7 @@
#include <gtest/gtest.h> #include <gtest/gtest.h>
#include <thread> #include <thread>
#include <folly/Memory.h>
#include <folly/io/async/EventBase.h> #include <folly/io/async/EventBase.h>
#include <folly/io/async/Request.h> #include <folly/io/async/Request.h>
...@@ -56,8 +57,7 @@ TEST(RequestContext, SimpleTest) { ...@@ -56,8 +57,7 @@ TEST(RequestContext, SimpleTest) {
EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test")); EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
RequestContext::get()->setContextData( RequestContext::get()->setContextData(
"test", "test", folly::make_unique<TestData>(10));
std::unique_ptr<TestData>(new TestData(10)));
base.runInEventBaseThread([&](){ base.runInEventBaseThread([&](){
EXPECT_TRUE(RequestContext::get() != nullptr); EXPECT_TRUE(RequestContext::get() != nullptr);
auto data = dynamic_cast<TestData*>( auto data = dynamic_cast<TestData*>(
...@@ -84,15 +84,15 @@ TEST(RequestContext, setIfAbsentTest) { ...@@ -84,15 +84,15 @@ TEST(RequestContext, setIfAbsentTest) {
EXPECT_TRUE(RequestContext::get() != nullptr); EXPECT_TRUE(RequestContext::get() != nullptr);
RequestContext::get()->setContextData( RequestContext::get()->setContextData(
"test", std::unique_ptr<TestData>(new TestData(10))); "test", folly::make_unique<TestData>(10));
EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent( EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
"test", std::unique_ptr<TestData>(new TestData(20)))); "test", folly::make_unique<TestData>(20)));
EXPECT_EQ(10, EXPECT_EQ(10,
dynamic_cast<TestData*>( dynamic_cast<TestData*>(
RequestContext::get()->getContextData("test"))->data_); RequestContext::get()->getContextData("test"))->data_);
EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent( EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
"test2", std::unique_ptr<TestData>(new TestData(20)))); "test2", folly::make_unique<TestData>(20)));
EXPECT_EQ(20, EXPECT_EQ(20,
dynamic_cast<TestData*>( dynamic_cast<TestData*>(
RequestContext::get()->getContextData("test2"))->data_); RequestContext::get()->getContextData("test2"))->data_);
...@@ -104,13 +104,13 @@ TEST(RequestContext, setIfAbsentTest) { ...@@ -104,13 +104,13 @@ TEST(RequestContext, setIfAbsentTest) {
TEST(RequestContext, testSetUnset) { TEST(RequestContext, testSetUnset) {
RequestContext::create(); RequestContext::create();
auto ctx1 = RequestContext::saveContext(); 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")); auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
// Override RequestContext // Override RequestContext
RequestContext::create(); RequestContext::create();
auto ctx2 = RequestContext::saveContext(); 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")); auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
// Check ctx1->onUnset was called // Check ctx1->onUnset was called
...@@ -129,11 +129,3 @@ TEST(RequestContext, testSetUnset) { ...@@ -129,11 +129,3 @@ TEST(RequestContext, testSetUnset) {
EXPECT_EQ(1, testData2->set_); EXPECT_EQ(1, testData2->set_);
EXPECT_EQ(1, testData2->unset_); 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();
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment