1#![cfg(feature = "macros")]
2#![allow(clippy::disallowed_names)]
3
4use std::sync::Arc;
5
6use tokio::sync::{oneshot, Semaphore};
7use tokio_test::{assert_pending, assert_ready, task};
8
9#[cfg(all(target_family = "wasm", not(target_os = "wasi")))]
10use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test;
11
12#[cfg(not(all(target_family = "wasm", not(target_os = "wasi"))))]
13use tokio::test as maybe_tokio_test;
14
15#[maybe_tokio_test]
16async fn sync_one_lit_expr_comma() {
17 let foo = tokio::try_join!(async { ok(1) },);
18
19 assert_eq!(foo, Ok((1,)));
20}
21
22#[maybe_tokio_test]
23async fn sync_one_lit_expr_no_comma() {
24 let foo = tokio::try_join!(async { ok(1) });
25
26 assert_eq!(foo, Ok((1,)));
27}
28
29#[maybe_tokio_test]
30async fn sync_two_lit_expr_comma() {
31 let foo = tokio::try_join!(async { ok(1) }, async { ok(2) },);
32
33 assert_eq!(foo, Ok((1, 2)));
34}
35
36#[maybe_tokio_test]
37async fn sync_two_lit_expr_no_comma() {
38 let foo = tokio::try_join!(async { ok(1) }, async { ok(2) });
39
40 assert_eq!(foo, Ok((1, 2)));
41}
42
43#[maybe_tokio_test]
44async fn two_await() {
45 let (tx1, rx1) = oneshot::channel::<&str>();
46 let (tx2, rx2) = oneshot::channel::<u32>();
47
48 let mut join =
49 task::spawn(async { tokio::try_join!(async { rx1.await }, async { rx2.await }) });
50
51 assert_pending!(join.poll());
52
53 tx2.send(123).unwrap();
54 assert!(join.is_woken());
55 assert_pending!(join.poll());
56
57 tx1.send("hello").unwrap();
58 assert!(join.is_woken());
59 let res: Result<(&str, u32), _> = assert_ready!(join.poll());
60
61 assert_eq!(Ok(("hello", 123)), res);
62}
63
64#[maybe_tokio_test]
65async fn err_abort_early() {
66 let (tx1, rx1) = oneshot::channel::<&str>();
67 let (tx2, rx2) = oneshot::channel::<u32>();
68 let (_tx3, rx3) = oneshot::channel::<u32>();
69
70 let mut join = task::spawn(async {
71 tokio::try_join!(async { rx1.await }, async { rx2.await }, async {
72 rx3.await
73 })
74 });
75
76 assert_pending!(join.poll());
77
78 tx2.send(123).unwrap();
79 assert!(join.is_woken());
80 assert_pending!(join.poll());
81
82 drop(tx1);
83 assert!(join.is_woken());
84
85 let res = assert_ready!(join.poll());
86
87 assert!(res.is_err());
88}
89
90#[test]
91#[cfg(target_pointer_width = "64")]
92fn join_size() {
93 use futures::future;
94 use std::mem;
95
96 let fut = async {
97 let ready = future::ready(ok(0i32));
98 tokio::try_join!(ready)
99 };
100 assert_eq!(mem::size_of_val(&fut), 32);
101
102 let fut = async {
103 let ready1 = future::ready(ok(0i32));
104 let ready2 = future::ready(ok(0i32));
105 tokio::try_join!(ready1, ready2)
106 };
107 assert_eq!(mem::size_of_val(&fut), 48);
108}
109
110fn ok<T>(val: T) -> Result<T, ()> {
111 Ok(val)
112}
113
114async fn non_cooperative_task(permits: Arc<Semaphore>) -> Result<usize, String> {
115 let mut exceeded_budget = 0;
116
117 for _ in 0..5 {
118 // Another task should run after this task uses its whole budget
119 for _ in 0..128 {
120 let _permit = permits.clone().acquire_owned().await.unwrap();
121 }
122
123 exceeded_budget += 1;
124 }
125
126 Ok(exceeded_budget)
127}
128
129async fn poor_little_task(permits: Arc<Semaphore>) -> Result<usize, String> {
130 let mut how_many_times_i_got_to_run = 0;
131
132 for _ in 0..5 {
133 let _permit = permits.clone().acquire_owned().await.unwrap();
134
135 how_many_times_i_got_to_run += 1;
136 }
137
138 Ok(how_many_times_i_got_to_run)
139}
140
141#[tokio::test]
142async fn try_join_does_not_allow_tasks_to_starve() {
143 let permits = Arc::new(Semaphore::new(10));
144
145 // non_cooperative_task should yield after its budget is exceeded and then poor_little_task should run.
146 let result = tokio::try_join!(
147 non_cooperative_task(Arc::clone(&permits)),
148 poor_little_task(permits)
149 );
150
151 let (non_cooperative_result, little_task_result) = result.unwrap();
152
153 assert_eq!(5, non_cooperative_result);
154 assert_eq!(5, little_task_result);
155}
156
157#[tokio::test]
158async fn a_different_future_is_polled_first_every_time_poll_fn_is_polled() {
159 let poll_order = Arc::new(std::sync::Mutex::new(vec![]));
160
161 let fut = |x, poll_order: Arc<std::sync::Mutex<Vec<i32>>>| async move {
162 for _ in 0..4 {
163 {
164 let mut guard = poll_order.lock().unwrap();
165
166 guard.push(x);
167 }
168
169 tokio::task::yield_now().await;
170 }
171 };
172
173 tokio::join!(
174 fut(1, Arc::clone(&poll_order)),
175 fut(2, Arc::clone(&poll_order)),
176 fut(3, Arc::clone(&poll_order)),
177 );
178
179 // Each time the future created by join! is polled, it should start
180 // by polling a different future first.
181 assert_eq!(
182 vec![1, 2, 3, 2, 3, 1, 3, 1, 2, 1, 2, 3],
183 *poll_order.lock().unwrap()
184 );
185}
186
187#[tokio::test]
188async fn empty_try_join() {
189 assert_eq!(tokio::try_join!() as Result<_, ()>, Ok(()));
190}
191