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/core/tests/memory_limit/repartition_mem_limit.rs b/datafusion/core/tests/memory_limit/repartition_mem_limit.rs index 27bcd33926e99..21a0eaec5473d 100644 --- a/datafusion/core/tests/memory_limit/repartition_mem_limit.rs +++ b/datafusion/core/tests/memory_limit/repartition_mem_limit.rs @@ -117,3 +117,157 @@ 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. +/// +/// This module is not compiled with `force_hash_collisions`. That feature makes +/// every hash 0, so the hash `RepartitionExec` sends all rows to one output +/// partition. The gate closes only when every channel holds data, so the +/// deadlock cannot occur and the test cannot test anything. The feature also +/// makes the group-by hash table degenerate, which made one attempt of this +/// query approximately 36 times slower and thus longer than the limit below. +#[cfg(not(feature = "force_hash_collisions"))] +mod spill_pool_deadlock { + use super::*; + + use std::time::Duration; + + use datafusion::execution::session_state::SessionStateBuilder; + use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion_physical_plan::collect; + + /// 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); + + /// Runs the query once. Returns the number of output rows and the number + /// of spills the `RepartitionExec`s did. + async fn run_once() -> datafusion::error::Result<(usize, usize)> { + 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 logical = ctx + .state() + .create_logical_plan( + "select g, count(distinct trace_id) as n from tv group by g", + ) + .await?; + let plan = ctx.state().create_physical_plan(&logical).await?; + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + + // Count the spills of the `RepartitionExec`s. The caller uses this to + // make sure the query still goes through the spill path: if a change + // stops it spilling, this test keeps passing but no longer covers the + // deadlock, and that must fail loudly instead. + let mut spills = 0; + plan.transform_down(|node| { + if node.is::() { + spills += node.metrics().and_then(|m| m.spill_count()).unwrap_or(0); + } + Ok(Transformed::no(node)) + })?; + + Ok((batches.iter().map(|b| b.num_rows()).sum(), spills)) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn repartition_spill_pool_does_not_deadlock() { + let mut attempts_that_spilled = 0; + 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, spills))) => { + assert_eq!(rows, 13, "attempt {attempt}: wrong row count"); + if spills > 0 { + attempts_that_spilled += 1; + } + } + // 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}"), + } + } + + // The memory limit must be tight enough that `RepartitionExec` spills, + // because the deadlock is in its spill pool. Timing decides how many + // attempts spill, but if none of them do, the limit no longer forces + // the spill path and this test has stopped testing anything. + assert!( + attempts_that_spilled > 0, + "no attempt reached the RepartitionExec spill path in {ATTEMPTS} \ + attempts; the memory limit no longer forces a spill" + ); + } +} diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 56d8d11201f46..a3366d6766170 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; } @@ -791,8 +833,20 @@ 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::{DataFusionError, exec_datafusion_err}; 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::sync::mpsc::{self, Sender}; + use std::time::Duration; fn create_test_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) @@ -919,8 +973,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 +1283,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 +1576,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 +1697,576 @@ mod tests { Ok(()) } + + type WriteHook = Arc; + /// What a read of a spill file does. The two failure modes reach the two + /// different error paths of [`SpillPoolFile::poll_file`]: one fails while it + /// builds the stream, the other fails while it polls a stream that is + /// already built. + enum ReadBehavior { + /// Read the real bytes, after this delay. + Delay(Duration), + /// Fail to open the file, so the stream is never built. + FailOpen(DataFusionError), + /// Open the file, then give this error as the first item of the stream. + FailFirstItem(DataFusionError), + } + + type ReadHook = Arc ReadBehavior + Send + Sync>; + + /// Test double for a spill file that runs `write_hook` before every disk + /// write or flush, and consults `read_hook` before every read. 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, and a read error fails the read. Data still goes to real + /// temporary files. + struct IoHookFactory { + inner: Arc, + write_hook: WriteHook, + read_hook: ReadHook, + } + + 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_hook: Arc::clone(&self.read_hook), + })) + } + } + + struct IoHookFile { + inner: Arc, + write_hook: WriteHook, + read_hook: ReadHook, + } + + 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 delay = match (self.read_hook)() { + ReadBehavior::Delay(delay) => delay, + ReadBehavior::FailOpen(e) => return Err(e), + ReadBehavior::FailFirstItem(e) => { + return Ok(Box::pin(futures::stream::once(async move { Err(e) }))); + } + }; + let mut inner = Some(self.inner.read_stream()?); + 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 + /// consult `read_hook` before every read, plus the `DiskManager` that owns + /// the files so tests can check disk usage. + fn spill_manager_with_io_hooks( + write_hook: WriteHook, + read_hook: ReadHook, + ) -> 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_hook, + })), + )) + .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, and + /// reports every other write so a test can wait for a second writer to + /// reach a file of its own instead of sleeping. + struct WriteGate { + taken: AtomicBool, + entered: Barrier, + release: Barrier, + other_writes: Sender<()>, + } + + impl WriteGate { + fn hold_if_first(&self) { + if self.taken.swap(true, Ordering::SeqCst) { + // A write only gets here once its writer holds a file of its + // own: a writer that had to wait for the held writer's file + // lock would not reach a write of its own at all. + let _ = self.other_writes.send(()); + } else { + 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 (other_writes_tx, other_writes) = mpsc::channel(); + let gate = Arc::new(WriteGate { + taken: AtomicBool::new(false), + entered: Barrier::new(2), + release: Barrier::new(2), + other_writes: other_writes_tx, + }); + 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(|| ReadBehavior::Delay(Duration::ZERO)), + )?; + + let (writer1, mut reader) = mpsc_channel(1024 * 1024, Arc::clone(&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 + }); + // Wait for writer 2 to reach a write of its own, so that writer 1 is + // only released once the pool really holds two open files. Sleeping + // instead would let a loaded runner schedule writer 2 after writer 1 + // was released, and writer 2 would then reuse writer 1's returned file: + // a single-file run that says nothing about the case under test. + other_writes + .recv_timeout(Duration::from_secs(30)) + .expect("writer 2 must write to a file of its own while writer 1 is held"); + gate.release.wait(); + let writer1 = writer1.join().unwrap(); + let writer2 = writer2.join().unwrap(); + assert_eq!( + spill_manager.metrics.spill_file_count.value(), + 2, + "the writers must have written to two different files" + ); + + // 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(()) + } + + /// A spill file that cannot be read must fail the reader rather than stall + /// it. `RepartitionExec` blocks on this stream after every spilled marker, + /// so an error that is swallowed here would hang the query exactly like the + /// deadlock this module guards against. + /// + /// Both error paths of `SpillPoolFile::poll_file` are covered: the file that + /// does not open, so the failure comes from the construction of the stream, + /// and the file that opens and then gives an error, so the failure comes + /// from a poll of a stream that the reader already holds. + #[tokio::test] + async fn test_read_errors_are_reported_to_the_reader() -> Result<()> { + async fn assert_reader_reports(behavior: fn() -> ReadBehavior) -> Result<()> { + let (spill_manager, _disk_manager) = + spill_manager_with_io_hooks(Arc::new(|| {}), Arc::new(behavior))?; + let (writer, mut reader) = spsc_channel(1024 * 1024, spill_manager); + + // Nothing reads the file until the reader does, so every read that + // the pool makes is a read of the batch below. + writer.push_batch(&create_test_batch(0, 10))?; + + let item = tokio::time::timeout(Duration::from_secs(5), reader.next()) + .await + .expect("reader must report the read error instead of waiting") + .expect("reader must report the read error instead of signalling EOF"); + let err = item.expect_err("a failed read must not produce a batch"); + assert!( + err.to_string().contains("injected spill read failure"), + "unexpected error: {err}" + ); + Ok(()) + } + + assert_reader_reports(|| { + ReadBehavior::FailOpen(exec_datafusion_err!("injected spill read failure")) + }) + .await?; + assert_reader_reports(|| { + ReadBehavior::FailFirstItem(exec_datafusion_err!( + "injected spill read failure" + )) + }) + .await?; + + 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_hook: ReadHook = { + 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) { + ReadBehavior::Delay(Duration::ZERO) + } else { + ReadBehavior::Delay(Duration::from_millis(rng.random_range(1..=20))) + } + }) + }; + let (spill_manager, disk_manager) = + spill_manager_with_io_hooks(write_hook, read_hook)?; + 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); + // `saturating_add` so a replay seed near `u64::MAX` cannot overflow. + for seed in first_seed..first_seed.saturating_add(iterations) { + run_spill_pool_scenario(seed).await?; + } + Ok(()) + } }