From b0752560ff47f36ce8896574829082e51f79088c Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 16 Sep 2026 16:16:41 +0000 Subject: [PATCH] fix: enforce fair memory limits across sibling reservations --- ...spilling_fuzz_in_memory_constrained_env.rs | 61 +++++- datafusion/execution/src/memory_pool/pool.rs | 200 ++++++++++++++++-- 2 files changed, 240 insertions(+), 21 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index 761c492d8dde3..aa9ca0efd06fc 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -26,7 +26,7 @@ use arrow::array::UInt64Array; use arrow::row::{RowConverter, SortField}; use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch}; use arrow_schema::{DataType, Field, Schema}; -use datafusion::common::Result; +use datafusion::common::{DataFusionError, Result}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::PhysicalSortExpr; @@ -624,7 +624,7 @@ async fn run_sort_test_with_limited_memory( let assert_output_batch_size = args.assert_all_output_batches_roughly_match_batch_size_conf; - let metrics = run_test(args, sort_exec, result).await?; + let metrics = run_test(args, sort_exec, result, false).await?; assert_baseline_metrics_for_non_empty_output( &metrics, @@ -764,15 +764,19 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ #[tokio::test] async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_sizes_of_record_batch_and_take_all_memory() -> Result<()> { + // Permanent non-spillable pressure must fail within the pool limit instead + // of letting replay overcommit memory after it has produced output. let record_batch_size = 8192; let pool_size = 2 * MB as usize; + let memory_pool = Arc::new(PeakRecordingPool::new(Arc::new(FairSpillPool::new( + pool_size, + )))); let task_ctx = { - let memory_pool = Arc::new(FairSpillPool::new(pool_size)); TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(record_batch_size)) .with_runtime(Arc::new( RuntimeEnvBuilder::new() - .with_memory_pool(memory_pool) + .with_memory_pool(Arc::clone(&memory_pool) as Arc) .build()?, )) }; @@ -793,6 +797,8 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ }) .await?; + assert!(memory_pool.peak_reserved() <= pool_size); + assert_eq!(memory_pool.reserved(), 0); Ok(()) } @@ -910,18 +916,29 @@ async fn run_test_aggregate_with_high_cardinality( let result = aggregate_final.execute(0, Arc::clone(&args.task_ctx))?; - run_test(args, aggregate_final, result).await + // A non-spilling competitor that permanently consumes every free byte can + // prevent later replay growth. Require a bounded error, not overcommit. + let expect_memory_exhaustion = matches!( + &args.memory_behavior, + MemoryBehavior::TakeAllMemoryAtTheBeginning + ); + run_test(args, aggregate_final, result, expect_memory_exhaustion).await } async fn run_test( args: RunTestWithLimitedMemoryArgs, plan: Arc, result_stream: SendableRecordBatchStream, + expect_memory_exhaustion: bool, ) -> Result { let number_of_record_batches = args.number_of_record_batches; - consume_stream_and_simulate_other_running_memory_consumers(args, result_stream) - .await?; + consume_stream_and_simulate_other_running_memory_consumers( + args, + result_stream, + expect_memory_exhaustion, + ) + .await?; let metrics = plan.metrics().expect("must have metrics"); let spill_count = assert_spill_count_metric(true, plan); @@ -938,6 +955,7 @@ async fn run_test( async fn consume_stream_and_simulate_other_running_memory_consumers( args: RunTestWithLimitedMemoryArgs, mut result_stream: SendableRecordBatchStream, + expect_memory_exhaustion: bool, ) -> Result<()> { let mut number_of_rows = 0; let record_batch_size = args.task_ctx.session_config().batch_size() as u64; @@ -950,6 +968,25 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( let mut memory_took = false; while let Some(batch) = result_stream.next().await { + let batch = match batch { + Ok(batch) => batch, + Err(DataFusionError::ResourcesExhausted(_)) if expect_memory_exhaustion => { + // Do not accept an early replay error before the mock actually + // takes memory away from an operator that has produced rows. + assert!(number_of_rows > 0); + assert!(memory_took && memory_reservation.size() > 0); + assert!(memory_pool.reserved() <= args.pool_size); + drop(result_stream); + drop(memory_reservation); + assert_eq!(memory_pool.reserved(), 0); + let progress = + args.task_ctx.runtime_env().disk_manager.spilling_progress(); + assert_eq!(progress.current_bytes, 0); + assert_eq!(progress.active_files_count, 0); + return Ok(()); + } + Err(error) => return Err(error), + }; match args.memory_behavior { MemoryBehavior::AsIs => { // Do nothing @@ -958,6 +995,11 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( if !memory_took { memory_took = true; grow_memory_as_much_as_possible(10, &mut memory_reservation)?; + if expect_memory_exhaustion { + assert!(memory_reservation.size() > 0); + assert!(memory_pool.reserved() <= args.pool_size); + assert!(args.pool_size - memory_pool.reserved() < 10); + } } } MemoryBehavior::TakeAllMemoryAndReleaseEveryNthBatch(n) => { @@ -974,12 +1016,15 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( } } - let batch = batch?; number_of_rows += batch.num_rows(); index += 1; } + assert!( + !expect_memory_exhaustion, + "expected memory exhaustion after external pressure" + ); assert_eq!( number_of_rows, args.number_of_record_batches * record_batch_size as usize diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs index d854cbd627cec..8f92847920535 100644 --- a/datafusion/execution/src/memory_pool/pool.rs +++ b/datafusion/execution/src/memory_pool/pool.rs @@ -144,7 +144,7 @@ impl Display for GreedyMemoryPool { /// A [`MemoryPool`] that prevents spillable reservations from using more than /// an even fraction of the available memory sans any unspillable reservations -/// (i.e. `(pool_size - unspillable_memory) / num_spillable_reservations`) +/// (i.e. `(pool_size - unspillable_memory) / num_spillable_consumers`) /// /// This pool works best when you know beforehand the query has /// multiple spillable operators that will likely all need to @@ -163,6 +163,11 @@ impl Display for GreedyMemoryPool { /// └───────────────────────z──────────────────────z───────────────┘ /// ``` /// +/// Reservations created with [`MemoryReservation::new_empty`], +/// [`MemoryReservation::split`], or [`MemoryReservation::take`] share their +/// consumer's allowance. Registering a new consumer does not revoke existing +/// reservations, but further fallible growth remains limited by total pool capacity. +/// /// Unspillable memory is allocated in a first-come, first-serve fashion #[derive(Debug)] pub struct FairSpillPool { @@ -174,12 +179,15 @@ pub struct FairSpillPool { #[derive(Debug)] struct FairSpillPoolState { - /// The number of consumers that can spill - num_spill: usize, - /// The total amount of memory reserved that can be spilled spillable: usize, + /// Total reservation across every sibling of each spillable consumer. + /// + /// `MemoryReservation::new_empty`, `split`, and `take` share a consumer + /// registration while maintaining separate reservation-size counters. + spillable_by_consumer: HashMap, + /// The total amount of memory reserved by consumers that cannot spill unspillable: usize, } @@ -191,8 +199,8 @@ impl FairSpillPool { Self { pool_size, state: Mutex::new(FairSpillPoolState { - num_spill: 0, spillable: 0, + spillable_by_consumer: HashMap::default(), unspillable: 0, }), } @@ -206,21 +214,30 @@ impl MemoryPool for FairSpillPool { fn register(&self, consumer: &MemoryConsumer) { if consumer.can_spill { - self.state.lock().num_spill += 1; + let mut state = self.state.lock(); + state.spillable_by_consumer.insert(consumer.id(), 0); } } fn unregister(&self, consumer: &MemoryConsumer) { if consumer.can_spill { let mut state = self.state.lock(); - state.num_spill = state.num_spill.checked_sub(1).unwrap(); + let released = state.spillable_by_consumer.remove(&consumer.id()); + debug_assert_eq!(released, Some(0)); } } fn grow(&self, reservation: &MemoryReservation, additional: usize) { let mut state = self.state.lock(); match reservation.registration.consumer.can_spill { - true => state.spillable += additional, + true => { + state.spillable += additional; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") += + additional; + } false => state.unspillable += additional, } } @@ -228,7 +245,13 @@ impl MemoryPool for FairSpillPool { fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { let mut state = self.state.lock(); match reservation.registration.consumer.can_spill { - true => state.spillable -= shrink, + true => { + state.spillable -= shrink; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") -= shrink; + } false => state.unspillable -= shrink, } } @@ -243,10 +266,18 @@ impl MemoryPool for FairSpillPool { // No spiller may use more than their fraction of the memory available let available = spill_available - .checked_div(state.num_spill) + .checked_div(state.spillable_by_consumer.len()) .unwrap_or(spill_available); - - if reservation.size() + additional > available { + let consumer_used = state + .spillable_by_consumer + .get(&reservation.consumer().id()) + .copied() + .expect("spillable memory consumer must remain registered"); + + if consumer_used + .checked_add(additional) + .is_none_or(|requested| requested > available) + { return Err(insufficient_capacity_err( reservation, additional, @@ -254,12 +285,28 @@ impl MemoryPool for FairSpillPool { self, )); } + let remaining = self + .pool_size + .saturating_sub(state.unspillable.saturating_add(state.spillable)); + if additional > remaining { + return Err(insufficient_capacity_err( + reservation, + additional, + remaining, + self, + )); + } state.spillable += additional; + *state + .spillable_by_consumer + .get_mut(&reservation.consumer().id()) + .expect("spillable memory consumer must remain registered") += + additional; } false => { let available = self .pool_size - .saturating_sub(state.unspillable + state.spillable); + .saturating_sub(state.unspillable.saturating_add(state.spillable)); if available < additional { return Err(insufficient_capacity_err( @@ -694,6 +741,133 @@ mod tests { assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 30.0 B for s4 with 0.0 B already allocated for this reservation - 20.0 B remain available for the total memory pool: fair(pool_size: 100.0 B)"); } + #[test] + fn test_fair_sibling_reservations_share_one_consumer_limit() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + let first_partition = parent.new_empty(); + let second_partition = parent.new_empty(); + + first_partition.try_grow(60).unwrap(); + second_partition.try_grow(40).unwrap(); + assert_eq!(pool.reserved(), 100); + assert!(parent.try_grow(1).is_err()); + assert!(second_partition.try_grow(1).is_err()); + assert_eq!(pool.reserved(), 100); + + drop(first_partition); + second_partition.try_grow(60).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(second_partition); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_siblings_respect_consumer_shares_and_global_capacity() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let unspillable = MemoryConsumer::new("fixed").register(&pool); + unspillable.try_grow(20).unwrap(); + + let first = MemoryConsumer::new("same name") + .with_can_spill(true) + .register(&pool); + let second = MemoryConsumer::new("same name") + .with_can_spill(true) + .register(&pool); + let sibling = first.new_empty(); + + first.try_grow(25).unwrap(); + sibling.try_grow(15).unwrap(); + assert!(sibling.try_grow(1).is_err()); + second.try_grow(40).unwrap(); + assert_eq!(pool.reserved(), 100); + + let split = first.split(10); + assert_eq!(pool.reserved(), 100); + assert!(split.try_grow(1).is_err()); + drop(split); + second.try_grow(1).unwrap_err(); + sibling.try_grow(10).unwrap(); + assert_eq!(pool.reserved(), 100); + } + + #[test] + fn test_fair_take_retains_usage_until_last_sibling_drops() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let mut parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + let other = MemoryConsumer::new("other") + .with_can_spill(true) + .register(&pool); + parent.try_grow(50).unwrap(); + let taken = parent.take(); + assert_eq!(parent.size(), 0); + assert_eq!(taken.size(), 50); + assert!(parent.try_grow(1).is_err()); + assert!(taken.try_grow(1).is_err()); + drop(parent); + assert_eq!(pool.reserved(), 50); + taken.shrink(10); + taken.try_grow(10).unwrap(); + assert!(other.try_grow(51).is_err()); + drop(taken); + other.try_grow(100).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(other); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_new_consumer_respects_existing_global_usage() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + first.try_grow(100).unwrap(); + let second = MemoryConsumer::new("second") + .with_can_spill(true) + .register(&pool); + assert!(second.try_grow(1).is_err()); + assert_eq!(pool.reserved(), 100); + first.shrink(50); + second.try_grow(50).unwrap(); + assert_eq!(pool.reserved(), 100); + } + + #[test] + fn test_fair_infallible_growth_is_charged_to_all_siblings() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + let sibling = first.new_empty(); + // Infallible growth remains permitted, including beyond the configured capacity. + first.grow(110); + assert_eq!(pool.reserved(), 110); + assert!(sibling.try_grow(1).is_err()); + first.shrink(20); + sibling.try_grow(10).unwrap(); + assert_eq!(pool.reserved(), 100); + drop(first); + drop(sibling); + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn test_fair_oversized_sibling_growth_does_not_overflow() { + let pool: Arc = Arc::new(FairSpillPool::new(100)); + let parent = MemoryConsumer::new("spilling operator") + .with_can_spill(true) + .register(&pool); + parent.try_grow(1).unwrap(); + let sibling = parent.new_empty(); + assert!(sibling.try_grow(usize::MAX).is_err()); + assert_eq!(pool.reserved(), 1); + } + #[test] fn test_tracked_consumers_pool() { let setting = make_settings();