1 use crate::sync::batch_semaphore::Semaphore;
2 use tokio_test::*;
3
4 #[test]
poll_acquire_one_available()5 fn poll_acquire_one_available() {
6 let s = Semaphore::new(100);
7 assert_eq!(s.available_permits(), 100);
8
9 // Polling for a permit succeeds immediately
10 assert_ready_ok!(task::spawn(s.acquire(1)).poll());
11 assert_eq!(s.available_permits(), 99);
12 }
13
14 #[test]
poll_acquire_many_available()15 fn poll_acquire_many_available() {
16 let s = Semaphore::new(100);
17 assert_eq!(s.available_permits(), 100);
18
19 // Polling for a permit succeeds immediately
20 assert_ready_ok!(task::spawn(s.acquire(5)).poll());
21 assert_eq!(s.available_permits(), 95);
22
23 assert_ready_ok!(task::spawn(s.acquire(5)).poll());
24 assert_eq!(s.available_permits(), 90);
25 }
26
27 #[test]
try_acquire_one_available()28 fn try_acquire_one_available() {
29 let s = Semaphore::new(100);
30 assert_eq!(s.available_permits(), 100);
31
32 assert_ok!(s.try_acquire(1));
33 assert_eq!(s.available_permits(), 99);
34
35 assert_ok!(s.try_acquire(1));
36 assert_eq!(s.available_permits(), 98);
37 }
38
39 #[test]
try_acquire_many_available()40 fn try_acquire_many_available() {
41 let s = Semaphore::new(100);
42 assert_eq!(s.available_permits(), 100);
43
44 assert_ok!(s.try_acquire(5));
45 assert_eq!(s.available_permits(), 95);
46
47 assert_ok!(s.try_acquire(5));
48 assert_eq!(s.available_permits(), 90);
49 }
50
51 #[test]
poll_acquire_one_unavailable()52 fn poll_acquire_one_unavailable() {
53 let s = Semaphore::new(1);
54
55 // Acquire the first permit
56 assert_ready_ok!(task::spawn(s.acquire(1)).poll());
57 assert_eq!(s.available_permits(), 0);
58
59 let mut acquire_2 = task::spawn(s.acquire(1));
60 // Try to acquire the second permit
61 assert_pending!(acquire_2.poll());
62 assert_eq!(s.available_permits(), 0);
63
64 s.release(1);
65
66 assert_eq!(s.available_permits(), 0);
67 assert!(acquire_2.is_woken());
68 assert_ready_ok!(acquire_2.poll());
69 assert_eq!(s.available_permits(), 0);
70
71 s.release(1);
72 assert_eq!(s.available_permits(), 1);
73 }
74
75 #[test]
poll_acquire_many_unavailable()76 fn poll_acquire_many_unavailable() {
77 let s = Semaphore::new(5);
78
79 // Acquire the first permit
80 assert_ready_ok!(task::spawn(s.acquire(1)).poll());
81 assert_eq!(s.available_permits(), 4);
82
83 // Try to acquire the second permit
84 let mut acquire_2 = task::spawn(s.acquire(5));
85 assert_pending!(acquire_2.poll());
86 assert_eq!(s.available_permits(), 0);
87
88 // Try to acquire the third permit
89 let mut acquire_3 = task::spawn(s.acquire(3));
90 assert_pending!(acquire_3.poll());
91 assert_eq!(s.available_permits(), 0);
92
93 s.release(1);
94
95 assert_eq!(s.available_permits(), 0);
96 assert!(acquire_2.is_woken());
97 assert_ready_ok!(acquire_2.poll());
98
99 assert!(!acquire_3.is_woken());
100 assert_eq!(s.available_permits(), 0);
101
102 s.release(1);
103 assert!(!acquire_3.is_woken());
104 assert_eq!(s.available_permits(), 0);
105
106 s.release(2);
107 assert!(acquire_3.is_woken());
108
109 assert_ready_ok!(acquire_3.poll());
110 }
111
112 #[test]
try_acquire_one_unavailable()113 fn try_acquire_one_unavailable() {
114 let s = Semaphore::new(1);
115
116 // Acquire the first permit
117 assert_ok!(s.try_acquire(1));
118 assert_eq!(s.available_permits(), 0);
119
120 assert_err!(s.try_acquire(1));
121
122 s.release(1);
123
124 assert_eq!(s.available_permits(), 1);
125 assert_ok!(s.try_acquire(1));
126
127 s.release(1);
128 assert_eq!(s.available_permits(), 1);
129 }
130
131 #[test]
try_acquire_many_unavailable()132 fn try_acquire_many_unavailable() {
133 let s = Semaphore::new(5);
134
135 // Acquire the first permit
136 assert_ok!(s.try_acquire(1));
137 assert_eq!(s.available_permits(), 4);
138
139 assert_err!(s.try_acquire(5));
140
141 s.release(1);
142 assert_eq!(s.available_permits(), 5);
143
144 assert_ok!(s.try_acquire(5));
145
146 s.release(1);
147 assert_eq!(s.available_permits(), 1);
148
149 s.release(1);
150 assert_eq!(s.available_permits(), 2);
151 }
152
153 #[test]
poll_acquire_one_zero_permits()154 fn poll_acquire_one_zero_permits() {
155 let s = Semaphore::new(0);
156 assert_eq!(s.available_permits(), 0);
157
158 // Try to acquire the permit
159 let mut acquire = task::spawn(s.acquire(1));
160 assert_pending!(acquire.poll());
161
162 s.release(1);
163
164 assert!(acquire.is_woken());
165 assert_ready_ok!(acquire.poll());
166 }
167
168 #[test]
169 #[should_panic]
validates_max_permits()170 fn validates_max_permits() {
171 use std::usize;
172 Semaphore::new((usize::MAX >> 2) + 1);
173 }
174
175 #[test]
close_semaphore_prevents_acquire()176 fn close_semaphore_prevents_acquire() {
177 let s = Semaphore::new(5);
178 s.close();
179
180 assert_eq!(5, s.available_permits());
181
182 assert_ready_err!(task::spawn(s.acquire(1)).poll());
183 assert_eq!(5, s.available_permits());
184
185 assert_ready_err!(task::spawn(s.acquire(1)).poll());
186 assert_eq!(5, s.available_permits());
187 }
188
189 #[test]
close_semaphore_notifies_permit1()190 fn close_semaphore_notifies_permit1() {
191 let s = Semaphore::new(0);
192 let mut acquire = task::spawn(s.acquire(1));
193
194 assert_pending!(acquire.poll());
195
196 s.close();
197
198 assert!(acquire.is_woken());
199 assert_ready_err!(acquire.poll());
200 }
201
202 #[test]
close_semaphore_notifies_permit2()203 fn close_semaphore_notifies_permit2() {
204 let s = Semaphore::new(2);
205
206 // Acquire a couple of permits
207 assert_ready_ok!(task::spawn(s.acquire(1)).poll());
208 assert_ready_ok!(task::spawn(s.acquire(1)).poll());
209
210 let mut acquire3 = task::spawn(s.acquire(1));
211 let mut acquire4 = task::spawn(s.acquire(1));
212 assert_pending!(acquire3.poll());
213 assert_pending!(acquire4.poll());
214
215 s.close();
216
217 assert!(acquire3.is_woken());
218 assert!(acquire4.is_woken());
219
220 assert_ready_err!(acquire3.poll());
221 assert_ready_err!(acquire4.poll());
222
223 assert_eq!(0, s.available_permits());
224
225 s.release(1);
226
227 assert_eq!(1, s.available_permits());
228
229 assert_ready_err!(task::spawn(s.acquire(1)).poll());
230
231 s.release(1);
232
233 assert_eq!(2, s.available_permits());
234 }
235
236 #[test]
cancel_acquire_releases_permits()237 fn cancel_acquire_releases_permits() {
238 let s = Semaphore::new(10);
239 s.try_acquire(4).expect("uncontended try_acquire succeeds");
240 assert_eq!(6, s.available_permits());
241
242 let mut acquire = task::spawn(s.acquire(8));
243 assert_pending!(acquire.poll());
244
245 assert_eq!(0, s.available_permits());
246 drop(acquire);
247
248 assert_eq!(6, s.available_permits());
249 assert_ok!(s.try_acquire(6));
250 }
251