1 use crate::runtime::Builder;
2 use crate::task::JoinSet;
3
4 #[test]
test_join_set()5 fn test_join_set() {
6 loom::model(|| {
7 let rt = Builder::new_multi_thread()
8 .worker_threads(1)
9 .build()
10 .unwrap();
11 let mut set = JoinSet::new();
12
13 rt.block_on(async {
14 assert_eq!(set.len(), 0);
15 set.spawn(async { () });
16 assert_eq!(set.len(), 1);
17 set.spawn(async { () });
18 assert_eq!(set.len(), 2);
19 let () = set.join_next().await.unwrap().unwrap();
20 assert_eq!(set.len(), 1);
21 set.spawn(async { () });
22 assert_eq!(set.len(), 2);
23 let () = set.join_next().await.unwrap().unwrap();
24 assert_eq!(set.len(), 1);
25 let () = set.join_next().await.unwrap().unwrap();
26 assert_eq!(set.len(), 0);
27 set.spawn(async { () });
28 assert_eq!(set.len(), 1);
29 });
30
31 drop(set);
32 drop(rt);
33 });
34 }
35
36 #[test]
abort_all_during_completion()37 fn abort_all_during_completion() {
38 use std::sync::{
39 atomic::{AtomicBool, Ordering::SeqCst},
40 Arc,
41 };
42
43 // These booleans assert that at least one execution had the task complete first, and that at
44 // least one execution had the task be cancelled before it completed.
45 let complete_happened = Arc::new(AtomicBool::new(false));
46 let cancel_happened = Arc::new(AtomicBool::new(false));
47
48 {
49 let complete_happened = complete_happened.clone();
50 let cancel_happened = cancel_happened.clone();
51 loom::model(move || {
52 let rt = Builder::new_multi_thread()
53 .worker_threads(1)
54 .build()
55 .unwrap();
56
57 let mut set = JoinSet::new();
58
59 rt.block_on(async {
60 set.spawn(async { () });
61 set.abort_all();
62
63 match set.join_next().await {
64 Some(Ok(())) => complete_happened.store(true, SeqCst),
65 Some(Err(err)) if err.is_cancelled() => cancel_happened.store(true, SeqCst),
66 Some(Err(err)) => panic!("fail: {}", err),
67 None => {
68 unreachable!("Aborting the task does not remove it from the JoinSet.")
69 }
70 }
71
72 assert!(matches!(set.join_next().await, None));
73 });
74
75 drop(set);
76 drop(rt);
77 });
78 }
79
80 assert!(complete_happened.load(SeqCst));
81 assert!(cancel_happened.load(SeqCst));
82 }
83