Thread-safe RequestContext putIfAbsent operation
[folly.git] / folly / io / async / test / RequestContextTest.cpp
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  *   http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 #include <gtest/gtest.h>
20 #include <thread>
21
22 #include <folly/io/async/EventBase.h>
23 #include <folly/io/async/Request.h>
24
25 using namespace folly;
26
27 class TestData : public RequestData {
28  public:
29   explicit TestData(int data) : data_(data) {}
30   ~TestData() override {}
31   int data_;
32 };
33
34 TEST(RequestContext, SimpleTest) {
35   EventBase base;
36
37
38   // There should always be a default context with get()
39   EXPECT_TRUE(RequestContext::get() != nullptr);
40
41
42   // but not with saveContext()
43   EXPECT_EQ(RequestContext::saveContext(), nullptr);
44   RequestContext::create();
45   EXPECT_NE(RequestContext::saveContext(), nullptr);
46   RequestContext::create();
47   EXPECT_NE(RequestContext::saveContext(), nullptr);
48
49   EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
50
51   RequestContext::get()->setContextData(
52     "test",
53     std::unique_ptr<TestData>(new TestData(10)));
54   base.runInEventBaseThread([&](){
55       EXPECT_TRUE(RequestContext::get() != nullptr);
56       auto data = dynamic_cast<TestData*>(
57         RequestContext::get()->getContextData("test"))->data_;
58       EXPECT_EQ(10, data);
59       base.terminateLoopSoon();
60     });
61   auto th = std::thread([&](){
62       base.loopForever();
63   });
64   th.join();
65   EXPECT_TRUE(RequestContext::get() != nullptr);
66   auto a = dynamic_cast<TestData*>(
67     RequestContext::get()->getContextData("test"));
68   auto data = a->data_;
69   EXPECT_EQ(10, data);
70
71   RequestContext::setContext(std::shared_ptr<RequestContext>());
72   // There should always be a default context
73   EXPECT_TRUE(nullptr != RequestContext::get());
74 }
75
76 TEST(RequestContext, setIfAbsentTest) {
77   EXPECT_TRUE(RequestContext::get() != nullptr);
78
79   RequestContext::get()->setContextData(
80       "test", std::unique_ptr<TestData>(new TestData(10)));
81   EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
82       "test", std::unique_ptr<TestData>(new TestData(20))));
83   EXPECT_EQ(10,
84             dynamic_cast<TestData*>(
85                 RequestContext::get()->getContextData("test"))->data_);
86
87   EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
88       "test2", std::unique_ptr<TestData>(new TestData(20))));
89   EXPECT_EQ(20,
90             dynamic_cast<TestData*>(
91                 RequestContext::get()->getContextData("test2"))->data_);
92
93   RequestContext::setContext(std::shared_ptr<RequestContext>());
94   EXPECT_TRUE(nullptr != RequestContext::get());
95 }
96
97 int main(int argc, char** argv) {
98   testing::InitGoogleTest(&argc, argv);
99   google::InitGoogleLogging(argv[0]);
100   google::ParseCommandLineFlags(&argc, &argv, true);
101
102   return RUN_ALL_TESTS();
103 }