adriangb opened a new pull request, #24891: URL: https://github.com/apache/datafusion/pull/24891
## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/24883 ## Rationale for this change Under a memory limit, a grouped aggregation on two or more partitions can stop and never complete. The query does not return and does not give an error. All worker threads wait. The issue has the full analysis. In short: when `RepartitionExec` cannot reserve memory for a batch, it writes the batch to a spill pool and sends a marker. The reader then blocks on the pool until it gets a batch. When two input tasks spill at the same time, the pool has two open files. The reader reads only the oldest file. When that file has no unread batch, the reader waits for it. But the batch that the reader needs is in the newer file. When each channel has data, the gate closes and the writers wait too. Nothing can wake the reader. ## What changes are included in this PR? The PR has three commits, in test-first order. 1. `test: reproduce the RepartitionExec spill pool deadlock with a SQL query`. This is the query from the issue as two tests. `datafusion/sqllogictest/test_files/repartition_memory_spill.slt` runs the query eight times and then makes sure that `RepartitionExec` spilled. On `main`, the file hangs in 8 of 8 runs, each stopped after 30 seconds. `repartition_spill_pool_does_not_deadlock` in `datafusion/core/tests/memory_limit/repartition_mem_limit.rs` runs the same query 12 times with a 20-second timeout per attempt. On `main`, attempt 0 does not complete and the test fails with a message. 2. `test: add a scenario fuzzer for the spill pool contract`. This adds a test double that can hold a writer inside its first disk write and can delay the reads of a file. It adds a regression test that holds writer 1 while writer 2 pushes a batch. It adds `spill_pool_scenario_fuzz`, which runs random scenarios from a seed and checks the pool contract: after overlapping pushes the reader can read each pushed batch without more writer activity, it does not signal EOF while a writer is alive, it gives exactly the pushed batches (in order for one writer), and it releases the files. Both tests are red on this commit. The extended CI job now runs 1000 seeds. 3. `fix: RepartitionExec deadlock when its spill pool holds two open files`. `SpillPoolReader` now keeps all the files that it received, oldest first, and gives the first available batch. It skips a file that has no unread batch. It waits for the oldest file that has an unread batch when the read of that file is not complete, so one writer gets FIFO order. The writer side does not change. Both tests are green on this commit. The fuzzer finds each known bug of this pool that a test can reach: | Pool version | Unit tests of that version | Fuzzer | |---|---|---| | Before #23522 (has #23447) | 16 of 16 pass | Seeds 0 and 1: stall with unread batches. Seed 4: stall after the drop of all writers (the #23447 signature). | | `main` (has #24883) | pass | Seed 0: stall with 5 of 13 batches unread. | | A draft of this fix that skips a file with a read in progress | 1 failure | Seed 2: FIFO violation with one writer, `[1, 3, 4, 2, ...]`. | | This fix | pass | 500 seeds pass. | ## What is the testing strategy for this PR? - The SLT file has no timeout, because sqllogictest has none. If the bug comes back, the file hangs. With the fix, 30 runs of the file passed. - The Rust reproducer runs 12 attempts with a timeout. With the fix, 156 runs and 3000 instrumented attempts had 0 hangs. One attempt in 3000 got a resources-exhausted error from the final aggregate. That error is a correct result of the 4 MB limit and the greedy pool, and the test accepts it. - The fuzzer runs 50 seeds by default, in approximately three seconds. `DATAFUSION_SPILL_POOL_FUZZ_ITERATIONS` sets the count. `DATAFUSION_SPILL_POOL_FUZZ_SEED` replays one seed. - The 33 memory-limit tests, the 109 spill and repartition unit tests, and the spill pool doctests pass. - `cargo fmt` and `cargo clippy --all-targets --all-features -- -D warnings` pass on the two changed crates. I also evaluated a second fix that keeps one open write file per pool. It passes the same tests. It serializes writers that spill at the same time on the file lock, which #23522 avoided on purpose. It is not part of this PR. ## Are there any user-facing changes? No. The `branch-54` line has the same bug, because #23654 backported the same file model. A backport of this fix is necessary. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9 -- 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]
