diff --git a/maitake/src/sync.rs b/maitake/src/sync.rs index 8081370f..975e49ea 100644 --- a/maitake/src/sync.rs +++ b/maitake/src/sync.rs @@ -32,6 +32,7 @@ //! [counting semaphore]: https://en.wikipedia.org/wiki/Semaphore_(programming) //! [`Waker`]: core::task::Waker #![warn(missing_docs, missing_debug_implementations)] +pub mod broadcast; pub mod mutex; pub mod rwlock; pub mod semaphore; diff --git a/maitake/src/sync/broadcast.rs b/maitake/src/sync/broadcast.rs new file mode 100644 index 00000000..f62a0ea4 --- /dev/null +++ b/maitake/src/sync/broadcast.rs @@ -0,0 +1,304 @@ +use super::{ + rwlock::{RwLock, RwLockReadGuard}, + WaitQueue, +}; +use crate::loom::{ + cell::UnsafeCell, + sync::atomic::{AtomicUsize, Ordering::*}, +}; +use alloc::sync::Arc; +use core::ops::Deref; +use mycelium_util::{fmt, sync::CachePadded}; + +#[cfg(test)] +mod tests; + +pub struct Publisher { + shared: Arc>, +} + +pub struct Subscriber { + shared: Arc>, + pos: usize, +} + +pub struct RecvRef<'a, T> { + slot: RwLockReadGuard<'a, Slot>, +} + +#[derive(Debug)] +pub enum TryRecvError { + Empty, + Closed, + Lagged(usize), +} + +pub fn channel(mut capacity: usize) -> (Publisher, Subscriber) { + capacity = capacity.next_power_of_two(); + + let shared = Arc::new(Shared { + mask: capacity - 1, + pubs: CachePadded::new(AtomicUsize::new(1)), + subs: CachePadded::new(AtomicUsize::new(1)), + tail: CachePadded::new(AtomicUsize::new(0)), + sub_wait: WaitQueue::new(), + slots: (0..capacity) + .map(|i| { + RwLock::new(Slot { + rem: AtomicUsize::new(0), + gen: i.wrapping_sub(capacity), + val: None, + }) + }) + .collect::<_>(), + }); + + let pub_ = Publisher { + shared: shared.clone(), + }; + let sub = Subscriber { shared, pos: 0 }; + + (pub_, sub) +} + +struct Shared { + tail: CachePadded, + /// Number of subscribers. + subs: CachePadded, + /// Number of publishers. + pubs: CachePadded, + /// Capapcity - 1 of the channel. + mask: usize, + sub_wait: WaitQueue, + slots: alloc::boxed::Box<[RwLock>]>, +} + +struct Slot { + /// Count of readers remaining to view this slot. + rem: AtomicUsize, + + gen: usize, + + /// The value broadcast at this position. + /// + /// The value is set when sending a value to this slot. + val: Option, +} + +// === impl Publisher === + +impl Publisher { + pub async fn send(&self, value: T) -> Result<(), T> { + test_debug!("Publisher::send"); + let tail = test_dbg!(self.shared.tail.fetch_add(1, AcqRel)); + let idx = tail & self.shared.mask; + { + let mut slot = self.shared.slots[test_dbg!(idx)].write().await; + + // load subscriber count + let subs = self.shared.subs.load(Acquire); + if test_dbg!(subs) == 0 { + return Err(value); + } + + // write to the slot + slot.gen = tail; + slot.rem.store(subs, Release); + slot.val = Some(value); + } + + // wake any waiting subscribers + self.shared.sub_wait.wake_all(); + test_debug!("wrote value to slot {idx}"); + + Ok(()) + } + + pub fn subscribe(&self) -> Subscriber { + self.shared.subs.fetch_add(1, Relaxed); + Subscriber { + shared: self.shared.clone(), + pos: self.shared.tail.load(Acquire), + } + } +} + +impl Clone for Publisher { + fn clone(&self) -> Self { + self.shared.pubs.fetch_add(1, Relaxed); + Self { + shared: self.shared.clone(), + } + } +} + +impl Drop for Publisher { + fn drop(&mut self) { + if self.shared.pubs.fetch_sub(1, AcqRel) == 1 { + self.shared.sub_wait.close(); + } + } +} + +impl fmt::Debug for Publisher { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let Self { shared } = self; + f.debug_struct("Publisher").field("shared", shared).finish() + } +} + +// === impl Subscriber === + +impl Subscriber { + pub async fn recv(&mut self) -> Result { + test_debug!("Subscriber::recv(pos: {})", self.pos); + self.recv_ref().await.map(|r| r.clone()) + } + + pub async fn try_recv(&mut self) -> Result { + test_debug!("Subscriber::try_recv(pos: {})", self.pos); + self.try_recv_ref().await.map(|r| r.clone()) + } +} + +impl Subscriber { + pub async fn recv_ref(&mut self) -> Result, TryRecvError> { + test_debug!("Subscriber::recv_ref(pos: {})", self.pos); + loop { + match Self::try_recv_ref2(&self.shared, &mut self.pos).await { + Ok(val) => { + test_debug!("Subscriber::recv_ref -> received!"); + return Ok(val); + } + Err(TryRecvError::Empty) => { + test_debug!("Subscriber::recv_ref -> empty; waiting..."); + // ignore errors here; the WaitQueue may close while there + // are still slots we have left to read, and the subsequent + // `try_recv_ref2` call will handle this. + let _ = test_dbg!(self.shared.sub_wait.wait().await); + } + Err(e) => { + test_debug!("Subscriber::recv_ref -> error {e:?}"); + return Err(e); + } + } + } + } + + pub async fn try_recv_ref(&mut self) -> Result, TryRecvError> { + test_debug!("Subscriber::recv(pos: {})", self.pos); + Self::try_recv_ref2(&self.shared, &mut self.pos).await + } + + async fn try_recv_ref2<'shared>( + shared: &'shared Shared, + pos: &mut usize, + ) -> Result, TryRecvError> { + let idx = test_dbg!(*pos) & shared.mask; + + let slot = shared.slots[test_dbg!(idx)].read().await; + + if test_dbg!(slot.gen != *pos) { + // we lagged behind, try to read the next slot. + let lap = slot.gen.wrapping_add(shared.slots.len()); + + // the channel is empty relative to this receiver. + if test_dbg!(lap == *pos) { + return if shared.sub_wait.is_closed() { + Err(TryRecvError::Closed) + } else { + Err(TryRecvError::Empty) + }; + } + + let tail = shared.tail.load(Acquire); + let next = tail.wrapping_sub(shared.slots.len()); + + let missed = next.wrapping_sub(*pos); + + // The receiver is slow but no values have been missed + if test_dbg!(missed) == 0 { + *pos = pos.wrapping_add(1); + + return Ok(RecvRef { slot }); + } + + *pos = next; + + return Err(TryRecvError::Lagged(missed)); + } + + *pos = pos.wrapping_add(1); + Ok(RecvRef { slot }) + } +} + +impl fmt::Debug for Subscriber { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let Self { pos, shared } = self; + f.debug_struct("Subscriber") + .field("shared", shared) + .field("pos", pos) + .finish() + } +} + +// === impl RecvRef === + +impl fmt::Debug for RecvRef<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.slot.val.as_ref().unwrap().fmt(f) + } +} + +impl Deref for RecvRef<'_, T> { + type Target = T; + + fn deref(&self) -> &Self::Target { + self.slot + .val + .as_ref() + .expect("a RecvRef is only returned if the slot is `Some`") + } +} + +// === impl Shared === + +impl fmt::Debug for Shared { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let Self { + tail, + subs, + pubs, + mask, + slots, + sub_wait, + } = self; + f.debug_struct("Shared") + .field("tail", tail) + .field("subs", subs) + .field("pubs", pubs) + .field("mask", &fmt::hex(mask)) + .field("slots", slots) + .field("sub_wait", sub_wait) + .finish() + } +} + +impl fmt::Debug for Slot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let Self { gen, rem, val } = self; + f.debug_struct("Slot") + .field("gen", gen) + .field("rem", rem) + .field( + "val", + &if val.is_some() { + format_args!("Some(...)") + } else { + format_args!("None") + }, + ) + .finish() + } +} diff --git a/maitake/src/sync/broadcast/tests.rs b/maitake/src/sync/broadcast/tests.rs new file mode 100644 index 00000000..40db7f08 --- /dev/null +++ b/maitake/src/sync/broadcast/tests.rs @@ -0,0 +1,215 @@ +use crate::sync::broadcast; +use crate::sync::broadcast::TryRecvError::{Closed, Empty, Lagged}; + +use crate::loom::{self, future::block_on, sync::Arc, thread}; +// use tokio_test::{assert_err, assert_ok}; + +#[test] +fn broadcast_send() { + loom::model(|| { + let (tx1, mut rx) = broadcast::channel(2); + let tx1 = Arc::new(tx1); + let tx2 = tx1.clone(); + + let th1 = thread::spawn(move || { + block_on(async { + test_dbg!(tx1.send("one").await).unwrap(); + test_dbg!(tx1.send("two").await).unwrap(); + test_dbg!(tx1.send("three").await).unwrap(); + }); + }); + + let th2 = thread::spawn(move || { + block_on(async { + test_dbg!(tx2.send("eins").await).unwrap(); + test_dbg!(tx2.send("zwei").await).unwrap(); + test_dbg!(tx2.send("drei").await).unwrap(); + }); + }); + + block_on(async { + let mut num = 0; + loop { + match test_dbg!(rx.recv().await) { + Ok(_) => num += 1, + Err(Closed) => break, + Err(Lagged(n)) => num += n, + Err(Empty) => panic!("unexpected empty"), + } + } + assert_eq!(num, 6); + }); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} + +// An `Arc` is used as the value in order to detect memory leaks. +#[test] +fn broadcast_two() { + loom::model(|| { + let (tx, mut rx1) = broadcast::channel::>(16); + let mut rx2 = tx.subscribe(); + + let th1 = thread::spawn(move || { + block_on(async move { + let v = test_dbg!(rx1.recv().await).unwrap(); + assert_eq!(*v, "hello"); + + let v = test_dbg!(rx1.recv().await).unwrap(); + assert_eq!(*v, "world"); + + match test_dbg!(rx1.recv().await).unwrap_err() { + Closed => {} + err => panic!("unexpected error: {err:?}"), + } + }); + }); + + let th2 = thread::spawn(move || { + block_on(async move { + let v = test_dbg!(rx2.recv().await).unwrap(); + assert_eq!(*v, "hello"); + + let v = test_dbg!(rx2.recv().await).unwrap(); + assert_eq!(*v, "world"); + + match test_dbg!(rx2.recv().await).unwrap_err() { + Closed => {} + _ => panic!(), + } + }); + }); + block_on(async move { + test_dbg!(tx.send(Arc::new("hello")).await).unwrap(); + test_dbg!(tx.send(Arc::new("world")).await).unwrap(); + }); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} + +#[test] +fn broadcast_wrap() { + loom::model(|| { + let (tx, mut rx1) = broadcast::channel(2); + let mut rx2 = tx.subscribe(); + + let th1 = thread::spawn(move || { + block_on(async { + let mut num = 0; + + loop { + match test_dbg!(rx1.recv().await) { + Ok(_) => num += 1, + Err(Closed) => break, + Err(Lagged(n)) => num += n, + Err(Empty) => panic!("unexpected empty"), + } + } + + assert_eq!(num, 3); + }); + }); + + let th2 = thread::spawn(move || { + block_on(async { + let mut num = 0; + + loop { + match test_dbg!(rx2.recv().await) { + Ok(_) => num += 1, + Err(Closed) => break, + Err(Lagged(n)) => num += n, + Err(Empty) => panic!("unexpected empty"), + } + } + + assert_eq!(num, 3); + }); + }); + + block_on(async move { + test_dbg!(tx.send("one").await).unwrap(); + test_dbg!(tx.send("two").await).unwrap(); + test_dbg!(tx.send("three").await).unwrap(); + }); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} + +#[test] +fn drop_rx() { + loom::model(|| { + let (tx, mut rx1) = broadcast::channel(16); + let rx2 = tx.subscribe(); + + let th1 = thread::spawn(move || { + block_on(async { + let v = test_dbg!(rx1.recv().await).unwrap(); + assert_eq!(v, "one"); + + let v = test_dbg!(rx1.recv().await).unwrap(); + assert_eq!(v, "two"); + + let v = test_dbg!(rx1.recv().await).unwrap(); + assert_eq!(v, "three"); + + match test_dbg!(rx1.recv().await).unwrap_err() { + Closed => {} + _ => panic!(), + } + }); + }); + + let th2 = thread::spawn(move || { + drop(rx2); + }); + + block_on(async move { + test_dbg!(tx.send("one").await).unwrap(); + test_dbg!(tx.send("two").await).unwrap(); + test_dbg!(tx.send("three").await).unwrap(); + }); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} + +#[test] +#[ignore] +fn drop_multiple_rx_with_overflow() { + loom::model(move || { + // It is essential to have multiple senders and receivers in this test case. + let (tx, mut rx) = broadcast::channel(1); + let _rx2 = tx.subscribe(); + + let tx = block_on(async { + let _ = test_dbg!(tx.send(()).await); + tx + }); + let tx2 = tx.clone(); + let th1 = thread::spawn(move || { + block_on(async { + for _ in 0..100 { + let _ = test_dbg!(tx2.send(()).await); + } + }); + }); + let tx = block_on(async { + let _ = test_dbg!(tx.send(()).await); + tx + }); + let th2 = thread::spawn(move || { + block_on(async { while let Ok(_) = test_dbg!(rx.recv().await) {} }); + }); + + th1.join().unwrap(); + th2.join().unwrap(); + }); +} diff --git a/maitake/src/sync/wait_queue.rs b/maitake/src/sync/wait_queue.rs index 597e6f85..c092e2b6 100644 --- a/maitake/src/sync/wait_queue.rs +++ b/maitake/src/sync/wait_queue.rs @@ -489,6 +489,10 @@ impl WaitQueue { } } + pub(crate) fn is_closed(&self) -> bool { + self.load().get(QueueState::STATE) == State::Closed + } + pub(crate) fn try_wait(&self) -> Poll> { let mut state = self.load(); let initial_wake_alls = state.get(QueueState::WAKE_ALLS);