Skip to content

fix: enforce fair memory limits across sibling reservations - #25172

Draft
sunchao wants to merge 1 commit into
apache:dev/chao/codex/pr25172-replay-headroomfrom
sunchao:codex/upstream-consumer-accounting
Draft

sunchao wants to merge 1 commit into
apache:dev/chao/codex/pr25172-replay-headroomfrom
sunchao:codex/upstream-consumer-accounting

Conversation

@sunchao

@sunchao sunchao commented Sep 10, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Fixes FairSpillPool admitting allocations beyond a consumer's fair share or the pool's remaining capacity.

Depends on #25383. This draft is based on that PR's branch; merge the prerequisite first, then rebase and retarget this PR to main.

Rationale for this change

An operator's memory allowance should cover all of its reservations together. FairSpillPool divides the memory available after non-spillable allocations equally among registered spillable consumers. An operator can use several reservations under one consumer to track different parts of its work. For example, during aggregate spill replay, the buffers used to merge spill files coexist with the group keys and accumulator state used to finish the aggregation. These reservations belong to the same consumer and need to share its allowance.

Today, the pool checks only the reservation requesting more memory. Consider a 200 MiB pool with two spillable consumers and no non-spillable allocations: each consumer has a 100 MiB share. One consumer holds 60 MiB in one reservation and 40 MiB in another, so it has used its entire share. Growing the second reservation by 10 MiB is nevertheless accepted because that reservation would hold only 50 MiB. The consumer now holds 110 MiB, exceeding its share even though each reservation individually fits. Splitting an operator's memory across more reservations can therefore let it take memory intended for other consumers.

There is a second problem when consumers register at different times. Suppose a single consumer fills a 100 MiB pool, then a second consumer registers. Their calculated shares become 50 MiB each, but registration does not release any of the first consumer's memory. The current check still lets the newcomer reserve 10 MiB because it is below its own share, taking total reservations to 110 MiB. A fair-share check alone cannot establish that the pool has space.

Both cases undermine the allocation failures that operators rely on to decide when to spill or stop. The pool needs to consider both who owns the memory and how much capacity actually remains, so allocations approved through try_grow respect these limits.

What changes are included in this PR?

FairSpillPool now checks both the consumer's shared allowance and the pool's remaining capacity before approving try_grow. It tracks the consumer's combined reservations and uses that total in the decision. Organizing memory into separate reservations, or transferring it between them, no longer creates additional allowance. Releasing memory reduces the shared total and makes that capacity available again.

In the first example, the extra 10 MiB is rejected because the consumer already holds its full 100 MiB share. In the second, the newcomer's request is rejected because the pool is full; after the first consumer releases 50 MiB, the newcomer can reserve its 50 MiB share. Existing reservations remain valid when another consumer joins, while subsequent growth is subject to the updated shares and remaining pool capacity.

Aggregate replay illustrates why this fix is paired with #25383. Once merge buffers and aggregate state are charged to the same allowance, a merge that uses all of that allowance leaves no room to process its output. The prerequisite makes aggregate spill merges leave room for that work. This PR then enforces the shared accounting in the pool. If another operator permanently takes the remaining memory, replay must report exhaustion when it cannot proceed within the limit.

Are these changes tested?

Six focused pool regressions cover shared reservations, competing consumers, reservation transfers and release, a newly registered consumer facing a full pool, infallible-growth accounting, and integer overflow. The aggregate pressure regression checks the end-to-end consequence: after output has begun and a competitor takes the remaining memory, replay reports ResourcesExhausted, reservations stay within the pool limit, and memory and spill files are released.

Validation results and environment

Previously validated on top of #25383:

  • 35 memory-pool tests, 2,263 physical-plan tests, and 2,281 core/CLI tests passed.
  • The permanent-pressure regression passed with the required resource error, bounded peak reservations, and complete memory/disk cleanup.
  • All 520 SQL logic test files passed.
  • cargo fmt --all, strict all-target/all-feature Clippy, and the full ./dev/rust_lint.sh passed.

The recorded local validation used Rust 1.98.1 and upstream revision 22651d24 with its unchanged dependency lockfile. It did not use newer main dependency versions, which were unavailable in the local registry. After the prerequisite merges, rebase this accounting change and rerun CI against main.

Are there any user-facing changes?

Fallible allocation requests through try_grow are rejected when the combined reservations would exceed the consumer's fair share or the requested growth would exceed the pool's remaining capacity. Queries that previously relied on those allocations being accepted may spill sooner or return ResourcesExhausted.

Infallible grow keeps its existing contract, including permitting growth beyond the configured limit; that usage is included in subsequent allocation checks. No public API is added.

@github-actions github-actions Bot added core Core DataFusion crate execution Related to the execution crate physical-plan Changes to the physical-plan crate labels Sep 10, 2026
@codecov-commenter

codecov-commenter commented Sep 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.94%. Comparing base (1ff4126) to head (b075256).

Additional details and impacted files
@@                            Coverage Diff                             @@
##           dev/chao/codex/pr25172-replay-headroom   #25172      +/-   ##
==========================================================================
- Coverage                                   82.02%   81.94%   -0.08%     
==========================================================================
  Files                                        1136     1136              
  Lines                                      439539   429614    -9925     
  Branches                                   439539   429614    -9925     
==========================================================================
- Hits                                       360522   352067    -8455     
+ Misses                                      58242    56491    -1751     
- Partials                                    20775    21056     +281     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@sunchao
sunchao force-pushed the codex/upstream-consumer-accounting branch from 9480bfa to b075256 Compare September 16, 2026 16:16
@sunchao
sunchao changed the base branch from main to dev/chao/codex/pr25172-replay-headroom September 16, 2026 16:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate execution Related to the execution crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants