1 use crate::runtime::scheduler::multi_thread::{queue, Stats};
2 use crate::runtime::tests::{unowned, NoopSchedule};
3
4 use loom::thread;
5 use std::cell::RefCell;
6
new_stats() -> Stats7 fn new_stats() -> Stats {
8 Stats::new(&crate::runtime::WorkerMetrics::new())
9 }
10
11 #[test]
basic()12 fn basic() {
13 loom::model(|| {
14 let (steal, mut local) = queue::local();
15 let inject = RefCell::new(vec![]);
16 let mut stats = new_stats();
17
18 let th = thread::spawn(move || {
19 let mut stats = new_stats();
20 let (_, mut local) = queue::local();
21 let mut n = 0;
22
23 for _ in 0..3 {
24 if steal.steal_into(&mut local, &mut stats).is_some() {
25 n += 1;
26 }
27
28 while local.pop().is_some() {
29 n += 1;
30 }
31 }
32
33 n
34 });
35
36 let mut n = 0;
37
38 for _ in 0..2 {
39 for _ in 0..2 {
40 let (task, _) = unowned(async {});
41 local.push_back_or_overflow(task, &inject, &mut stats);
42 }
43
44 if local.pop().is_some() {
45 n += 1;
46 }
47
48 // Push another task
49 let (task, _) = unowned(async {});
50 local.push_back_or_overflow(task, &inject, &mut stats);
51
52 while local.pop().is_some() {
53 n += 1;
54 }
55 }
56
57 n += inject.borrow_mut().drain(..).count();
58
59 n += th.join().unwrap();
60
61 assert_eq!(6, n);
62 });
63 }
64
65 #[test]
steal_overflow()66 fn steal_overflow() {
67 loom::model(|| {
68 let (steal, mut local) = queue::local();
69 let inject = RefCell::new(vec![]);
70 let mut stats = new_stats();
71
72 let th = thread::spawn(move || {
73 let mut stats = new_stats();
74 let (_, mut local) = queue::local();
75 let mut n = 0;
76
77 if steal.steal_into(&mut local, &mut stats).is_some() {
78 n += 1;
79 }
80
81 while local.pop().is_some() {
82 n += 1;
83 }
84
85 n
86 });
87
88 let mut n = 0;
89
90 // push a task, pop a task
91 let (task, _) = unowned(async {});
92 local.push_back_or_overflow(task, &inject, &mut stats);
93
94 if local.pop().is_some() {
95 n += 1;
96 }
97
98 for _ in 0..6 {
99 let (task, _) = unowned(async {});
100 local.push_back_or_overflow(task, &inject, &mut stats);
101 }
102
103 n += th.join().unwrap();
104
105 while local.pop().is_some() {
106 n += 1;
107 }
108
109 n += inject.borrow_mut().drain(..).count();
110
111 assert_eq!(7, n);
112 });
113 }
114
115 #[test]
multi_stealer()116 fn multi_stealer() {
117 const NUM_TASKS: usize = 5;
118
119 fn steal_tasks(steal: queue::Steal<NoopSchedule>) -> usize {
120 let mut stats = new_stats();
121 let (_, mut local) = queue::local();
122
123 if steal.steal_into(&mut local, &mut stats).is_none() {
124 return 0;
125 }
126
127 let mut n = 1;
128
129 while local.pop().is_some() {
130 n += 1;
131 }
132
133 n
134 }
135
136 loom::model(|| {
137 let (steal, mut local) = queue::local();
138 let inject = RefCell::new(vec![]);
139 let mut stats = new_stats();
140
141 // Push work
142 for _ in 0..NUM_TASKS {
143 let (task, _) = unowned(async {});
144 local.push_back_or_overflow(task, &inject, &mut stats);
145 }
146
147 let th1 = {
148 let steal = steal.clone();
149 thread::spawn(move || steal_tasks(steal))
150 };
151
152 let th2 = thread::spawn(move || steal_tasks(steal));
153
154 let mut n = 0;
155
156 while local.pop().is_some() {
157 n += 1;
158 }
159
160 n += inject.borrow_mut().drain(..).count();
161
162 n += th1.join().unwrap();
163 n += th2.join().unwrap();
164
165 assert_eq!(n, NUM_TASKS);
166 });
167 }
168
169 #[test]
chained_steal()170 fn chained_steal() {
171 loom::model(|| {
172 let mut stats = new_stats();
173 let (s1, mut l1) = queue::local();
174 let (s2, mut l2) = queue::local();
175 let inject = RefCell::new(vec![]);
176
177 // Load up some tasks
178 for _ in 0..4 {
179 let (task, _) = unowned(async {});
180 l1.push_back_or_overflow(task, &inject, &mut stats);
181
182 let (task, _) = unowned(async {});
183 l2.push_back_or_overflow(task, &inject, &mut stats);
184 }
185
186 // Spawn a task to steal from **our** queue
187 let th = thread::spawn(move || {
188 let mut stats = new_stats();
189 let (_, mut local) = queue::local();
190 s1.steal_into(&mut local, &mut stats);
191
192 while local.pop().is_some() {}
193 });
194
195 // Drain our tasks, then attempt to steal
196 while l1.pop().is_some() {}
197
198 s2.steal_into(&mut l1, &mut stats);
199
200 th.join().unwrap();
201
202 while l1.pop().is_some() {}
203 while l2.pop().is_some() {}
204 });
205 }
206