1 #![cfg(feature = "async")]
2 
3 use core::{future, mem, pin, task};
4 
5 #[cfg(loom)]
6 pub use loom::sync::{Arc, Mutex};
7 #[cfg(not(loom))]
8 pub use std::sync::{Arc, Mutex};
9 
10 mod helpers;
11 use helpers::maybe_loom_model;
12 
13 #[test]
multiple_receiver_polls_keeps_only_latest_waker()14 fn multiple_receiver_polls_keeps_only_latest_waker() {
15     #[derive(Default)]
16     struct MockWaker {
17         cloned: usize,
18         dropped: usize,
19     }
20 
21     fn clone_mock_waker(waker: *const ()) -> task::RawWaker {
22         let mock_waker = unsafe { Arc::from_raw(waker as *const Mutex<MockWaker>) };
23         mock_waker.lock().unwrap().cloned += 1;
24         let new_waker =
25             task::RawWaker::new(Arc::into_raw(mock_waker.clone()) as *const (), &VTABLE);
26         mem::forget(mock_waker);
27         new_waker
28     }
29 
30     fn drop_mock_waker(waker: *const ()) {
31         let mock_waker = unsafe { Arc::from_raw(waker as *const Mutex<MockWaker>) };
32         mock_waker.lock().unwrap().dropped += 1;
33     }
34 
35     const VTABLE: task::RawWakerVTable =
36         task::RawWakerVTable::new(clone_mock_waker, |_| (), |_| (), drop_mock_waker);
37 
38     maybe_loom_model(|| {
39         let mock_waker1 = Arc::new(Mutex::new(MockWaker::default()));
40         let raw_waker1 =
41             task::RawWaker::new(Arc::into_raw(mock_waker1.clone()) as *const (), &VTABLE);
42         let waker1 = unsafe { task::Waker::from_raw(raw_waker1) };
43         let mut context1 = task::Context::from_waker(&waker1);
44 
45         let (_sender, mut receiver) = oneshot::channel::<()>();
46 
47         let poll_result = future::Future::poll(pin::Pin::new(&mut receiver), &mut context1);
48         assert_eq!(poll_result, task::Poll::Pending);
49         assert_eq!(mock_waker1.lock().unwrap().cloned, 1);
50         assert_eq!(mock_waker1.lock().unwrap().dropped, 0);
51 
52         let mock_waker2 = Arc::new(Mutex::new(MockWaker::default()));
53         let raw_waker2 =
54             task::RawWaker::new(Arc::into_raw(mock_waker2.clone()) as *const (), &VTABLE);
55         let waker2 = unsafe { task::Waker::from_raw(raw_waker2) };
56         let mut context2 = task::Context::from_waker(&waker2);
57 
58         let poll_result = future::Future::poll(pin::Pin::new(&mut receiver), &mut context2);
59         assert_eq!(poll_result, task::Poll::Pending);
60         assert_eq!(mock_waker2.lock().unwrap().cloned, 1);
61         assert_eq!(mock_waker2.lock().unwrap().dropped, 0);
62         assert_eq!(mock_waker1.lock().unwrap().cloned, 1);
63         assert_eq!(mock_waker1.lock().unwrap().dropped, 1);
64     });
65 }
66