1 #![allow(unknown_lints, unexpected_cfgs)]
2 #![warn(rust_2018_idioms)]
3 #![cfg(all(feature = "full", tokio_unstable, target_has_atomic = "64"))]
4 
5 use std::collections::HashSet;
6 use std::sync::atomic::{AtomicUsize, Ordering};
7 use std::sync::{Arc, Mutex};
8 
9 use tokio::runtime::Builder;
10 
11 const TASKS: usize = 8;
12 const ITERATIONS: usize = 64;
13 /// Assert that the spawn task hook always fires when set.
14 #[test]
spawn_task_hook_fires()15 fn spawn_task_hook_fires() {
16     let count = Arc::new(AtomicUsize::new(0));
17     let count2 = Arc::clone(&count);
18 
19     let ids = Arc::new(Mutex::new(HashSet::new()));
20     let ids2 = Arc::clone(&ids);
21 
22     let runtime = Builder::new_current_thread()
23         .on_task_spawn(move |data| {
24             ids2.lock().unwrap().insert(data.id());
25 
26             count2.fetch_add(1, Ordering::SeqCst);
27         })
28         .build()
29         .unwrap();
30 
31     for _ in 0..TASKS {
32         runtime.spawn(std::future::pending::<()>());
33     }
34 
35     let count_realized = count.load(Ordering::SeqCst);
36     assert_eq!(
37         TASKS, count_realized,
38         "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}",
39         count_realized
40     );
41 
42     let count_ids_realized = ids.lock().unwrap().len();
43 
44     assert_eq!(
45         TASKS, count_ids_realized,
46         "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}",
47         count_realized
48     );
49 }
50 
51 /// Assert that the terminate task hook always fires when set.
52 #[test]
terminate_task_hook_fires()53 fn terminate_task_hook_fires() {
54     let count = Arc::new(AtomicUsize::new(0));
55     let count2 = Arc::clone(&count);
56 
57     let runtime = Builder::new_current_thread()
58         .on_task_terminate(move |_data| {
59             count2.fetch_add(1, Ordering::SeqCst);
60         })
61         .build()
62         .unwrap();
63 
64     for _ in 0..TASKS {
65         runtime.spawn(std::future::ready(()));
66     }
67 
68     runtime.block_on(async {
69         // tick the runtime a bunch to close out tasks
70         for _ in 0..ITERATIONS {
71             tokio::task::yield_now().await;
72         }
73     });
74 
75     assert_eq!(TASKS, count.load(Ordering::SeqCst));
76 }
77