fix: RepartitionExec deadlock when its spill pool holds two open files

.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 \

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 <https://github.com/apache/datafusion/issues/24883>.
+///
+/// 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<usize> {
+        let pool: Arc<dyn MemoryPool> = 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}"),
+            }
+        }
+    }
+}