RepartitionExec deadlocks under a memory limit when its spill pool holds two open files #24883

Describe the bug

Under a memory limit, a grouped aggregation running on two or more partitions can
deadlock permanently. Every tokio worker parks, CPU drops to 0%, and the query neither
returns nor errors. It reproduces from a stock datafusion-cli session with three SET
statements and no data files.

To Reproduce

Build and run datafusion-cli from a checkout, then paste the block below:

$ cargo build -p datafusion-cli
$ ./target/debug/datafusion-cli -f deadlock.sql

deadlock.sql:

set datafusion.execution.target_partitions = 2;
set datafusion.execution.batch_size = 64;

create table t as
select v % 13 as g, md5(cast(v % 337 as varchar)) as trace_id
from generate_series(1, 40000) as s(v);

set datafusion.runtime.memory_limit = '2M';

-- Each of these takes about 0.1s when it completes. One of them never returns.
select g, count(distinct trace_id) as n from t group by g;
select g, count(distinct trace_id) as n from t group by g;
select g, count(distinct trace_id) as n from t group by g;
select g, count(distinct trace_id) as n from t group by g;
select g, count(distinct trace_id) as n from t group by g;

No CLI flags are needed. datafusion-cli already gives you a multi-threaded tokio
runtime, and its default RuntimeEnv has the disk manager enabled
(DiskManagerMode::OsTmpDirectory), so SET covers the rest.

Reproduction rate: 20 of 20 runs on an unpatched tree (each run capped at 25s;
apache/main at ad5820c84ed00d31125842e62a5124303a39c8d2, macOS arm64, 12 cores,
cargo build -p datafusion-cli). It also hangs 10 of 10 on an optimized build
(cargo build --profile release-nonlto -p datafusion-cli).

Time to hang: 0.11s to 2.46s from process start, mean 0.29s, so a few seconds of no
output is enough to conclude it is stuck. The wedged session has printed only the
0 row(s) fetched lines for the SET statements and the CREATE TABLE, or those plus a
few completed SELECTs, and then nothing.

A single SELECT hangs 10 of 12 runs. The statement is repeated five times only to make
one paste reliable; the first SELECT is usually the one that wedges.

Negative control, same session, same memory limit:

set datafusion.execution.target_partitions = 1;   -- instead of 2

0 hangs in 10 runs, and all five SELECTs return the correct 13 rows every time, in
about 0.25s for the whole script.

The memory limit is not delicate. With the block above, 256K, 512K, 1M, 2M, 3M
and 4M all hang 6 of 6; 8M hangs 2 of 6.

Note that datafusion.execution.target_partitions must be set explicitly. Left at the
default (one per core, 12 on this machine) the same query needs more memory before
RepartitionExec reaches its spill path, and 2M does not hang.

Where the two sides are parked

lldb -p <pid> -o "thread backtrace all" on the wedged process is not enough on its own,
and it is worth saying why: both deadlocked participants are suspended async tasks,
not blocked threads. Their poll_next returned Poll::Pending and the stack unwound, so
the state machines sit on the heap inside their tokio tasks and appear on no thread stack.
The thread dump shows only the scheduler parking (details below).

Everything below comes from one wedged process of a plain debug build
(cargo build -p datafusion-cli). Frames resolved fully without
-C force-frame-pointers or a custom profile.

The frames that matter are the ones at the moment each side last returned Pending.
Capturing std::backtrace::Backtrace::force_capture() at the two park sites, with a
global sequence number, gives the last four events of the query and then silence:

seq=331  SpillPoolReader parked  pool=0x...61e90  files_len=2 open_write_len=2 writers=2
                                 head_file=0x...50c50 batches_read=44 batches_written=44 writer_finished=false
seq=332  SpillPoolReader parked  pool=0x...61d90  files_len=2 open_write_len=2 writers=2
                                 head_file=0x...70c50 batches_read=29 batches_written=29 writer_finished=false
seq=333  DistributionSender parked on the distributor gate  output_channel=0
seq=334  DistributionSender parked on the distributor gate  output_channel=1
<nothing further, 0% CPU>

Both readers park on a head file they have fully drained (batches_read == batches_written)
that is not finished, while their pool holds two files and two open write files. Then both
producers park on the gate. The two stacks, DataFusion frames only:

consumer (parked, seq=332)                         producer (parked, seq=333 and 334)
  SpillPoolReader::poll_next                         DistributionSender::poll
    spill_pool.rs:744  Poll::Pending                   distributor_channels.rs:224  Poll::Pending
  PerPartitionStream::poll_next_inner                OutputChannel::send
    repartition/mod.rs:2434  ReadingSpilled             repartition/mod.rs:239  sender.send(..).await
  ...                                                RepartitionExec::pull_from_input
  RepartitionExec::pull_from_input                     repartition/mod.rs:2179  output_channel.send(batch).await
    repartition/mod.rs:2158  stream.next().await

Same RepartitionExec on both sides: its PerPartitionStream is blocked in
StreamState::ReadingSpilled on its own SpillPoolReader, and its input tasks are blocked
in OutputChannel::send on its own distributor gate.

DataSourceExec (memory, 40000 rows)
  AggregateExec Partial            group_by: g, trace_id as alias1
    RepartitionExec Hash([g, alias1], 2)      <-- this one deadlocks
      AggregateExec FinalPartitioned  group_by: g, alias1
        AggregateExec Partial           group_by: g, aggr: count(alias1)
          RepartitionExec Hash([g], 2)
            AggregateExec FinalPartitioned
              ProjectionExec

 3  SpillPoolReader::poll_next          spill/spill_pool.rs:744
        the `Poll::Pending` arm: head file drained, writer not finished, waker registered
 6  PerPartitionStream::poll_next_inner repartition/mod.rs:2434
        StreamState::ReadingSpilled, polling self.spill_stream, blocks to preserve ordering
 7  PerPartitionStream::poll_next       repartition/mod.rs:2469
11  RecordBatchStreamAdapter::poll_next stream.rs:469
14  GroupedHashAggregateStream::handle_reading_input   aggregates/hash_stream.rs:1139
15  GroupedHashAggregateStream::poll_next              aggregates/hash_stream.rs:1511
        AggregateExec FinalPartitioned (group_by: g, alias1), pulling the repartition output
18  GroupedHashAggregateStream::handle_reading_input   aggregates/hash_stream.rs:565
19  GroupedHashAggregateStream::poll_next              aggregates/hash_stream.rs:956
        AggregateExec Partial (aggr: count(alias1))
23  RepartitionExec::pull_from_input    repartition/mod.rs:2158
        `stream.next().await` in the *upper* RepartitionExec's input task
24  spawned task body                   common-runtime/src/trace_utils.rs:137

 3  DistributionSender::poll            repartition/distributor_channels.rs:224
        the `empty_channels == 0` gate: every output channel is non-empty, so senders park
 4  OutputChannel::send                 repartition/mod.rs:239
        `self.sender.send(Some(payload)).await`, after `push_batch` spilled the batch
 5  RepartitionExec::pull_from_input    repartition/mod.rs:2179
        `output_channel.send(batch).await`
 6  spawned task body                   common-runtime/src/trace_utils.rs:137

* thread #1, name = 'main'
    frame  #0: __psynch_cvwait
    frame  #9: park at park.rs:111
    frame #15: block_on<...datafusion_cli::main::{async_block_env#0}...> at park.rs:288
    frame #22: datafusion-cli`main at main.rs:196:15

  thread #2, name = 'tokio-rt-worker'          (and 10 more identical)
    frame  #0: __psynch_cvwait
    frame  #9: park_condvar at park.rs:203
    frame #12: park_internal at worker.rs:887
    frame #14: run at worker.rs:614

Root cause

The deadlock is in RepartitionExec's spill path.

When RepartitionExec cannot reserve memory for an output batch it writes the batch to
the output partition's SpillPool and sends a RepartitionBatch::Spilled marker down
the distributor channel instead of the batch (OutputChannel::send,
datafusion/physical-plan/src/repartition/mod.rs). On the read side a Spilled marker
moves the output stream into StreamState::ReadingSpilled, where it blocks on the
spill stream to preserve ordering:

Poll::Pending => {
    // Spilled batch not ready yet, must wait
    // This preserves ordering by blocking until spill data arrives
    return Poll::Pending;
}

In non-preserve_order mode all input tasks share one pool per output partition via
SpillPoolWriter::new_sink (PartitionSpillWriters::Shared,
repartition/mod.rs:206). SpillPoolSink::push_batch pops the front of
SpillPoolShared::open_write_files, releases the pool lock for the file I/O, and pushes
the file back afterwards (spill/spill_pool.rs:216-283). Two input tasks that overlap
inside that window leave the pool with two open spill files: the second finds the
deque empty and creates a new file. This multiple-open-file model was introduced
deliberately in #23522.

SpillPoolReader however drains SpillPoolShared::files strictly head first and only
advances past the head when that file yields Ready(None), which requires
writer_finished == true on the head (spill_pool.rs:715-760). Below
max_spill_file_size_bytes (128 MB by default) nothing rotates, so a file is finished
only when the last sink is dropped (SpillPoolSink::drop, spill_pool.rs:143-176).

Batches written to the second open file are therefore unreachable while the first is
still open, which closes this cycle:

1. An output partition's reader pops a Spilled marker whose batch went into the second
   open file. Its head file has batches_read == batches_written and
   writer_finished == false, so the reader parks and stops draining that partition's
   memory channel.
2. Every distributor channel ends up non-empty, so the global gate in
   distributor_channels.rs closes (empty_channels == 0, line 220).
3. Both input tasks park in DistributionSender::send. No sink is dropped and no further
   batch is written to the head file.
4. Nothing can wake anybody.

#23522 fixed the drop-time variant of this (finish every open file when the last sink is
dropped, closing #23447). The steady-state variant above is not covered, because the
writers never reach their drop: they are blocked on the gate that only the parked reader
could open.

A separate instrumented run, viewed chronologically, shows the second open file being
created while the first is still open, and then the gate closing on both channels:

DBG new_spill_file pool=0x20000d10510 file=0x200081a0c50 files_len=2 open_write_len=1
DBG send_gate_closed chan=1
DBG file_pending file=0x20008030710 read=12 written=15 finished=false
DBG   pool=0x20000d10410 files_len=2 open_write_len=2 writers=2
DBG send_gate_closed chan=0
DBG file_pending file=0x200080301d0 read=18 written=18 finished=false
DBG   pool=0x20000d10510 files_len=2 open_write_len=2 writers=2
DBG send_gate_closed chan=0
DBG file_pending file=0x20008030710 read=16 written=16 finished=false
DBG   pool=0x20000d10410 files_len=2 open_write_len=2 writers=2
<nothing further, 0% CPU>

Both output-partition pools end with files_len=2 open_write_len=2, both readers park on
a head file whose read == written and finished=false, and both producers park on the
distributor gate, reproducing the same end state as the run whose stacks are shown
above.

Ingredients that matter

Each verified by turning it off, all with the SQL block above unless stated:

• More than one tokio worker thread. datafusion-cli supplies this. In the Rust
  reproducer below, a current_thread runtime (plain #[tokio::test]) passes every time
  (3 runs x 12 attempts), and so does worker_threads = 1 (3 runs). With
  worker_threads = 2 or the multi_thread default it deadlocks on the first attempt
  every time (8 runs). push_batch has no await point, so on a single worker the two
  input tasks can never overlap inside it and a second open file is never created.
• target_partitions >= 2. set datafusion.execution.target_partitions = 1: 0 hangs
  in 10 runs from the CLI, and 0 hangs in 120 standalone Rust runs across 6 memory limits
  and 5 column types.
• A memory limit low enough for RepartitionExec to reach its spill path, plus a disk
  manager that permits spilling. With DiskManagerMode::Disabled the query fails with
  a resource error instead: 0 hangs in 180 Rust runs across the whole risky band.
• A spill file that does not rotate.
  set datafusion.execution.max_spill_file_size_bytes = 16384 removes the hang entirely
  (0 hangs in 10 CLI runs, 0 in 120 Rust runs in the worst cells), because the head file
  is finished on rotation and the reader can advance.
• A small batch_size. Dropping set datafusion.execution.batch_size = 64 and using
  the default 8192 removes the hang: fewer, larger batches, and the query no longer
  reaches the repartition spill path at this limit.
• set datafusion.optimizer.repartition_aggregations = false also removes it (0 of 10),
  while set datafusion.optimizer.enable_round_robin_repartition = false does not
  (9 of 10 still hang), so it is the hash RepartitionExec between the partial and final
  aggregate.

The column type is not special. Utf8View and BinaryView hang at a 4 MB limit,
plain Utf8 (used above) hangs from 256 KB to 8 MB. The type only decides how much
accounted memory the query needs and therefore whether the repartition spill path is
reached at a given limit. COUNT(DISTINCT ...) is likewise incidental; it is a
convenient way to build enough aggregate state to reach that path at a small limit.

Expected behavior

The query completes, or fails with a resource error. It should not park forever.

Rust reproducer

Equivalent regression test, if a test in the tree is preferred over the SQL above. It
needs a multi-threaded runtime and fails on the first attempt on an unfixed tree.

use std::sync::Arc;
use std::time::Duration;

use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool};

/// 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(_) => panic!(
                "attempt {attempt}: grouped COUNT(DISTINCT Utf8View) never completed \
                 within {PER_ATTEMPT_TIMEOUT:?}; RepartitionExec spill pool deadlock"
            ),
            Ok(Ok(rows)) => assert_eq!(rows, 13, "attempt {attempt}: wrong row count"),
            Ok(Err(e)) => panic!("attempt {attempt}: query failed: {e}"),
        }
    }
}

Additional context

• Reproduced on apache/main at ad5820c84ed00d31125842e62a5124303a39c8d2
  (workspace version 55.0.0), macOS arm64, Rust 1.97.0 per rust-toolchain.toml.
• Related: SMJ + aggregate finish with one partition but stuck forever when number of partitions is big #23447 and its fix Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters #23522, which introduced the multiple-open-file model and
  fixed the drop-time lost wakeup. The case above is the steady-state one that a
  drop-time fix cannot reach.
• Two directions a fix could take, for discussion: let SpillPoolReader read from any
  open file in the queue rather than only the head, or make push_batch keep a single
  open file per pool so the head is always the file being written.