1#![warn(rust_2018_idioms)]
2
3use tokio::pin;
4use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
5
6use core::future::Future;
7use core::task::{Context, Poll};
8use futures_test::task::new_count_waker;
9
10#[test]
11fn cancel_token() {
12 let (waker, wake_counter) = new_count_waker();
13 let token = CancellationToken::new();
14 assert!(!token.is_cancelled());
15
16 let wait_fut = token.cancelled();
17 pin!(wait_fut);
18
19 assert_eq!(
20 Poll::Pending,
21 wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
22 );
23 assert_eq!(wake_counter, 0);
24
25 let wait_fut_2 = token.cancelled();
26 pin!(wait_fut_2);
27
28 token.cancel();
29 assert_eq!(wake_counter, 1);
30 assert!(token.is_cancelled());
31
32 assert_eq!(
33 Poll::Ready(()),
34 wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
35 );
36 assert_eq!(
37 Poll::Ready(()),
38 wait_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
39 );
40}
41
42#[test]
43fn cancel_token_owned() {
44 let (waker, wake_counter) = new_count_waker();
45 let token = CancellationToken::new();
46 assert!(!token.is_cancelled());
47
48 let wait_fut = token.clone().cancelled_owned();
49 pin!(wait_fut);
50
51 assert_eq!(
52 Poll::Pending,
53 wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
54 );
55 assert_eq!(wake_counter, 0);
56
57 let wait_fut_2 = token.clone().cancelled_owned();
58 pin!(wait_fut_2);
59
60 token.cancel();
61 assert_eq!(wake_counter, 1);
62 assert!(token.is_cancelled());
63
64 assert_eq!(
65 Poll::Ready(()),
66 wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
67 );
68 assert_eq!(
69 Poll::Ready(()),
70 wait_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
71 );
72}
73
74#[test]
75fn cancel_token_owned_drop_test() {
76 let (waker, wake_counter) = new_count_waker();
77 let token = CancellationToken::new();
78
79 let future = token.cancelled_owned();
80 pin!(future);
81
82 assert_eq!(
83 Poll::Pending,
84 future.as_mut().poll(&mut Context::from_waker(&waker))
85 );
86 assert_eq!(wake_counter, 0);
87
88 // let future be dropped while pinned and under pending state to
89 // find potential memory related bugs.
90}
91
92#[test]
93fn cancel_child_token_through_parent() {
94 let (waker, wake_counter) = new_count_waker();
95 let token = CancellationToken::new();
96
97 let child_token = token.child_token();
98 assert!(!child_token.is_cancelled());
99
100 let child_fut = child_token.cancelled();
101 pin!(child_fut);
102 let parent_fut = token.cancelled();
103 pin!(parent_fut);
104
105 assert_eq!(
106 Poll::Pending,
107 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
108 );
109 assert_eq!(
110 Poll::Pending,
111 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
112 );
113 assert_eq!(wake_counter, 0);
114
115 token.cancel();
116 assert_eq!(wake_counter, 2);
117 assert!(token.is_cancelled());
118 assert!(child_token.is_cancelled());
119
120 assert_eq!(
121 Poll::Ready(()),
122 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
123 );
124 assert_eq!(
125 Poll::Ready(()),
126 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
127 );
128}
129
130#[test]
131fn cancel_grandchild_token_through_parent_if_child_was_dropped() {
132 let (waker, wake_counter) = new_count_waker();
133 let token = CancellationToken::new();
134
135 let intermediate_token = token.child_token();
136 let child_token = intermediate_token.child_token();
137 drop(intermediate_token);
138 assert!(!child_token.is_cancelled());
139
140 let child_fut = child_token.cancelled();
141 pin!(child_fut);
142 let parent_fut = token.cancelled();
143 pin!(parent_fut);
144
145 assert_eq!(
146 Poll::Pending,
147 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
148 );
149 assert_eq!(
150 Poll::Pending,
151 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
152 );
153 assert_eq!(wake_counter, 0);
154
155 token.cancel();
156 assert_eq!(wake_counter, 2);
157 assert!(token.is_cancelled());
158 assert!(child_token.is_cancelled());
159
160 assert_eq!(
161 Poll::Ready(()),
162 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
163 );
164 assert_eq!(
165 Poll::Ready(()),
166 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
167 );
168}
169
170#[test]
171fn cancel_child_token_without_parent() {
172 let (waker, wake_counter) = new_count_waker();
173 let token = CancellationToken::new();
174
175 let child_token_1 = token.child_token();
176
177 let child_fut = child_token_1.cancelled();
178 pin!(child_fut);
179 let parent_fut = token.cancelled();
180 pin!(parent_fut);
181
182 assert_eq!(
183 Poll::Pending,
184 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
185 );
186 assert_eq!(
187 Poll::Pending,
188 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
189 );
190 assert_eq!(wake_counter, 0);
191
192 child_token_1.cancel();
193 assert_eq!(wake_counter, 1);
194 assert!(!token.is_cancelled());
195 assert!(child_token_1.is_cancelled());
196
197 assert_eq!(
198 Poll::Ready(()),
199 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
200 );
201 assert_eq!(
202 Poll::Pending,
203 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
204 );
205
206 let child_token_2 = token.child_token();
207 let child_fut_2 = child_token_2.cancelled();
208 pin!(child_fut_2);
209
210 assert_eq!(
211 Poll::Pending,
212 child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
213 );
214 assert_eq!(
215 Poll::Pending,
216 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
217 );
218
219 token.cancel();
220 assert_eq!(wake_counter, 3);
221 assert!(token.is_cancelled());
222 assert!(child_token_2.is_cancelled());
223
224 assert_eq!(
225 Poll::Ready(()),
226 child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
227 );
228 assert_eq!(
229 Poll::Ready(()),
230 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
231 );
232}
233
234#[test]
235fn create_child_token_after_parent_was_cancelled() {
236 for drop_child_first in [true, false].iter().cloned() {
237 let (waker, wake_counter) = new_count_waker();
238 let token = CancellationToken::new();
239 token.cancel();
240
241 let child_token = token.child_token();
242 assert!(child_token.is_cancelled());
243
244 {
245 let child_fut = child_token.cancelled();
246 pin!(child_fut);
247 let parent_fut = token.cancelled();
248 pin!(parent_fut);
249
250 assert_eq!(
251 Poll::Ready(()),
252 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
253 );
254 assert_eq!(
255 Poll::Ready(()),
256 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
257 );
258 assert_eq!(wake_counter, 0);
259 }
260
261 if drop_child_first {
262 drop(child_token);
263 drop(token);
264 } else {
265 drop(token);
266 drop(child_token);
267 }
268 }
269}
270
271#[test]
272fn drop_multiple_child_tokens() {
273 for drop_first_child_first in &[true, false] {
274 let token = CancellationToken::new();
275 let mut child_tokens = [None, None, None];
276 for child in &mut child_tokens {
277 *child = Some(token.child_token());
278 }
279
280 assert!(!token.is_cancelled());
281 assert!(!child_tokens[0].as_ref().unwrap().is_cancelled());
282
283 for i in 0..child_tokens.len() {
284 if *drop_first_child_first {
285 child_tokens[i] = None;
286 } else {
287 child_tokens[child_tokens.len() - 1 - i] = None;
288 }
289 assert!(!token.is_cancelled());
290 }
291
292 drop(token);
293 }
294}
295
296#[test]
297fn cancel_only_all_descendants() {
298 // ARRANGE
299 let (waker, wake_counter) = new_count_waker();
300
301 let parent_token = CancellationToken::new();
302 let token = parent_token.child_token();
303 let sibling_token = parent_token.child_token();
304 let child1_token = token.child_token();
305 let child2_token = token.child_token();
306 let grandchild_token = child1_token.child_token();
307 let grandchild2_token = child1_token.child_token();
308 let great_grandchild_token = grandchild_token.child_token();
309
310 assert!(!parent_token.is_cancelled());
311 assert!(!token.is_cancelled());
312 assert!(!sibling_token.is_cancelled());
313 assert!(!child1_token.is_cancelled());
314 assert!(!child2_token.is_cancelled());
315 assert!(!grandchild_token.is_cancelled());
316 assert!(!grandchild2_token.is_cancelled());
317 assert!(!great_grandchild_token.is_cancelled());
318
319 let parent_fut = parent_token.cancelled();
320 let fut = token.cancelled();
321 let sibling_fut = sibling_token.cancelled();
322 let child1_fut = child1_token.cancelled();
323 let child2_fut = child2_token.cancelled();
324 let grandchild_fut = grandchild_token.cancelled();
325 let grandchild2_fut = grandchild2_token.cancelled();
326 let great_grandchild_fut = great_grandchild_token.cancelled();
327
328 pin!(parent_fut);
329 pin!(fut);
330 pin!(sibling_fut);
331 pin!(child1_fut);
332 pin!(child2_fut);
333 pin!(grandchild_fut);
334 pin!(grandchild2_fut);
335 pin!(great_grandchild_fut);
336
337 assert_eq!(
338 Poll::Pending,
339 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
340 );
341 assert_eq!(
342 Poll::Pending,
343 fut.as_mut().poll(&mut Context::from_waker(&waker))
344 );
345 assert_eq!(
346 Poll::Pending,
347 sibling_fut.as_mut().poll(&mut Context::from_waker(&waker))
348 );
349 assert_eq!(
350 Poll::Pending,
351 child1_fut.as_mut().poll(&mut Context::from_waker(&waker))
352 );
353 assert_eq!(
354 Poll::Pending,
355 child2_fut.as_mut().poll(&mut Context::from_waker(&waker))
356 );
357 assert_eq!(
358 Poll::Pending,
359 grandchild_fut
360 .as_mut()
361 .poll(&mut Context::from_waker(&waker))
362 );
363 assert_eq!(
364 Poll::Pending,
365 grandchild2_fut
366 .as_mut()
367 .poll(&mut Context::from_waker(&waker))
368 );
369 assert_eq!(
370 Poll::Pending,
371 great_grandchild_fut
372 .as_mut()
373 .poll(&mut Context::from_waker(&waker))
374 );
375 assert_eq!(wake_counter, 0);
376
377 // ACT
378 token.cancel();
379
380 // ASSERT
381 assert_eq!(wake_counter, 6);
382 assert!(!parent_token.is_cancelled());
383 assert!(token.is_cancelled());
384 assert!(!sibling_token.is_cancelled());
385 assert!(child1_token.is_cancelled());
386 assert!(child2_token.is_cancelled());
387 assert!(grandchild_token.is_cancelled());
388 assert!(grandchild2_token.is_cancelled());
389 assert!(great_grandchild_token.is_cancelled());
390
391 assert_eq!(
392 Poll::Ready(()),
393 fut.as_mut().poll(&mut Context::from_waker(&waker))
394 );
395 assert_eq!(
396 Poll::Ready(()),
397 child1_fut.as_mut().poll(&mut Context::from_waker(&waker))
398 );
399 assert_eq!(
400 Poll::Ready(()),
401 child2_fut.as_mut().poll(&mut Context::from_waker(&waker))
402 );
403 assert_eq!(
404 Poll::Ready(()),
405 grandchild_fut
406 .as_mut()
407 .poll(&mut Context::from_waker(&waker))
408 );
409 assert_eq!(
410 Poll::Ready(()),
411 grandchild2_fut
412 .as_mut()
413 .poll(&mut Context::from_waker(&waker))
414 );
415 assert_eq!(
416 Poll::Ready(()),
417 great_grandchild_fut
418 .as_mut()
419 .poll(&mut Context::from_waker(&waker))
420 );
421 assert_eq!(wake_counter, 6);
422}
423
424#[test]
425fn drop_parent_before_child_tokens() {
426 let token = CancellationToken::new();
427 let child1 = token.child_token();
428 let child2 = token.child_token();
429
430 drop(token);
431 assert!(!child1.is_cancelled());
432
433 drop(child1);
434 drop(child2);
435}
436
437#[test]
438fn derives_send_sync() {
439 fn assert_send<T: Send>() {}
440 fn assert_sync<T: Sync>() {}
441
442 assert_send::<CancellationToken>();
443 assert_sync::<CancellationToken>();
444
445 assert_send::<WaitForCancellationFuture<'static>>();
446 assert_sync::<WaitForCancellationFuture<'static>>();
447}
448