//! # Single producer, multi-consumer notify //! //! See [`Master`] for more info on it. use crate::spsc; /// Single producer, multi-consumer notify. Essentially just collection of /// [`spsc`] masters. #[derive(Default)] pub struct Master { masters: Vec, } impl Master { /// Create master. pub const fn new() -> Self { Self { masters: Vec::new(), } } /// Preallocate storage for `cap` slaves. pub fn with_capacity(cap: usize) -> Self { Self { masters: Vec::with_capacity(cap), } } } impl Master { /// Register slave and return handle to it. Dropping slave /// does nothing, as for now, space reclamation is not implemented. pub fn make_slave(&mut self) -> spsc::Slave { let (master, slave) = spsc::make(); self.masters.push(master); slave } /// Get number of currently registered slaves. pub fn len(&self) -> usize { self.masters.len() } /// if [`Master::len`] == 0. pub fn is_empty(&self) -> bool { self.masters.is_empty() } /// Notify every registered slave. In tokio terms, this method stores /// permit for every slave, thus next await of [`spsc::Slave`] will wake-up /// immediately. Only one permit is stored. pub fn notify_waiters(&mut self) { for master in self.masters.iter_mut() { master.notify(); } } } #[cfg(test)] mod tests { use std::time::Duration; use tokio::{sync::mpsc, time::timeout}; use super::*; type Rx = mpsc::UnboundedReceiver; async fn task(mut slave: spsc::Slave, tx: impl Fn() + Send + 'static) { loop { (&mut slave).await; tx(); } } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_wakeups() { const WORKERS: usize = 2000; const TRIES: usize = 4; const TIMEOUT: Duration = Duration::from_millis(500); let correct: Vec = (0..WORKERS).collect(); let mut buf = Vec::with_capacity(WORKERS); async fn case(rx: &mut Rx, correct: &[usize], buf: &mut Vec) { unsafe { buf.set_len(0) }; for _ in 0..correct.len() { let item = timeout(TIMEOUT, rx.recv()) .await .expect("timed out") .unwrap(); buf.push(item); } buf.sort_unstable(); assert_eq!(buf, correct); } let mut master = Master::default(); let (tx, mut rx) = mpsc::unbounded_channel(); for worker_id in 0..WORKERS { let slave = master.make_slave(); let tx = tx.clone(); tokio::spawn(task(slave, move || tx.send(worker_id).unwrap())); } for try_no in 0..TRIES { master.notify_waiters(); eprint!("try#{try_no}..."); case(&mut rx, &correct, &mut buf).await; eprintln!(" success"); } assert!(timeout(Duration::from_millis(100), rx.recv()) .await .is_err()); } }