use std::cell::Cell; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::thread; use std::time::Duration; use async_task_ffi::Runnable; use atomic_waker::AtomicWaker; // Creates a future with event counters. // // Usage: `future!(f, get_waker, POLL, DROP)` // // The future `f` always sleeps for 200 ms, and returns `Poll::Ready` the second // time it is polled. When it gets polled, `POLL` is incremented. // When it gets dropped, `DROP` is incremented. // // Every time the future is run, it stores the waker into a global variable. // This waker can be extracted using the `get_waker()` function. macro_rules! future { ($name:pat, $get_waker:pat, $poll:ident, $drop:ident) => { static $poll: AtomicUsize = AtomicUsize::new(0); static $drop: AtomicUsize = AtomicUsize::new(0); static WAKER: AtomicWaker = AtomicWaker::new(); let ($name, $get_waker) = { struct Fut(Cell, Box); impl Future for Fut { type Output = Box; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { WAKER.register(cx.waker()); $poll.fetch_add(1, Ordering::SeqCst); thread::sleep(ms(200)); if self.0.get() { Poll::Ready(Box::new(0)) } else { self.0.set(true); Poll::Pending } } } impl Drop for Fut { fn drop(&mut self) { $drop.fetch_add(1, Ordering::SeqCst); } } (Fut(Cell::new(false), Box::new(0)), || WAKER.take().unwrap()) }; }; } // Creates a schedule function with event counters. // // Usage: `schedule!(s, d, chan, SCHED, DATA, DROP)` // // The schedule function `s` pushes the task into `chan`. // When it gets invoked, `SCHED` and `DATA` are incremented. // When it gets dropped, `DROP` is incremented. // When it gets dropped, `DROP` is incremented. // The user data `d` references `DATA`. // // Receiver `chan` extracts the task when it is scheduled. macro_rules! schedule { ($sched_name:pat, $data_name:pat, $chan:pat, $sched:ident, $data:ident, $drop:ident) => { static $drop: AtomicUsize = AtomicUsize::new(0); static $sched: AtomicUsize = AtomicUsize::new(0); static $data: AtomicUsize = AtomicUsize::new(0); let ($sched_name, $data_name, $chan) = { let (s, r) = flume::unbounded(); struct Guard(Box); impl Drop for Guard { fn drop(&mut self) { $drop.fetch_add(1, Ordering::SeqCst); } } struct Data(Box<&'static AtomicUsize>); impl Drop for Data { fn drop(&mut self) { $data.fetch_add(1, Ordering::SeqCst); } } let guard = Guard(Box::new(0)); let sched = move |runnable: Runnable| { &guard; $sched.fetch_add(1, Ordering::SeqCst); runnable.data().0.fetch_add(1, Ordering::SeqCst); s.send(runnable).unwrap(); }; let data = Data(Box::new(&$data)); (sched, data, r) }; }; } fn ms(ms: u64) -> Duration { Duration::from_millis(ms) } #[test] fn wake() { future!(f, get_waker, POLL, DROP_F); schedule!(s, d, chan, SCHEDULE, DATA, DROP_S); let (mut runnable, task) = async_task_ffi::spawn_with(f, s, d); task.detach(); assert!(chan.is_empty()); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); assert_eq!(DATA.load(Ordering::SeqCst), 0); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); get_waker().wake(); runnable = chan.recv().unwrap(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); get_waker().wake(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 2); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 1); assert_eq!(chan.len(), 0); } #[test] fn wake_by_ref() { future!(f, get_waker, POLL, DROP_F); schedule!(s, d, chan, SCHEDULE, DATA, DROP_S); let (mut runnable, task) = async_task_ffi::spawn_with(f, s, d); task.detach(); assert!(chan.is_empty()); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); assert_eq!(DATA.load(Ordering::SeqCst), 0); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); get_waker().wake_by_ref(); runnable = chan.recv().unwrap(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); get_waker().wake_by_ref(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 2); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 1); assert_eq!(chan.len(), 0); } #[test] fn clone() { future!(f, get_waker, POLL, DROP_F); schedule!(s, d, chan, SCHEDULE, DATA, DROP_S); let (mut runnable, task) = async_task_ffi::spawn_with(f, s, d); task.detach(); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); assert_eq!(DATA.load(Ordering::SeqCst), 0); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); let w2 = get_waker(); let w3 = w2.clone(); let w4 = w3.clone(); w4.wake(); runnable = chan.recv().unwrap(); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); w3.wake(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); drop(w2); drop(get_waker()); assert_eq!(DATA.load(Ordering::SeqCst), 2); assert_eq!(DROP_S.load(Ordering::SeqCst), 1); } #[test] fn wake_dropped() { future!(f, get_waker, POLL, DROP_F); schedule!(s, d, chan, SCHEDULE, DATA, DROP_S); let (runnable, task) = async_task_ffi::spawn_with(f, s, d); task.detach(); runnable.run(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); assert_eq!(DATA.load(Ordering::SeqCst), 0); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); let waker = get_waker(); waker.wake_by_ref(); drop(chan.recv().unwrap()); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); waker.wake(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 2); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 1); assert_eq!(chan.len(), 0); } #[test] fn wake_completed() { future!(f, get_waker, POLL, DROP_F); schedule!(s, d, chan, SCHEDULE, DATA, DROP_S); let (runnable, task) = async_task_ffi::spawn_with(f, s, d); task.detach(); runnable.run(); let waker = get_waker(); assert_eq!(POLL.load(Ordering::SeqCst), 1); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 0); assert_eq!(DATA.load(Ordering::SeqCst), 0); assert_eq!(DROP_F.load(Ordering::SeqCst), 0); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); waker.wake(); chan.recv().unwrap().run(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 1); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 0); assert_eq!(chan.len(), 0); get_waker().wake(); assert_eq!(POLL.load(Ordering::SeqCst), 2); assert_eq!(SCHEDULE.load(Ordering::SeqCst), 1); assert_eq!(DATA.load(Ordering::SeqCst), 2); assert_eq!(DROP_F.load(Ordering::SeqCst), 1); assert_eq!(DROP_S.load(Ordering::SeqCst), 1); assert_eq!(chan.len(), 0); }