1 use futures::channel::{mpsc, oneshot};
2 use futures::executor::{block_on, block_on_stream};
3 use futures::future::{poll_fn, FutureExt};
4 use futures::pin_mut;
5 use futures::sink::{Sink, SinkExt};
6 use futures::stream::{Stream, StreamExt};
7 use futures::task::{Context, Poll};
8 use futures_test::task::{new_count_waker, noop_context};
9 use std::sync::atomic::{AtomicUsize, Ordering};
10 use std::sync::{Arc, Mutex};
11 use std::thread;
12 
13 trait AssertSend: Send {}
14 impl AssertSend for mpsc::Sender<i32> {}
15 impl AssertSend for mpsc::Receiver<i32> {}
16 
17 #[test]
send_recv()18 fn send_recv() {
19     let (mut tx, rx) = mpsc::channel::<i32>(16);
20 
21     block_on(tx.send(1)).unwrap();
22     drop(tx);
23     let v: Vec<_> = block_on(rx.collect());
24     assert_eq!(v, vec![1]);
25 }
26 
27 #[test]
send_recv_no_buffer()28 fn send_recv_no_buffer() {
29     // Run on a task context
30     block_on(poll_fn(move |cx| {
31         let (tx, rx) = mpsc::channel::<i32>(0);
32         pin_mut!(tx, rx);
33 
34         assert!(tx.as_mut().poll_flush(cx).is_ready());
35         assert!(tx.as_mut().poll_ready(cx).is_ready());
36 
37         // Send first message
38         assert!(tx.as_mut().start_send(1).is_ok());
39         assert!(tx.as_mut().poll_ready(cx).is_pending());
40 
41         // poll_ready said Pending, so no room in buffer, therefore new sends
42         // should get rejected with is_full.
43         assert!(tx.as_mut().start_send(0).unwrap_err().is_full());
44         assert!(tx.as_mut().poll_ready(cx).is_pending());
45 
46         // Take the value
47         assert_eq!(rx.as_mut().poll_next(cx), Poll::Ready(Some(1)));
48         assert!(tx.as_mut().poll_ready(cx).is_ready());
49 
50         // Send second message
51         assert!(tx.as_mut().poll_ready(cx).is_ready());
52         assert!(tx.as_mut().start_send(2).is_ok());
53         assert!(tx.as_mut().poll_ready(cx).is_pending());
54 
55         // Take the value
56         assert_eq!(rx.as_mut().poll_next(cx), Poll::Ready(Some(2)));
57         assert!(tx.as_mut().poll_ready(cx).is_ready());
58 
59         Poll::Ready(())
60     }));
61 }
62 
63 #[test]
send_shared_recv()64 fn send_shared_recv() {
65     let (mut tx1, rx) = mpsc::channel::<i32>(16);
66     let mut rx = block_on_stream(rx);
67     let mut tx2 = tx1.clone();
68 
69     block_on(tx1.send(1)).unwrap();
70     assert_eq!(rx.next(), Some(1));
71 
72     block_on(tx2.send(2)).unwrap();
73     assert_eq!(rx.next(), Some(2));
74 }
75 
76 #[test]
send_recv_threads()77 fn send_recv_threads() {
78     let (mut tx, rx) = mpsc::channel::<i32>(16);
79 
80     let t = thread::spawn(move || {
81         block_on(tx.send(1)).unwrap();
82     });
83 
84     let v: Vec<_> = block_on(rx.take(1).collect());
85     assert_eq!(v, vec![1]);
86 
87     t.join().unwrap();
88 }
89 
90 #[test]
send_recv_threads_no_capacity()91 fn send_recv_threads_no_capacity() {
92     let (mut tx, rx) = mpsc::channel::<i32>(0);
93 
94     let t = thread::spawn(move || {
95         block_on(tx.send(1)).unwrap();
96         block_on(tx.send(2)).unwrap();
97     });
98 
99     let v: Vec<_> = block_on(rx.collect());
100     assert_eq!(v, vec![1, 2]);
101 
102     t.join().unwrap();
103 }
104 
105 #[test]
recv_close_gets_none()106 fn recv_close_gets_none() {
107     let (mut tx, mut rx) = mpsc::channel::<i32>(10);
108 
109     // Run on a task context
110     block_on(poll_fn(move |cx| {
111         rx.close();
112 
113         assert_eq!(rx.poll_next_unpin(cx), Poll::Ready(None));
114         match tx.poll_ready(cx) {
115             Poll::Pending | Poll::Ready(Ok(_)) => panic!(),
116             Poll::Ready(Err(e)) => assert!(e.is_disconnected()),
117         };
118 
119         Poll::Ready(())
120     }));
121 }
122 
123 #[test]
tx_close_gets_none()124 fn tx_close_gets_none() {
125     let (_, mut rx) = mpsc::channel::<i32>(10);
126 
127     // Run on a task context
128     block_on(poll_fn(move |cx| {
129         assert_eq!(rx.poll_next_unpin(cx), Poll::Ready(None));
130         Poll::Ready(())
131     }));
132 }
133 
134 // #[test]
135 // fn spawn_sends_items() {
136 //     let core = local_executor::Core::new();
137 //     let stream = unfold(0, |i| Some(ok::<_,u8>((i, i + 1))));
138 //     let rx = mpsc::spawn(stream, &core, 1);
139 //     assert_eq!(core.run(rx.take(4).collect()).unwrap(),
140 //                [0, 1, 2, 3]);
141 // }
142 
143 // #[test]
144 // fn spawn_kill_dead_stream() {
145 //     use std::thread;
146 //     use std::time::Duration;
147 //     use futures::future::Either;
148 //     use futures::sync::oneshot;
149 //
150 //     // a stream which never returns anything (maybe a remote end isn't
151 //     // responding), but dropping it leads to observable side effects
152 //     // (like closing connections, releasing limited resources, ...)
153 //     #[derive(Debug)]
154 //     struct Dead {
155 //         // when dropped you should get Err(oneshot::Canceled) on the
156 //         // receiving end
157 //         done: oneshot::Sender<()>,
158 //     }
159 //     impl Stream for Dead {
160 //         type Item = ();
161 //         type Error = ();
162 //
163 //         fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
164 //             Ok(Poll::Pending)
165 //         }
166 //     }
167 //
168 //     // need to implement a timeout for the test, as it would hang
169 //     // forever right now
170 //     let (timeout_tx, timeout_rx) = oneshot::channel();
171 //     thread::spawn(move || {
172 //         thread::sleep(Duration::from_millis(1000));
173 //         let _ = timeout_tx.send(());
174 //     });
175 //
176 //     let core = local_executor::Core::new();
177 //     let (done_tx, done_rx) = oneshot::channel();
178 //     let stream = Dead{done: done_tx};
179 //     let rx = mpsc::spawn(stream, &core, 1);
180 //     let res = core.run(
181 //         Ok::<_, ()>(())
182 //         .into_future()
183 //         .then(move |_| {
184 //             // now drop the spawned stream: maybe some timeout exceeded,
185 //             // or some connection on this end was closed by the remote
186 //             // end.
187 //             drop(rx);
188 //             // and wait for the spawned stream to release its resources
189 //             done_rx
190 //         })
191 //         .select2(timeout_rx)
192 //     );
193 //     match res {
194 //         Err(Either::A((oneshot::Canceled, _))) => (),
195 //         _ => {
196 //             panic!("dead stream wasn't canceled");
197 //         },
198 //     }
199 // }
200 
201 #[test]
stress_shared_unbounded()202 fn stress_shared_unbounded() {
203     const AMT: u32 = 10000;
204     const NTHREADS: u32 = 8;
205     let (tx, rx) = mpsc::unbounded::<i32>();
206 
207     let t = thread::spawn(move || {
208         let result: Vec<_> = block_on(rx.collect());
209         assert_eq!(result.len(), (AMT * NTHREADS) as usize);
210         for item in result {
211             assert_eq!(item, 1);
212         }
213     });
214 
215     for _ in 0..NTHREADS {
216         let tx = tx.clone();
217 
218         thread::spawn(move || {
219             for _ in 0..AMT {
220                 tx.unbounded_send(1).unwrap();
221             }
222         });
223     }
224 
225     drop(tx);
226 
227     t.join().ok().unwrap();
228 }
229 
230 #[test]
stress_shared_bounded_hard()231 fn stress_shared_bounded_hard() {
232     const AMT: u32 = 10000;
233     const NTHREADS: u32 = 8;
234     let (tx, rx) = mpsc::channel::<i32>(0);
235 
236     let t = thread::spawn(move || {
237         let result: Vec<_> = block_on(rx.collect());
238         assert_eq!(result.len(), (AMT * NTHREADS) as usize);
239         for item in result {
240             assert_eq!(item, 1);
241         }
242     });
243 
244     for _ in 0..NTHREADS {
245         let mut tx = tx.clone();
246 
247         thread::spawn(move || {
248             for _ in 0..AMT {
249                 block_on(tx.send(1)).unwrap();
250             }
251         });
252     }
253 
254     drop(tx);
255 
256     t.join().unwrap();
257 }
258 
259 #[allow(clippy::same_item_push)]
260 #[test]
stress_receiver_multi_task_bounded_hard()261 fn stress_receiver_multi_task_bounded_hard() {
262     const AMT: usize = 10_000;
263     const NTHREADS: u32 = 2;
264 
265     let (mut tx, rx) = mpsc::channel::<usize>(0);
266     let rx = Arc::new(Mutex::new(Some(rx)));
267     let n = Arc::new(AtomicUsize::new(0));
268 
269     let mut th = vec![];
270 
271     for _ in 0..NTHREADS {
272         let rx = rx.clone();
273         let n = n.clone();
274 
275         let t = thread::spawn(move || {
276             let mut i = 0;
277 
278             loop {
279                 i += 1;
280                 let mut rx_opt = rx.lock().unwrap();
281                 if let Some(rx) = &mut *rx_opt {
282                     if i % 5 == 0 {
283                         let item = block_on(rx.next());
284 
285                         if item.is_none() {
286                             *rx_opt = None;
287                             break;
288                         }
289 
290                         n.fetch_add(1, Ordering::Relaxed);
291                     } else {
292                         // Just poll
293                         let n = n.clone();
294                         match rx.poll_next_unpin(&mut noop_context()) {
295                             Poll::Ready(Some(_)) => {
296                                 n.fetch_add(1, Ordering::Relaxed);
297                             }
298                             Poll::Ready(None) => {
299                                 *rx_opt = None;
300                                 break;
301                             }
302                             Poll::Pending => {}
303                         }
304                     }
305                 } else {
306                     break;
307                 }
308             }
309         });
310 
311         th.push(t);
312     }
313 
314     for i in 0..AMT {
315         block_on(tx.send(i)).unwrap();
316     }
317     drop(tx);
318 
319     for t in th {
320         t.join().unwrap();
321     }
322 
323     assert_eq!(AMT, n.load(Ordering::Relaxed));
324 }
325 
326 /// Stress test that receiver properly receives all the messages
327 /// after sender dropped.
328 #[test]
stress_drop_sender()329 fn stress_drop_sender() {
330     fn list() -> impl Stream<Item = i32> {
331         let (tx, rx) = mpsc::channel(1);
332         thread::spawn(move || {
333             block_on(send_one_two_three(tx));
334         });
335         rx
336     }
337 
338     for _ in 0..10000 {
339         let v: Vec<_> = block_on(list().collect());
340         assert_eq!(v, vec![1, 2, 3]);
341     }
342 }
343 
send_one_two_three(mut tx: mpsc::Sender<i32>)344 async fn send_one_two_three(mut tx: mpsc::Sender<i32>) {
345     for i in 1..=3 {
346         tx.send(i).await.unwrap();
347     }
348 }
349 
350 /// Stress test that after receiver dropped,
351 /// no messages are lost.
stress_close_receiver_iter()352 fn stress_close_receiver_iter() {
353     let (tx, rx) = mpsc::unbounded();
354     let mut rx = block_on_stream(rx);
355     let (unwritten_tx, unwritten_rx) = std::sync::mpsc::channel();
356     let th = thread::spawn(move || {
357         for i in 1.. {
358             if tx.unbounded_send(i).is_err() {
359                 unwritten_tx.send(i).expect("unwritten_tx");
360                 return;
361             }
362         }
363     });
364 
365     // Read one message to make sure thread effectively started
366     assert_eq!(Some(1), rx.next());
367 
368     rx.close();
369 
370     for i in 2.. {
371         match rx.next() {
372             Some(r) => assert!(i == r),
373             None => {
374                 let unwritten = unwritten_rx.recv().expect("unwritten_rx");
375                 assert_eq!(unwritten, i);
376                 th.join().unwrap();
377                 return;
378             }
379         }
380     }
381 }
382 
383 #[test]
stress_close_receiver()384 fn stress_close_receiver() {
385     for _ in 0..10000 {
386         stress_close_receiver_iter();
387     }
388 }
389 
stress_poll_ready_sender(mut sender: mpsc::Sender<u32>, count: u32)390 async fn stress_poll_ready_sender(mut sender: mpsc::Sender<u32>, count: u32) {
391     for i in (1..=count).rev() {
392         sender.send(i).await.unwrap();
393     }
394 }
395 
396 /// Tests that after `poll_ready` indicates capacity a channel can always send without waiting.
397 #[allow(clippy::same_item_push)]
398 #[test]
stress_poll_ready()399 fn stress_poll_ready() {
400     const AMT: u32 = 1000;
401     const NTHREADS: u32 = 8;
402 
403     /// Run a stress test using the specified channel capacity.
404     fn stress(capacity: usize) {
405         let (tx, rx) = mpsc::channel(capacity);
406         let mut threads = Vec::new();
407         for _ in 0..NTHREADS {
408             let sender = tx.clone();
409             threads.push(thread::spawn(move || block_on(stress_poll_ready_sender(sender, AMT))));
410         }
411         drop(tx);
412 
413         let result: Vec<_> = block_on(rx.collect());
414         assert_eq!(result.len() as u32, AMT * NTHREADS);
415 
416         for thread in threads {
417             thread.join().unwrap();
418         }
419     }
420 
421     stress(0);
422     stress(1);
423     stress(8);
424     stress(16);
425 }
426 
427 #[test]
try_send_1()428 fn try_send_1() {
429     const N: usize = 3000;
430     let (mut tx, rx) = mpsc::channel(0);
431 
432     let t = thread::spawn(move || {
433         for i in 0..N {
434             loop {
435                 if tx.try_send(i).is_ok() {
436                     break;
437                 }
438             }
439         }
440     });
441 
442     let result: Vec<_> = block_on(rx.collect());
443     for (i, j) in result.into_iter().enumerate() {
444         assert_eq!(i, j);
445     }
446 
447     t.join().unwrap();
448 }
449 
450 #[test]
try_send_2()451 fn try_send_2() {
452     let (mut tx, rx) = mpsc::channel(0);
453     let mut rx = block_on_stream(rx);
454 
455     tx.try_send("hello").unwrap();
456 
457     let (readytx, readyrx) = oneshot::channel::<()>();
458 
459     let th = thread::spawn(move || {
460         block_on(poll_fn(|cx| {
461             assert!(tx.poll_ready(cx).is_pending());
462             Poll::Ready(())
463         }));
464 
465         drop(readytx);
466         block_on(tx.send("goodbye")).unwrap();
467     });
468 
469     let _ = block_on(readyrx);
470     assert_eq!(rx.next(), Some("hello"));
471     assert_eq!(rx.next(), Some("goodbye"));
472     assert_eq!(rx.next(), None);
473 
474     th.join().unwrap();
475 }
476 
477 #[test]
try_send_fail()478 fn try_send_fail() {
479     let (mut tx, rx) = mpsc::channel(0);
480     let mut rx = block_on_stream(rx);
481 
482     tx.try_send("hello").unwrap();
483 
484     // This should fail
485     assert!(tx.try_send("fail").is_err());
486 
487     assert_eq!(rx.next(), Some("hello"));
488 
489     tx.try_send("goodbye").unwrap();
490     drop(tx);
491 
492     assert_eq!(rx.next(), Some("goodbye"));
493     assert_eq!(rx.next(), None);
494 }
495 
496 #[test]
try_send_recv()497 fn try_send_recv() {
498     let (mut tx, mut rx) = mpsc::channel(1);
499     tx.try_send("hello").unwrap();
500     tx.try_send("hello").unwrap();
501     tx.try_send("hello").unwrap_err(); // should be full
502     rx.try_next().unwrap();
503     rx.try_next().unwrap();
504     rx.try_next().unwrap_err(); // should be empty
505     tx.try_send("hello").unwrap();
506     rx.try_next().unwrap();
507     rx.try_next().unwrap_err(); // should be empty
508 }
509 
510 #[test]
same_receiver()511 fn same_receiver() {
512     let (mut txa1, _) = mpsc::channel::<i32>(1);
513     let txa2 = txa1.clone();
514 
515     let (mut txb1, _) = mpsc::channel::<i32>(1);
516     let txb2 = txb1.clone();
517 
518     assert!(txa1.same_receiver(&txa2));
519     assert!(txb1.same_receiver(&txb2));
520     assert!(!txa1.same_receiver(&txb1));
521 
522     txa1.disconnect();
523     txb1.close_channel();
524 
525     assert!(!txa1.same_receiver(&txa2));
526     assert!(txb1.same_receiver(&txb2));
527 }
528 
529 #[test]
is_connected_to()530 fn is_connected_to() {
531     let (txa, rxa) = mpsc::channel::<i32>(1);
532     let (txb, rxb) = mpsc::channel::<i32>(1);
533 
534     assert!(txa.is_connected_to(&rxa));
535     assert!(txb.is_connected_to(&rxb));
536     assert!(!txa.is_connected_to(&rxb));
537     assert!(!txb.is_connected_to(&rxa));
538 }
539 
540 #[test]
hash_receiver()541 fn hash_receiver() {
542     use std::collections::hash_map::DefaultHasher;
543     use std::hash::Hasher;
544 
545     let mut hasher_a1 = DefaultHasher::new();
546     let mut hasher_a2 = DefaultHasher::new();
547     let mut hasher_b1 = DefaultHasher::new();
548     let mut hasher_b2 = DefaultHasher::new();
549     let (mut txa1, _) = mpsc::channel::<i32>(1);
550     let txa2 = txa1.clone();
551 
552     let (mut txb1, _) = mpsc::channel::<i32>(1);
553     let txb2 = txb1.clone();
554 
555     txa1.hash_receiver(&mut hasher_a1);
556     let hash_a1 = hasher_a1.finish();
557     txa2.hash_receiver(&mut hasher_a2);
558     let hash_a2 = hasher_a2.finish();
559     txb1.hash_receiver(&mut hasher_b1);
560     let hash_b1 = hasher_b1.finish();
561     txb2.hash_receiver(&mut hasher_b2);
562     let hash_b2 = hasher_b2.finish();
563 
564     assert_eq!(hash_a1, hash_a2);
565     assert_eq!(hash_b1, hash_b2);
566     assert!(hash_a1 != hash_b1);
567 
568     txa1.disconnect();
569     txb1.close_channel();
570 
571     let mut hasher_a1 = DefaultHasher::new();
572     let mut hasher_a2 = DefaultHasher::new();
573     let mut hasher_b1 = DefaultHasher::new();
574     let mut hasher_b2 = DefaultHasher::new();
575 
576     txa1.hash_receiver(&mut hasher_a1);
577     let hash_a1 = hasher_a1.finish();
578     txa2.hash_receiver(&mut hasher_a2);
579     let hash_a2 = hasher_a2.finish();
580     txb1.hash_receiver(&mut hasher_b1);
581     let hash_b1 = hasher_b1.finish();
582     txb2.hash_receiver(&mut hasher_b2);
583     let hash_b2 = hasher_b2.finish();
584 
585     assert!(hash_a1 != hash_a2);
586     assert_eq!(hash_b1, hash_b2);
587 }
588 
589 #[test]
send_backpressure()590 fn send_backpressure() {
591     let (waker, counter) = new_count_waker();
592     let mut cx = Context::from_waker(&waker);
593 
594     let (mut tx, mut rx) = mpsc::channel(1);
595     block_on(tx.send(1)).unwrap();
596 
597     let mut task = tx.send(2);
598     assert_eq!(task.poll_unpin(&mut cx), Poll::Pending);
599     assert_eq!(counter, 0);
600 
601     let item = block_on(rx.next()).unwrap();
602     assert_eq!(item, 1);
603     assert_eq!(counter, 1);
604     assert_eq!(task.poll_unpin(&mut cx), Poll::Ready(Ok(())));
605 
606     let item = block_on(rx.next()).unwrap();
607     assert_eq!(item, 2);
608 }
609 
610 #[test]
send_backpressure_multi_senders()611 fn send_backpressure_multi_senders() {
612     let (waker, counter) = new_count_waker();
613     let mut cx = Context::from_waker(&waker);
614 
615     let (mut tx1, mut rx) = mpsc::channel(1);
616     let mut tx2 = tx1.clone();
617     block_on(tx1.send(1)).unwrap();
618 
619     let mut task = tx2.send(2);
620     assert_eq!(task.poll_unpin(&mut cx), Poll::Pending);
621     assert_eq!(counter, 0);
622 
623     let item = block_on(rx.next()).unwrap();
624     assert_eq!(item, 1);
625     assert_eq!(counter, 1);
626     assert_eq!(task.poll_unpin(&mut cx), Poll::Ready(Ok(())));
627 
628     let item = block_on(rx.next()).unwrap();
629     assert_eq!(item, 2);
630 }
631