2 * Copyright 2004-present Facebook, Inc.
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
8 * http://www.apache.org/licenses/LICENSE-2.0
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
19 #include <folly/Memory.h>
20 #include <folly/io/async/EventBase.h>
21 #include <folly/io/async/Request.h>
22 #include <folly/portability/GTest.h>
24 using namespace folly;
26 class TestData : public RequestData {
28 explicit TestData(int data) : data_(data) {}
29 ~TestData() override {}
30 void onSet() override {
33 void onUnset() override {
36 int set_ = 0, unset_ = 0;
40 TEST(RequestContext, SimpleTest) {
44 // There should always be a default context with get()
45 EXPECT_TRUE(RequestContext::get() != nullptr);
48 // but not with saveContext()
49 EXPECT_EQ(RequestContext::saveContext(), nullptr);
50 RequestContext::create();
51 EXPECT_NE(RequestContext::saveContext(), nullptr);
52 RequestContext::create();
53 EXPECT_NE(RequestContext::saveContext(), nullptr);
55 EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
57 RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
58 base.runInEventBaseThread([&](){
59 EXPECT_TRUE(RequestContext::get() != nullptr);
60 auto data = dynamic_cast<TestData*>(
61 RequestContext::get()->getContextData("test"))->data_;
63 base.terminateLoopSoon();
65 auto th = std::thread([&](){
69 EXPECT_TRUE(RequestContext::get() != nullptr);
70 auto a = dynamic_cast<TestData*>(
71 RequestContext::get()->getContextData("test"));
75 RequestContext::setContext(std::shared_ptr<RequestContext>());
76 // There should always be a default context
77 EXPECT_TRUE(nullptr != RequestContext::get());
80 TEST(RequestContext, nonDefaultContextsAreThreadLocal) {
81 RequestContext* ctx1 = nullptr;
82 RequestContext* ctx2 = nullptr;
84 std::vector<std::thread> ts;
85 for (size_t i = 0; i < 2; ++i) {
86 auto*& ctx = (i == 0 ? ctx1 : ctx2);
87 ts.emplace_back([&ctx]() {
88 RequestContext::create();
89 ctx = RequestContext::get();
96 EXPECT_NE(nullptr, ctx1);
97 EXPECT_NE(nullptr, ctx2);
98 EXPECT_NE(ctx1, ctx2);
101 TEST(RequestContext, customRequestContextProvider) {
102 auto customContext = std::make_shared<RequestContext>();
103 auto customProvider = [&customContext]() -> std::shared_ptr<RequestContext>& {
104 return customContext;
107 auto* const originalContext = RequestContext::get();
108 EXPECT_NE(nullptr, originalContext);
110 // Install new RequestContext provider
111 auto originalProvider =
112 RequestContext::setRequestContextProvider(std::move(customProvider));
114 auto* const newContext = RequestContext::get();
115 EXPECT_EQ(customContext.get(), newContext);
116 EXPECT_NE(originalContext, newContext);
118 // Restore original RequestContext provider
119 RequestContext::setRequestContextProvider(std::move(originalProvider));
120 EXPECT_EQ(originalContext, RequestContext::get());
123 TEST(RequestContext, setIfAbsentTest) {
124 EXPECT_TRUE(RequestContext::get() != nullptr);
126 RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
127 EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
128 "test", std::make_unique<TestData>(20)));
130 dynamic_cast<TestData*>(
131 RequestContext::get()->getContextData("test"))->data_);
133 EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
134 "test2", std::make_unique<TestData>(20)));
136 dynamic_cast<TestData*>(
137 RequestContext::get()->getContextData("test2"))->data_);
139 RequestContext::setContext(std::shared_ptr<RequestContext>());
140 EXPECT_TRUE(nullptr != RequestContext::get());
143 TEST(RequestContext, testSetUnset) {
144 RequestContext::create();
145 auto ctx1 = RequestContext::saveContext();
146 ctx1->setContextData("test", std::make_unique<TestData>(10));
147 auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
149 // Override RequestContext
150 RequestContext::create();
151 auto ctx2 = RequestContext::saveContext();
152 ctx2->setContextData("test", std::make_unique<TestData>(20));
153 auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
155 // Check ctx1->onUnset was called
156 EXPECT_EQ(0, testData1->set_);
157 EXPECT_EQ(1, testData1->unset_);
159 RequestContext::setContext(ctx1);
160 EXPECT_EQ(1, testData1->set_);
161 EXPECT_EQ(1, testData1->unset_);
162 EXPECT_EQ(0, testData2->set_);
163 EXPECT_EQ(1, testData2->unset_);
165 RequestContext::setContext(ctx2);
166 EXPECT_EQ(1, testData1->set_);
167 EXPECT_EQ(2, testData1->unset_);
168 EXPECT_EQ(1, testData2->set_);
169 EXPECT_EQ(1, testData2->unset_);
172 TEST(RequestContext, deadlockTest) {
173 class DeadlockTestData : public RequestData {
175 explicit DeadlockTestData(const std::string& val) : val_(val) {}
177 ~DeadlockTestData() override {
178 RequestContext::get()->setContextData(
179 val_, std::make_unique<TestData>(1));
182 void onSet() override {}
184 void onUnset() override {}
189 RequestContext::get()->setContextData(
190 "test", std::make_unique<DeadlockTestData>("test2"));
191 RequestContext::get()->clearContextData("test");