diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a5d85a..0335243 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ All notable changes to this project will be documented in this file. ### Bug fixes * Release cancelled wait registrations promptly and reclaim fulfilled `Semaphore::forget_exact` debt nodes. +* Serialize broadcast publication so receivers cannot observe reserved slots or messages overwritten out of sequence. ### Improvements diff --git a/asyncband/src/broadcast/overflow/mod.rs b/asyncband/src/broadcast/overflow/mod.rs index 4c996f2..176379a 100644 --- a/asyncband/src/broadcast/overflow/mod.rs +++ b/asyncband/src/broadcast/overflow/mod.rs @@ -72,6 +72,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::Waker; use crate::internal::mutex::Mutex; use crate::internal::rwlock::RwLock; @@ -117,9 +118,11 @@ pub fn channel(capacity: usize) -> (Sender, Receiver) { buffer: buffer.into_boxed_slice(), capacity, mask, - tail_cnt: AtomicU64::new(0), + tail: AtomicU64::new(0), + state: Mutex::new(State { + waiters: WaitSet::new(), + }), senders: AtomicUsize::new(1), - waiters: Mutex::new(WaitSet::new()), }); let sender = Sender { shared: shared.clone(), @@ -190,25 +193,22 @@ struct Shared { buffer: Box<[RwLock>]>, capacity: usize, mask: usize, - /// The global tail cursor. Points to the next slot to write. - /// Strictly monotonically increasing. - tail_cnt: AtomicU64, + /// The next sequence after the contiguous prefix of fully published slots. + tail: AtomicU64, + /// Serializes senders and makes publishing atomic with draining or registering waiters. + state: Mutex, /// Number of active senders. senders: AtomicUsize, - /// Waiters (receivers) waiting for new messages. - waiters: Mutex, } -impl Shared { - fn wake_waiters(&self) { - let wakers = { - let mut waiters = self.waiters.lock(); - waiters.take_wakers() - }; +struct State { + /// Receivers waiting for a new message. + waiters: WaitSet, +} - for waker in wakers { - waker.wake(); - } +fn wake_waiters(wakers: impl IntoIterator) { + for waker in wakers { + waker.wake(); } } @@ -235,7 +235,8 @@ impl Drop for Sender { 1 => { // If this is the last sender, we need to wake up the receiver so it can // observe the disconnected state. - self.shared.wake_waiters(); + let wakers = self.shared.state.lock().waiters.take_wakers(); + wake_waiters(wakers); } _ => { // there are still other senders left, do nothing @@ -261,17 +262,28 @@ impl Sender { /// assert_eq!(rx.try_recv(), Ok(10)); /// ``` pub fn send(&self, msg: T) { - let tail = self.shared.tail_cnt.fetch_add(1, Ordering::SeqCst); - let idx = (tail as usize) & self.shared.mask; + let wakers = { + let mut state = self.shared.state.lock(); + let tail = self.shared.tail.load(Ordering::Relaxed); + let idx = (tail as usize) & self.shared.mask; - { let mut slot = self.shared.buffer[idx].write(); slot.msg = Some(msg); slot.version = tail; - } + + // Publish the completed slot before releasing its write lock. A receiver that sees the + // new tail either held the old slot lock first or waits until the new value is + // complete. + self.shared + .tail + .store(tail.wrapping_add(1), Ordering::Release); + drop(slot); + + state.waiters.take_wakers() + }; // Notify all waiting receivers. - self.shared.wake_waiters(); + wake_waiters(wakers); } /// Creates a new receiver that starts receiving messages from the current tail of the channel. @@ -295,7 +307,7 @@ impl Sender { /// ``` pub fn subscribe(&self) -> Receiver { // Receiver starts at the current tail. - let head = self.shared.tail_cnt.load(Ordering::SeqCst); + let head = self.shared.tail.load(Ordering::Acquire); let shared = self.shared.clone(); Receiver { shared, head } } @@ -375,49 +387,63 @@ impl Receiver { let shared = &self.shared; let cap = shared.capacity as u64; - let tail = shared.tail_cnt.load(Ordering::SeqCst); - let head = self.head; - - // diff represents how far behind the head is from the tail. - let diff = tail.wrapping_sub(head); + loop { + let tail = shared.tail.load(Ordering::Acquire); + let head = self.head; - // 1. Check for Lag - if diff > cap { - let missed = diff - cap; - self.head = tail.wrapping_sub(cap); - return Err(TryRecvError::Lagged(missed)); - } + // diff represents how far behind the head is from the tail. + let diff = tail.wrapping_sub(head); - // 2. Check if a message is available - if diff > 0 { - let idx = (head as usize) & shared.mask; - let slot = shared.buffer[idx].read(); - - if slot.version == head { - return if let Some(msg) = &slot.msg { - self.head = head.wrapping_add(1); - Ok(msg.clone()) - } else { - Err(TryRecvError::Empty) - }; + // 1. Check for Lag + if diff > cap { + let missed = diff - cap; + self.head = tail.wrapping_sub(cap); + return Err(TryRecvError::Lagged(missed)); } - drop(slot); + // 2. Check if a message is available + if diff > 0 { + let idx = (head as usize) & shared.mask; + let slot = shared.buffer[idx].read(); + + if slot.version == head { + return if let Some(msg) = &slot.msg { + self.head = head.wrapping_add(1); + Ok(msg.clone()) + } else { + Err(TryRecvError::Empty) + }; + } - // If version != head, the slot was overwritten. - // This means we lagged, but the `diff > cap` check missed it (likely due to overflow - // wrapping). We treat this as a lag. - let missed = tail.wrapping_sub(self.head).wrapping_sub(cap); - self.head = tail.wrapping_sub(cap); - return Err(TryRecvError::Lagged(missed)); - } + drop(slot); + + // The slot may have been overwritten after the first tail snapshot. Publication + // happens while holding the slot write lock, so a fresh tail now includes that + // overwrite and produces an accurate lag count. + let tail = shared.tail.load(Ordering::Acquire); + let diff = tail.wrapping_sub(head); + if diff > cap { + let missed = diff - cap; + self.head = tail.wrapping_sub(cap); + return Err(TryRecvError::Lagged(missed)); + } - // 3. No message available (diff == 0). Check for Closed. - if shared.senders.load(Ordering::Acquire) == 0 { - return Err(TryRecvError::Disconnected); - } + return Err(TryRecvError::Empty); + } + + // 3. No message available (diff == 0). Check for Closed. + if shared.senders.load(Ordering::Acquire) == 0 { + // Observing the final sender drop synchronizes with all preceding sends, but the + // first tail snapshot predates that acquire. Reload it before declaring the + // channel drained so a published final message cannot be hidden by closure. + if shared.tail.load(Ordering::Acquire) != head { + continue; + } + return Err(TryRecvError::Disconnected); + } - Err(TryRecvError::Empty) + return Err(TryRecvError::Empty); + } } } @@ -444,7 +470,7 @@ impl Receiver { /// ``` pub fn resubscribe(&self) -> Self { // Resubscribe starts at the current tail. - let head = self.shared.tail_cnt.load(Ordering::SeqCst); + let head = self.shared.tail.load(Ordering::Acquire); let shared = self.shared.clone(); Self { shared, head } } @@ -484,13 +510,12 @@ impl Future for Recv<'_, T> { } let shared = &receiver.shared; - let mut waiters = shared.waiters.lock(); + let mut state = shared.state.lock(); // Double check tail to avoid race conditions. - let tail_now = shared.tail_cnt.load(Ordering::SeqCst); - if tail_now != receiver.head { + if shared.tail.load(Ordering::Acquire) != receiver.head { // New message arrived while acquiring the lock. Retry. - drop(waiters); + drop(state); continue; } @@ -502,8 +527,8 @@ impl Future for Recv<'_, T> { } // Register Waker - let replaced_waker = waiters.register_waker(registration, cx); - drop(waiters); + let replaced_waker = state.waiters.register_waker(registration, cx); + drop(state); drop(replaced_waker); return Poll::Pending; } @@ -514,8 +539,8 @@ impl Drop for Recv<'_, T> { fn drop(&mut self) { if self.registration.is_some() { let removed_waker = { - let mut waiters = self.receiver.shared.waiters.lock(); - waiters.unregister_waker(&mut self.registration) + let mut state = self.receiver.shared.state.lock(); + state.waiters.unregister_waker(&mut self.registration) }; drop(removed_waker); } diff --git a/asyncband/src/broadcast/overflow/tests.rs b/asyncband/src/broadcast/overflow/tests.rs index 6c2b951..da49b0c 100644 --- a/asyncband/src/broadcast/overflow/tests.rs +++ b/asyncband/src/broadcast/overflow/tests.rs @@ -27,7 +27,7 @@ async fn sequence_number_wraparound() { let mut rx2 = rx.clone(); let boundary = u64::MAX - 2; - tx.shared.tail_cnt.store(boundary, Ordering::SeqCst); + tx.shared.tail.store(boundary, Ordering::Release); rx.head = boundary; tx.send(1); @@ -54,7 +54,7 @@ async fn sequence_number_wraparound_exactly_overwritten() { let mut rx2 = rx.clone(); let boundary = u64::MAX - 2; - tx.shared.tail_cnt.store(boundary, Ordering::SeqCst); + tx.shared.tail.store(boundary, Ordering::Release); rx.head = boundary; tx.send(1); @@ -84,14 +84,3 @@ fn capacity_is_rounded_to_a_power_of_two() { assert_eq!(tx.shared.capacity, 8); assert_eq!(tx.shared.mask, 7); } - -#[test] -fn try_recv_treats_an_unwritten_slot_as_empty() { - let (tx, mut rx) = channel::(2); - drop(tx); - - rx.shared.tail_cnt.store(1, Ordering::SeqCst); - - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - assert_eq!(rx.head, 0); -} diff --git a/benchmarks/broadcast.rs b/benchmarks/broadcast.rs index aec7c72..87d9870 100644 --- a/benchmarks/broadcast.rs +++ b/benchmarks/broadcast.rs @@ -16,6 +16,10 @@ // under the License. use std::pin::pin; +use std::sync::Arc; +use std::sync::Barrier; +use std::thread; +use std::thread::JoinHandle; use asyncband::broadcast::overflow; use divan::Bencher; @@ -26,6 +30,175 @@ use super::support::poll_pending; use super::support::poll_pinned_ready; const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; +const CONCURRENCY_COUNTS: &[usize] = &[1, 2, 4, 8]; +const CONCURRENT_BATCH_SIZE: usize = 4096; + +struct ConcurrentSend { + _receiver: overflow::Receiver, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentSend { + fn new(sender_count: usize) -> Self { + let (sender, receiver) = overflow::channel(CONCURRENT_BATCH_SIZE); + let ready = Arc::new(Barrier::new(sender_count + 1)); + let start = Arc::new(Barrier::new(sender_count + 1)); + let done = Arc::new(Barrier::new(sender_count + 1)); + let sends_per_worker = CONCURRENT_BATCH_SIZE / sender_count; + let mut workers = Vec::with_capacity(sender_count); + + for worker_index in 0..sender_count { + let sender = sender.clone(); + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let first = worker_index * sends_per_worker; + for value in first..first + sends_per_worker { + sender.send(black_box(value)); + } + done.wait(); + })); + } + drop(sender); + ready.wait(); + + Self { + _receiver: receiver, + start, + done, + workers, + } + } + + fn run(&mut self) { + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for ConcurrentSend { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +struct ConcurrentFanout { + sender: overflow::Sender, + start: Arc, + done: Arc, + workers: Vec>, +} + +impl ConcurrentFanout { + fn new(receiver_count: usize) -> Self { + let (sender, receiver) = overflow::channel(CONCURRENT_BATCH_SIZE); + let mut receivers = Vec::with_capacity(receiver_count); + receivers.push(receiver); + for _ in 1..receiver_count { + receivers.push(receivers[0].clone()); + } + + let ready = Arc::new(Barrier::new(receiver_count + 1)); + let start = Arc::new(Barrier::new(receiver_count + 1)); + let done = Arc::new(Barrier::new(receiver_count + 1)); + let mut workers = Vec::with_capacity(receiver_count); + + for mut receiver in receivers { + let ready = ready.clone(); + let start = start.clone(); + let done = done.clone(); + workers.push(thread::spawn(move || { + ready.wait(); + start.wait(); + let result = (0..CONCURRENT_BATCH_SIZE).try_for_each(|_| { + receiver.try_recv().map(|value| { + black_box(value); + }) + }); + done.wait(); + result.unwrap(); + })); + } + ready.wait(); + + Self { + sender, + start, + done, + workers, + } + } + + fn run(&mut self) { + for value in 0..CONCURRENT_BATCH_SIZE { + self.sender.send(black_box(value)); + } + self.start.wait(); + self.done.wait(); + } +} + +impl Drop for ConcurrentFanout { + fn drop(&mut self) { + for worker in self.workers.drain(..) { + worker.join().unwrap(); + } + } +} + +#[divan::bench] +fn send_overwrite(bencher: Bencher) { + let (sender, receiver) = overflow::channel::(1); + bencher.bench_local(|| sender.send(black_box(1))); + black_box(receiver); +} + +#[divan::bench] +fn try_recv_empty(bencher: Bencher) { + let (sender, mut receiver) = overflow::channel::(1); + bencher.bench_local(|| black_box(receiver.try_recv())); + black_box(sender); +} + +#[divan::bench] +fn send_and_try_recv(bencher: Bencher) { + let (sender, mut receiver) = overflow::channel(1); + bencher.bench_local(|| { + sender.send(black_box(1)); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_send(bencher: Bencher, sender_count: usize) { + bencher + .with_inputs(|| ConcurrentSend::new(sender_count)) + .bench_local_refs(ConcurrentSend::run); +} + +#[divan::bench( + args = CONCURRENCY_COUNTS, + sample_count = 50, + sample_size = 1, + counters = [CONCURRENT_BATCH_SIZE] +)] +fn concurrent_fanout(bencher: Bencher, receiver_count: usize) { + bencher + .with_inputs(|| ConcurrentFanout::new(receiver_count)) + .bench_local_refs(ConcurrentFanout::run); +} #[divan::bench] fn cancel_pending(bencher: Bencher) { diff --git a/tests-integration/tests/broadcast_test.rs b/tests-integration/tests/broadcast_test.rs index fd77e2d..56ab69d 100644 --- a/tests-integration/tests/broadcast_test.rs +++ b/tests-integration/tests/broadcast_test.rs @@ -17,11 +17,14 @@ use std::future::Future; use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Wake; use std::task::Waker; +use std::thread; use asyncband::broadcast::overflow::*; @@ -33,6 +36,31 @@ impl Wake for TrackWake { } } +#[derive(Debug)] +struct PanicOnDrop { + value: u64, + panic: bool, + panicked: Arc, +} + +impl Clone for PanicOnDrop { + fn clone(&self) -> Self { + Self { + value: self.value, + panic: false, + panicked: self.panicked.clone(), + } + } +} + +impl Drop for PanicOnDrop { + fn drop(&mut self) { + if self.panic && !self.panicked.swap(true, Ordering::Relaxed) { + panic!("panic while replacing a broadcast slot"); + } + } +} + #[tokio::test] async fn test_broadcast_basic() { let (tx, mut rx1) = channel(10); @@ -174,6 +202,119 @@ async fn test_try_recv_lagged() { assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } +#[test] +fn panicking_send_does_not_publish_an_unwritten_slot() { + let panicked = Arc::new(AtomicBool::new(false)); + let (tx, mut rx) = channel(1); + tx.send(PanicOnDrop { + value: 1, + panic: true, + panicked: panicked.clone(), + }); + + let received = rx.try_recv().unwrap(); + assert_eq!(received.value, 1); + drop(received); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + tx.send(PanicOnDrop { + value: 2, + panic: false, + panicked: panicked.clone(), + }); + })); + assert!(result.is_err()); + assert!(panicked.load(Ordering::Relaxed)); + + assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty))); + tx.send(PanicOnDrop { + value: 3, + panic: false, + panicked: panicked.clone(), + }); + assert_eq!(rx.try_recv().unwrap().value, 3); + + drop(tx); + assert!(matches!(rx.try_recv(), Err(TryRecvError::Disconnected))); +} + +#[test] +fn concurrent_overwrite_preserves_sequence_and_lag_count() { + const MESSAGE_COUNT: u64 = 200_000; + + let (tx, mut rx) = channel(2); + let producer = thread::spawn(move || { + for value in 0..MESSAGE_COUNT { + tx.send(value); + } + }); + + let mut next = 0_u64; + loop { + match rx.try_recv() { + Ok(value) => { + assert_eq!(value, next); + next = next.wrapping_add(1); + } + Err(TryRecvError::Lagged(missed)) => { + assert!(missed > 0); + next = next.wrapping_add(missed); + } + Err(TryRecvError::Empty) => thread::yield_now(), + Err(TryRecvError::Disconnected) => break, + } + } + + producer.join().unwrap(); + assert_eq!(next, MESSAGE_COUNT); +} + +#[test] +fn concurrent_receivers_observe_the_same_sequence() { + const MESSAGE_COUNT: usize = 4096; + const RECEIVER_COUNT: usize = 8; + + let (tx, receiver) = channel(MESSAGE_COUNT); + let mut receivers = Vec::with_capacity(RECEIVER_COUNT); + receivers.push(receiver); + for _ in 1..RECEIVER_COUNT { + receivers.push(receivers[0].clone()); + } + + let ready = Arc::new(Barrier::new(RECEIVER_COUNT + 1)); + let workers = receivers + .into_iter() + .map(|mut receiver| { + let ready = ready.clone(); + thread::spawn(move || { + ready.wait(); + let mut received = Vec::with_capacity(MESSAGE_COUNT); + loop { + match receiver.try_recv() { + Ok(value) => received.push(value), + Err(TryRecvError::Empty) => thread::yield_now(), + Err(TryRecvError::Disconnected) => return received, + Err(TryRecvError::Lagged(missed)) => { + panic!("receiver unexpectedly lagged by {missed}") + } + } + } + }) + }) + .collect::>(); + + ready.wait(); + for value in 0..MESSAGE_COUNT { + tx.send(value); + } + drop(tx); + + let expected = (0..MESSAGE_COUNT).collect::>(); + for worker in workers { + assert_eq!(worker.join().unwrap(), expected); + } +} + #[tokio::test] async fn test_multi_senders_concurrent() { let (tx, mut rx) = channel(100);