1 #![cfg(tokio_unstable)]
2
3 use tokio::pin;
4 use tokio::sync::CancellationToken;
5
6 use core::future::Future;
7 use core::task::{Context, Poll};
8 use futures_test::task::new_count_waker;
9
10 #[test]
cancel_token()11 fn cancel_token() {
12 let (waker, wake_counter) = new_count_waker();
13 let token = CancellationToken::new();
14 assert_eq!(false, 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_eq!(true, 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]
cancel_child_token_through_parent()43 fn cancel_child_token_through_parent() {
44 let (waker, wake_counter) = new_count_waker();
45 let token = CancellationToken::new();
46
47 let child_token = token.child_token();
48 assert!(!child_token.is_cancelled());
49
50 let child_fut = child_token.cancelled();
51 pin!(child_fut);
52 let parent_fut = token.cancelled();
53 pin!(parent_fut);
54
55 assert_eq!(
56 Poll::Pending,
57 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
58 );
59 assert_eq!(
60 Poll::Pending,
61 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
62 );
63 assert_eq!(wake_counter, 0);
64
65 token.cancel();
66 assert_eq!(wake_counter, 2);
67 assert_eq!(true, token.is_cancelled());
68 assert_eq!(true, child_token.is_cancelled());
69
70 assert_eq!(
71 Poll::Ready(()),
72 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
73 );
74 assert_eq!(
75 Poll::Ready(()),
76 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
77 );
78 }
79
80 #[test]
cancel_child_token_without_parent()81 fn cancel_child_token_without_parent() {
82 let (waker, wake_counter) = new_count_waker();
83 let token = CancellationToken::new();
84
85 let child_token_1 = token.child_token();
86
87 let child_fut = child_token_1.cancelled();
88 pin!(child_fut);
89 let parent_fut = token.cancelled();
90 pin!(parent_fut);
91
92 assert_eq!(
93 Poll::Pending,
94 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
95 );
96 assert_eq!(
97 Poll::Pending,
98 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
99 );
100 assert_eq!(wake_counter, 0);
101
102 child_token_1.cancel();
103 assert_eq!(wake_counter, 1);
104 assert_eq!(false, token.is_cancelled());
105 assert_eq!(true, child_token_1.is_cancelled());
106
107 assert_eq!(
108 Poll::Ready(()),
109 child_fut.as_mut().poll(&mut Context::from_waker(&waker))
110 );
111 assert_eq!(
112 Poll::Pending,
113 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
114 );
115
116 let child_token_2 = token.child_token();
117 let child_fut_2 = child_token_2.cancelled();
118 pin!(child_fut_2);
119
120 assert_eq!(
121 Poll::Pending,
122 child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
123 );
124 assert_eq!(
125 Poll::Pending,
126 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
127 );
128
129 token.cancel();
130 assert_eq!(wake_counter, 3);
131 assert_eq!(true, token.is_cancelled());
132 assert_eq!(true, child_token_2.is_cancelled());
133
134 assert_eq!(
135 Poll::Ready(()),
136 child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
137 );
138 assert_eq!(
139 Poll::Ready(()),
140 parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
141 );
142 }
143
144 #[test]
create_child_token_after_parent_was_cancelled()145 fn create_child_token_after_parent_was_cancelled() {
146 for drop_child_first in [true, false].iter().cloned() {
147 let (waker, wake_counter) = new_count_waker();
148 let token = CancellationToken::new();
149 token.cancel();
150
151 let child_token = token.child_token();
152 assert!(child_token.is_cancelled());
153
154 {
155 let child_fut = child_token.cancelled();
156 pin!(child_fut);
157 let parent_fut = token.cancelled();
158 pin!(parent_fut);
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 assert_eq!(wake_counter, 0);
169
170 drop(child_fut);
171 drop(parent_fut);
172 }
173
174 if drop_child_first {
175 drop(child_token);
176 drop(token);
177 } else {
178 drop(token);
179 drop(child_token);
180 }
181 }
182 }
183
184 #[test]
drop_multiple_child_tokens()185 fn drop_multiple_child_tokens() {
186 for drop_first_child_first in &[true, false] {
187 let token = CancellationToken::new();
188 let mut child_tokens = [None, None, None];
189 for i in 0..child_tokens.len() {
190 child_tokens[i] = Some(token.child_token());
191 }
192
193 assert!(!token.is_cancelled());
194 assert!(!child_tokens[0].as_ref().unwrap().is_cancelled());
195
196 for i in 0..child_tokens.len() {
197 if *drop_first_child_first {
198 child_tokens[i] = None;
199 } else {
200 child_tokens[child_tokens.len() - 1 - i] = None;
201 }
202 assert!(!token.is_cancelled());
203 }
204
205 drop(token);
206 }
207 }
208
209 #[test]
drop_parent_before_child_tokens()210 fn drop_parent_before_child_tokens() {
211 let token = CancellationToken::new();
212 let child1 = token.child_token();
213 let child2 = token.child_token();
214
215 drop(token);
216 assert!(!child1.is_cancelled());
217
218 drop(child1);
219 drop(child2);
220 }
221