1 #![allow(clippy::blacklisted_name)]
2 use tokio::sync::{mpsc, oneshot};
3 use tokio::task;
4 use tokio_test::{assert_ok, assert_pending, assert_ready};
5 
6 use futures::future::poll_fn;
7 use std::task::Poll::Ready;
8 
9 #[tokio::test]
sync_one_lit_expr_comma()10 async fn sync_one_lit_expr_comma() {
11     let foo = tokio::select! {
12         foo = async { 1 } => foo,
13     };
14 
15     assert_eq!(foo, 1);
16 }
17 
18 #[tokio::test]
nested_one()19 async fn nested_one() {
20     let foo = tokio::select! {
21         foo = async { 1 } => tokio::select! {
22             bar = async { foo } => bar,
23         },
24     };
25 
26     assert_eq!(foo, 1);
27 }
28 
29 #[tokio::test]
sync_one_lit_expr_no_comma()30 async fn sync_one_lit_expr_no_comma() {
31     let foo = tokio::select! {
32         foo = async { 1 } => foo
33     };
34 
35     assert_eq!(foo, 1);
36 }
37 
38 #[tokio::test]
sync_one_lit_expr_block()39 async fn sync_one_lit_expr_block() {
40     let foo = tokio::select! {
41         foo = async { 1 } => { foo }
42     };
43 
44     assert_eq!(foo, 1);
45 }
46 
47 #[tokio::test]
sync_one_await()48 async fn sync_one_await() {
49     let foo = tokio::select! {
50         foo = one() => foo,
51     };
52 
53     assert_eq!(foo, 1);
54 }
55 
56 #[tokio::test]
sync_one_ident()57 async fn sync_one_ident() {
58     let one = one();
59 
60     let foo = tokio::select! {
61         foo = one => foo,
62     };
63 
64     assert_eq!(foo, 1);
65 }
66 
67 #[tokio::test]
sync_two()68 async fn sync_two() {
69     use std::cell::Cell;
70 
71     let cnt = Cell::new(0);
72 
73     let res = tokio::select! {
74         foo = async {
75             cnt.set(cnt.get() + 1);
76             1
77         } => foo,
78         bar = async {
79             cnt.set(cnt.get() + 1);
80             2
81         } => bar,
82     };
83 
84     assert_eq!(1, cnt.get());
85     assert!(res == 1 || res == 2);
86 }
87 
88 #[tokio::test]
drop_in_fut()89 async fn drop_in_fut() {
90     let s = "hello".to_string();
91 
92     let res = tokio::select! {
93         foo = async {
94             let v = one().await;
95             drop(s);
96             v
97         } => foo
98     };
99 
100     assert_eq!(res, 1);
101 }
102 
103 #[tokio::test]
one_ready()104 async fn one_ready() {
105     let (tx1, rx1) = oneshot::channel::<i32>();
106     let (_tx2, rx2) = oneshot::channel::<i32>();
107 
108     tx1.send(1).unwrap();
109 
110     let v = tokio::select! {
111         res = rx1 => {
112             assert_ok!(res)
113         },
114         _ = rx2 => unreachable!(),
115     };
116 
117     assert_eq!(1, v);
118 }
119 
120 #[tokio::test]
select_streams()121 async fn select_streams() {
122     let (tx1, mut rx1) = mpsc::unbounded_channel::<i32>();
123     let (tx2, mut rx2) = mpsc::unbounded_channel::<i32>();
124 
125     tokio::spawn(async move {
126         assert_ok!(tx2.send(1));
127         task::yield_now().await;
128 
129         assert_ok!(tx1.send(2));
130         task::yield_now().await;
131 
132         assert_ok!(tx2.send(3));
133         task::yield_now().await;
134 
135         drop((tx1, tx2));
136     });
137 
138     let mut rem = true;
139     let mut msgs = vec![];
140 
141     while rem {
142         tokio::select! {
143             Some(x) = rx1.recv() => {
144                 msgs.push(x);
145             }
146             Some(y) = rx2.recv() => {
147                 msgs.push(y);
148             }
149             else => {
150                 rem = false;
151             }
152         }
153     }
154 
155     msgs.sort_unstable();
156     assert_eq!(&msgs[..], &[1, 2, 3]);
157 }
158 
159 #[tokio::test]
move_uncompleted_futures()160 async fn move_uncompleted_futures() {
161     let (tx1, mut rx1) = oneshot::channel::<i32>();
162     let (tx2, mut rx2) = oneshot::channel::<i32>();
163 
164     tx1.send(1).unwrap();
165     tx2.send(2).unwrap();
166 
167     let ran;
168 
169     tokio::select! {
170         res = &mut rx1 => {
171             assert_eq!(1, assert_ok!(res));
172             assert_eq!(2, assert_ok!(rx2.await));
173             ran = true;
174         },
175         res = &mut rx2 => {
176             assert_eq!(2, assert_ok!(res));
177             assert_eq!(1, assert_ok!(rx1.await));
178             ran = true;
179         },
180     }
181 
182     assert!(ran);
183 }
184 
185 #[tokio::test]
nested()186 async fn nested() {
187     let res = tokio::select! {
188         x = async { 1 } => {
189             tokio::select! {
190                 y = async { 2 } => x + y,
191             }
192         }
193     };
194 
195     assert_eq!(res, 3);
196 }
197 
198 #[tokio::test]
struct_size()199 async fn struct_size() {
200     use futures::future;
201     use std::mem;
202 
203     let fut = async {
204         let ready = future::ready(0i32);
205 
206         tokio::select! {
207             _ = ready => {},
208         }
209     };
210 
211     assert!(mem::size_of_val(&fut) <= 32);
212 
213     let fut = async {
214         let ready1 = future::ready(0i32);
215         let ready2 = future::ready(0i32);
216 
217         tokio::select! {
218             _ = ready1 => {},
219             _ = ready2 => {},
220         }
221     };
222 
223     assert!(mem::size_of_val(&fut) <= 40);
224 
225     let fut = async {
226         let ready1 = future::ready(0i32);
227         let ready2 = future::ready(0i32);
228         let ready3 = future::ready(0i32);
229 
230         tokio::select! {
231             _ = ready1 => {},
232             _ = ready2 => {},
233             _ = ready3 => {},
234         }
235     };
236 
237     assert!(mem::size_of_val(&fut) <= 48);
238 }
239 
240 #[tokio::test]
mutable_borrowing_future_with_same_borrow_in_block()241 async fn mutable_borrowing_future_with_same_borrow_in_block() {
242     let mut value = 234;
243 
244     tokio::select! {
245         _ = require_mutable(&mut value) => { },
246         _ = async_noop() => {
247             value += 5;
248         },
249     }
250 
251     assert!(value >= 234);
252 }
253 
254 #[tokio::test]
mutable_borrowing_future_with_same_borrow_in_block_and_else()255 async fn mutable_borrowing_future_with_same_borrow_in_block_and_else() {
256     let mut value = 234;
257 
258     tokio::select! {
259         _ = require_mutable(&mut value) => { },
260         _ = async_noop() => {
261             value += 5;
262         },
263         else => {
264             value += 27;
265         },
266     }
267 
268     assert!(value >= 234);
269 }
270 
271 #[tokio::test]
future_panics_after_poll()272 async fn future_panics_after_poll() {
273     use tokio_test::task;
274 
275     let (tx, rx) = oneshot::channel();
276 
277     let mut polled = false;
278 
279     let f = poll_fn(|_| {
280         assert!(!polled);
281         polled = true;
282         Ready(None::<()>)
283     });
284 
285     let mut f = task::spawn(async {
286         tokio::select! {
287             Some(_) = f => unreachable!(),
288             ret = rx => ret.unwrap(),
289         }
290     });
291 
292     assert_pending!(f.poll());
293     assert_pending!(f.poll());
294 
295     assert_ok!(tx.send(1));
296 
297     let res = assert_ready!(f.poll());
298     assert_eq!(1, res);
299 }
300 
301 #[tokio::test]
disable_with_if()302 async fn disable_with_if() {
303     use tokio_test::task;
304 
305     let f = poll_fn(|_| panic!());
306     let (tx, rx) = oneshot::channel();
307 
308     let mut f = task::spawn(async {
309         tokio::select! {
310             _ = f, if false => unreachable!(),
311             _ = rx => (),
312         }
313     });
314 
315     assert_pending!(f.poll());
316 
317     assert_ok!(tx.send(()));
318     assert!(f.is_woken());
319 
320     assert_ready!(f.poll());
321 }
322 
323 #[tokio::test]
join_with_select()324 async fn join_with_select() {
325     use tokio_test::task;
326 
327     let (tx1, mut rx1) = oneshot::channel();
328     let (tx2, mut rx2) = oneshot::channel();
329 
330     let mut f = task::spawn(async {
331         let mut a = None;
332         let mut b = None;
333 
334         while a.is_none() || b.is_none() {
335             tokio::select! {
336                 v1 = &mut rx1, if a.is_none() => a = Some(assert_ok!(v1)),
337                 v2 = &mut rx2, if b.is_none() => b = Some(assert_ok!(v2))
338             }
339         }
340 
341         (a.unwrap(), b.unwrap())
342     });
343 
344     assert_pending!(f.poll());
345 
346     assert_ok!(tx1.send(123));
347     assert!(f.is_woken());
348     assert_pending!(f.poll());
349 
350     assert_ok!(tx2.send(456));
351     assert!(f.is_woken());
352     let (a, b) = assert_ready!(f.poll());
353 
354     assert_eq!(a, 123);
355     assert_eq!(b, 456);
356 }
357 
358 #[tokio::test]
use_future_in_if_condition()359 async fn use_future_in_if_condition() {
360     use tokio::time::{self, Duration};
361 
362     tokio::select! {
363         _ = time::sleep(Duration::from_millis(10)), if false => {
364             panic!("if condition ignored")
365         }
366         _ = async { 1u32 } => {
367         }
368     }
369 }
370 
371 #[tokio::test]
use_future_in_if_condition_biased()372 async fn use_future_in_if_condition_biased() {
373     use tokio::time::{self, Duration};
374 
375     tokio::select! {
376         biased;
377         _ = time::sleep(Duration::from_millis(10)), if false => {
378             panic!("if condition ignored")
379         }
380         _ = async { 1u32 } => {
381         }
382     }
383 }
384 
385 #[tokio::test]
many_branches()386 async fn many_branches() {
387     let num = tokio::select! {
388         x = async { 1 } => x,
389         x = async { 1 } => x,
390         x = async { 1 } => x,
391         x = async { 1 } => x,
392         x = async { 1 } => x,
393         x = async { 1 } => x,
394         x = async { 1 } => x,
395         x = async { 1 } => x,
396         x = async { 1 } => x,
397         x = async { 1 } => x,
398         x = async { 1 } => x,
399         x = async { 1 } => x,
400         x = async { 1 } => x,
401         x = async { 1 } => x,
402         x = async { 1 } => x,
403         x = async { 1 } => x,
404         x = async { 1 } => x,
405         x = async { 1 } => x,
406         x = async { 1 } => x,
407         x = async { 1 } => x,
408         x = async { 1 } => x,
409         x = async { 1 } => x,
410         x = async { 1 } => x,
411         x = async { 1 } => x,
412         x = async { 1 } => x,
413         x = async { 1 } => x,
414         x = async { 1 } => x,
415         x = async { 1 } => x,
416         x = async { 1 } => x,
417         x = async { 1 } => x,
418         x = async { 1 } => x,
419         x = async { 1 } => x,
420         x = async { 1 } => x,
421         x = async { 1 } => x,
422         x = async { 1 } => x,
423         x = async { 1 } => x,
424         x = async { 1 } => x,
425         x = async { 1 } => x,
426         x = async { 1 } => x,
427         x = async { 1 } => x,
428         x = async { 1 } => x,
429         x = async { 1 } => x,
430         x = async { 1 } => x,
431         x = async { 1 } => x,
432         x = async { 1 } => x,
433         x = async { 1 } => x,
434         x = async { 1 } => x,
435         x = async { 1 } => x,
436         x = async { 1 } => x,
437         x = async { 1 } => x,
438         x = async { 1 } => x,
439         x = async { 1 } => x,
440         x = async { 1 } => x,
441         x = async { 1 } => x,
442         x = async { 1 } => x,
443         x = async { 1 } => x,
444         x = async { 1 } => x,
445         x = async { 1 } => x,
446         x = async { 1 } => x,
447         x = async { 1 } => x,
448         x = async { 1 } => x,
449         x = async { 1 } => x,
450         x = async { 1 } => x,
451     };
452 
453     assert_eq!(1, num);
454 }
455 
456 #[tokio::test]
never_branch_no_warnings()457 async fn never_branch_no_warnings() {
458     let t = tokio::select! {
459         _ = async_never() => 0,
460         one_async_ready = one() => one_async_ready,
461     };
462     assert_eq!(t, 1);
463 }
464 
one() -> usize465 async fn one() -> usize {
466     1
467 }
468 
require_mutable(_: &mut i32)469 async fn require_mutable(_: &mut i32) {}
async_noop()470 async fn async_noop() {}
471 
async_never() -> !472 async fn async_never() -> ! {
473     futures::future::pending().await
474 }
475 
476 // From https://github.com/tokio-rs/tokio/issues/2857
477 #[tokio::test]
mut_on_left_hand_side()478 async fn mut_on_left_hand_side() {
479     let v = async move {
480         let ok = async { 1 };
481         tokio::pin!(ok);
482         tokio::select! {
483             mut a = &mut ok => {
484                 a += 1;
485                 a
486             }
487         }
488     }
489     .await;
490     assert_eq!(v, 2);
491 }
492 
493 #[tokio::test]
biased_one_not_ready()494 async fn biased_one_not_ready() {
495     let (_tx1, rx1) = oneshot::channel::<i32>();
496     let (tx2, rx2) = oneshot::channel::<i32>();
497     let (tx3, rx3) = oneshot::channel::<i32>();
498 
499     tx2.send(2).unwrap();
500     tx3.send(3).unwrap();
501 
502     let v = tokio::select! {
503         biased;
504 
505         _ = rx1 => unreachable!(),
506         res = rx2 => {
507             assert_ok!(res)
508         },
509         _ = rx3 => {
510             panic!("This branch should never be activated because `rx2` should be polled before `rx3` due to `biased;`.")
511         }
512     };
513 
514     assert_eq!(2, v);
515 }
516 
517 #[tokio::test]
biased_eventually_ready()518 async fn biased_eventually_ready() {
519     use tokio::task::yield_now;
520 
521     let one = async {};
522     let two = async { yield_now().await };
523     let three = async { yield_now().await };
524 
525     let mut count = 0u8;
526 
527     tokio::pin!(one, two, three);
528 
529     loop {
530         tokio::select! {
531             biased;
532 
533             _ = &mut two, if count < 2 => {
534                 count += 1;
535                 assert_eq!(count, 2);
536             }
537             _ = &mut three, if count < 3 => {
538                 count += 1;
539                 assert_eq!(count, 3);
540             }
541             _ = &mut one, if count < 1 => {
542                 count += 1;
543                 assert_eq!(count, 1);
544             }
545             else => break,
546         }
547     }
548 
549     assert_eq!(count, 3);
550 }
551 
552 // https://github.com/tokio-rs/tokio/issues/3830
553 // https://github.com/rust-lang/rust-clippy/issues/7304
554 #[warn(clippy::default_numeric_fallback)]
default_numeric_fallback()555 pub async fn default_numeric_fallback() {
556     tokio::select! {
557         _ = async {} => (),
558         else => (),
559     }
560 }
561