From 15a8ebc62590e7c919dcf2a6b78e6998ba6a2b3a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 2 Sep 2026 20:36:54 +0000 Subject: [PATCH 1/3] test: reproduce the RepartitionExec spill pool deadlock with a SQL query Add a test that runs the query from issue #24883 under a 4 MB memory limit, with two partitions and 64-row batches. The query is a grouped COUNT(DISTINCT) on a Utf8View column. It sends `RepartitionExec` into its spill path. On the current code, the first attempt does not complete. The test stops the attempt after 20 seconds and fails. The test does 12 attempts. One attempt that completes is not sufficient for a pass. The test uses a memory limit that is too small for the query. This is necessary to make `RepartitionExec` spill. Once in a few thousand attempts, the greedy memory pool refuses an allocation of an aggregate, because the repartition reservations hold the pool. The test accepts a resources-exhausted error. That error is a correct result of the limit and not the deadlock. Ref: https://github.com/apache/datafusion/issues/24883 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9 --- .../memory_limit/repartition_mem_limit.rs | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/datafusion/core/tests/memory_limit/repartition_mem_limit.rs b/datafusion/core/tests/memory_limit/repartition_mem_limit.rs index 27bcd33926e99..15ec0ebfcf4eb 100644 --- a/datafusion/core/tests/memory_limit/repartition_mem_limit.rs +++ b/datafusion/core/tests/memory_limit/repartition_mem_limit.rs @@ -16,14 +16,18 @@ // under the License. use std::sync::Arc; +use std::time::Duration; use arrow::array::{ArrayRef, Int32Array, RecordBatch}; use datafusion::{ assert_batches_sorted_eq, + execution::session_state::SessionStateBuilder, prelude::{SessionConfig, SessionContext}, }; use datafusion_catalog::MemTable; use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_plan::{ExecutionPlanProperties, repartition::RepartitionExec}; use futures::TryStreamExt; @@ -117,3 +121,111 @@ async fn test_repartition_memory_limit() { ]; assert_batches_sorted_eq!(expected, &all_batches); } + +/// Regression test for . +/// +/// In non-preserve-order mode every input task of a `RepartitionExec` shares one +/// spill pool per output partition. When two input tasks spilled concurrently the +/// pool could end up with two open spill files while the reader only ever drained +/// the head file: it parked on a drained-but-unfinished head file even though the +/// batch it was waiting for had been written to the second file. Once every +/// distributor channel was non-empty the gate closed, both input tasks parked in +/// `send`, no sink was dropped, and nothing could wake anybody. +mod spill_pool_deadlock { + use super::*; + + /// Memory limit that puts `RepartitionExec` into its spilling path for this + /// query without failing the aggregation outright. + const MEMORY_LIMIT: usize = 4 * 1024 * 1024; + + /// Number of attempts. On an unfixed tree the first attempt deadlocks; the + /// budget is only there so a fix cannot pass by luck. + const ATTEMPTS: usize = 12; + + /// Generous per-attempt budget: a healthy run of this query takes well under + /// a second. + const PER_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(20); + + async fn run_once() -> datafusion::error::Result { + let pool: Arc = Arc::new(GreedyMemoryPool::new(MEMORY_LIMIT)); + let runtime = RuntimeEnvBuilder::new() + // Spilling must be possible: with the disk manager disabled the query + // fails with a resource error instead of deadlocking. + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .with_memory_pool(pool) + .build_arc()?; + + let config = SessionConfig::new() + // Two input partitions feeding one hash RepartitionExec is the + // smallest configuration with two concurrent writers per spill pool. + .with_target_partitions(2) + // Small batches so many small batches reach the repartition spill + // path while the spill file stays far below + // `max_spill_file_size_bytes` and so never rotates. + .with_batch_size(64); + + let state = SessionStateBuilder::new() + .with_config(config) + .with_runtime_env(runtime) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + + ctx.sql( + "create table trace_events as + select v % 13 as g, + case when v % 29 = 0 then null + else md5(cast(v % 337 as varchar)) end as trace_id + from generate_series(1, 40000) as t(v)", + ) + .await? + .collect() + .await?; + + ctx.sql( + "create view tv as + select g, arrow_cast(trace_id, 'Utf8View') as trace_id from trace_events", + ) + .await? + .collect() + .await?; + + let batches = ctx + .sql("select g, count(distinct trace_id) as n from tv group by g") + .await? + .collect() + .await?; + + Ok(batches.iter().map(|b| b.num_rows()).sum()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn repartition_spill_pool_does_not_deadlock() { + for attempt in 0..ATTEMPTS { + match tokio::time::timeout(PER_ATTEMPT_TIMEOUT, run_once()).await { + Err(elapsed) => panic!( + "attempt {attempt}: grouped COUNT(DISTINCT Utf8View) never completed \ + within {PER_ATTEMPT_TIMEOUT:?} ({elapsed}); RepartitionExec spill \ + pool deadlock" + ), + Ok(Ok(rows)) => { + assert_eq!(rows, 13, "attempt {attempt}: wrong row count") + } + // The budget is deliberately far too small for the query, and the + // greedy pool hands memory out first come first served, so once in + // a few thousand attempts an aggregate's allocation is refused + // while the repartition reservations hold the pool. That is a + // legitimate outcome of the limit, not the hang this test guards + // against. + Ok(Err(e)) + if matches!( + e.find_root(), + datafusion::error::DataFusionError::ResourcesExhausted(_) + ) => {} + Ok(Err(e)) => panic!("attempt {attempt}: query failed: {e}"), + } + } + } +} From f943030d2e486768f94a0219d30b31229cc857bf Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 2 Sep 2026 20:37:05 +0000 Subject: [PATCH 2/3] test: add a scenario fuzzer for the spill pool contract The spill pool had three bugs since its unbuffered reads: #20683, #23447 and #24883. In each bug, two writers were in `push_batch` at the same time. No unit test could cause that overlap. This commit adds a test double and two tests: - `IoHookFactory` is a `TempFileFactory` that runs a hook before each disk write. A writer writes while it holds the lock of its file. Thus a hook that blocks can hold a writer at the point where it has a file checked out. The factory can also delay the reads of a file. - `test_reader_does_not_wait_on_drained_file_while_another_has_data` holds writer 1 inside its first write while writer 2 pushes a batch. Then the reader must give both batches while both writers are alive. - `spill_pool_scenario_fuzz` runs random scenarios from a seed: one to four writers, file rotation after each batch, after some batches or never, and one to six phases. In each phase, a random set of writers pushes one to three batches each. A pauser holds each writer of the phase inside its first write, so the pushes overlap. Then the reader reads a random number of the available batches. Some phases drop a writer. Each scenario checks the contract that `RepartitionExec` depends on: - When no push is in progress, the reader can read each pushed batch without more writer activity. - The reader does not signal EOF while a writer is alive. - After the last writer is dropped, the reader gives the backlog and then EOF. - The reader gives exactly the pushed batches. With one writer, it gives them in push order. - Disk usage is zero after the reader is dropped. The two tests are red on this commit. The fuzzer also fails on these versions of the pool: - The code before #23522. Seed 4 shows the #23447 stall after the drop of all writers. Seeds 0 and 1 show the #24883 stall. All 16 pool tests of that time pass on that code. - A draft of the fix in the next commit that skips a file with a read in progress. Seed 2 shows a FIFO violation with one writer. The default of 50 seeds takes approximately three seconds. The extended CI job runs 1000 seeds. `DATAFUSION_SPILL_POOL_FUZZ_SEED` replays one seed. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9 --- .github/workflows/extended.yml | 3 + .../physical-plan/src/spill/spill_pool.rs | 505 +++++++++++++++++- 2 files changed, 501 insertions(+), 7 deletions(-) diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index a6e303e3d6ff4..ae7408ae771fa 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -88,6 +88,9 @@ jobs: - name: Run tests (excluding doctests) env: RUST_BACKTRACE: 1 + # Run more random scenarios of the spill pool fuzzer than the + # default suite does (about a minute). + DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS: 1000 run: | cargo test \ --profile ci \ diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 56d8d11201f46..43575388edea9 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -791,8 +791,18 @@ mod tests { use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::instant::Instant; use datafusion_common_runtime::{JoinSet, SpawnedTask}; + use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, + }; use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::{SpillFile, SpillWriter, TempFileFactory}; + use std::pin::Pin; + use std::sync::Barrier; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; fn create_test_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) @@ -919,8 +929,7 @@ mod tests { // Reader should pend since no batches were written let mut reader = reader; let result = - tokio::time::timeout(std::time::Duration::from_millis(100), reader.next()) - .await; + tokio::time::timeout(Duration::from_millis(100), reader.next()).await; assert!(result.is_err(), "Reader should timeout on empty writer"); @@ -1230,7 +1239,7 @@ mod tests { let batch = create_test_batch(i * 10, 10); writer.push_batch(&batch).unwrap(); // Small delay to simulate real concurrent work - tokio::time::sleep(std::time::Duration::from_millis(5)).await; + tokio::time::sleep(Duration::from_millis(5)).await; } }); @@ -1523,10 +1532,9 @@ mod tests { proceed_tx.send(()).unwrap(); // The reader should wait (Pending) for writer2's data, not EOF. - let batch2 = - tokio::time::timeout(std::time::Duration::from_secs(5), reader.next()) - .await - .expect("Reader timed out — should not hang"); + let batch2 = tokio::time::timeout(Duration::from_secs(5), reader.next()) + .await + .expect("Reader timed out — should not hang"); assert!( batch2.is_some(), @@ -1645,4 +1653,487 @@ mod tests { Ok(()) } + + type WriteHook = Arc; + type ReadDelay = Arc Duration + Send + Sync>; + + /// Test double for a spill file that runs `write_hook` before every disk + /// write or flush, and delays the first bytes of every read stream by + /// `read_delay()`. Writers write while holding their file's lock, so a hook + /// that blocks pauses a `push_batch` at exactly the point where it has a + /// file checked out for writing; a read delay makes an older file's bytes + /// arrive after a newer file's. Data still goes to real temporary files. + struct IoHookFactory { + inner: Arc, + write_hook: WriteHook, + read_delay: ReadDelay, + } + + impl TempFileFactory for IoHookFactory { + fn create_temp_file(&self, description: &str) -> Result> { + Ok(Arc::new(IoHookFile { + inner: self.inner.create_tmp_file(description)?, + write_hook: Arc::clone(&self.write_hook), + read_delay: Arc::clone(&self.read_delay), + })) + } + } + + struct IoHookFile { + inner: Arc, + write_hook: WriteHook, + read_delay: ReadDelay, + } + + impl SpillFile for IoHookFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option { + self.inner.size() + } + + fn read_stream( + &self, + ) -> Result> + Send>>> { + let mut inner = Some(self.inner.read_stream()?); + let delay = (self.read_delay)(); + Ok(Box::pin( + futures::stream::once(tokio::time::sleep(delay)) + .flat_map(move |_| inner.take().expect("polled once")), + )) + } + + fn open_writer(&self) -> Result> { + Ok(Box::new(IoHookWriter { + inner: self.inner.open_writer()?, + hook: Arc::clone(&self.write_hook), + })) + } + } + + struct IoHookWriter { + inner: Box, + hook: WriteHook, + } + + impl std::io::Write for IoHookWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + (self.hook)(); + self.inner.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + (self.hook)(); + self.inner.flush() + } + } + + impl SpillWriter for IoHookWriter { + fn finish(&mut self) -> Result<()> { + self.inner.finish() + } + } + + /// A `SpillManager` whose files run `write_hook` before every disk write and + /// delay reads by `read_delay()`, plus the `DiskManager` that owns the files + /// so tests can check disk usage. + fn spill_manager_with_io_hooks( + write_hook: WriteHook, + read_delay: ReadDelay, + ) -> Result<(Arc, Arc)> { + let disk_manager = Arc::new(DiskManagerBuilder::default().build()?); + let runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder(DiskManagerBuilder::default().with_mode( + DiskManagerMode::Custom(Arc::new(IoHookFactory { + inner: Arc::clone(&disk_manager), + write_hook, + read_delay, + })), + )) + .build_arc()?; + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = + Arc::new(SpillManager::new(runtime, metrics, create_test_schema())); + Ok((spill_manager, disk_manager)) + } + + /// Holds the first disk write on the pool until the test releases it. + struct WriteGate { + taken: AtomicBool, + entered: Barrier, + release: Barrier, + } + + impl WriteGate { + fn hold_if_first(&self) { + if !self.taken.swap(true, Ordering::SeqCst) { + self.entered.wait(); + self.release.wait(); + } + } + } + + /// Regression test for . + /// + /// Two writers push concurrently. The first writer is held inside its first + /// write to disk (holding its file's lock) while the second writer pushes. + /// On the unfixed pool that leaves two open files with one batch each, and + /// the reader, which only ever drained the oldest file, parked on that file + /// once it had read its single batch even though the second batch was on + /// disk. `RepartitionExec` blocks on this stream once per pushed batch and, + /// with its channel gate closed, cannot push anything else to wake it. + /// + /// The reader must yield both batches while both writers are still alive. + #[tokio::test] + async fn test_reader_does_not_wait_on_drained_file_while_another_has_data() + -> Result<()> { + let gate = Arc::new(WriteGate { + taken: AtomicBool::new(false), + entered: Barrier::new(2), + release: Barrier::new(2), + }); + let write_hook: WriteHook = { + let gate = Arc::clone(&gate); + Arc::new(move || gate.hold_if_first()) + }; + let (spill_manager, _disk_manager) = + spill_manager_with_io_hooks(write_hook, Arc::new(|| Duration::ZERO))?; + + let (writer1, mut reader) = mpsc_channel(1024 * 1024, spill_manager); + let writer2 = writer1.clone(); + + // Writer 1 pushes and is held inside its first disk write, holding its + // file's lock. The writers run on plain threads because `push_batch` + // blocks. + let writer1 = std::thread::spawn(move || { + writer1.push_batch(&create_test_batch(0, 10)).unwrap(); + writer1 + }); + gate.entered.wait(); + + // Writer 2 pushes while writer 1 is held. + let writer2 = std::thread::spawn(move || { + writer2.push_batch(&create_test_batch(10, 10)).unwrap(); + writer2 + }); + // Let writer 2 either complete (into a second file) or block on writer + // 1's file lock, depending on the pool implementation. + std::thread::sleep(Duration::from_millis(200)); + gate.release.wait(); + let writer1 = writer1.join().unwrap(); + let writer2 = writer2.join().unwrap(); + + // Both batches are on disk and both writers are still alive, so no file + // is finished. The reader must yield both batches without waiting for a + // writer. + let mut values = vec![]; + for _ in 0..2 { + let batch = tokio::time::timeout(Duration::from_secs(5), reader.next()) + .await + .expect( + "reader waited on a drained file while another file had an unread batch", + ) + .expect("reader must not signal EOF while writers are alive")?; + assert_eq!(batch.num_rows(), 10); + values.push(id_of(&batch)); + } + // Multiple writers: the order between their batches is not guaranteed. + values.sort_unstable(); + assert_eq!(values, vec![0, 10]); + + // Only once every writer is gone does the reader signal EOF. + drop(writer1); + drop(writer2); + assert!(reader.next().await.is_none()); + + Ok(()) + } + + /// Pauses every writer of a "phase" inside its first disk write, so that + /// several `push_batch` calls overlap at the point where each holds a file + /// checked out for writing. A writer that blocks earlier (for example on + /// another writer's file lock) never arrives; `release_phase` stops waiting + /// for it after a short while. + #[derive(Default)] + struct PhasePauser { + state: Mutex, + changed: parking_lot::Condvar, + } + + #[derive(Default)] + struct PhaseState { + phase: u64, + expected: usize, + held: usize, + released: bool, + } + + thread_local! { + /// The phase in which the current writer thread has already been held. + static HELD_IN_PHASE: std::cell::Cell = const { std::cell::Cell::new(0) }; + } + + impl PhasePauser { + fn begin_phase(&self, expected: usize) { + let mut state = self.state.lock(); + state.phase += 1; + state.expected = expected; + state.held = 0; + state.released = false; + } + + /// Called by writer threads before every disk write. + fn on_io(&self) { + let mut state = self.state.lock(); + if state.released + || state.held >= state.expected + || HELD_IN_PHASE.get() == state.phase + { + return; + } + HELD_IN_PHASE.set(state.phase); + state.held += 1; + self.changed.notify_all(); + while !state.released { + self.changed.wait(&mut state); + } + } + + /// Waits until every writer of the phase is held (or 300 ms have + /// passed), then releases them all at once. + fn release_phase(&self) { + let mut state = self.state.lock(); + let deadline = Instant::now() + Duration::from_millis(300); + while state.held < state.expected { + if self.changed.wait_until(&mut state, deadline).timed_out() { + break; + } + } + state.released = true; + self.changed.notify_all(); + } + } + + fn batch_with_id(id: i32, rows: usize) -> RecordBatch { + let a: ArrayRef = Arc::new(Int32Array::from(vec![id; rows])); + RecordBatch::try_new(create_test_schema(), vec![a]).unwrap() + } + + fn id_of(batch: &RecordBatch) -> i32 { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + } + + /// One randomly generated scenario, see [`spill_pool_scenario_fuzz`]. + async fn run_spill_pool_scenario(seed: u64) -> Result<()> { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + const READ_TIMEOUT: Duration = Duration::from_secs(5); + + let mut rng = StdRng::seed_from_u64(seed); + let n_writers = rng.random_range(1..=4); + let batch_bytes = batch_with_id(0, 10).get_array_memory_size(); + // Rotate after every batch, after a few batches, or never. + let max_file_size = match rng.random_range(0..3) { + 0 => batch_bytes / 2, + 1 => batch_bytes * 3, + _ => usize::MAX / 2, + }; + let n_phases = rng.random_range(1..=6); + let context = format!( + "seed {seed}: {n_writers} writers, max_file_size {max_file_size}, {n_phases} phases" + ); + + let pauser = Arc::new(PhasePauser::default()); + let write_hook: WriteHook = { + let pauser = Arc::clone(&pauser); + Arc::new(move || pauser.on_io()) + }; + // Some files are slow to read back, so an older file's bytes can arrive + // after a newer file's. Batches must still come back in file order for a + // single writer. + let read_delay: ReadDelay = { + let rng = Mutex::new(StdRng::seed_from_u64(seed ^ 0x5EED_4EAD)); + Arc::new(move || { + let mut rng = rng.lock(); + if rng.random_bool(0.5) { + Duration::ZERO + } else { + Duration::from_millis(rng.random_range(1..=20)) + } + }) + }; + let (spill_manager, disk_manager) = + spill_manager_with_io_hooks(write_hook, read_delay)?; + let (writer, mut reader) = mpsc_channel(max_file_size, spill_manager); + let mut sinks: Vec> = + (0..n_writers).map(|_| Some(writer.new_sink())).collect(); + drop(writer); + + let mut pushed: Vec = vec![]; + let mut read: Vec = vec![]; + + for phase in 0..n_phases { + let alive: Vec = + (0..n_writers).filter(|w| sinks[*w].is_some()).collect(); + if alive.is_empty() { + break; + } + + // A random non-empty subset of the live writers each pushes a few + // batches, concurrently, held so that their pushes overlap. + let mut pushing: Vec = alive + .iter() + .copied() + .filter(|_| rng.random_bool(0.6)) + .collect(); + if pushing.is_empty() { + pushing.push(alive[rng.random_range(0..alive.len())]); + } + pauser.begin_phase(pushing.len()); + let mut threads = Vec::with_capacity(pushing.len()); + for w in pushing { + let sink = sinks[w].take().unwrap(); + let batches: Vec = (0..rng.random_range(1..=3)) + .map(|_| { + let id = pushed.len() as i32 + 1; + pushed.push(id); + batch_with_id(id, rng.random_range(1..=40)) + }) + .collect(); + threads.push(( + w, + std::thread::spawn(move || { + for batch in &batches { + sink.push_batch(batch).unwrap(); + } + sink + }), + )); + } + pauser.release_phase(); + for (w, thread) in threads { + sinks[w] = Some(thread.join().unwrap()); + } + + // Quiescent point: no push is in progress and none will happen until + // the reader is done. Every batch pushed so far must be readable now, + // in any amount. `RepartitionExec` relies on exactly this: it blocks + // on the reader once per pushed batch while its channel gate keeps + // the writers parked. Sometimes leave a backlog for later phases. + let read_now = rng.random_range(0..=pushed.len() - read.len()); + for _ in 0..read_now { + let next = tokio::time::timeout(READ_TIMEOUT, reader.next()) + .await + .unwrap_or_else(|_| { + panic!( + "{context}, phase {phase}: reader stalled with {} of {} \ + batches unread and {} writers alive", + pushed.len() - read.len(), + pushed.len(), + alive.len() + ) + }); + let batch = next.unwrap_or_else(|| { + panic!("{context}, phase {phase}: EOF while writers are alive") + })?; + read.push(id_of(&batch)); + } + + // With nothing left to read the reader must wait, not signal EOF. + if read.len() == pushed.len() && rng.random_bool(0.3) { + let probe = + tokio::time::timeout(Duration::from_millis(50), reader.next()).await; + assert!( + probe.is_err(), + "{context}, phase {phase}: reader yielded {probe:?} with nothing \ + pushed and writers alive" + ); + } + + // Dropping one of several writers must not disturb anything. + if alive.len() > 1 && rng.random_bool(0.3) { + sinks[alive[rng.random_range(0..alive.len())]] = None; + } + } + + // Once the last writer is gone the reader must hand out the backlog and + // then signal EOF. + sinks.clear(); + while read.len() < pushed.len() { + let next = tokio::time::timeout(READ_TIMEOUT, reader.next()) + .await + .unwrap_or_else(|_| { + panic!( + "{context}: reader stalled with {} of {} batches unread after \ + all writers were dropped", + pushed.len() - read.len(), + pushed.len() + ) + }); + let batch = next.unwrap_or_else(|| { + panic!( + "{context}: EOF with {} batches unread", + pushed.len() - read.len() + ) + })?; + read.push(id_of(&batch)); + } + let eof = tokio::time::timeout(READ_TIMEOUT, reader.next()) + .await + .unwrap_or_else(|_| { + panic!("{context}: reader stalled instead of signalling EOF after all writers were dropped") + }); + assert!(eof.is_none(), "{context}: batch after everything was read"); + + let mut expected = pushed.clone(); + expected.sort_unstable(); + let mut actual = read.clone(); + actual.sort_unstable(); + assert_eq!(actual, expected, "{context}: batches lost or duplicated"); + if n_writers == 1 { + assert_eq!(read, pushed, "{context}: single-writer pool must be FIFO"); + } + + drop(reader); + assert_eq!( + disk_manager.used_disk_space(), + 0, + "{context}: spill files not released" + ); + Ok(()) + } + + /// Randomized scenarios against the pool's contract: after any set of + /// overlapping pushes the reader can drain everything without further + /// writer activity (no lost wakeups, no parking on one file while another + /// has data), it never signals EOF while a writer is alive, it returns + /// exactly the pushed batches (in order for a single writer), and the + /// files are released when it is dropped. + /// + /// `DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS` and + /// `DATAFUSION_SPILL_POOL_FUZZ_SEED` select how many scenarios run and + /// from which seed; a failure names its seed so it can be replayed. + #[tokio::test] + async fn spill_pool_scenario_fuzz() -> Result<()> { + let env_u64 = |name: &str, default: u64| { + std::env::var(name) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + }; + let first_seed = env_u64("DATAFUSION_SPILL_POOL_FUZZ_SEED", 0); + let iterations = env_u64("DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS", 50); + for seed in first_seed..first_seed + iterations { + run_spill_pool_scenario(seed).await?; + } + Ok(()) + } } From fe8f3a4aa54b7291e5be25a8dcbb1d684f53c6a2 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 2 Sep 2026 20:37:06 +0000 Subject: [PATCH 3/3] fix: RepartitionExec deadlock when its spill pool holds two open files In non-preserve-order mode, all input tasks of a `RepartitionExec` share one spill pool per output partition. When two input tasks spilled at the same time, the pool had two open spill files. The `SpillPoolReader` read only the oldest file. When that file had no unread batch and was not finished, the reader waited for it. But the batch that the reader needed was in the newer file. When each distributor channel had data, the gate closed. Then both input tasks stopped in `send`, no sink was dropped, and nothing could wake the reader. The reader now keeps all the files that it received, oldest first. It gives the first batch that is available in one of them. It skips a file that has no unread batch. It waits for the oldest file that has an unread batch when the read of that file is not complete. Thus one writer gets its batches in FIFO order. This commit makes the tests from the two commits before it green. The fuzzer passes 500 seeds. Ref: https://github.com/apache/datafusion/issues/24883 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9 --- .../physical-plan/src/spill/spill_pool.rs | 214 +++++++++++------- 1 file changed, 128 insertions(+), 86 deletions(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 43575388edea9..ecb0dc8648fc4 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -48,10 +48,11 @@ use super::spill_manager::SpillManager; /// **Lock ordering discipline**: Never hold both locks simultaneously to prevent deadlock. /// Always: acquire outer lock → release outer lock → acquire inner lock (if needed). struct SpillPoolShared { - /// Queue of ALL files (including the current write files if any exist). - /// Readers always read from the front of this queue (FIFO). - /// Each file has its own lock to enable concurrent reader/writer access. - files: VecDeque>>, + /// Files created by writers that the reader has not picked up yet, in creation + /// order. The reader moves them into its own list on every poll (see + /// [`SpillPoolReader`]), so this queue only ever holds files the reader has not + /// seen. Each file has its own lock to enable concurrent reader/writer access. + new_files: VecDeque>>, /// SpillManager for creating files and tracking metrics spill_manager: Arc, /// Pool-level waker to notify when new files are available (single reader) @@ -70,7 +71,7 @@ impl SpillPoolShared { /// Creates a new shared pool state fn new(spill_manager: Arc) -> Self { Self { - files: VecDeque::new(), + new_files: VecDeque::new(), spill_manager, waker: None, open_write_files: VecDeque::new(), @@ -237,7 +238,7 @@ impl SpillPoolSink { // Re-acquire lock and push to shared queue shared = self.shared.lock(); - shared.files.push_back(Arc::clone(&file_shared)); + shared.new_files.push_back(Arc::clone(&file_shared)); shared.wake(); // Wake readers waiting for new files file_shared }; @@ -570,17 +571,34 @@ struct SpillPoolFile { spill_manager: Arc, } -impl Stream for SpillPoolFile { - type Item = Result; +/// Outcome of polling a single file of the pool, see [`SpillPoolFile::poll_file`]. +enum FilePoll { + /// A batch was read from the file, or reading it failed. + Item(Result), + /// The file's stream ended: every batch written to it has been read and the + /// writer has finished it (the caller verifies the latter). + Done, + /// Every batch written so far has been read, but the writer has not finished + /// the file. The file's waker is registered, so the task is woken when the + /// writer appends another batch or finishes the file. + CaughtUp, + /// The file has unread batches, but the underlying I/O is not ready yet. The + /// I/O has registered the task's waker. + Pending, +} - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { +impl SpillPoolFile { + /// Polls the file for its next batch. + /// + /// Unlike a plain stream poll this tells the caller *why* no batch is + /// available, so [`SpillPoolReader`] can move on to a newer file when this + /// one is only waiting for its writer, but keeps waiting on it while its + /// data is still being read from disk (which keeps batches in file order). + fn poll_file(&mut self, cx: &mut std::task::Context<'_>) -> FilePoll { use std::task::Poll; // Step 1: Lock shared state and check coordination - let (should_read, file) = { + let file = { let mut shared = self.shared.lock(); // Determine if we can read @@ -588,24 +606,23 @@ impl Stream for SpillPoolFile { if batches_read < shared.batches_written { // More data available to read - take the file if we don't have a reader yet - let file = if self.reader.is_none() { + if self.reader.is_none() { shared.file.take() } else { None - }; - (true, file) + } } else if shared.writer_finished { // No more data and writer is done - EOF - return Poll::Ready(None); + return FilePoll::Done; } else { // Caught up to writer, but writer still active - register waker and wait shared.register_waker(cx.waker().clone()); - return Poll::Pending; + return FilePoll::CaughtUp; } }; // Lock released here // Step 2: Lazy-create reader stream if needed - if self.reader.is_none() && should_read { + if self.reader.is_none() { if let Some(file) = file { // we want this unbuffered because files are actively being written to match self @@ -618,36 +635,35 @@ impl Stream for SpillPoolFile { batches_read: 0, }); } - Err(e) => return Poll::Ready(Some(Err(e))), + Err(e) => return FilePoll::Item(Err(e)), } } else { // File not available yet (writer hasn't finished or already taken) // Register waker and wait for file to be ready let mut shared = self.shared.lock(); shared.register_waker(cx.waker().clone()); - return Poll::Pending; + return FilePoll::CaughtUp; } } // Step 3: Poll the reader stream (no lock held) - if let Some(reader) = &mut self.reader { - match reader.stream.poll_next_unpin(cx) { - Poll::Ready(Some(Ok(batch))) => { - // Successfully read a batch - increment counter - reader.batches_read += 1; - Poll::Ready(Some(Ok(batch))) - } - Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), - Poll::Ready(None) => { - // Stream exhausted unexpectedly - // This shouldn't happen if coordination is correct, but handle gracefully - Poll::Ready(None) - } - Poll::Pending => Poll::Pending, - } - } else { + let Some(reader) = &mut self.reader else { // Should not reach here, but handle gracefully - Poll::Ready(None) + return FilePoll::Done; + }; + match reader.stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + // Successfully read a batch - increment counter + reader.batches_read += 1; + FilePoll::Item(Ok(batch)) + } + Poll::Ready(Some(Err(e))) => FilePoll::Item(Err(e)), + Poll::Ready(None) => { + // Stream exhausted unexpectedly + // This shouldn't happen if coordination is correct, but handle gracefully + FilePoll::Done + } + Poll::Pending => FilePoll::Pending, } } } @@ -669,11 +685,31 @@ impl Stream for SpillPoolFile { /// /// This makes it suitable for continuous streaming scenarios where the writer may /// produce data intermittently. +/// +/// # Reading from several open files +/// +/// With a single writer ([`spsc_channel`]) at most one file is open for writing at +/// a time and every earlier file is finished, so reading the files oldest first is +/// exactly FIFO. With several writers ([`mpsc_channel`]) concurrent pushes can +/// leave more than one file open, and a batch may land in a newer file while the +/// oldest file is drained but not finished. The reader therefore polls every file +/// it knows about, oldest first, and returns the first batch that is available. +/// A file that is merely caught up with its writer is skipped; the reader only +/// waits on a file while that file has unread batches whose bytes are still being +/// read from disk, which keeps batches in file order (strict FIFO for one writer). +/// +/// This matters for [`RepartitionExec`](crate::repartition::RepartitionExec), which +/// sends a "spilled" marker through its channel after every `push_batch` and then +/// blocks on this stream until it yields a batch. If the reader only ever drained the +/// oldest file, it could wait for a batch that had been written to a newer file; with +/// the channel gate closed the writers could not push anything to wake it, so the +/// query would deadlock (). pub struct SpillPoolReader { /// Shared reference to the spill pool shared: Arc>, - /// Current SpillPoolFile we're reading from - current_file: Option, + /// Files this reader has picked up from the pool and not fully consumed yet, in + /// creation order. Each carries its own lazily-created stream and read position. + files: VecDeque, /// Schema of the spilled data schema: SchemaRef, } @@ -689,7 +725,7 @@ impl SpillPoolReader { fn new(shared: Arc>, schema: SchemaRef) -> Self { Self { shared, - current_file: None, + files: VecDeque::new(), schema, } } @@ -704,75 +740,81 @@ impl Stream for SpillPoolReader { ) -> std::task::Poll> { use std::task::Poll; + // `Self: Unpin`; reborrow once so `files` and `shared` can be used together + let this = &mut *self; + loop { - // If we have a current file, try to read from it - if let Some(ref mut file) = self.current_file { - match file.poll_next_unpin(cx) { - Poll::Ready(Some(Ok(batch))) => { - // Got a batch, return it - return Poll::Ready(Some(Ok(batch))); - } - Poll::Ready(Some(Err(e))) => { - // Error reading batch - return Poll::Ready(Some(Err(e))); + // Pick up the files writers have created since the last poll, in + // creation order. + { + let mut shared = this.shared.lock(); + while let Some(file_shared) = shared.new_files.pop_front() { + let spill_manager = Arc::clone(&shared.spill_manager); + this.files.push_back(SpillPoolFile { + shared: file_shared, + reader: None, + spill_manager, + }); + } + } // Lock released here + + // Poll the files oldest first and return the first available batch. + // A file that is finished and fully read is dropped, which releases + // its disk space. A file that is only waiting for its writer has + // registered its own waker, so we move on to the next file instead + // of waiting on it. + let mut idx = 0; + while idx < this.files.len() { + match this.files[idx].poll_file(cx) { + FilePoll::Item(item) => { + // Got a batch (or an error), return it + return Poll::Ready(Some(item)); } - Poll::Ready(None) => { - // Current file stream exhausted + FilePoll::Done => { + // File stream exhausted // Check if this file is marked as writer_finished - let writer_finished = { file.shared.lock().writer_finished }; + let writer_finished = + { this.files[idx].shared.lock().writer_finished }; if writer_finished { - // File is complete, pop it from the queue and move to next - let mut shared = self.shared.lock(); - shared.files.pop_front(); - drop(shared); // Release lock - - // Clear current file and continue loop to get next file - self.current_file = None; - continue; + // File is complete, drop it and move on to the next + this.files.remove(idx); } else { // Stream exhausted but writer not finished - unexpected // This shouldn't happen with proper coordination return Poll::Ready(None); } } - Poll::Pending => { - // File not ready yet (waiting for writer) - // Register waker so we get notified when writer adds more batches - let mut shared = self.shared.lock(); - shared.register_waker(cx.waker().clone()); + FilePoll::Pending => { + // The oldest file with unread batches is still being read + // from disk. Wait for it rather than skipping ahead, so + // batches come back in file order (FIFO for one writer). return Poll::Pending; } + FilePoll::CaughtUp => { + // Nothing to read here until its writer appends more; a + // newer file may have unread batches, so try the next one + idx += 1; + } } } - // No current file, need to get the next one - let mut shared = self.shared.lock(); - - // Peek at the front of the queue (don't pop yet) - if let Some(file_shared) = shared.files.front() { - // Create a SpillPoolFile from the shared state - let spill_manager = Arc::clone(&shared.spill_manager); - let file_shared = Arc::clone(file_shared); - drop(shared); // Release lock before creating SpillPoolFile - - self.current_file = Some(SpillPoolFile { - shared: file_shared, - reader: None, - spill_manager, - }); + // No known file has an unread batch + let mut shared = this.shared.lock(); - // Continue loop to poll the new file + if !shared.new_files.is_empty() { + // A writer created a file while we were polling; pick it up continue; } - // No files in queue - check if writer is done - if shared.remaining_writer_count == 0 { + if this.files.is_empty() && shared.remaining_writer_count == 0 { // Writer is done and no more files will be added - EOF return Poll::Ready(None); } - // Writer still active, register waker that will get notified when new files are added + // Writer still active: register the pool-level waker so we are + // notified when a new file is created or the last writer is dropped. + // Each pending file has registered its own waker for new batches. shared.register_waker(cx.waker().clone()); return Poll::Pending; }