diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index 681596a52fc60..341d64e68d84d 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -69,10 +69,14 @@ jobs: - name: Show available disk space run: | df -h - - name: Run tests (excluding doctests) + - name: Run extended tests run: cargo xtask ci step test extended - - name: Run fuzzers with more iterations + - name: Run fuzz tests run: cargo xtask ci step test fuzz + # Migrate it to the above unified fuzz runner over time + # Tracker: https://github.com/apache/datafusion/issues/27082 + - name: Run unmigrated fuzzers with more iterations + run: cargo xtask ci step test fuzz-legacy - parallel: - name: Verify Working Directory Clean run: git diff --exit-code diff --git a/Cargo.lock b/Cargo.lock index 6fd56360c54f2..792ea91080d76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2411,6 +2411,22 @@ dependencies = [ "datafusion-physical-expr-common", ] +[[package]] +name = "datafusion-fuzz" +version = "0.1.0" +dependencies = [ + "arrow", + "bytes", + "datafusion-common", + "datafusion-execution", + "datafusion-physical-plan", + "futures", + "parking_lot 0.13.0", + "rand 0.9.4", + "test-utils", + "tokio", +] + [[package]] name = "datafusion-macros" version = "57.0.0-dev" diff --git a/Cargo.toml b/Cargo.toml index a96d3cbf37e58..c08d7af6d7096 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ members = [ "datafusion/execution", "datafusion/ffi", "datafusion/ffi/integration-tests", + "datafusion/fuzz", "datafusion/functions", "datafusion/functions-aggregate", "datafusion/functions-aggregate-common", diff --git a/datafusion/fuzz/Cargo.toml b/datafusion/fuzz/Cargo.toml new file mode 100644 index 0000000000000..661c4be58b276 --- /dev/null +++ b/datafusion/fuzz/Cargo.toml @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +# A private crate providing a single entry point for all fuzz tests. + +[package] +name = "datafusion-fuzz" +version = "0.1.0" +edition = { workspace = true } +# Ensure this crate is private +publish = false +autotests = false + +[lints] +workspace = true + +[features] +extended_tests = [] + +[dev-dependencies] +arrow = { workspace = true } +bytes = { workspace = true } +datafusion-common = { workspace = true } +datafusion-execution = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } +parking_lot = { workspace = true } +rand = { workspace = true } +test-utils = { path = "../../test-utils" } +tokio = { workspace = true, features = ["rt-multi-thread", "time"] } + +[[test]] +name = "fuzz" +path = "tests/fuzz.rs" +required-features = ["extended_tests"] diff --git a/datafusion/fuzz/tests/fuzz.rs b/datafusion/fuzz/tests/fuzz.rs new file mode 100644 index 0000000000000..6eae5f9127eb5 --- /dev/null +++ b/datafusion/fuzz/tests/fuzz.rs @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! One integration entrypoint for DataFusion's registered fuzz suites. + +use std::future::Future; +use std::pin::Pin; + +use datafusion_common::Result; +use test_utils::fuzz::{FuzzTest, RunConfig, run}; + +mod suites; + +/// The single registration point for project-wide fuzz campaigns. +/// +/// Migration of the remaining fuzz tests is tracked in [#27082]. +/// +/// [#27082]: https://github.com/apache/datafusion/issues/27082 +#[tokio::test(flavor = "multi_thread")] +async fn fuzz_tests() -> Result<()> { + let spill_pool = suites::spill_pool::SpillPoolFuzzer; + let suites = [Suite::new( + &spill_pool, + RunConfig { + run_seed: Some(0), + max_cases: Some(1000), + time_limit: None, + cases_per_dataset: 1000, + }, + )]; + run_suites(suites).await +} + +async fn run_suites<'a>(suites: impl IntoIterator>) -> Result<()> { + for suite in suites { + eprintln!("Running fuzz suite: {}", suite.name); + suite.run.await?; + } + Ok(()) +} + +/// Wrapper for one fuzzing test suite +struct Suite<'a> { + name: &'a str, + /// Run the test suite. Report an error on failure or execution statistics on success. + run: Pin> + 'a>>, +} + +impl<'a> Suite<'a> { + fn new(test: &'a T, config: RunConfig) -> Self { + Self { + name: test.name(), + run: Box::pin(async move { + let stats = run(test, config).await?; + eprintln!( + "Fuzz suite {}: {} cases, {} datasets, {:?}", + test.name(), + stats.cases_passed, + stats.datasets_prepared, + stats.elapsed + ); + Ok(()) + }), + } + } +} diff --git a/datafusion/fuzz/tests/suites/mod.rs b/datafusion/fuzz/tests/suites/mod.rs new file mode 100644 index 0000000000000..a443d433bd089 --- /dev/null +++ b/datafusion/fuzz/tests/suites/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Implementation for fuzz test cases. + +pub(crate) mod spill_pool; diff --git a/datafusion/fuzz/tests/suites/spill_pool.rs b/datafusion/fuzz/tests/suites/spill_pool.rs new file mode 100644 index 0000000000000..d6669bfe8e265 --- /dev/null +++ b/datafusion/fuzz/tests/suites/spill_pool.rs @@ -0,0 +1,410 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Seeded scenarios checking concurrent writes, reader progress, ordering, and +//! spill-file cleanup. Seeds reproduce the workload and delay policy; thread +//! scheduling may still vary. + +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int32Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::instant::Instant; +use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::{SpillFile, SpillWriter, TempFileFactory}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; +use datafusion_physical_plan::spill::SpillManager; +use datafusion_physical_plan::spill::spill_pool::{SpillPoolSink, mpsc_channel}; +use futures::{Stream, StreamExt}; +use parking_lot::Mutex; +use test_utils::fuzz::FuzzTest; + +pub(crate) struct SpillPoolFuzzer; + +impl FuzzTest for SpillPoolFuzzer { + type Dataset = (); + + fn name(&self) -> &str { + "spill_pool_scenario" + } + + fn prepare_dataset(&self, _: u64) -> impl Future> { + std::future::ready(Ok(())) + } + + async fn run_case(&self, _: &(), 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) { + Duration::ZERO + } else { + 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(()) + } +} + +fn create_test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) +} + +type WriteHook = Arc; +type ReadHook = Arc Duration + 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. 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 = (self.read_hook)(); + 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)) +} + +/// 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) +} diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 48e67e72d2e75..f1a426ccd59bd 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -820,7 +820,6 @@ 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::{ @@ -1970,76 +1969,6 @@ mod tests { 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) @@ -2048,212 +1977,4 @@ mod tests { .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(()) - } } diff --git a/test-utils/src/fuzz.rs b/test-utils/src/fuzz.rs new file mode 100644 index 0000000000000..f932f75f4fb34 --- /dev/null +++ b/test-utils/src/fuzz.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Shared execution and replay for [`FuzzTest`] experiments. + +use std::future::Future; +use std::time::Duration; + +use datafusion_common::instant::Instant; +use datafusion_common::{Result, exec_err}; + +/// Implement this trait to add a fuzzing unit test to DataFusion. +/// +/// The fuzz implementation prepares the data and checks one case. The [centralized driver] +/// repeats cases, manages the time budget, and reports failures. +/// +/// See the [centralized driver] for how to use it. At a high level: +/// 1. Generate one dataset with `prepare_dataset()`. +/// 2. Run several cases on that dataset with `run_case()`. +/// 3. Repeat steps 1 and 2 until the time budget is reached. +/// +/// (All configuration, such as the total time budget, lives in [`RunConfig`].) +/// +/// [centralized driver]: https://github.com/apache/datafusion/blob/main/datafusion/fuzz/tests/fuzz.rs +#[expect(async_fn_in_trait)] // Serial execution deliberately does not require Send futures. +pub trait FuzzTest { + /// Reusable fixtures; `()` is valid when no preparation can be shared. + type Dataset; + + /// Test identity, including any static variant needed for replay. + fn name(&self) -> &str; + + /// Build all fixtures from this seed and the test's fixed configuration. + async fn prepare_dataset(&self, dataset_seed: u64) -> Result; + + /// Execute one complete, independent validation and clean up its resources. + /// + /// The implementation should ensure: + /// - The test case is reproducible given the same dataset and seed. + /// - If the test fails, the error message includes all the information needed + /// to reproduce it (by providing the data seed and case seed). + async fn run_case(&self, dataset: &Self::Dataset, case_seed: u64) -> Result<()>; + + /// Reproduce the failing case deterministically. + /// + /// The input seed should be printed in the error message from [`Self::run_case`] + /// when the test fails. + async fn reproduce(&self, dataset_seed: u64, case_seed: u64) -> Result<()> { + let dataset = with_context(self.name(), None, dataset_seed, None, || { + self.prepare_dataset(dataset_seed) + }) + .await?; + with_context(self.name(), None, dataset_seed, Some(case_seed), || { + self.run_case(&dataset, case_seed) + }) + .await + } +} + +/// Execution limits for [`run`]. At least one stopping condition is required. +#[derive(Debug, Clone)] +pub struct RunConfig { + /// Initial seed, or `None` to choose a fresh seed once per run. + pub run_seed: Option, + /// Maximum successful cases; must be positive when supplied. + pub max_cases: Option, + /// Soft budget including preparation. The first complete case always runs, + /// even with a zero budget. Active preparation and cases are never cancelled. + pub time_limit: Option, + /// Maximum cases sharing one fixture; must be positive. + pub cases_per_dataset: usize, +} + +impl RunConfig { + fn validate(&self) -> Result<()> { + if self.max_cases.is_none() && self.time_limit.is_none() { + return exec_err!("Fuzz run requires max_cases or time_limit"); + } + if self.max_cases == Some(0) { + return exec_err!("Fuzz run max_cases must be greater than zero"); + } + if self.cases_per_dataset == 0 { + return exec_err!("Fuzz run cases_per_dataset must be greater than zero"); + } + Ok(()) + } + + // Keep admission separate from the clock so deadline boundaries can be + // tested without sleeping. Only the first case may start past the deadline. + fn should_stop(&self, cases_passed: usize, elapsed: Duration) -> bool { + self.max_cases.is_some_and(|limit| cases_passed >= limit) + || (cases_passed > 0 && self.time_limit.is_some_and(|limit| elapsed >= limit)) + } +} + +/// Counters returned by a successful campaign. Failures stop the run immediately. +#[derive(Debug)] +pub struct RunStats { + /// Resolved seed, including when the caller requested a fresh one. + pub run_seed: u64, + /// Number of successfully prepared fixtures. + pub datasets_prepared: usize, + /// Number of complete validations that returned successfully. + pub cases_passed: usize, + /// Total execution time, including fixture preparation and cleanup. + pub elapsed: Duration, +} + +/// Run cases serially until either limit is reached or a failure occurs. +/// +/// Case seeds are `run_seed + case_index`, wrapping at `u64::MAX`. Dataset seeds +/// independently use `run_seed + dataset_index`. Neither sequence depends on +/// randomness consumed by the test. This also lets existing consecutive-seed +/// fuzzers retain their coverage. A timed run may finish different case counts +/// on different machines; use a fixed count for a repeatable seed sequence. +pub async fn run(test: &T, config: RunConfig) -> Result { + config.validate()?; + let run_seed = config.run_seed.unwrap_or_else(rand::random); + let start = Instant::now(); + let mut stats = RunStats { + run_seed, + datasets_prepared: 0, + cases_passed: 0, + elapsed: Duration::ZERO, + }; + while !config.should_stop(stats.cases_passed, start.elapsed()) { + let dataset_seed = run_seed.wrapping_add(stats.datasets_prepared as u64); + let dataset = + with_context(test.name(), Some(run_seed), dataset_seed, None, || { + test.prepare_dataset(dataset_seed) + }) + .await?; + stats.datasets_prepared += 1; + for _ in 0..config.cases_per_dataset { + if config.should_stop(stats.cases_passed, start.elapsed()) { + break; + } + let case_seed = run_seed.wrapping_add(stats.cases_passed as u64); + with_context( + test.name(), + Some(run_seed), + dataset_seed, + Some(case_seed), + || test.run_case(&dataset, case_seed), + ) + .await?; + stats.cases_passed += 1; + } + } + stats.elapsed = start.elapsed(); + Ok(stats) +} + +/// Print before invoking the operation so libtest captures context even for +/// assertion panics. Returned errors retain their original source; panics unwind +/// normally with their original payload and hook, without process-global changes. +async fn with_context>>( + name: &str, + run_seed: Option, + dataset_seed: u64, + case_seed: Option, + operation: impl FnOnce() -> F, +) -> Result { + let phase = if case_seed.is_some() { + "run_case" + } else { + "prepare_dataset" + }; + let run_context = run_seed + .map(|seed| format!("Run seed: {seed}\n")) + .unwrap_or_default(); + let instruction = if let Some(case_seed) = case_seed { + format!( + "Case seed: {case_seed}\nReproduce with this test's configuration: \ + test.reproduce({dataset_seed}, {case_seed}).await" + ) + } else { + format!( + "Retry with this test's configuration: \ + test.prepare_dataset({dataset_seed}).await" + ) + }; + let context = format!( + "Fuzz test: {name}\nPhase: {phase}\n{run_context}Dataset seed: {dataset_seed}\n{instruction}" + ); + eprintln!("{context}"); + operation().await.map_err(|error| error.context(context)) +} diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs index 55d801e3b5ba2..be0f4ff806f3f 100644 --- a/test-utils/src/lib.rs +++ b/test-utils/src/lib.rs @@ -12,6 +12,7 @@ use rand::{Rng, SeedableRng}; pub mod array_gen; mod data_gen; +pub mod fuzz; mod string_gen; pub mod tpcds; pub mod tpch; diff --git a/xtask/src/ci_steps.rs b/xtask/src/ci_steps.rs index 4ece399a79a2e..a4cfc8f831fc8 100644 --- a/xtask/src/ci_steps.rs +++ b/xtask/src/ci_steps.rs @@ -15,9 +15,8 @@ type StepRunner = fn(&StepContext, &[String]) -> StepResult; const CI_COMMAND: &str = "cargo xtask ci step"; const CI_SHORTCUT: &str = "cargo ci-step"; const TPCH_DATA_DIR: &str = "datafusion/sqllogictest/test_files/tpch/data"; -/// The `datafusion-physical-plan` unit tests that the `fuzz` step runs. -const FUZZ_TESTS: &[&str] = &[ - "spill::spill_pool::tests::spill_pool_scenario_fuzz", +/// Physical-plan fuzzers awaiting migration to the centralized driver. +const LEGACY_FUZZ_TESTS: &[&str] = &[ "joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz", "joins::key_range_bitmap::tests::union_fuzz", ]; @@ -44,7 +43,7 @@ static CI_STEPS: &[StepInfo] = &[ }, StepInfo { command: "test", - help_usage: "test [--explain]", + help_usage: "test [--explain]", help_examples: &["test cli", "test workspace"], help_description: "Run a CI test suite", runner: StepContext::run_test, @@ -444,10 +443,7 @@ impl StepContext { } } "extended" => { - cmd_builder.doc( - "Run tests behind the `extended_tests` feature. The `workspace`\n\ - step runs the remaining tests.", - ); + cmd_builder.doc("Run the core extended tests."); cmd_builder .program("cargo") .args([ @@ -456,6 +452,8 @@ impl StepContext { "ci", "-p", "datafusion", + "-p", + "datafusion-fuzz", "--test", "extended_tests", "--features", @@ -464,10 +462,28 @@ impl StepContext { .env("RUST_BACKTRACE", "1"); } "fuzz" => { - cmd_builder.doc( - "Run more random fuzzer scenarios than the `workspace` step. This\n\ - takes about a minute for the spill pool, seconds for the others.", - ); + cmd_builder.doc("Run the centralized fuzz tests."); + // Match extended's packages and features to reuse compiled dependencies. + cmd_builder + .program("cargo") + .args([ + "test", + "--profile", + "ci", + "-p", + "datafusion", + "-p", + "datafusion-fuzz", + "--test", + "fuzz", + "--features", + "avro,backtrace,extended_tests,parquet_encryption", + ]) + .env("RUST_BACKTRACE", "1"); + } + "fuzz-legacy" => { + cmd_builder + .doc("Run unmigrated physical-plan fuzz tests with more iterations."); cmd_builder .program("cargo") .args([ @@ -480,9 +496,8 @@ impl StepContext { "--", "--exact", ]) - .args(FUZZ_TESTS) + .args(LEGACY_FUZZ_TESTS) .env("RUST_BACKTRACE", "1") - .env("DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS", "1000") .env("DATAFUSION_ROUTED_LOOKUP_FUZZ_ITERATIONS", "10000") .env("DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS", "10000"); } @@ -565,7 +580,7 @@ mod tests { use super::*; use std::ffi::OsStr; - const TEST_USAGE: &str = "usage: cargo xtask ci step test [--explain]"; + const TEST_USAGE: &str = "usage: cargo xtask ci step test [--explain]"; fn args(args: &[&str]) -> Vec { args.iter().map(|arg| (*arg).to_string()).collect() @@ -617,8 +632,8 @@ mod tests { /// A filter that matches no test passes without running anything, so a /// renamed or moved fuzzer would silently stop running in CI. #[test] - fn fuzz_tests_exist() { - for path in FUZZ_TESTS { + fn legacy_fuzz_tests_exist() { + for path in LEGACY_FUZZ_TESTS { let (module, test) = path .rsplit_once("::tests::") .expect("test path has a `tests` module"); @@ -629,15 +644,15 @@ mod tests { .unwrap_or_else(|e| panic!("cannot read {}: {e}", file.display())); assert!( source.contains("mod tests {") && source.contains(&format!("fn {test}(")), - "{} has no test `{test}` in `mod tests`; update FUZZ_TESTS", + "{} has no test `{test}` in `mod tests`; update LEGACY_FUZZ_TESTS", file.display() ); } } - /// The `extended` step runs only the `extended_tests` test target, which - /// requires the `extended_tests` feature. A test gated on the feature with a - /// `cfg` in any other target would never run in CI. + /// Extended workflows explicitly select targets gated by Cargo's + /// `required-features`. A test hidden behind a `cfg` in another target + /// would never run in those workflows. #[test] fn extended_tests_feature_only_gates_whole_targets() { fn visit(dir: &Path, found: &mut Vec) { @@ -661,8 +676,9 @@ mod tests { ); assert!( found.is_empty(), - "put tests behind `extended_tests` in the `extended_tests` test target \ - (see datafusion/core/Cargo.toml), not behind a `cfg` in {found:?}" + "gate complete test targets with `required-features = [\"extended_tests\"]` \ + (see datafusion/core/Cargo.toml and datafusion/fuzz/Cargo.toml), \ + not with a `cfg` in {found:?}" ); } @@ -937,14 +953,14 @@ cargo test \ command: &["test", "extended"], expected: |actual| { insta::assert_snapshot!(actual, @r" -# Run tests behind the `extended_tests` feature. The `workspace` -# step runs the remaining tests. +# Run the core extended tests. cd /workspace && \ RUST_BACKTRACE=1 \ cargo test \ --profile ci \ -p datafusion \ +-p datafusion-fuzz \ --test extended_tests \ --features avro,backtrace,extended_tests,parquet_encryption ") @@ -954,19 +970,34 @@ cargo test \ command: &["test", "fuzz"], expected: |actual| { insta::assert_snapshot!(actual, @r" -# Run more random fuzzer scenarios than the `workspace` step. This -# takes about a minute for the spill pool, seconds for the others. +# Run the centralized fuzz tests. + +cd /workspace && \ +RUST_BACKTRACE=1 \ +cargo test \ +--profile ci \ +-p datafusion \ +-p datafusion-fuzz \ +--test fuzz \ +--features avro,backtrace,extended_tests,parquet_encryption +") + }, + }, + Cmd { + command: &["test", "fuzz-legacy"], + expected: |actual| { + insta::assert_snapshot!(actual, @r" +# Run unmigrated physical-plan fuzz tests with more iterations. cd /workspace && \ DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS=10000 \ DATAFUSION_ROUTED_LOOKUP_FUZZ_ITERATIONS=10000 \ -DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS=1000 \ RUST_BACKTRACE=1 \ cargo test \ --profile ci \ -p datafusion-physical-plan \ --lib \ --- --exact spill::spill_pool::tests::spill_pool_scenario_fuzz joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz joins::key_range_bitmap::tests::union_fuzz +-- --exact joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz joins::key_range_bitmap::tests::union_fuzz ") }, },