From 7c24cfdff2493991b2538554fae9e43a9346e310 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Sun, 13 Sep 2026 20:40:22 -0400 Subject: [PATCH 01/10] evict broadcast entries after all consumers read --- src/execution_plans/broadcast.rs | 192 ++++++++++++++++++++++++++----- 1 file changed, 163 insertions(+), 29 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 4d06d4d09..63183a072 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -12,6 +12,7 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, internal_err, }; use futures::{Stream, StreamExt}; +use std::collections::VecDeque; use std::fmt::Formatter; use std::pin::Pin; use std::sync::{Arc, Mutex, OnceLock}; @@ -165,7 +166,7 @@ impl ExecutionPlan for BroadcastExec { let input = Arc::clone(&self.input); let queue_or_err = self.queues[real_partition].get_or_init(|| { - let queue = BroadcastQueue::new(); + let queue = BroadcastQueue::new(self.consumer_task_count); let consumers = SegQueue::new(); for _ in 0..self.consumer_task_count { consumers.push(Box::pin(RecordBatchStreamAdapter::new( @@ -219,26 +220,45 @@ impl ExecutionPlan for BroadcastExec { } } +#[derive(Debug)] +struct Entry { + value: T, + remaining_readers: usize, +} + #[derive(Debug, Clone, Copy)] struct BroadcastState { - len: usize, + tail_sequence: usize, closed: bool, } +#[derive(Debug)] +struct QueueState { + entries: VecDeque>, + base_sequence: usize, + tail_sequence: usize, + active_readers: usize, +} + #[derive(Debug)] struct BroadcastQueue { - entries: Arc>>, + state: Arc>>, notify: tokio::sync::watch::Sender, } impl BroadcastQueue { - fn new() -> Self { + fn new(expected_readers: usize) -> Self { let (notify, _rx) = tokio::sync::watch::channel(BroadcastState { - len: 0, + tail_sequence: 0, closed: false, }); Self { - entries: Arc::new(Mutex::new(vec![])), + state: Arc::new(Mutex::new(QueueState { + entries: VecDeque::new(), + base_sequence: 0, + tail_sequence: 0, + active_readers: expected_readers, + })), notify, } } @@ -247,22 +267,42 @@ impl BroadcastQueue { let rx = self.notify.subscribe(); let state = *rx.borrow(); BroadcastConsumer { - index: 0, - entries: Arc::clone(&self.entries), + next_sequence: 0, + state: Arc::clone(&self.state), notify: WatchStream::new(rx), - state, + notification: state, + registered: true, } } - fn push(&self, entry: T) { - let len = { - let mut entries = self.entries.lock().unwrap(); - entries.push(entry); - entries.len() + fn push(&self, value: T) { + let tail_sequence = { + let mut queue_state = self.state.lock().unwrap(); + + // Once every consumer has been dropped, no future consumer can execute this + // queue, so don't add values that cannot be read. + if queue_state.active_readers == 0 { + return; + } + + let tail_sequence = queue_state.tail_sequence; + let remaining_readers = queue_state.active_readers; + queue_state.entries.push_back(Entry { + value, + remaining_readers, + }); + queue_state.tail_sequence += 1; + tail_sequence + 1 }; - let mut state = *self.notify.borrow(); - state.len = len; - let _ = self.notify.send(state); + + let mut broadcast_state = *self.notify.borrow(); + broadcast_state.tail_sequence = tail_sequence; + let _ = self.notify.send(broadcast_state); + } + + #[cfg(test)] + fn buffered_len(&self) -> usize { + self.state.lock().unwrap().entries.len() } } @@ -275,11 +315,12 @@ impl Drop for BroadcastQueue { } /// A consumer stream that reads from the broadcast queue. -struct BroadcastConsumer { - index: usize, - entries: Arc>>, +struct BroadcastConsumer { + next_sequence: usize, + state: Arc>>, notify: WatchStream, - state: BroadcastState, + notification: BroadcastState, + registered: bool, } impl Stream for BroadcastConsumer { @@ -287,24 +328,51 @@ impl Stream for BroadcastConsumer { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - if self.index < self.state.len { - let entry = self.entries.lock().unwrap().get(self.index).cloned(); - if let Some(v) = entry { - self.index += 1; - return Poll::Ready(Some(v)); + let value = { + let mut queue_state = self.state.lock().unwrap(); + if self.next_sequence < queue_state.tail_sequence { + let offset = self + .next_sequence + .checked_sub(queue_state.base_sequence) + .expect("broadcast consumer fell behind evicted entries"); + let entry = queue_state + .entries + .get_mut(offset) + .expect("broadcast entry sequence was not retained"); + debug_assert!(entry.remaining_readers > 0); + let value = entry.value.clone(); + entry.remaining_readers -= 1; + + while queue_state + .entries + .front() + .is_some_and(|entry| entry.remaining_readers == 0) + { + queue_state.entries.pop_front(); + queue_state.base_sequence += 1; + } + + Some(value) + } else { + None } + }; + + if let Some(value) = value { + self.next_sequence += 1; + return Poll::Ready(Some(value)); } - if self.state.closed { + if self.notification.closed { return Poll::Ready(None); } match Pin::new(&mut self.notify).poll_next(cx) { Poll::Ready(Some(state)) => { - self.state = state; + self.notification = state; } Poll::Ready(None) => { - self.state.closed = true; + self.notification.closed = true; } Poll::Pending => return Poll::Pending, } @@ -312,6 +380,31 @@ impl Stream for BroadcastConsumer { } } +impl Drop for BroadcastConsumer { + fn drop(&mut self) { + if !self.registered { + return; + } + + let mut state = self.state.lock().unwrap(); + debug_assert!(self.next_sequence >= state.base_sequence); + let offset = self.next_sequence.saturating_sub(state.base_sequence); + for entry in state.entries.iter_mut().skip(offset) { + entry.remaining_readers = entry.remaining_readers.saturating_sub(1); + } + state.active_readers = state.active_readers.saturating_sub(1); + + while state + .entries + .front() + .is_some_and(|entry| entry.remaining_readers == 0) + { + state.entries.pop_front(); + state.base_sequence += 1; + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -337,6 +430,47 @@ mod tests { } } + #[tokio::test] + async fn broadcast_queue_evicts_consumed_prefix() { + let queue = BroadcastQueue::new(2); + let mut consumer0 = queue.new_consumer(); + let mut consumer1 = queue.new_consumer(); + + queue.push(10); + queue.push(20); + assert_eq!(queue.buffered_len(), 2); + + assert_eq!(consumer0.next().await, Some(10)); + assert_eq!(queue.buffered_len(), 2); + assert_eq!(consumer1.next().await, Some(10)); + assert_eq!(queue.buffered_len(), 1); + + assert_eq!(consumer0.next().await, Some(20)); + assert_eq!(queue.buffered_len(), 1); + assert_eq!(consumer1.next().await, Some(20)); + assert_eq!(queue.buffered_len(), 0); + } + + #[tokio::test] + async fn broadcast_queue_drop_releases_unread_entries() { + let queue = BroadcastQueue::new(2); + let mut consumer0 = queue.new_consumer(); + let mut consumer1 = queue.new_consumer(); + + queue.push(10); + queue.push(20); + assert_eq!(consumer0.next().await, Some(10)); + drop(consumer0); + + // The dropped consumer must no longer pin the unread suffix, including entries produced + // after cancellation. + queue.push(30); + assert_eq!(consumer1.next().await, Some(10)); + assert_eq!(consumer1.next().await, Some(20)); + assert_eq!(consumer1.next().await, Some(30)); + assert_eq!(queue.buffered_len(), 0); + } + #[tokio::test] async fn broadcast_exec_reuses_queue_for_virtual_partitions() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); From f5d2eb91a4cc35b7d1869b15a7290c3fe97a5634 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Sun, 13 Sep 2026 20:56:06 -0400 Subject: [PATCH 02/10] get rid of the verbose RSS mode and just use a summary --- .../benches/broadcast_cache_scenarios.rs | 52 +++++++------------ 1 file changed, 18 insertions(+), 34 deletions(-) diff --git a/benchmarks/benches/broadcast_cache_scenarios.rs b/benchmarks/benches/broadcast_cache_scenarios.rs index d18791c85..3b6771d16 100644 --- a/benchmarks/benches/broadcast_cache_scenarios.rs +++ b/benchmarks/benches/broadcast_cache_scenarios.rs @@ -326,16 +326,6 @@ fn scenario_matrix() -> Vec { scenarios } -fn verbose_enabled() -> bool { - match std::env::var("BROADCAST_BENCH_VERBOSE") { - Ok(val) => { - let val = val.to_ascii_lowercase(); - val == "1" || val == "true" || val == "yes" - } - Err(_) => false, - } -} - fn rss_enabled() -> bool { match std::env::var("BROADCAST_BENCH_RSS") { Ok(val) => { @@ -362,7 +352,6 @@ fn bench_broadcast_cache(c: &mut Criterion) { let mut group = c.benchmark_group("broadcast_cache_scenarios"); group.sample_size(10); - let verbose = verbose_enabled(); let sample_rss = rss_enabled(); let task_ctx = SessionContext::new().task_ctx(); @@ -383,11 +372,12 @@ fn bench_broadcast_cache(c: &mut Criterion) { .collect::>(); let batches = Arc::new(batches); + let mut rss_peaks = Vec::new(); group.bench_function(BenchmarkId::new("scenario", scenario.name), |b| { b.iter_custom(|iters| { let mut total = Duration::ZERO; let mut peaks = Vec::with_capacity(iters as usize); - for i in 0..iters { + for _ in 0..iters { let (elapsed, peak_kb) = rt .block_on(run_scenario( &scenario, @@ -397,35 +387,29 @@ fn bench_broadcast_cache(c: &mut Criterion) { sample_rss, )) .expect("scenario"); - if verbose || sample_rss { - eprintln!( - "scenario={} iter={} peak_rss_kb={} elapsed_ms={}", - scenario.name, - i, - peak_kb, - elapsed.as_millis() - ); - } peaks.push(peak_kb); total += elapsed; } - if sample_rss && !peaks.is_empty() { - peaks.sort_unstable(); - let min = peaks[0]; - let max = peaks[peaks.len() - 1]; - let median = peaks[peaks.len() / 2]; - eprintln!( - "scenario={} peak_rss_kb[min/median/max]={}/{}/{} runs={}", - scenario.name, - min, - median, - max, - peaks.len() - ); + if sample_rss { + rss_peaks.extend(peaks); } total }); }); + if sample_rss && !rss_peaks.is_empty() { + rss_peaks.sort_unstable(); + let min = rss_peaks[0]; + let max = rss_peaks[rss_peaks.len() - 1]; + let median = rss_peaks[rss_peaks.len() / 2]; + eprintln!( + "scenario={} peak_rss_kb[min/median/max]={}/{}/{} runs={}", + scenario.name, + min, + median, + max, + rss_peaks.len() + ); + } } group.finish(); From 5baba4604a95e7cdff4422b7bff2b1eb20b34daf Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Mon, 14 Sep 2026 05:58:16 -0400 Subject: [PATCH 03/10] take the non cleanup path out the lock --- src/execution_plans/broadcast.rs | 50 +++++++++++++++----------------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 63183a072..f9218da2a 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -328,34 +328,32 @@ impl Stream for BroadcastConsumer { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - let value = { + let value = if self.next_sequence < self.notification.tail_sequence { let mut queue_state = self.state.lock().unwrap(); - if self.next_sequence < queue_state.tail_sequence { - let offset = self - .next_sequence - .checked_sub(queue_state.base_sequence) - .expect("broadcast consumer fell behind evicted entries"); - let entry = queue_state - .entries - .get_mut(offset) - .expect("broadcast entry sequence was not retained"); - debug_assert!(entry.remaining_readers > 0); - let value = entry.value.clone(); - entry.remaining_readers -= 1; - - while queue_state - .entries - .front() - .is_some_and(|entry| entry.remaining_readers == 0) - { - queue_state.entries.pop_front(); - queue_state.base_sequence += 1; - } - - Some(value) - } else { - None + let offset = self + .next_sequence + .checked_sub(queue_state.base_sequence) + .expect("broadcast consumer fell behind evicted entries"); + let entry = queue_state + .entries + .get_mut(offset) + .expect("broadcast entry sequence was not retained"); + debug_assert!(entry.remaining_readers > 0); + let value = entry.value.clone(); + entry.remaining_readers -= 1; + + while queue_state + .entries + .front() + .is_some_and(|entry| entry.remaining_readers == 0) + { + queue_state.entries.pop_front(); + queue_state.base_sequence += 1; } + + Some(value) + } else { + None }; if let Some(value) = value { From e2ce4a319f63f83f04890455da8299878dd1dcf5 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Mon, 14 Sep 2026 09:28:16 -0400 Subject: [PATCH 04/10] only reclaim memory periodically --- src/execution_plans/broadcast.rs | 213 +++++++++++++++++++------------ 1 file changed, 131 insertions(+), 82 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index f9218da2a..68a922518 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -15,10 +15,13 @@ use futures::{Stream, StreamExt}; use std::collections::VecDeque; use std::fmt::Formatter; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; use std::task::{Context, Poll}; use tokio_stream::wrappers::WatchStream; +const RECLAIM_INTERVAL: usize = 8; + /// [ExecutionPlan] that scales up partitions for network broadcasting. /// /// This plan takes N input partitions and exposes N*M output partitions, @@ -169,9 +172,10 @@ impl ExecutionPlan for BroadcastExec { let queue = BroadcastQueue::new(self.consumer_task_count); let consumers = SegQueue::new(); for _ in 0..self.consumer_task_count { + let consumer = queue.new_consumer().map_err(Arc::new)?; consumers.push(Box::pin(RecordBatchStreamAdapter::new( self.schema(), - queue.new_consumer().map(|msg| match msg { + consumer.map(|msg| match msg { Ok((batch, _reservation)) => Ok(batch), Err(e) => Err(DataFusionError::Shared(e)), }), @@ -221,9 +225,10 @@ impl ExecutionPlan for BroadcastExec { } #[derive(Debug)] -struct Entry { - value: T, - remaining_readers: usize, +struct QueueState { + entries: VecDeque, + base_sequence: usize, + tail_sequence: usize, } #[derive(Debug, Clone, Copy)] @@ -233,16 +238,16 @@ struct BroadcastState { } #[derive(Debug)] -struct QueueState { - entries: VecDeque>, - base_sequence: usize, - tail_sequence: usize, - active_readers: usize, +struct ConsumerRegistry { + consumers: Vec>, + next_consumer: AtomicUsize, + active_count: AtomicUsize, } #[derive(Debug)] struct BroadcastQueue { - state: Arc>>, + queue_state: Arc>>, + consumers: Arc, notify: tokio::sync::watch::Sender, } @@ -253,46 +258,63 @@ impl BroadcastQueue { closed: false, }); Self { - state: Arc::new(Mutex::new(QueueState { + queue_state: Arc::new(Mutex::new(QueueState { entries: VecDeque::new(), base_sequence: 0, tail_sequence: 0, - active_readers: expected_readers, })), + consumers: Arc::new(ConsumerRegistry { + consumers: (0..expected_readers) + .map(|_| { + Arc::new(ConsumerProgress { + sequence: AtomicUsize::new(0), + active: AtomicBool::new(true), + }) + }) + .collect(), + next_consumer: AtomicUsize::new(0), + active_count: AtomicUsize::new(expected_readers), + }), notify, } } - fn new_consumer(&self) -> BroadcastConsumer { + fn new_consumer(&self) -> Result> { let rx = self.notify.subscribe(); let state = *rx.borrow(); - BroadcastConsumer { + let consumer_index = self.consumers.next_consumer.fetch_add(1, Ordering::Relaxed); + let Some(progress) = self.consumers.consumers.get(consumer_index) else { + return internal_err!( + "broadcast queue created more consumers than expected (index {consumer_index})" + ); + }; + let progress = Arc::clone(progress); + Ok(BroadcastConsumer { next_sequence: 0, - state: Arc::clone(&self.state), + progress, + consumers: Arc::clone(&self.consumers), + queue_state: Arc::clone(&self.queue_state), notify: WatchStream::new(rx), notification: state, - registered: true, - } + }) } fn push(&self, value: T) { let tail_sequence = { - let mut queue_state = self.state.lock().unwrap(); + let mut queue_state = self.queue_state.lock().unwrap(); // Once every consumer has been dropped, no future consumer can execute this // queue, so don't add values that cannot be read. - if queue_state.active_readers == 0 { + if self.consumers.active_count.load(Ordering::Acquire) == 0 { return; } - let tail_sequence = queue_state.tail_sequence; - let remaining_readers = queue_state.active_readers; - queue_state.entries.push_back(Entry { - value, - remaining_readers, - }); + queue_state.entries.push_back(value); queue_state.tail_sequence += 1; - tail_sequence + 1 + if queue_state.tail_sequence.is_multiple_of(RECLAIM_INTERVAL) { + Self::reclaim_locked(&self.consumers, &mut queue_state); + } + queue_state.tail_sequence }; let mut broadcast_state = *self.notify.borrow(); @@ -300,9 +322,22 @@ impl BroadcastQueue { let _ = self.notify.send(broadcast_state); } - #[cfg(test)] - fn buffered_len(&self) -> usize { - self.state.lock().unwrap().entries.len() + fn reclaim_locked(consumers: &ConsumerRegistry, queue_state: &mut QueueState) { + let minimum_sequence = consumers + .consumers + .iter() + .filter(|consumer| consumer.active.load(Ordering::Acquire)) + .map(|consumer| consumer.sequence.load(Ordering::Acquire)) + .min() + .unwrap_or(queue_state.tail_sequence); + + while queue_state.base_sequence < minimum_sequence { + queue_state + .entries + .pop_front() + .expect("broadcast sequence bounds were inconsistent"); + queue_state.base_sequence += 1; + } } } @@ -314,13 +349,20 @@ impl Drop for BroadcastQueue { } } +#[derive(Debug)] +struct ConsumerProgress { + sequence: AtomicUsize, + active: AtomicBool, +} + /// A consumer stream that reads from the broadcast queue. struct BroadcastConsumer { next_sequence: usize, - state: Arc>>, + progress: Arc, + consumers: Arc, + queue_state: Arc>>, notify: WatchStream, notification: BroadcastState, - registered: bool, } impl Stream for BroadcastConsumer { @@ -329,39 +371,34 @@ impl Stream for BroadcastConsumer { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { let value = if self.next_sequence < self.notification.tail_sequence { - let mut queue_state = self.state.lock().unwrap(); - let offset = self - .next_sequence - .checked_sub(queue_state.base_sequence) - .expect("broadcast consumer fell behind evicted entries"); - let entry = queue_state - .entries - .get_mut(offset) - .expect("broadcast entry sequence was not retained"); - debug_assert!(entry.remaining_readers > 0); - let value = entry.value.clone(); - entry.remaining_readers -= 1; - - while queue_state - .entries - .front() - .is_some_and(|entry| entry.remaining_readers == 0) - { - queue_state.entries.pop_front(); - queue_state.base_sequence += 1; + let queue_state = self.queue_state.lock().unwrap(); + if self.next_sequence >= queue_state.tail_sequence { + None + } else { + let offset = self + .next_sequence + .checked_sub(queue_state.base_sequence) + .expect("broadcast consumer fell behind evicted entries"); + let entry = queue_state + .entries + .get(offset) + .expect("broadcast entry sequence was not retained"); + Some(entry.clone()) } - - Some(value) } else { None }; if let Some(value) = value { self.next_sequence += 1; + self.progress + .sequence + .store(self.next_sequence, Ordering::Release); return Poll::Ready(Some(value)); } if self.notification.closed { + self.reclaim(); return Poll::Ready(None); } @@ -378,28 +415,19 @@ impl Stream for BroadcastConsumer { } } +impl BroadcastConsumer { + fn reclaim(&self) { + let mut queue_state = self.queue_state.lock().unwrap(); + BroadcastQueue::::reclaim_locked(&self.consumers, &mut queue_state); + } +} + impl Drop for BroadcastConsumer { fn drop(&mut self) { - if !self.registered { - return; - } - - let mut state = self.state.lock().unwrap(); - debug_assert!(self.next_sequence >= state.base_sequence); - let offset = self.next_sequence.saturating_sub(state.base_sequence); - for entry in state.entries.iter_mut().skip(offset) { - entry.remaining_readers = entry.remaining_readers.saturating_sub(1); - } - state.active_readers = state.active_readers.saturating_sub(1); - - while state - .entries - .front() - .is_some_and(|entry| entry.remaining_readers == 0) - { - state.entries.pop_front(); - state.base_sequence += 1; + if self.progress.active.swap(false, Ordering::AcqRel) { + self.consumers.active_count.fetch_sub(1, Ordering::AcqRel); } + self.reclaim(); } } @@ -428,32 +456,52 @@ mod tests { } } + fn buffered_len(queue: &BroadcastQueue) -> usize { + queue.queue_state.lock().unwrap().entries.len() + } + + fn sequence_bounds(queue: &BroadcastQueue) -> (usize, usize) { + let state = queue.queue_state.lock().unwrap(); + (state.base_sequence, state.tail_sequence) + } + #[tokio::test] async fn broadcast_queue_evicts_consumed_prefix() { let queue = BroadcastQueue::new(2); - let mut consumer0 = queue.new_consumer(); - let mut consumer1 = queue.new_consumer(); + let mut consumer0 = queue.new_consumer().expect("consumer 0 registration"); + let mut consumer1 = queue.new_consumer().expect("consumer 1 registration"); queue.push(10); queue.push(20); - assert_eq!(queue.buffered_len(), 2); + assert_eq!(buffered_len(&queue), 2); assert_eq!(consumer0.next().await, Some(10)); - assert_eq!(queue.buffered_len(), 2); + assert_eq!(buffered_len(&queue), 2); assert_eq!(consumer1.next().await, Some(10)); - assert_eq!(queue.buffered_len(), 1); + assert_eq!(sequence_bounds(&queue), (0, 2)); assert_eq!(consumer0.next().await, Some(20)); - assert_eq!(queue.buffered_len(), 1); + assert_eq!(buffered_len(&queue), 2); assert_eq!(consumer1.next().await, Some(20)); - assert_eq!(queue.buffered_len(), 0); + assert_eq!(buffered_len(&queue), 2); + + // The eighth append observes that consumers have advanced and removes the consumed prefix. + for value in [30, 40, 50, 60, 70, 80] { + queue.push(value); + } + assert_eq!(sequence_bounds(&queue), (2, 8)); + assert_eq!(buffered_len(&queue), 6); + + drop(consumer0); + drop(consumer1); + assert_eq!(buffered_len(&queue), 0); } #[tokio::test] async fn broadcast_queue_drop_releases_unread_entries() { let queue = BroadcastQueue::new(2); - let mut consumer0 = queue.new_consumer(); - let mut consumer1 = queue.new_consumer(); + let mut consumer0 = queue.new_consumer().expect("consumer 0 registration"); + let mut consumer1 = queue.new_consumer().expect("consumer 1 registration"); queue.push(10); queue.push(20); @@ -466,7 +514,8 @@ mod tests { assert_eq!(consumer1.next().await, Some(10)); assert_eq!(consumer1.next().await, Some(20)); assert_eq!(consumer1.next().await, Some(30)); - assert_eq!(queue.buffered_len(), 0); + drop(consumer1); + assert_eq!(buffered_len(&queue), 0); } #[tokio::test] From 3718cab7ae8d88c3950d08411664942d89787d6a Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Mon, 14 Sep 2026 11:29:39 -0400 Subject: [PATCH 05/10] separte criterion and memory bench paths using peak recroding pool --- .../benches/broadcast_cache_scenarios.rs | 208 ++++++++++++++---- 1 file changed, 161 insertions(+), 47 deletions(-) diff --git a/benchmarks/benches/broadcast_cache_scenarios.rs b/benchmarks/benches/broadcast_cache_scenarios.rs index 3b6771d16..33dc1aa34 100644 --- a/benchmarks/benches/broadcast_cache_scenarios.rs +++ b/benchmarks/benches/broadcast_cache_scenarios.rs @@ -5,6 +5,8 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Statistics; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::Result; +use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -12,7 +14,7 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, }; -use datafusion::prelude::SessionContext; +use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_distributed::BroadcastExec; use futures::{StreamExt, stream}; use std::sync::{ @@ -49,31 +51,30 @@ struct Scenario { struct PeakRssSampler { stop: Arc, - peak_kb: Arc, + peak_bytes: Arc, handle: Option>, } impl PeakRssSampler { fn start(interval: Duration) -> Self { let stop = Arc::new(AtomicBool::new(false)); - let peak_kb = Arc::new(AtomicU64::new(0)); + let peak_bytes = Arc::new(AtomicU64::new(0)); let stop_clone = Arc::clone(&stop); - let peak_clone = Arc::clone(&peak_kb); + let peak_clone = Arc::clone(&peak_bytes); let handle = thread::spawn(move || { let mut sys = System::new(); let pid = get_current_pid().expect("pid"); while !stop_clone.load(Ordering::Relaxed) { sys.refresh_process(pid); if let Some(proc) = sys.process(pid) { - let mem_kb = proc.memory(); - peak_clone.fetch_max(mem_kb, Ordering::Relaxed); + peak_clone.fetch_max(proc.memory(), Ordering::Relaxed); } thread::sleep(interval); } }); Self { stop, - peak_kb, + peak_bytes, handle: Some(handle), } } @@ -83,7 +84,7 @@ impl PeakRssSampler { if let Some(handle) = self.handle.take() { let _ = handle.join(); } - self.peak_kb.load(Ordering::Relaxed) + self.peak_bytes.load(Ordering::Relaxed) } } @@ -92,11 +93,17 @@ struct SyntheticExec { schema: SchemaRef, partitions: usize, batches: Arc>>, + interval: Option, properties: Arc, } impl SyntheticExec { - fn new(schema: SchemaRef, partitions: usize, batches: Arc>>) -> Self { + fn new( + schema: SchemaRef, + partitions: usize, + batches: Arc>>, + interval: Option, + ) -> Self { let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(Arc::clone(&schema)), Partitioning::UnknownPartitioning(partitions), @@ -107,6 +114,7 @@ impl SyntheticExec { schema, partitions, batches, + interval, properties, } } @@ -158,12 +166,29 @@ impl ExecutionPlan for SyntheticExec { let batches = Arc::clone(&self.batches); let len = batches.len(); - let stream = stream::iter((0..len).map(move |idx| { - let batch = &batches[idx]; - Ok(batch.as_ref().clone()) - })); - - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + match self.interval { + None => { + let stream = stream::iter((0..len).map(move |idx| { + let batch = &batches[idx]; + Ok(batch.as_ref().clone()) + })); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + Some(interval) => { + let stream = futures::stream::unfold(0usize, move |idx| { + let batches = Arc::clone(&batches); + async move { + if idx >= len { + return None; + } + tokio::time::sleep(interval).await; + let batch = batches[idx].as_ref().clone(); + Some((Ok(batch), idx + 1)) + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + } } fn partition_statistics(&self, _partition: Option) -> Result> { @@ -197,17 +222,11 @@ async fn consume_partition( async fn run_scenario( scenario: &Scenario, - schema: Arc, - batches: Arc>>, + input: Arc, task_ctx: Arc, sample_rss: bool, -) -> Result<(Duration, u64)> { - let input: Arc = Arc::new(SyntheticExec::new( - Arc::clone(&schema), - scenario.input_partitions, - batches, - )); - + recording_pool: Option>, +) -> Result<(Duration, u64, Option<(usize, usize)>)> { let broadcast = Arc::new(BroadcastExec::new( Arc::clone(&input), scenario.consumer_tasks, @@ -232,8 +251,51 @@ async fn run_scenario( } let elapsed = start.elapsed(); - let peak_kb = sampler.map(|s| s.stop()).unwrap_or(0); - Ok((elapsed, peak_kb)) + let peak_rss_bytes = sampler.map(|s| s.stop()).unwrap_or(0); + let memory = recording_pool.map(|pool| (pool.peak_reserved(), pool.reserved())); + Ok((elapsed, peak_rss_bytes, memory)) +} + +fn memory_context() -> Result<(Arc, Arc)> { + let inner_pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let recording_pool = Arc::new(PeakRecordingPool::new(inner_pool)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&recording_pool) as Arc) + .build()?; + let task_ctx = + SessionContext::new_with_config_rt(SessionConfig::new(), Arc::new(runtime)).task_ctx(); + Ok((task_ctx, recording_pool)) +} + +fn synthetic_input( + scenario: &Scenario, + schema: &SchemaRef, + interval: Option, +) -> Arc { + let batches = (0..scenario.num_batches) + .map(|_| { + let array = UInt8Array::from(vec![0u8; scenario.rows_per_batch]); + let batch = + RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(array)]).expect("batch"); + Arc::new(batch) + }) + .collect::>(); + Arc::new(SyntheticExec::new( + Arc::clone(schema), + scenario.input_partitions, + Arc::new(batches), + interval, + )) +} + +fn prebuilt_input(scenario: &Scenario, schema: &SchemaRef) -> Arc { + synthetic_input(scenario, schema, None) +} + +fn paced_input(scenario: &Scenario, schema: &SchemaRef) -> Arc { + // Consumers are started before the first batch and receive a scheduling window between + // batches. This is intentionally outside the timing measurement. + synthetic_input(scenario, schema, Some(Duration::from_millis(1))) } fn all_fast_consumers(output_partitions: usize) -> Vec { @@ -336,6 +398,16 @@ fn rss_enabled() -> bool { } } +fn memory_enabled() -> bool { + match std::env::var("BROADCAST_BENCH_MEMORY") { + Ok(val) => { + let val = val.to_ascii_lowercase(); + val == "1" || val == "true" || val == "yes" + } + Err(_) => false, + } +} + fn runtime_threads() -> Option { std::env::var("BROADCAST_BENCH_THREADS") .ok() @@ -343,6 +415,8 @@ fn runtime_threads() -> Option { .filter(|threads| *threads > 0) } +const DIAGNOSTIC_RUNS: usize = 5; + fn bench_broadcast_cache(c: &mut Criterion) { let mut rt_builder = RuntimeBuilder::new_multi_thread(); if let Some(threads) = runtime_threads() { @@ -353,6 +427,7 @@ fn bench_broadcast_cache(c: &mut Criterion) { let mut group = c.benchmark_group("broadcast_cache_scenarios"); group.sample_size(10); let sample_rss = rss_enabled(); + let sample_memory = memory_enabled(); let task_ctx = SessionContext::new().task_ctx(); for scenario in scenario_matrix() { @@ -361,48 +436,63 @@ fn bench_broadcast_cache(c: &mut Criterion) { DataType::UInt8, false, )])); - let batches = (0..scenario.num_batches) - .map(|_| { - let data = vec![0u8; scenario.rows_per_batch]; - let array = UInt8Array::from(data); - let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]) - .expect("batch"); - Arc::new(batch) - }) - .collect::>(); - let batches = Arc::new(batches); - + let input = prebuilt_input(&scenario, &schema); let mut rss_peaks = Vec::new(); + let mut memory_peaks = Vec::new(); + let mut memory_residuals = Vec::new(); group.bench_function(BenchmarkId::new("scenario", scenario.name), |b| { b.iter_custom(|iters| { let mut total = Duration::ZERO; - let mut peaks = Vec::with_capacity(iters as usize); for _ in 0..iters { - let (elapsed, peak_kb) = rt + let (elapsed, _, _) = rt .block_on(run_scenario( &scenario, - Arc::clone(&schema), - Arc::clone(&batches), + Arc::clone(&input), Arc::clone(&task_ctx), - sample_rss, + false, + None, )) .expect("scenario"); - peaks.push(peak_kb); total += elapsed; } - if sample_rss { - rss_peaks.extend(peaks); - } total }); }); + + if sample_rss || sample_memory { + for _ in 0..DIAGNOSTIC_RUNS { + let (diagnostic_task_ctx, recording_pool) = if sample_memory { + let (task_ctx, pool) = memory_context().expect("memory context"); + (task_ctx, Some(pool)) + } else { + (Arc::clone(&task_ctx), None) + }; + let (_, peak_rss_bytes, memory) = rt + .block_on(run_scenario( + &scenario, + paced_input(&scenario, &schema), + diagnostic_task_ctx, + sample_rss, + recording_pool, + )) + .expect("diagnostic scenario"); + if sample_rss { + rss_peaks.push(peak_rss_bytes); + } + if let Some((peak_reserved, residual_reserved)) = memory { + memory_peaks.push(peak_reserved); + memory_residuals.push(residual_reserved); + } + } + } + if sample_rss && !rss_peaks.is_empty() { rss_peaks.sort_unstable(); let min = rss_peaks[0]; let max = rss_peaks[rss_peaks.len() - 1]; let median = rss_peaks[rss_peaks.len() / 2]; eprintln!( - "scenario={} peak_rss_kb[min/median/max]={}/{}/{} runs={}", + "scenario={} peak_rss_bytes[min/median/max]={}/{}/{} runs={}", scenario.name, min, median, @@ -410,6 +500,30 @@ fn bench_broadcast_cache(c: &mut Criterion) { rss_peaks.len() ); } + if sample_memory && !memory_peaks.is_empty() { + memory_peaks.sort_unstable(); + memory_residuals.sort_unstable(); + let peak_min = memory_peaks[0]; + let peak_median = memory_peaks[memory_peaks.len() / 2]; + let peak_max = memory_peaks[memory_peaks.len() - 1]; + let residual_min = memory_residuals[0]; + let residual_median = memory_residuals[memory_residuals.len() / 2]; + let residual_max = memory_residuals[memory_residuals.len() - 1]; + eprintln!( + concat!( + "scenario={} reserved_bytes[peak_min/median/max]={}/{}/{} ", + "end_min/median/max={}/{}/{} runs={}" + ), + scenario.name, + peak_min, + peak_median, + peak_max, + residual_min, + residual_median, + residual_max, + memory_peaks.len() + ); + } } group.finish(); From bc88d3b32317b8fe1c5af0f2384e8f6e58de067f Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 15 Sep 2026 06:06:06 -0400 Subject: [PATCH 06/10] simplify the amount of state --- src/execution_plans/broadcast.rs | 406 ++++++++++++++++++++----------- 1 file changed, 261 insertions(+), 145 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 68a922518..0f252c2d7 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -1,10 +1,10 @@ use crate::common::{OnceLockResult, require_one_child}; -use crossbeam_queue::SegQueue; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::runtime::SpawnedTask; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError, Result}; -use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; @@ -15,10 +15,10 @@ use futures::{Stream, StreamExt}; use std::collections::VecDeque; use std::fmt::Formatter; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; use std::task::{Context, Poll}; use tokio_stream::wrappers::WatchStream; +use tokio_util::sync::CancellationToken; const RECLAIM_INTERVAL: usize = 8; @@ -81,7 +81,9 @@ pub struct BroadcastExec { queues: Vec>, } -type StreamAndTask = (SegQueue, Arc>); +type BroadcastMessage = + std::result::Result<(RecordBatch, Arc), Arc>; +type StreamAndTask = (BroadcastReaders, Arc>); impl BroadcastExec { pub fn new(input: Arc, consumer_task_count: usize) -> Self { @@ -116,6 +118,18 @@ impl BroadcastExec { } } +impl Drop for BroadcastExec { + fn drop(&mut self) { + // The last plan reference is gone, so its unclaimed virtual partitions cannot + // execute. Active streams may still hold their producer task and queue alive. + for queue in &self.queues { + if let Some(Ok((readers, _task))) = queue.get() { + readers.release_pending(); + } + } + } +} + impl DisplayAs for BroadcastExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { let input_partition_count = self.input_partition_count(); @@ -164,58 +178,60 @@ impl ExecutionPlan for BroadcastExec { partition: usize, context: Arc, ) -> Result { - let real_partition = partition % self.input_partition_count(); + let input_partition_count = self.input_partition_count(); + let real_partition = partition % input_partition_count; + let consumer_id = partition / input_partition_count; let input = Arc::clone(&self.input); let queue_or_err = self.queues[real_partition].get_or_init(|| { let queue = BroadcastQueue::new(self.consumer_task_count); - let consumers = SegQueue::new(); - for _ in 0..self.consumer_task_count { - let consumer = queue.new_consumer().map_err(Arc::new)?; - consumers.push(Box::pin(RecordBatchStreamAdapter::new( - self.schema(), - consumer.map(|msg| match msg { - Ok((batch, _reservation)) => Ok(batch), - Err(e) => Err(DataFusionError::Shared(e)), - }), - )) as SendableRecordBatchStream); - } + let readers = queue.readers(); let pool = Arc::clone(context.memory_pool()); let mut stream = input.execute(real_partition, context).map_err(Arc::new)?; + let cancel = queue.shared.cancel.clone(); let task = SpawnedTask::spawn(async move { let mem_consumer = MemoryConsumer::new(format!("BroadcastExec[{real_partition}]")); - while let Some(msg) = stream.next().await { + loop { + let msg = tokio::select! { + _ = cancel.cancelled() => break, + msg = stream.next() => msg, + }; + let Some(msg) = msg else { break }; match msg { Ok(record_batch) => { let reservation = mem_consumer.clone_with_new_id().register(&pool); reservation.grow(record_batch.get_array_memory_size()); - queue.push(Ok((record_batch, Arc::new(reservation)))); + if !queue.push(Ok((record_batch, Arc::new(reservation)))) { + break; + } } Err(err) => { - queue.push(Err(Arc::new(err))); + let _ = queue.push(Err(Arc::new(err))); break; } } } }); - Ok::<_, Arc>((consumers, Arc::new(task))) + Ok::<_, Arc>((readers, Arc::new(task))) }); let (consumer, task) = match queue_or_err { - Ok((consumers, task)) => (consumers.pop(), Arc::clone(task)), + Ok((readers, task)) => (readers.claim(consumer_id)?, Arc::clone(task)), Err(err) => return Err(DataFusionError::Shared(Arc::clone(err))), }; - let Some(consumer) = consumer else { - return internal_err!("Too many consumers for real partition {real_partition}"); - }; Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), - consumer.inspect(move |_| { - let _ = &task; - }), + consumer + .map(|msg| match msg { + Ok((batch, _reservation)) => Ok(batch), + Err(e) => Err(DataFusionError::Shared(e)), + }) + .inspect(move |_| { + let _ = &task; + }), ))) } @@ -229,105 +245,88 @@ struct QueueState { entries: VecDeque, base_sequence: usize, tail_sequence: usize, + readers: Box<[ReaderSlot]>, + retaining_readers: usize, + closed: bool, } -#[derive(Debug, Clone, Copy)] -struct BroadcastState { - tail_sequence: usize, - closed: bool, +#[derive(Debug)] +enum ReaderSlot { + Pending, + Reading { next_sequence: usize }, + Released, } #[derive(Debug)] -struct ConsumerRegistry { - consumers: Vec>, - next_consumer: AtomicUsize, - active_count: AtomicUsize, +struct BroadcastShared { + queue_state: Mutex>, + notify: tokio::sync::watch::Sender<()>, + cancel: CancellationToken, } #[derive(Debug)] struct BroadcastQueue { - queue_state: Arc>>, - consumers: Arc, - notify: tokio::sync::watch::Sender, + shared: Arc>, +} + +#[derive(Debug)] +struct BroadcastReaders { + shared: Arc>, } impl BroadcastQueue { fn new(expected_readers: usize) -> Self { - let (notify, _rx) = tokio::sync::watch::channel(BroadcastState { - tail_sequence: 0, - closed: false, - }); + let (notify, _rx) = tokio::sync::watch::channel(()); Self { - queue_state: Arc::new(Mutex::new(QueueState { - entries: VecDeque::new(), - base_sequence: 0, - tail_sequence: 0, - })), - consumers: Arc::new(ConsumerRegistry { - consumers: (0..expected_readers) - .map(|_| { - Arc::new(ConsumerProgress { - sequence: AtomicUsize::new(0), - active: AtomicBool::new(true), - }) - }) - .collect(), - next_consumer: AtomicUsize::new(0), - active_count: AtomicUsize::new(expected_readers), + shared: Arc::new(BroadcastShared { + queue_state: Mutex::new(QueueState { + entries: VecDeque::new(), + base_sequence: 0, + tail_sequence: 0, + readers: (0..expected_readers).map(|_| ReaderSlot::Pending).collect(), + retaining_readers: expected_readers, + closed: false, + }), + notify, + cancel: CancellationToken::new(), }), - notify, } } - fn new_consumer(&self) -> Result> { - let rx = self.notify.subscribe(); - let state = *rx.borrow(); - let consumer_index = self.consumers.next_consumer.fetch_add(1, Ordering::Relaxed); - let Some(progress) = self.consumers.consumers.get(consumer_index) else { - return internal_err!( - "broadcast queue created more consumers than expected (index {consumer_index})" - ); - }; - let progress = Arc::clone(progress); - Ok(BroadcastConsumer { - next_sequence: 0, - progress, - consumers: Arc::clone(&self.consumers), - queue_state: Arc::clone(&self.queue_state), - notify: WatchStream::new(rx), - notification: state, - }) + fn readers(&self) -> BroadcastReaders { + BroadcastReaders { + shared: Arc::clone(&self.shared), + } } - fn push(&self, value: T) { - let tail_sequence = { - let mut queue_state = self.queue_state.lock().unwrap(); + fn push(&self, value: T) -> bool { + { + let mut queue_state = self.shared.queue_state.lock().unwrap(); - // Once every consumer has been dropped, no future consumer can execute this - // queue, so don't add values that cannot be read. - if self.consumers.active_count.load(Ordering::Acquire) == 0 { - return; + if queue_state.retaining_readers == 0 { + return false; } queue_state.entries.push_back(value); queue_state.tail_sequence += 1; if queue_state.tail_sequence.is_multiple_of(RECLAIM_INTERVAL) { - Self::reclaim_locked(&self.consumers, &mut queue_state); + Self::reclaim_locked(&mut queue_state); } - queue_state.tail_sequence - }; + } - let mut broadcast_state = *self.notify.borrow(); - broadcast_state.tail_sequence = tail_sequence; - let _ = self.notify.send(broadcast_state); + self.shared.notify.send_replace(()); + true } - fn reclaim_locked(consumers: &ConsumerRegistry, queue_state: &mut QueueState) { - let minimum_sequence = consumers - .consumers + fn reclaim_locked(queue_state: &mut QueueState) { + let minimum_sequence = queue_state + .readers .iter() - .filter(|consumer| consumer.active.load(Ordering::Acquire)) - .map(|consumer| consumer.sequence.load(Ordering::Acquire)) + .filter_map(|reader| match reader { + ReaderSlot::Pending => Some(0), + ReaderSlot::Reading { next_sequence } => Some(*next_sequence), + ReaderSlot::Released => None, + }) .min() .unwrap_or(queue_state.tail_sequence); @@ -343,26 +342,59 @@ impl BroadcastQueue { impl Drop for BroadcastQueue { fn drop(&mut self) { - let mut state = *self.notify.borrow(); - state.closed = true; - let _ = self.notify.send(state); + { + let mut state = self.shared.queue_state.lock().unwrap(); + state.closed = true; + } + self.shared.notify.send_replace(()); } } -#[derive(Debug)] -struct ConsumerProgress { - sequence: AtomicUsize, - active: AtomicBool, +impl BroadcastReaders { + fn claim(&self, consumer_id: usize) -> Result> { + let rx = self.shared.notify.subscribe(); + let mut state = self.shared.queue_state.lock().unwrap(); + let Some(reader) = state.readers.get_mut(consumer_id) else { + return internal_err!("broadcast consumer {consumer_id} is out of range"); + }; + match reader { + ReaderSlot::Pending => *reader = ReaderSlot::Reading { next_sequence: 0 }, + ReaderSlot::Reading { .. } | ReaderSlot::Released => { + return internal_err!("broadcast consumer {consumer_id} cannot execute twice"); + } + } + Ok(BroadcastConsumer { + consumer_id, + shared: Arc::clone(&self.shared), + notify: WatchStream::new(rx), + }) + } + + fn release_pending(&self) { + let no_readers_remain = { + let mut state = self.shared.queue_state.lock().unwrap(); + let mut released = 0; + for reader in &mut state.readers { + if matches!(reader, ReaderSlot::Pending) { + *reader = ReaderSlot::Released; + released += 1; + } + } + state.retaining_readers -= released; + BroadcastQueue::::reclaim_locked(&mut state); + state.retaining_readers == 0 + }; + if no_readers_remain { + self.shared.cancel.cancel(); + } + } } /// A consumer stream that reads from the broadcast queue. struct BroadcastConsumer { - next_sequence: usize, - progress: Arc, - consumers: Arc, - queue_state: Arc>>, - notify: WatchStream, - notification: BroadcastState, + consumer_id: usize, + shared: Arc>, + notify: WatchStream<()>, } impl Stream for BroadcastConsumer { @@ -370,44 +402,45 @@ impl Stream for BroadcastConsumer { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - let value = if self.next_sequence < self.notification.tail_sequence { - let queue_state = self.queue_state.lock().unwrap(); - if self.next_sequence >= queue_state.tail_sequence { - None - } else { - let offset = self - .next_sequence - .checked_sub(queue_state.base_sequence) + let (value, closed) = { + let mut state = self.shared.queue_state.lock().unwrap(); + let next_sequence = match state.readers[self.consumer_id] { + ReaderSlot::Reading { next_sequence } => next_sequence, + ReaderSlot::Released => return Poll::Ready(None), + ReaderSlot::Pending => unreachable!("an unclaimed consumer was polled"), + }; + if next_sequence < state.tail_sequence { + let offset = next_sequence + .checked_sub(state.base_sequence) .expect("broadcast consumer fell behind evicted entries"); - let entry = queue_state + let value = state .entries .get(offset) - .expect("broadcast entry sequence was not retained"); - Some(entry.clone()) + .expect("broadcast entry sequence was not retained") + .clone(); + state.readers[self.consumer_id] = ReaderSlot::Reading { + next_sequence: next_sequence + 1, + }; + (Some(value), false) + } else { + (None, state.closed) } - } else { - None }; if let Some(value) = value { - self.next_sequence += 1; - self.progress - .sequence - .store(self.next_sequence, Ordering::Release); return Poll::Ready(Some(value)); } - if self.notification.closed { - self.reclaim(); + if closed { + self.shared.release(self.consumer_id); return Poll::Ready(None); } match Pin::new(&mut self.notify).poll_next(cx) { - Poll::Ready(Some(state)) => { - self.notification = state; - } + Poll::Ready(Some(_)) => continue, Poll::Ready(None) => { - self.notification.closed = true; + self.shared.release(self.consumer_id); + return Poll::Ready(None); } Poll::Pending => return Poll::Pending, } @@ -415,19 +448,27 @@ impl Stream for BroadcastConsumer { } } -impl BroadcastConsumer { - fn reclaim(&self) { - let mut queue_state = self.queue_state.lock().unwrap(); - BroadcastQueue::::reclaim_locked(&self.consumers, &mut queue_state); +impl BroadcastShared { + fn release(&self, consumer_id: usize) { + let no_readers_remain = { + let mut state = self.queue_state.lock().unwrap(); + if matches!(state.readers[consumer_id], ReaderSlot::Released) { + return; + } + state.readers[consumer_id] = ReaderSlot::Released; + state.retaining_readers -= 1; + BroadcastQueue::::reclaim_locked(&mut state); + state.retaining_readers == 0 + }; + if no_readers_remain { + self.cancel.cancel(); + } } } impl Drop for BroadcastConsumer { fn drop(&mut self) { - if self.progress.active.swap(false, Ordering::AcqRel) { - self.consumers.active_count.fetch_sub(1, Ordering::AcqRel); - } - self.reclaim(); + self.shared.release(self.consumer_id); } } @@ -457,19 +498,19 @@ mod tests { } fn buffered_len(queue: &BroadcastQueue) -> usize { - queue.queue_state.lock().unwrap().entries.len() + queue.shared.queue_state.lock().unwrap().entries.len() } fn sequence_bounds(queue: &BroadcastQueue) -> (usize, usize) { - let state = queue.queue_state.lock().unwrap(); + let state = queue.shared.queue_state.lock().unwrap(); (state.base_sequence, state.tail_sequence) } #[tokio::test] async fn broadcast_queue_evicts_consumed_prefix() { let queue = BroadcastQueue::new(2); - let mut consumer0 = queue.new_consumer().expect("consumer 0 registration"); - let mut consumer1 = queue.new_consumer().expect("consumer 1 registration"); + let mut consumer0 = queue.readers().claim(0).expect("consumer 0 registration"); + let mut consumer1 = queue.readers().claim(1).expect("consumer 1 registration"); queue.push(10); queue.push(20); @@ -500,8 +541,8 @@ mod tests { #[tokio::test] async fn broadcast_queue_drop_releases_unread_entries() { let queue = BroadcastQueue::new(2); - let mut consumer0 = queue.new_consumer().expect("consumer 0 registration"); - let mut consumer1 = queue.new_consumer().expect("consumer 1 registration"); + let mut consumer0 = queue.readers().claim(0).expect("consumer 0 registration"); + let mut consumer1 = queue.readers().claim(1).expect("consumer 1 registration"); queue.push(10); queue.push(20); @@ -518,6 +559,64 @@ mod tests { assert_eq!(buffered_len(&queue), 0); } + #[tokio::test] + async fn broadcast_queue_releases_reader_at_eof() { + let queue = BroadcastQueue::new(1); + let shared = Arc::clone(&queue.shared); + let mut consumer = queue.readers().claim(0).expect("consumer registration"); + + assert!(queue.push(10)); + assert_eq!(consumer.next().await, Some(10)); + drop(queue); + + assert_eq!(consumer.next().await, None); + assert_eq!(shared.queue_state.lock().unwrap().entries.len(), 0); + assert!(shared.cancel.is_cancelled()); + assert_eq!(consumer.next().await, None); + } + + #[tokio::test] + async fn broadcast_queue_cancels_when_all_readers_drop() { + let queue = BroadcastQueue::new(2); + let reader0 = queue.readers().claim(0).expect("consumer 0 registration"); + let reader1 = queue.readers().claim(1).expect("consumer 1 registration"); + + assert!(queue.push(10)); + drop(reader0); + assert!(!queue.shared.cancel.is_cancelled()); + drop(reader1); + assert!(queue.shared.cancel.is_cancelled()); + assert_eq!(buffered_len(&queue), 0); + assert!(!queue.push(20)); + } + + #[tokio::test] + async fn broadcast_exec_releases_pending_readers_when_plan_drops() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input = Arc::new(MockExec::new_partitioned(vec![vec![]], Arc::clone(&schema))); + let broadcast = BroadcastExec::new(input, 2); + let task_ctx = SessionContext::new().task_ctx(); + + let stream = broadcast.execute(0, task_ctx)?; + let shared = Arc::clone( + &broadcast.queues[0] + .get() + .expect("initialized queue") + .as_ref() + .expect("queue initialization") + .0 + .shared, + ); + + drop(broadcast); + assert_eq!(shared.queue_state.lock().unwrap().retaining_readers, 1); + drop(stream); + assert_eq!(shared.queue_state.lock().unwrap().retaining_readers, 0); + assert!(shared.cancel.is_cancelled()); + + Ok(()) + } + #[tokio::test] async fn broadcast_exec_reuses_queue_for_virtual_partitions() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -608,6 +707,23 @@ mod tests { Ok(()) } + #[tokio::test] + async fn broadcast_exec_claims_virtual_partition_once() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input = Arc::new(MockExec::new_partitioned(vec![vec![]], Arc::clone(&schema))); + let broadcast = BroadcastExec::new(input, 2); + let task_ctx = SessionContext::new().task_ctx(); + + let _stream = broadcast.execute(0, Arc::clone(&task_ctx))?; + let err = broadcast + .execute(0, task_ctx) + .err() + .expect("duplicate execute"); + assert!(err.to_string().contains("consumer 0 cannot execute twice")); + + Ok(()) + } + #[tokio::test] async fn broadcast_exec_queue_survives_cancellation() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); From 3b9211250ac3896ff15246f7e6643ecf38b31f68 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 15 Sep 2026 13:27:24 -0400 Subject: [PATCH 07/10] add some sick diagrams --- src/execution_plans/broadcast.rs | 154 ++++++++++++++++++++++++++++--- 1 file changed, 143 insertions(+), 11 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 0f252c2d7..8ee0a3b15 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -180,7 +180,7 @@ impl ExecutionPlan for BroadcastExec { ) -> Result { let input_partition_count = self.input_partition_count(); let real_partition = partition % input_partition_count; - let consumer_id = partition / input_partition_count; + let consumer_task = partition / input_partition_count; let input = Arc::clone(&self.input); @@ -205,6 +205,7 @@ impl ExecutionPlan for BroadcastExec { let reservation = mem_consumer.clone_with_new_id().register(&pool); reservation.grow(record_batch.get_array_memory_size()); if !queue.push(Ok((record_batch, Arc::new(reservation)))) { + // If there are no remaining readers, short-circuit. break; } } @@ -219,7 +220,7 @@ impl ExecutionPlan for BroadcastExec { Ok::<_, Arc>((readers, Arc::new(task))) }); let (consumer, task) = match queue_or_err { - Ok((readers, task)) => (readers.claim(consumer_id)?, Arc::clone(task)), + Ok((readers, task)) => (readers.claim(consumer_task)?, Arc::clone(task)), Err(err) => return Err(DataFusionError::Shared(Arc::clone(err))), }; Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -240,6 +241,49 @@ impl ExecutionPlan for BroadcastExec { } } +/// Represents the queue for a single real partition, which multiple virtual partitions may share. +/// Assume we have 4 consumer tasks each represented as a reader, this could create a situation as +/// such: +/// +/// ```text +/// +/// base_sequence tail_sequence +/// │ │ +/// ▼ ▼ +/// ┌ ─ ─┌────┬────┬────┐ ┌────┬────┐ +/// entries: e0 │ e1 │ e2 │ e3 │ ... │eN-1│ eN │ +/// └ ─ ─└────┴────┴────┘ └────┴────┘ +/// ▲ ▲ +/// ┌──────┘ ┌────────────────┴────┐ +/// │ │ │ +/// ┌──────────┬──────────┬──────────┬──────────┐ +/// readers: │ r0: │ r1: │ r2: │ r3: │ +/// │Reading(2)│Reading(N)│ Released │Reading(N)│ +/// └──────────┴──────────┴──────────┴──────────┘ +/// ``` +/// +/// The `base_sequence` represents the first retained entry while `tail_sequence` indicates the next +/// append position. In this example, entry `e0` is shown using a dashed line because it has already +/// been evicted and `base_sequence` now points at `e1`. Every retaining reader has advanced past +/// `e1`, so `e1` is reclaimable but remains temporarily retained until an append reaches the next +/// `RECLAIM_INTERVAL` boundary. At that boundary, the queue removes entries before the smallest +/// `next_sequence` (the most lagging retaining reader, `r0` here) and advances `base_sequence`: +/// +/// ```text +/// base_sequence tail_sequence +/// │ │ +/// ▼ ▼ +/// ┌ ─ ─┌ ─ ─┌────┬────┐ ┌────┬────┬────┐ +/// entries: e0 │ e1 │ e2 │ e3 │ ... │eN-1│ eN │eN+1│ +/// └ ─ ─└ ─ ─└────┴────┘ └────┴────┴────┘ +/// ▲ ▲ ▲ +/// ┌──────┘ ┌────────────────┘ │ +/// │ │ │ +/// ┌──────────┬──────────┬──────────┬────────────┐ +/// readers: │ r0: │ r1: │ r2: │ r3: │ +/// │Reading(2)│Reading(N)│ Released │Reading(N+1)│ +/// └──────────┴──────────┴──────────┴────────────┘ +/// ``` #[derive(Debug)] struct QueueState { entries: VecDeque, @@ -257,6 +301,77 @@ enum ReaderSlot { Released, } +/// Shared state and signals for one real input partition. +/// +/// The producer task owns the input stream and queue handle while every consumer stream shares the +/// same queue state but has its own reader cursor and notification receiver. There are two flows +/// this is responsible for: queue updates and cancellation. +/// +/// ## Queue Push Flow: +/// +/// ```text +/// ┌───────────────────reads──────────────────┐ +/// │ ┌─────────(each acquire lock)─────────┐ │ +/// │ │ ┌─────────────────────────────────┐ │ │ +/// │ │ │ │ │ │ +/// │ │ │ ┌────────────────┐ │ │ │ +/// ┌─────────────────────────┐ │ │ │ │ │ │ │ │ +/// │BroadcastShared │ │ │ │ ┌────▶│ Consumer 0 │──┘ │ │ +/// │ ┌─────────────────────┐◀┼────────────────────┘ │ │ │ │ │ │ │ +/// ┌───────┼▶│ Mutex(QueueState) │◀┼───────────────────────┘ │ │ └────────────────┘ │ │ +/// push │ └─────────────────────┘◀┼─────────────────────────┘ │ │ │ +/// (acquires lock)│ │notify │ │ │ │ +/// ┌────────────────┐ │ │ ▼ │ │ ┌────────────────┐ │ │ +/// │ │ │ │ ┌─────────────────────┐ │ notify ┌───────────────┐ │ │ │ │ │ +/// │ Producer │────────┘ │ │ Sender │ ├──signals──▶│ Watch Channel │─notify──▶│ Consumer 1 │────┘ │ +/// │ │ │ └─────────────────────┘ │ └───────────────┘ signal │ │ │ +/// └────────────────┘ │ ┌─────────────────────┐ │ │ └────────────────┘ │ +/// │ │ CancellationToken │ │ │ │ +/// │ └─────────────────────┘ │ │ ... │ +/// └─────────────────────────┘ │ │ +/// │ ┌────────────────┐ │ +/// │ │ │ │ +/// └────▶│ Consumer N │──────┘ +/// │ │ +/// └────────────────┘ +/// ``` +/// +/// A consumer stream has access to the shared `Arc`, but only `queue_state` is locked so its +/// `poll_next` holds that mutex while it reads an entry. The producer follows the same rule, +/// it appends under the mutex, then sends the notification after unlocking. +/// +/// ## Cancellation flow: +/// +/// ```text +/// ┌────────────────┐ +/// │ │───────────────────────┐ ┌─────────────────────────┐ +/// │ Consumer 0 │─────────────┐ │ │BroadcastShared │ ┌────────────close ───────┐ +/// │ │◀───────┐ │ │ release reader │ ┌─────────────────────┐ │ │ (acquires lock) │ +/// └────────────────┘ │ │ └────(acquires lock)───┼▶│ Mutex(QueueState) │◀┼───────┘ │ +/// │ │ │ └─────────────────────┘ │ ┌────────────────┐ │ +/// │ │ │ │ notify on │ │ │ │ +/// ┌────────────────┐ │ │ │ ▼ close │ ┌──────▶│ Producer │──┘ +/// │ │ │ │ ┌───────────────┐ notify │ ┌─────────────────────┐ │ │ │ │ +/// │ Consumer 1 │◀────notify──┼───│ Watch Channel │◀──signals──┼─│ Sender │ │ cancel └────────────────┘ +/// │ │ signal │ └───────────────┘ │ └─────────────────────┘ │ signal +/// └────────────────┘ │ │ │ ┌─────────────────────┐ │ │ +/// │ └────────────cancel──────────────┼▶│ CancellationToken │─┼─────┘ +/// ... │ │ └─────────────────────┘ │ +/// │ └─────────────────────────┘ +/// ┌────────────────┐ │ +/// │ │ │ +/// │ Consumer N │◀───────┘ +/// │ │ +/// └────────────────┘ +/// ``` +/// +/// In this case, the consumer initiates the action by mutating the queue state to release itself. +/// Only in the case that the last reader has dropped the consumer will also set the +/// `CancellationToken` to tell the producer to close the queue. +/// +/// Also, a consumer stream may outlive the `BroadcastExec`. In this case, the plan's `Drop` releases +/// only still `Pending` readers while active `BroadcastConsumer`s keep the shared state alive, and +/// keeps the producer task alive until those streams finish or are dropped. #[derive(Debug)] struct BroadcastShared { queue_state: Mutex>, @@ -299,6 +414,10 @@ impl BroadcastQueue { } } + /// Appends a value to the entry queue and increments `tail_sequence`. Every `RECLAIM_INTERVAL` + /// calls, this checks for reclaimable entries in the queue. + /// + /// This method will not append the value and returns `false` if no retaining readers remain. fn push(&self, value: T) -> bool { { let mut queue_state = self.shared.queue_state.lock().unwrap(); @@ -310,7 +429,7 @@ impl BroadcastQueue { queue_state.entries.push_back(value); queue_state.tail_sequence += 1; if queue_state.tail_sequence.is_multiple_of(RECLAIM_INTERVAL) { - Self::reclaim_locked(&mut queue_state); + Self::reclaim_processed_entries(&mut queue_state); } } @@ -318,7 +437,9 @@ impl BroadcastQueue { true } - fn reclaim_locked(queue_state: &mut QueueState) { + /// Frees all entries in the queue that have been processed and updates `base_sequence` to point + /// at the first non-freeable position. + fn reclaim_processed_entries(queue_state: &mut QueueState) { let minimum_sequence = queue_state .readers .iter() @@ -351,25 +472,32 @@ impl Drop for BroadcastQueue { } impl BroadcastReaders { - fn claim(&self, consumer_id: usize) -> Result> { + /// Creates a new `BroadcastConsumer` for a given consumer task and claims its reader slot, + /// starting at sequence zero. + /// + /// Returns an error if the consumer task is out of range or if the same reader is claimed more + /// than once. + fn claim(&self, consumer_task: usize) -> Result> { let rx = self.shared.notify.subscribe(); let mut state = self.shared.queue_state.lock().unwrap(); - let Some(reader) = state.readers.get_mut(consumer_id) else { - return internal_err!("broadcast consumer {consumer_id} is out of range"); + let Some(reader) = state.readers.get_mut(consumer_task) else { + return internal_err!("broadcast consumer {consumer_task} is out of range"); }; match reader { ReaderSlot::Pending => *reader = ReaderSlot::Reading { next_sequence: 0 }, ReaderSlot::Reading { .. } | ReaderSlot::Released => { - return internal_err!("broadcast consumer {consumer_id} cannot execute twice"); + return internal_err!("broadcast consumer {consumer_task} cannot execute twice"); } } Ok(BroadcastConsumer { - consumer_id, + consumer_id: consumer_task, shared: Arc::clone(&self.shared), notify: WatchStream::new(rx), }) } + /// Sets all `ReaderSlot::Pending` readers to `ReaderSlot::Released`. This also cleans up newly + /// freeable entries and cancels the producer if all readers are released. fn release_pending(&self) { let no_readers_remain = { let mut state = self.shared.queue_state.lock().unwrap(); @@ -381,7 +509,7 @@ impl BroadcastReaders { } } state.retaining_readers -= released; - BroadcastQueue::::reclaim_locked(&mut state); + BroadcastQueue::::reclaim_processed_entries(&mut state); state.retaining_readers == 0 }; if no_readers_remain { @@ -400,6 +528,10 @@ struct BroadcastConsumer { impl Stream for BroadcastConsumer { type Item = T; + /// Poll the next value from the stream reading from the shared entry queue. + /// + /// TODO: Profile lock contention and inspect if a lock free implementation has better + /// performance. fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { let (value, closed) = { @@ -457,7 +589,7 @@ impl BroadcastShared { } state.readers[consumer_id] = ReaderSlot::Released; state.retaining_readers -= 1; - BroadcastQueue::::reclaim_locked(&mut state); + BroadcastQueue::::reclaim_processed_entries(&mut state); state.retaining_readers == 0 }; if no_readers_remain { From 73dcdaae1d3565a9627d3a98cf632f44a3c6197c Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 15 Sep 2026 20:14:01 -0400 Subject: [PATCH 08/10] put the drop outside the lock --- src/execution_plans/broadcast.rs | 40 ++++++++++++++++++-------------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 8ee0a3b15..4fbff03f8 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -419,7 +419,7 @@ impl BroadcastQueue { /// /// This method will not append the value and returns `false` if no retaining readers remain. fn push(&self, value: T) -> bool { - { + let reclaimed = { let mut queue_state = self.shared.queue_state.lock().unwrap(); if queue_state.retaining_readers == 0 { @@ -429,9 +429,13 @@ impl BroadcastQueue { queue_state.entries.push_back(value); queue_state.tail_sequence += 1; if queue_state.tail_sequence.is_multiple_of(RECLAIM_INTERVAL) { - Self::reclaim_processed_entries(&mut queue_state); + Self::reclaim_processed_entries(&mut queue_state) + } else { + Vec::new() } - } + }; + + drop(reclaimed); self.shared.notify.send_replace(()); true @@ -439,7 +443,7 @@ impl BroadcastQueue { /// Frees all entries in the queue that have been processed and updates `base_sequence` to point /// at the first non-freeable position. - fn reclaim_processed_entries(queue_state: &mut QueueState) { + fn reclaim_processed_entries(queue_state: &mut QueueState) -> Vec { let minimum_sequence = queue_state .readers .iter() @@ -451,13 +455,9 @@ impl BroadcastQueue { .min() .unwrap_or(queue_state.tail_sequence); - while queue_state.base_sequence < minimum_sequence { - queue_state - .entries - .pop_front() - .expect("broadcast sequence bounds were inconsistent"); - queue_state.base_sequence += 1; - } + let reclaim_count = minimum_sequence - queue_state.base_sequence; + queue_state.base_sequence = minimum_sequence; + queue_state.entries.drain(..reclaim_count).collect() } } @@ -499,7 +499,7 @@ impl BroadcastReaders { /// Sets all `ReaderSlot::Pending` readers to `ReaderSlot::Released`. This also cleans up newly /// freeable entries and cancels the producer if all readers are released. fn release_pending(&self) { - let no_readers_remain = { + let (no_readers_remain, reclaimed) = { let mut state = self.shared.queue_state.lock().unwrap(); let mut released = 0; for reader in &mut state.readers { @@ -509,9 +509,12 @@ impl BroadcastReaders { } } state.retaining_readers -= released; - BroadcastQueue::::reclaim_processed_entries(&mut state); - state.retaining_readers == 0 + let reclaimed = BroadcastQueue::::reclaim_processed_entries(&mut state); + (state.retaining_readers == 0, reclaimed) }; + + drop(reclaimed); + if no_readers_remain { self.shared.cancel.cancel(); } @@ -582,16 +585,19 @@ impl Stream for BroadcastConsumer { impl BroadcastShared { fn release(&self, consumer_id: usize) { - let no_readers_remain = { + let (no_readers_remain, reclaimed) = { let mut state = self.queue_state.lock().unwrap(); if matches!(state.readers[consumer_id], ReaderSlot::Released) { return; } state.readers[consumer_id] = ReaderSlot::Released; state.retaining_readers -= 1; - BroadcastQueue::::reclaim_processed_entries(&mut state); - state.retaining_readers == 0 + let reclaimed = BroadcastQueue::::reclaim_processed_entries(&mut state); + (state.retaining_readers == 0, reclaimed) }; + + drop(reclaimed); + if no_readers_remain { self.cancel.cancel(); } From 6fed81b12f91d09e972159bfb6a75409473a392f Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Thu, 24 Sep 2026 20:18:06 -0400 Subject: [PATCH 09/10] clean up bench --- Cargo.lock | 1 - benchmarks/Cargo.toml | 1 - .../benches/broadcast_cache_scenarios.rs | 99 ++----------------- 3 files changed, 8 insertions(+), 93 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8fa95485e..4c499b697 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2431,7 +2431,6 @@ dependencies = [ "serde_json", "sketches-ddsketch", "structopt", - "sysinfo", "tempfile", "tokio", "tonic", diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 7b19f0ea9..00d7ac359 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -34,7 +34,6 @@ tempfile = "3" [dev-dependencies] criterion = "0.5" -sysinfo = "0.30" [build-dependencies] built = { version = "0.8", features = ["git2", "chrono"] } diff --git a/benchmarks/benches/broadcast_cache_scenarios.rs b/benchmarks/benches/broadcast_cache_scenarios.rs index 33dc1aa34..dcf89fcab 100644 --- a/benchmarks/benches/broadcast_cache_scenarios.rs +++ b/benchmarks/benches/broadcast_cache_scenarios.rs @@ -17,13 +17,8 @@ use datafusion::physical_plan::{ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_distributed::BroadcastExec; use futures::{StreamExt, stream}; -use std::sync::{ - Arc, - atomic::{AtomicBool, AtomicU64, Ordering}, -}; -use std::thread; +use std::sync::Arc; use std::time::{Duration, Instant}; -use sysinfo::{System, get_current_pid}; use tokio::runtime::Builder as RuntimeBuilder; #[derive(Clone, Copy)] @@ -49,45 +44,6 @@ struct Scenario { consumers: Vec, } -struct PeakRssSampler { - stop: Arc, - peak_bytes: Arc, - handle: Option>, -} - -impl PeakRssSampler { - fn start(interval: Duration) -> Self { - let stop = Arc::new(AtomicBool::new(false)); - let peak_bytes = Arc::new(AtomicU64::new(0)); - let stop_clone = Arc::clone(&stop); - let peak_clone = Arc::clone(&peak_bytes); - let handle = thread::spawn(move || { - let mut sys = System::new(); - let pid = get_current_pid().expect("pid"); - while !stop_clone.load(Ordering::Relaxed) { - sys.refresh_process(pid); - if let Some(proc) = sys.process(pid) { - peak_clone.fetch_max(proc.memory(), Ordering::Relaxed); - } - thread::sleep(interval); - } - }); - Self { - stop, - peak_bytes, - handle: Some(handle), - } - } - - fn stop(mut self) -> u64 { - self.stop.store(true, Ordering::Relaxed); - if let Some(handle) = self.handle.take() { - let _ = handle.join(); - } - self.peak_bytes.load(Ordering::Relaxed) - } -} - #[derive(Debug)] struct SyntheticExec { schema: SchemaRef, @@ -224,15 +180,13 @@ async fn run_scenario( scenario: &Scenario, input: Arc, task_ctx: Arc, - sample_rss: bool, recording_pool: Option>, -) -> Result<(Duration, u64, Option<(usize, usize)>)> { +) -> Result<(Duration, Option<(usize, usize)>)> { let broadcast = Arc::new(BroadcastExec::new( Arc::clone(&input), scenario.consumer_tasks, )); - let sampler = sample_rss.then(|| PeakRssSampler::start(Duration::from_millis(25))); let start = Instant::now(); let mut join_set = tokio::task::JoinSet::new(); @@ -251,9 +205,8 @@ async fn run_scenario( } let elapsed = start.elapsed(); - let peak_rss_bytes = sampler.map(|s| s.stop()).unwrap_or(0); let memory = recording_pool.map(|pool| (pool.peak_reserved(), pool.reserved())); - Ok((elapsed, peak_rss_bytes, memory)) + Ok((elapsed, memory)) } fn memory_context() -> Result<(Arc, Arc)> { @@ -388,16 +341,6 @@ fn scenario_matrix() -> Vec { scenarios } -fn rss_enabled() -> bool { - match std::env::var("BROADCAST_BENCH_RSS") { - Ok(val) => { - let val = val.to_ascii_lowercase(); - val == "1" || val == "true" || val == "yes" - } - Err(_) => false, - } -} - fn memory_enabled() -> bool { match std::env::var("BROADCAST_BENCH_MEMORY") { Ok(val) => { @@ -426,7 +369,6 @@ fn bench_broadcast_cache(c: &mut Criterion) { let mut group = c.benchmark_group("broadcast_cache_scenarios"); group.sample_size(10); - let sample_rss = rss_enabled(); let sample_memory = memory_enabled(); let task_ctx = SessionContext::new().task_ctx(); @@ -437,19 +379,17 @@ fn bench_broadcast_cache(c: &mut Criterion) { false, )])); let input = prebuilt_input(&scenario, &schema); - let mut rss_peaks = Vec::new(); let mut memory_peaks = Vec::new(); let mut memory_residuals = Vec::new(); group.bench_function(BenchmarkId::new("scenario", scenario.name), |b| { b.iter_custom(|iters| { let mut total = Duration::ZERO; for _ in 0..iters { - let (elapsed, _, _) = rt + let (elapsed, _) = rt .block_on(run_scenario( &scenario, Arc::clone(&input), Arc::clone(&task_ctx), - false, None, )) .expect("scenario"); @@ -459,26 +399,17 @@ fn bench_broadcast_cache(c: &mut Criterion) { }); }); - if sample_rss || sample_memory { + if sample_memory { for _ in 0..DIAGNOSTIC_RUNS { - let (diagnostic_task_ctx, recording_pool) = if sample_memory { - let (task_ctx, pool) = memory_context().expect("memory context"); - (task_ctx, Some(pool)) - } else { - (Arc::clone(&task_ctx), None) - }; - let (_, peak_rss_bytes, memory) = rt + let (diagnostic_task_ctx, pool) = memory_context().expect("memory context"); + let (_, memory) = rt .block_on(run_scenario( &scenario, paced_input(&scenario, &schema), diagnostic_task_ctx, - sample_rss, - recording_pool, + Some(pool), )) .expect("diagnostic scenario"); - if sample_rss { - rss_peaks.push(peak_rss_bytes); - } if let Some((peak_reserved, residual_reserved)) = memory { memory_peaks.push(peak_reserved); memory_residuals.push(residual_reserved); @@ -486,20 +417,6 @@ fn bench_broadcast_cache(c: &mut Criterion) { } } - if sample_rss && !rss_peaks.is_empty() { - rss_peaks.sort_unstable(); - let min = rss_peaks[0]; - let max = rss_peaks[rss_peaks.len() - 1]; - let median = rss_peaks[rss_peaks.len() / 2]; - eprintln!( - "scenario={} peak_rss_bytes[min/median/max]={}/{}/{} runs={}", - scenario.name, - min, - median, - max, - rss_peaks.len() - ); - } if sample_memory && !memory_peaks.is_empty() { memory_peaks.sort_unstable(); memory_residuals.sort_unstable(); From 89d3022cfd3de3f593d7c59a8d213447b1d9ca45 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 25 Sep 2026 14:33:05 -0400 Subject: [PATCH 10/10] clean up tests --- src/execution_plans/broadcast.rs | 23 ++++------------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/src/execution_plans/broadcast.rs b/src/execution_plans/broadcast.rs index 4fbff03f8..c5f278778 100644 --- a/src/execution_plans/broadcast.rs +++ b/src/execution_plans/broadcast.rs @@ -532,9 +532,6 @@ impl Stream for BroadcastConsumer { type Item = T; /// Poll the next value from the stream reading from the shared entry queue. - /// - /// TODO: Profile lock contention and inspect if a lock free implementation has better - /// performance. fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { let (value, closed) = { @@ -686,15 +683,18 @@ mod tests { queue.push(20); assert_eq!(consumer0.next().await, Some(10)); drop(consumer0); + assert!(!queue.shared.cancel.is_cancelled()); // The dropped consumer must no longer pin the unread suffix, including entries produced - // after cancellation. + // after the first reader is released. queue.push(30); assert_eq!(consumer1.next().await, Some(10)); assert_eq!(consumer1.next().await, Some(20)); assert_eq!(consumer1.next().await, Some(30)); drop(consumer1); + assert!(queue.shared.cancel.is_cancelled()); assert_eq!(buffered_len(&queue), 0); + assert!(!queue.push(20)); } #[tokio::test] @@ -713,21 +713,6 @@ mod tests { assert_eq!(consumer.next().await, None); } - #[tokio::test] - async fn broadcast_queue_cancels_when_all_readers_drop() { - let queue = BroadcastQueue::new(2); - let reader0 = queue.readers().claim(0).expect("consumer 0 registration"); - let reader1 = queue.readers().claim(1).expect("consumer 1 registration"); - - assert!(queue.push(10)); - drop(reader0); - assert!(!queue.shared.cancel.is_cancelled()); - drop(reader1); - assert!(queue.shared.cancel.is_cancelled()); - assert_eq!(buffered_len(&queue), 0); - assert!(!queue.push(20)); - } - #[tokio::test] async fn broadcast_exec_releases_pending_readers_when_plan_drops() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));