fix: RepartitionExec deadlock when its spill pool holds two open files - #73
Closed
adriangb wants to merge 3 commits into
Closed
fix: RepartitionExec deadlock when its spill pool holds two open files#73adriangb wants to merge 3 commits into
adriangb wants to merge 3 commits into
Conversation
Add a test that runs the query from issue apache#24883 under a 4 MB memory limit, with two partitions and 64-row batches. The query is a grouped COUNT(DISTINCT) on a Utf8View column. It sends `RepartitionExec` into its spill path. On the current code, the first attempt does not complete. The test stops the attempt after 20 seconds and fails. The test does 12 attempts. One attempt that completes is not sufficient for a pass. The test uses a memory limit that is too small for the query. This is necessary to make `RepartitionExec` spill. Once in a few thousand attempts, the greedy memory pool refuses an allocation of an aggregate, because the repartition reservations hold the pool. The test accepts a resources-exhausted error. That error is a correct result of the limit and not the deadlock. Ref: apache#24883 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9
The spill pool had three bugs since its unbuffered reads: apache#20683, apache#23447 and apache#24883. In each bug, two writers were in `push_batch` at the same time. No unit test could cause that overlap. This commit adds a test double and two tests: - `IoHookFactory` is a `TempFileFactory` that runs a hook before each disk write. A writer writes while it holds the lock of its file. Thus a hook that blocks can hold a writer at the point where it has a file checked out. The factory can also delay the reads of a file. - `test_reader_does_not_wait_on_drained_file_while_another_has_data` holds writer 1 inside its first write while writer 2 pushes a batch. Then the reader must give both batches while both writers are alive. - `spill_pool_scenario_fuzz` runs random scenarios from a seed: one to four writers, file rotation after each batch, after some batches or never, and one to six phases. In each phase, a random set of writers pushes one to three batches each. A pauser holds each writer of the phase inside its first write, so the pushes overlap. Then the reader reads a random number of the available batches. Some phases drop a writer. Each scenario checks the contract that `RepartitionExec` depends on: - When no push is in progress, the reader can read each pushed batch without more writer activity. - The reader does not signal EOF while a writer is alive. - After the last writer is dropped, the reader gives the backlog and then EOF. - The reader gives exactly the pushed batches. With one writer, it gives them in push order. - Disk usage is zero after the reader is dropped. The two tests are red on this commit. The fuzzer also fails on these versions of the pool: - The code before apache#23522. Seed 4 shows the apache#23447 stall after the drop of all writers. Seeds 0 and 1 show the apache#24883 stall. All 16 pool tests of that time pass on that code. - A draft of the fix in the next commit that skips a file with a read in progress. Seed 2 shows a FIFO violation with one writer. The default of 50 seeds takes approximately three seconds. The extended CI job runs 1000 seeds. `DATAFUSION_SPILL_POOL_FUZZ_SEED` replays one seed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9
In non-preserve-order mode, all input tasks of a `RepartitionExec` share one spill pool per output partition. When two input tasks spilled at the same time, the pool had two open spill files. The `SpillPoolReader` read only the oldest file. When that file had no unread batch and was not finished, the reader waited for it. But the batch that the reader needed was in the newer file. When each distributor channel had data, the gate closed. Then both input tasks stopped in `send`, no sink was dropped, and nothing could wake the reader. The reader now keeps all the files that it received, oldest first. It gives the first batch that is available in one of them. 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. Thus one writer gets its batches in FIFO order. This commit makes the tests from the two commits before it green. The fuzzer passes 500 seeds. Ref: apache#24883 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
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
RepartitionExeccannot 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.
test: reproduce the RepartitionExec spill pool deadlock with a SQL query. This is the query from the issue as a test indatafusion/core/tests/memory_limit/repartition_mem_limit.rs. Onmain, attempt 0 does not complete and the test stops it after 20 seconds.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 addsspill_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.fix: RepartitionExec deadlock when its spill pool holds two open files.SpillPoolReadernow 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:
main(has apache#24883)[1, 3, 4, 2, ...].What is the testing strategy for this PR?
DATAFUSION_SPILL_POOL_FUZZ_ITERATIONSsets the count.DATAFUSION_SPILL_POOL_FUZZ_SEEDreplays one seed.cargo fmtandcargo clippy --all-targets --all-features -- -D warningspass 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 apache#23522 avoided on purpose. It is not part of this PR.
Are there any user-facing changes?
No.
The
branch-54line has the same bug, because apache#23654 backported the same file model. A backport of this fix is necessary.🤖 Generated with Claude Code
https://claude.ai/code/session_01RauD5PqPCp4y92RUZUD6C9
Generated by Claude Code