adriangb opened a new issue, #24883:
URL: https://github.com/apache/datafusion/issues/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:
```console
$ cargo build -p datafusion-cli
$ ./target/debug/datafusion-cli -f deadlock.sql
```
`deadlock.sql`:
```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 `SELECT`s, 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:**
```sql
set datafusion.execution.target_partitions = 1; -- instead of 2
```
0 hangs in 10 runs, and all five `SELECT`s 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.
<details>
<summary>Full annotated stacks, and the lldb thread dump</summary>
The physical plan for the MRE, for orientation (bottom up):
```
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
```
**Consumer side.** Line numbers are from the unpatched tree; the capture
itself reported
`spill_pool.rs:770` because the instrumentation added lines above it.
```
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
```
10 DataFusion frames of 86; the other 76 are `futures`, `std` and tokio
task-harness
boilerplate. Both of the query's two spill-pool readers produced this
identical signature.
**Producer side.**
```
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
```
4 DataFusion frames of 68. All 8 sender captures from this run share this
one signature.
**lldb thread dump of the same wedged process.** 15 threads, 4 distinct
stacks, and
**zero DataFusion frames anywhere in the dump**:
| threads | stack | leaf |
| --- | --- | --- |
| 11 `tokio-rt-worker` | `worker::run` -> `park_internal` (worker.rs:887) |
`park_condvar` -> `__psynch_cvwait` |
| 1 `tokio-rt-worker` | same, holding the I/O driver | `kevent` |
| 2 `tokio-rt-worker` (blocking pool) | `blocking::pool::run` ->
`wait_timeout` (pool.rs:529), idle | `__psynch_cvwait` |
| 1 `main` | `datafusion_cli::main` (main.rs:196) -> `Runtime::block_on` ->
`park` | `__psynch_cvwait` |
```
* 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
```
Nothing is spinning and nothing is waiting on I/O, so this is a deadlock
rather than a
slow query. But no thread stack can name the DataFusion code that is stuck,
which is why
the `Pending`-site captures above are the useful artifact.
`lsof` on the same process independently corroborates the pool state:
several spill files
open under the OS temp directory, some carrying two write handles and a read
handle on the
same file.
</details>
### 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:
```rust
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.
<details>
<summary>tests/repartition_spill_deadlock.rs</summary>
```rust
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}"),
}
}
}
```
</details>
### Additional context
- Reproduced on `apache/main` at `ad5820c84ed00d31125842e62a5124303a39c8d2`
(workspace version 55.0.0), macOS arm64, Rust 1.97.0 per
`rust-toolchain.toml`.
- Related: #23447 and its fix #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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
