use super::ThreadPool;
use crate::rayon::{ThreadPoolBuilder, sleep::with_sleep_notification};
use commonware_utils::sync::Mutex;
use futures::{channel::oneshot, future::poll_fn};
use std::{
panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
pin::Pin,
sync::{Arc, mpsc},
task::Poll,
thread,
};
fn with_pool<T>(threads: usize, run: impl FnOnce(&Arc<ThreadPool>) -> T) -> T {
let pool = Arc::new(
ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap(),
);
let registry = Arc::clone(&pool.registry);
let result = catch_unwind(AssertUnwindSafe(|| run(&pool)));
drop(pool);
registry.wait_until_stopped();
result.unwrap_or_else(|panic| resume_unwind(panic))
}
#[allow(
clippy::disallowed_types,
reason = "Exercise the channel used by Strategy::spawn."
)]
fn sleeping_receiver(
release: impl FnOnce(&Arc<ThreadPool>, oneshot::Sender<u32>),
) -> Result<u32, oneshot::Canceled> {
with_pool(1, |pool| {
let (tx, mut rx) = oneshot::channel();
let (asleep_tx, asleep_rx) = mpsc::channel();
thread::scope(|scope| {
let waiter = scope.spawn(|| {
pool.install(|| {
with_sleep_notification(asleep_tx, || {
pool.block_on_worker(Pin::new(&mut rx)).unwrap()
})
})
});
asleep_rx.recv().unwrap();
release(pool, tx);
waiter.join().unwrap()
})
})
}
#[commonware_macros::test_group("miri")]
#[test]
fn result_wakes_sleeping_worker() {
assert_eq!(sleeping_receiver(|_, tx| tx.send(42).unwrap()), Ok(42));
}
#[commonware_macros::test_group("miri")]
#[test]
fn canceled_sender_wakes_sleeping_worker() {
assert!(sleeping_receiver(|_, tx| drop(tx)).is_err());
}
#[commonware_macros::test_group("miri")]
#[test]
fn queued_work_wakes_sleeping_worker() {
assert_eq!(
sleeping_receiver(|pool, tx| pool.spawn(move || tx.send(42).unwrap())),
Ok(42)
);
}
#[commonware_macros::test_group("miri")]
#[test]
fn outside_and_foreign_workers_do_not_poll() {
with_pool(1, |pool| {
let mut future = poll_fn(|_| -> Poll<()> { panic!("future must not be polled") });
assert!(pool.block_on_worker(Pin::new(&mut future)).is_none());
with_pool(1, |other| {
other.install(|| {
assert!(pool.block_on_worker(Pin::new(&mut future)).is_none());
});
});
});
}
#[commonware_macros::test_group("miri")]
#[test]
fn duplicate_wakes_during_poll_are_retained() {
with_pool(1, |pool| {
pool.install(|| {
let mut polls = 0;
let mut future = poll_fn(|cx| {
polls += 1;
if polls == 1 {
cx.waker().wake_by_ref();
cx.waker().wake_by_ref();
Poll::Pending
} else {
Poll::Ready(42)
}
});
assert_eq!(pool.block_on_worker(Pin::new(&mut future)), Some(42));
assert_eq!(polls, 2);
});
});
}
#[commonware_macros::test_group("miri")]
#[test]
fn retained_wakers_outlive_ready_panic_and_pool() {
let wakers = with_pool(1, |pool| {
pool.install(|| {
let mut ready_wakers = Vec::new();
let mut ready = poll_fn(|cx| {
ready_wakers.push(cx.waker().clone());
if ready_wakers.len() == 1 {
cx.waker().wake_by_ref();
Poll::Pending
} else {
Poll::Ready(())
}
});
assert_eq!(pool.block_on_worker(Pin::new(&mut ready)), Some(()));
let mut panic_waker = None;
let mut panicking = poll_fn(|cx| -> Poll<()> {
panic_waker = Some(cx.waker().clone());
panic!("future poll panicked");
});
assert!(
catch_unwind(AssertUnwindSafe(|| {
pool.block_on_worker(Pin::new(&mut panicking));
}))
.is_err()
);
ready_wakers.push(panic_waker.unwrap());
ready_wakers
})
});
for waker in wakers {
waker.wake_by_ref();
waker.wake_by_ref();
waker.wake();
}
}
#[commonware_macros::test_group("miri")]
#[test]
fn nested_waits_service_later_work() {
with_pool(1, |pool| {
let (outer_tx, mut outer_rx) = oneshot::channel();
let (inner_tx, mut inner_rx) = oneshot::channel();
let (outer_asleep_tx, outer_asleep_rx) = mpsc::channel();
let (inner_asleep_tx, inner_asleep_rx) = mpsc::channel();
thread::scope(|scope| {
let waiter = scope.spawn(|| {
pool.install(|| {
with_sleep_notification(outer_asleep_tx, || {
pool.block_on_worker(Pin::new(&mut outer_rx)).unwrap()
})
})
});
outer_asleep_rx.recv().unwrap();
let inner_pool = Arc::clone(pool);
pool.spawn(move || {
let value = with_sleep_notification(inner_asleep_tx, || {
inner_pool
.block_on_worker(Pin::new(&mut inner_rx))
.unwrap()
.unwrap()
});
outer_tx.send(value).unwrap();
});
inner_asleep_rx.recv().unwrap();
pool.spawn(move || inner_tx.send(42).unwrap());
assert_eq!(waiter.join().unwrap(), Ok(42));
});
});
}
#[commonware_macros::test_group("miri")]
#[test]
fn workers_wait_independently() {
with_pool(2, |pool| {
let (asleep_tx, asleep_rx) = mpsc::channel();
let (done_tx, done_rx) = mpsc::channel();
let (senders, receivers): (Vec<_>, Vec<_>) = (0..2)
.map(|_| {
let (tx, rx) = oneshot::channel();
(tx, Mutex::new(Some(rx)))
})
.unzip();
let waiting_pool = Arc::clone(pool);
pool.spawn_broadcast(move |context| {
let mut rx = receivers[context.index()].lock().take().unwrap();
let result = with_sleep_notification(asleep_tx.clone(), || {
waiting_pool.block_on_worker(Pin::new(&mut rx)).unwrap()
});
done_tx.send(result.unwrap()).unwrap();
});
asleep_rx.recv().unwrap();
asleep_rx.recv().unwrap();
let mut senders = senders.into_iter();
let worker0 = senders.next().unwrap();
let worker1 = senders.next().unwrap();
worker1.send(1).unwrap();
assert_eq!(done_rx.recv().unwrap(), 1);
worker0.send(0).unwrap();
assert_eq!(done_rx.recv().unwrap(), 0);
});
}