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 @@
#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(
......
......@@ -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 {
......
......@@ -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();
}
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