1 use std::future::Future;
2 use std::pin::Pin;
3 use std::task::{Context, Poll};
4 use std::thread;
5 use std::time::Duration;
6 
7 use async_task::Task;
8 use crossbeam::atomic::AtomicCell;
9 use futures::executor::block_on;
10 use futures::future;
11 use lazy_static::lazy_static;
12 
13 // Creates a future with event counters.
14 //
15 // Usage: `future!(f, POLL, DROP_F, DROP_O)`
16 //
17 // The future `f` sleeps for 200 ms and outputs `Poll::Ready`.
18 // When it gets polled, `POLL` is incremented.
19 // When it gets dropped, `DROP_F` is incremented.
20 // When the output gets dropped, `DROP_O` is incremented.
21 macro_rules! future {
22     ($name:pat, $poll:ident, $drop_f:ident, $drop_o:ident) => {
23         lazy_static! {
24             static ref $poll: AtomicCell<usize> = AtomicCell::new(0);
25             static ref $drop_f: AtomicCell<usize> = AtomicCell::new(0);
26             static ref $drop_o: AtomicCell<usize> = AtomicCell::new(0);
27         }
28 
29         let $name = {
30             struct Fut(Box<i32>);
31 
32             impl Future for Fut {
33                 type Output = Out;
34 
35                 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
36                     $poll.fetch_add(1);
37                     thread::sleep(ms(400));
38                     Poll::Ready(Out(Box::new(0)))
39                 }
40             }
41 
42             impl Drop for Fut {
43                 fn drop(&mut self) {
44                     $drop_f.fetch_add(1);
45                 }
46             }
47 
48             struct Out(Box<i32>);
49 
50             impl Drop for Out {
51                 fn drop(&mut self) {
52                     $drop_o.fetch_add(1);
53                 }
54             }
55 
56             Fut(Box::new(0))
57         };
58     };
59 }
60 
61 // Creates a schedule function with event counters.
62 //
63 // Usage: `schedule!(s, SCHED, DROP)`
64 //
65 // The schedule function `s` does nothing.
66 // When it gets invoked, `SCHED` is incremented.
67 // When it gets dropped, `DROP` is incremented.
68 macro_rules! schedule {
69     ($name:pat, $sched:ident, $drop:ident) => {
70         lazy_static! {
71             static ref $sched: AtomicCell<usize> = AtomicCell::new(0);
72             static ref $drop: AtomicCell<usize> = AtomicCell::new(0);
73         }
74 
75         let $name = {
76             struct Guard(Box<i32>);
77 
78             impl Drop for Guard {
79                 fn drop(&mut self) {
80                     $drop.fetch_add(1);
81                 }
82             }
83 
84             let guard = Guard(Box::new(0));
85             move |_task: Task<_>| {
86                 &guard;
87                 $sched.fetch_add(1);
88             }
89         };
90     };
91 }
92 
93 // Creates a task with event counters.
94 //
95 // Usage: `task!(task, handle f, s, DROP)`
96 //
97 // A task with future `f` and schedule function `s` is created.
98 // The `Task` and `JoinHandle` are bound to `task` and `handle`, respectively.
99 // When the tag inside the task gets dropped, `DROP` is incremented.
100 macro_rules! task {
101     ($task:pat, $handle: pat, $future:expr, $schedule:expr, $drop:ident) => {
102         lazy_static! {
103             static ref $drop: AtomicCell<usize> = AtomicCell::new(0);
104         }
105 
106         let ($task, $handle) = {
107             struct Tag(Box<i32>);
108 
109             impl Drop for Tag {
110                 fn drop(&mut self) {
111                     $drop.fetch_add(1);
112                 }
113             }
114 
115             async_task::spawn($future, $schedule, Tag(Box::new(0)))
116         };
117     };
118 }
119 
ms(ms: u64) -> Duration120 fn ms(ms: u64) -> Duration {
121     Duration::from_millis(ms)
122 }
123 
124 #[test]
cancel_during_run()125 fn cancel_during_run() {
126     future!(f, POLL, DROP_F, DROP_O);
127     schedule!(s, SCHEDULE, DROP_S);
128     task!(task, handle, f, s, DROP_T);
129 
130     crossbeam::scope(|scope| {
131         scope.spawn(|_| {
132             task.run();
133             assert_eq!(POLL.load(), 1);
134             assert_eq!(SCHEDULE.load(), 0);
135             assert_eq!(DROP_F.load(), 1);
136             assert_eq!(DROP_S.load(), 0);
137             assert_eq!(DROP_T.load(), 0);
138             assert_eq!(DROP_O.load(), 1);
139         });
140 
141         thread::sleep(ms(200));
142 
143         handle.cancel();
144         assert_eq!(POLL.load(), 1);
145         assert_eq!(SCHEDULE.load(), 0);
146         assert_eq!(DROP_F.load(), 0);
147         assert_eq!(DROP_S.load(), 0);
148         assert_eq!(DROP_T.load(), 0);
149         assert_eq!(DROP_O.load(), 0);
150 
151         thread::sleep(ms(400));
152 
153         assert_eq!(POLL.load(), 1);
154         assert_eq!(SCHEDULE.load(), 0);
155         assert_eq!(DROP_F.load(), 1);
156         assert_eq!(DROP_S.load(), 0);
157         assert_eq!(DROP_T.load(), 0);
158         assert_eq!(DROP_O.load(), 1);
159 
160         drop(handle);
161         assert_eq!(POLL.load(), 1);
162         assert_eq!(SCHEDULE.load(), 0);
163         assert_eq!(DROP_F.load(), 1);
164         assert_eq!(DROP_S.load(), 1);
165         assert_eq!(DROP_T.load(), 1);
166         assert_eq!(DROP_O.load(), 1);
167     })
168     .unwrap();
169 }
170 
171 #[test]
join_during_run()172 fn join_during_run() {
173     future!(f, POLL, DROP_F, DROP_O);
174     schedule!(s, SCHEDULE, DROP_S);
175     task!(task, handle, f, s, DROP_T);
176 
177     crossbeam::scope(|scope| {
178         scope.spawn(|_| {
179             task.run();
180             assert_eq!(POLL.load(), 1);
181             assert_eq!(SCHEDULE.load(), 0);
182             assert_eq!(DROP_F.load(), 1);
183 
184             thread::sleep(ms(200));
185             assert_eq!(DROP_S.load(), 1);
186             assert_eq!(DROP_T.load(), 1);
187         });
188 
189         thread::sleep(ms(200));
190 
191         assert!(block_on(handle).is_some());
192         assert_eq!(POLL.load(), 1);
193         assert_eq!(SCHEDULE.load(), 0);
194         assert_eq!(DROP_F.load(), 1);
195         assert_eq!(DROP_O.load(), 1);
196 
197         thread::sleep(ms(200));
198         assert_eq!(DROP_S.load(), 1);
199         assert_eq!(DROP_T.load(), 1);
200     })
201     .unwrap();
202 }
203 
204 #[test]
try_join_during_run()205 fn try_join_during_run() {
206     future!(f, POLL, DROP_F, DROP_O);
207     schedule!(s, SCHEDULE, DROP_S);
208     task!(task, mut handle, f, s, DROP_T);
209 
210     crossbeam::scope(|scope| {
211         scope.spawn(|_| {
212             task.run();
213             assert_eq!(POLL.load(), 1);
214             assert_eq!(SCHEDULE.load(), 0);
215             assert_eq!(DROP_F.load(), 1);
216             assert_eq!(DROP_S.load(), 1);
217             assert_eq!(DROP_T.load(), 1);
218             assert_eq!(DROP_O.load(), 1);
219         });
220 
221         thread::sleep(ms(200));
222 
223         block_on(future::select(&mut handle, future::ready(())));
224         assert_eq!(POLL.load(), 1);
225         assert_eq!(SCHEDULE.load(), 0);
226         assert_eq!(DROP_F.load(), 0);
227         assert_eq!(DROP_S.load(), 0);
228         assert_eq!(DROP_T.load(), 0);
229         assert_eq!(DROP_O.load(), 0);
230         drop(handle);
231     })
232     .unwrap();
233 }
234 
235 #[test]
drop_handle_during_run()236 fn drop_handle_during_run() {
237     future!(f, POLL, DROP_F, DROP_O);
238     schedule!(s, SCHEDULE, DROP_S);
239     task!(task, handle, f, s, DROP_T);
240 
241     crossbeam::scope(|scope| {
242         scope.spawn(|_| {
243             task.run();
244             assert_eq!(POLL.load(), 1);
245             assert_eq!(SCHEDULE.load(), 0);
246             assert_eq!(DROP_F.load(), 1);
247             assert_eq!(DROP_S.load(), 1);
248             assert_eq!(DROP_T.load(), 1);
249             assert_eq!(DROP_O.load(), 1);
250         });
251 
252         thread::sleep(ms(200));
253 
254         drop(handle);
255         assert_eq!(POLL.load(), 1);
256         assert_eq!(SCHEDULE.load(), 0);
257         assert_eq!(DROP_F.load(), 0);
258         assert_eq!(DROP_S.load(), 0);
259         assert_eq!(DROP_T.load(), 0);
260         assert_eq!(DROP_O.load(), 0);
261     })
262     .unwrap();
263 }
264