use core::task::{Poll, Waker};
use std::sync::Arc as StdArc;
use loom::sync::atomic::{AtomicUsize, Ordering};
use loom::sync::{Arc, Notify};
use loom::thread;
use super::{After, Hooks, Side};
struct Unpark(Notify);
impl std::task::Wake for Unpark {
fn wake(self: StdArc<Self>) {
self.0.notify();
}
fn wake_by_ref(self: &StdArc<Self>) {
self.0.notify();
}
}
#[derive(Default)]
struct Model {
ready: bool,
got: usize,
reader: Option<Waker>,
}
fn read(m: &mut Model, kernel: &AtomicUsize) {
match kernel.swap(0, Ordering::SeqCst) {
0 => m.ready = false,
n => m.got += n,
}
}
struct Serve<'a> {
kernel: &'a AtomicUsize,
wake: bool,
}
impl Hooks<Model> for Serve<'_> {
fn serve(&self, m: &mut Model, after: &mut After) {
m.ready = true;
read(m, self.kernel);
if m.got > 0
&& self.wake
&& let Some(w) = m.reader.take()
{
after.wake_one(w);
}
}
}
fn step(m: &mut Model, kernel: &AtomicUsize, waker: &Waker) -> Poll<()> {
if m.ready {
read(m, kernel);
}
if m.got > 0 {
return Poll::Ready(());
}
m.reader = Some(waker.clone());
Poll::Pending
}
fn poll_until_done(side: &Side<Model>, kernel: &AtomicUsize, hooks: &Serve<'_>) {
let park = StdArc::new(Unpark(Notify::new()));
let waker = Waker::from(StdArc::clone(&park));
loop {
let mut after = After::default();
let polled = side.poll(&mut after, &waker, hooks, |m, _| step(m, kernel, &waker));
after.run();
if polled.is_ready() {
return;
}
park.0.wait();
}
}
fn edge_model(wake: bool) {
let side = Arc::new(Side::new(Model::default()));
let kernel = Arc::new(AtomicUsize::new(0));
let driver = {
let (side, kernel) = (side.clone(), kernel.clone());
thread::spawn(move || {
kernel.fetch_add(1, Ordering::SeqCst);
let mut after = After::default();
side.drive(
&mut after,
&Serve {
kernel: &kernel,
wake,
},
0,
);
after.run();
})
};
poll_until_done(
&side,
&kernel,
&Serve {
kernel: &kernel,
wake,
},
);
driver.join().unwrap();
assert_eq!(
kernel.load(Ordering::SeqCst),
0,
"the byte was left in the kernel"
);
}
#[test]
fn an_edge_during_a_poll_is_never_lost() {
loom::model(|| edge_model(true));
}
#[test]
#[should_panic]
fn calibration_a_driver_part_that_forgets_the_wake_strands_the_reader() {
loom::model(|| edge_model(false));
}
struct Count<'a>(&'a AtomicUsize);
impl Hooks<usize> for Count<'_> {
fn serve(&self, _st: &mut usize, _after: &mut After) {}
}
fn contended_model(wakes: bool) {
let side = Arc::new(Side::new(0usize));
let steps = Arc::new(AtomicUsize::new(0));
let other = {
let (side, steps) = (side.clone(), steps.clone());
thread::spawn(move || poll_once_to_completion(&side, &steps, true))
};
poll_once_to_completion(&side, &steps, wakes);
other.join().unwrap();
assert_eq!(steps.load(Ordering::SeqCst), 2, "a poll never ran its step");
let mut after = After::default();
let mut held = side.try_hold().expect("the side was left held");
assert_eq!(*held.get(), 2);
held.release(&mut after, &Count(&steps));
}
fn poll_once_to_completion(side: &Side<usize>, steps: &AtomicUsize, wakes: bool) {
let park = StdArc::new(Unpark(Notify::new()));
let waker = if wakes {
Waker::from(StdArc::clone(&park))
} else {
Waker::noop().clone()
};
loop {
let mut after = After::default();
let polled = side.poll(&mut after, &waker, &Count(steps), |n, _| {
*n += 1;
steps.fetch_add(1, Ordering::SeqCst);
Poll::Ready(())
});
after.run();
if polled.is_ready() {
return;
}
park.0.wait();
}
}
#[test]
fn a_contended_poll_is_woken_by_the_holder() {
loom::model(|| contended_model(true));
}
#[test]
#[should_panic]
fn calibration_a_contended_poll_that_is_never_woken_stays_parked() {
loom::model(|| contended_model(false));
}