1 /*
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
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
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
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.
15 */
16
17 #include <folly/CancellationToken.h>
18
19 #include <chrono>
20 #include <thread>
21
22 #include <folly/Optional.h>
23 #include <folly/portability/GTest.h>
24 #include <folly/synchronization/Baton.h>
25
26 using namespace folly;
27 using namespace std::literals::chrono_literals;
28
TEST(CancellationTokenTest,DefaultCancellationTokenIsNotCancellable)29 TEST(CancellationTokenTest, DefaultCancellationTokenIsNotCancellable) {
30 CancellationToken t;
31 EXPECT_FALSE(t.isCancellationRequested());
32 EXPECT_FALSE(t.canBeCancelled());
33
34 CancellationToken tCopy = t;
35 EXPECT_FALSE(tCopy.isCancellationRequested());
36 EXPECT_FALSE(tCopy.canBeCancelled());
37
38 CancellationToken tMoved = std::move(t);
39 EXPECT_FALSE(tMoved.isCancellationRequested());
40 EXPECT_FALSE(tMoved.canBeCancelled());
41 }
42
TEST(CancellationTokenTest,Polling)43 TEST(CancellationTokenTest, Polling) {
44 CancellationSource src;
45 EXPECT_FALSE(src.isCancellationRequested());
46 EXPECT_TRUE(src.canBeCancelled());
47
48 CancellationToken token = src.getToken();
49 EXPECT_FALSE(token.isCancellationRequested());
50 EXPECT_TRUE(token.canBeCancelled());
51
52 CancellationToken tokenCopy = token;
53 EXPECT_FALSE(tokenCopy.isCancellationRequested());
54 EXPECT_TRUE(tokenCopy.canBeCancelled());
55
56 src.requestCancellation();
57 EXPECT_TRUE(token.isCancellationRequested());
58 EXPECT_TRUE(tokenCopy.isCancellationRequested());
59 }
60
TEST(CancellationTokenTest,MultiThreadedPolling)61 TEST(CancellationTokenTest, MultiThreadedPolling) {
62 CancellationSource src;
63
64 std::thread t1{[t = src.getToken()] {
65 while (!t.isCancellationRequested()) {
66 std::this_thread::yield();
67 }
68 }};
69
70 src.requestCancellation();
71
72 t1.join();
73 }
74
TEST(CancellationTokenTest,TokenIsNotCancellableOnceLastSourceIsDestroyed)75 TEST(CancellationTokenTest, TokenIsNotCancellableOnceLastSourceIsDestroyed) {
76 CancellationToken token;
77 {
78 CancellationSource src;
79 token = src.getToken();
80 {
81 CancellationSource srcCopy1;
82 CancellationSource srcCopy2;
83 EXPECT_TRUE(token.canBeCancelled());
84 }
85 EXPECT_TRUE(token.canBeCancelled());
86 }
87 EXPECT_FALSE(token.canBeCancelled());
88 }
89
TEST(CancellationTokenTest,TokenRemainsCancellableEvenOnceLastSourceIsDestroyed)90 TEST(
91 CancellationTokenTest,
92 TokenRemainsCancellableEvenOnceLastSourceIsDestroyed) {
93 CancellationToken token;
94 {
95 CancellationSource src;
96 token = src.getToken();
97 {
98 CancellationSource srcCopy1;
99 CancellationSource srcCopy2;
100 EXPECT_TRUE(token.canBeCancelled());
101 }
102 EXPECT_TRUE(token.canBeCancelled());
103 src.requestCancellation();
104 }
105 EXPECT_TRUE(token.canBeCancelled());
106 EXPECT_TRUE(token.isCancellationRequested());
107 }
108
TEST(CancellationTokenTest,CallbackRegistration)109 TEST(CancellationTokenTest, CallbackRegistration) {
110 CancellationSource src;
111
112 bool callbackExecuted = false;
113 CancellationCallback cb{src.getToken(), [&] { callbackExecuted = true; }};
114
115 EXPECT_FALSE(callbackExecuted);
116
117 src.requestCancellation();
118
119 EXPECT_TRUE(callbackExecuted);
120 }
121
TEST(CancellationTokenTest,CallbackExecutesImmediatelyIfAlreadyCancelled)122 TEST(CancellationTokenTest, CallbackExecutesImmediatelyIfAlreadyCancelled) {
123 CancellationSource src;
124 src.requestCancellation();
125
126 bool callbackExecuted = false;
127 CancellationCallback cb{src.getToken(), [&] { callbackExecuted = true; }};
128
129 EXPECT_TRUE(callbackExecuted);
130 }
131
TEST(CancellationTokenTest,CallbackShouldNotBeExecutedMultipleTimes)132 TEST(CancellationTokenTest, CallbackShouldNotBeExecutedMultipleTimes) {
133 CancellationSource src;
134
135 int callbackExecutionCount = 0;
136 CancellationCallback cb{src.getToken(), [&] { ++callbackExecutionCount; }};
137
138 src.requestCancellation();
139
140 EXPECT_EQ(1, callbackExecutionCount);
141
142 src.requestCancellation();
143
144 EXPECT_EQ(1, callbackExecutionCount);
145 }
146
TEST(CancellationTokenTest,RegisterMultipleCallbacks)147 TEST(CancellationTokenTest, RegisterMultipleCallbacks) {
148 CancellationSource src;
149
150 bool executed1 = false;
151 CancellationCallback cb1{src.getToken(), [&] { executed1 = true; }};
152
153 bool executed2 = false;
154 CancellationCallback cb2{src.getToken(), [&] { executed2 = true; }};
155
156 EXPECT_FALSE(executed1);
157 EXPECT_FALSE(executed2);
158
159 src.requestCancellation();
160
161 EXPECT_TRUE(executed1);
162 EXPECT_TRUE(executed2);
163 }
164
TEST(CancellationTokenTest,DeregisteredCallbacksDontExecute)165 TEST(CancellationTokenTest, DeregisteredCallbacksDontExecute) {
166 CancellationSource src;
167
168 bool executed1 = false;
169 bool executed2 = false;
170
171 CancellationCallback cb1{src.getToken(), [&] { executed1 = true; }};
172 {
173 CancellationCallback cb2{src.getToken(), [&] { executed2 = true; }};
174 }
175
176 src.requestCancellation();
177
178 EXPECT_TRUE(executed1);
179 EXPECT_FALSE(executed2);
180 }
181
TEST(CancellationTokenTest,CallbackThatDeregistersItself)182 TEST(CancellationTokenTest, CallbackThatDeregistersItself) {
183 CancellationSource src;
184 // Check that this doesn't deadlock when a callback tries to deregister
185 // itself from within the callback.
186 folly::Optional<CancellationCallback> callback;
187 callback.emplace(src.getToken(), [&] { callback.reset(); });
188 src.requestCancellation();
189 }
TEST(CancellationTokenTest,ManyCallbacks)190 TEST(CancellationTokenTest, ManyCallbacks) {
191 // This test checks that the CancellationSource internal state is able to
192 // grow to accommodate a large number of callbacks and that there are no
193 // memory leaks when it's all eventually destroyed.
194 CancellationSource src;
195 auto addLotsOfCallbacksAndWait = [t = src.getToken()] {
196 int counter = 0;
197 std::vector<std::unique_ptr<CancellationCallback>> callbacks;
198 for (int i = 0; i < 100; ++i) {
199 callbacks.push_back(
200 std::make_unique<CancellationCallback>(t, [&] { ++counter; }));
201 }
202
203 Baton<> baton;
204 CancellationCallback cb{t, [&] { baton.post(); }};
205 baton.wait();
206 };
207
208 std::thread t1{addLotsOfCallbacksAndWait};
209 std::thread t2{addLotsOfCallbacksAndWait};
210 std::thread t3{addLotsOfCallbacksAndWait};
211 std::thread t4{addLotsOfCallbacksAndWait};
212
213 src.requestCancellation();
214
215 t1.join();
216 t2.join();
217 t3.join();
218 t4.join();
219 }
220
TEST(CancellationTokenTest,ManyConcurrentCallbackAddRemove)221 TEST(CancellationTokenTest, ManyConcurrentCallbackAddRemove) {
222 auto runTest = [](CancellationToken ct) {
223 auto cb = [] { std::this_thread::sleep_for(1ms); };
224 while (!ct.isCancellationRequested()) {
225 CancellationCallback cb1{ct, cb};
226 CancellationCallback cb2{ct, cb};
227 CancellationCallback cb3{ct, cb};
228 CancellationCallback cb5{ct, cb};
229 CancellationCallback cb6{ct, cb};
230 CancellationCallback cb7{ct, cb};
231 CancellationCallback cb8{ct, cb};
232 }
233 };
234
235 CancellationSource src;
236
237 std::vector<std::thread> threads;
238 for (int i = 0; i < 10; ++i) {
239 threads.emplace_back([&, t = src.getToken()] { runTest(t); });
240 }
241
242 std::this_thread::sleep_for(1s);
243
244 src.requestCancellation();
245
246 for (auto& t : threads) {
247 t.join();
248 }
249 }
250
TEST(CancellationTokenTest,NonCancellableSource)251 TEST(CancellationTokenTest, NonCancellableSource) {
252 CancellationSource src = CancellationSource::invalid();
253 CHECK(!src.canBeCancelled());
254 CHECK(!src.isCancellationRequested());
255 CHECK(!src.requestCancellation());
256 CHECK(!src.isCancellationRequested());
257 CHECK(!src.canBeCancelled());
258
259 auto token = src.getToken();
260 CHECK(!src.canBeCancelled());
261 CHECK(!src.isCancellationRequested());
262 CHECK(token == CancellationToken{});
263 }
264
TEST(CancellationTokenTest,MergedToken)265 TEST(CancellationTokenTest, MergedToken) {
266 CancellationSource src1, src2;
267
268 auto token = CancellationToken::merge(src1.getToken(), src2.getToken());
269
270 EXPECT_TRUE(token.canBeCancelled());
271 EXPECT_FALSE(token.isCancellationRequested());
272
273 bool callbackExecuted = false;
274 CancellationCallback cb{token, [&] { callbackExecuted = true; }};
275
276 EXPECT_FALSE(callbackExecuted);
277 EXPECT_FALSE(token.isCancellationRequested());
278
279 src1.requestCancellation();
280
281 EXPECT_TRUE(callbackExecuted);
282 EXPECT_TRUE(token.isCancellationRequested());
283
284 src2.requestCancellation();
285
286 EXPECT_TRUE(callbackExecuted);
287 EXPECT_TRUE(token.isCancellationRequested());
288
289 token = CancellationToken::merge();
290 EXPECT_FALSE(token.canBeCancelled());
291
292 token = CancellationToken::merge(CancellationToken());
293 EXPECT_FALSE(token.canBeCancelled());
294 }
295
TEST(CancellationTokenTest,TokenWithData)296 TEST(CancellationTokenTest, TokenWithData) {
297 struct Guard {
298 int& counter;
299 explicit Guard(int& c) : counter(c) {}
300 ~Guard() { ++counter; }
301 };
302 int counter = 0;
303
304 {
305 CancellationToken token;
306 {
307 auto [source, data] =
308 CancellationSource::create(detail::WithDataTag<Guard>{}, counter);
309 EXPECT_EQ(counter, 0);
310 token = source.getToken();
311 EXPECT_EQ(counter, 0);
312 }
313 EXPECT_EQ(counter, 0);
314 }
315 EXPECT_EQ(counter, 1);
316 }
317