fix: enforce fair memory share per reservation - #5466
Conversation
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Summary
Reviewed the full single-file diff at head a658c1960ea2ae27cfa4584e77800163fab16645 against base bd7dc601aa9a99289530598d4778c40e6ded7772. The reservation-based fair-share calculation fixes the reported premature rejection, but one P2 remains: the change also removes enforcement of Comet's aggregate configured memory budget. That finding is reported inline.
Prior state and problem
The previous implementation divided pool_size by the registered-consumer count, then compared that per-consumer allowance with state.used, which is the total across the pool. With a 32 MiB pool, two consumers holding 10 MiB and 6 MiB exhaust the old 16 MiB comparison even though the requester still has 10 MiB of its own allowance available. DataFusion updates a reservation's size after each successful grow, so subsequent sequential calls do include previous allocations; the removed comment incorrectly implied otherwise. The pool also inherited MemoryLimit::Unknown despite having a configured size.
Design approach
The new helper compares the requesting reservation's current size plus the additional request against pool_size / num. Saturating addition avoids wraparound in that comparison, and memory_limit() now exposes the configured size through MemoryLimit::Finite. The existing policy of counting every registered consumer is retained, leaving the spillability policy in #5465 separate from this correction.
Correctness / compatibility analysis
The sequential reservation calculation agrees with DataFusion 54.1's API ordering, and the finite-limit API is supported by the pinned dependency. Registration, aggregate bookkeeping, JNI acquisition/release, and partial-grant rollback are otherwise unchanged. However, individual reservation bounds do not imply an aggregate bound: registrations can arrive after earlier allocations, and new_empty()/split() create separate reservation sizes while sharing one consumer registration. Spark's backing memory manager does not enforce Comet's smaller spark.comet.exec.memoryPool.fraction budget, so the inline P2 needs an independent aggregate check.
Validation separates actual CI from the focused harness. The Rust CI job tested the merge of these pins and passed 901 tests with 4 skipped, including the added regression; the current PR checks are 64 successful and 9 skipped. A separate Linux harness used the unchanged head/base pool implementations and checksum-verified DataFusion reservation source: all six checks passed in both debug and optimized modes, including assertions that the head accepts the excess reservations while the base rejects them. That harness mocked JNI grants/releases and supporting mutex/error adapters; it was not a native Spark integration test.
Key design decisions
Keeping the fair-share calculation tied to the requester is the right correction for the original 10+6 example. Keeping state.used as the total remains necessary for bookkeeping and for enforcing the configured aggregate budget independently of fairness. The unchanged spillability policy and unchanged JNI rollback behavior keep this PR's intended scope narrow; neither needs redesign to address the P2.
Implementation sketch
fair_limit_exceeded returns the requesting reservation's usage and its calculated allowance when growth would exceed that allowance. try_grow calls this helper while holding the existing state lock, then follows the existing JNI acquisition and accounting path. The new test exercises the helper with real DataFusion reservations backed by an unbounded pool, plus the finite-limit accessor on a Comet pool with a null JVM handle. It does not call the Comet pool's allocating try_grow path, so it does not cover aggregate admission or JNI behavior.
Behavioral changes worth calling out
The requesting consumer can now use its remaining allowance even when another consumer already holds memory, which addresses the original premature-spill condition. Callers can also observe the configured finite limit instead of an unknown limit. Without the additional aggregate guard, some valid reservation sequences can consume the headroom deliberately reserved by a fraction below 1.0; the inline example reaches 40 MiB against a configured 32 MiB pool. The reported workload timings were not independently rerun, so this review makes no performance claim from them.
Suggested improvements
Address the P2 by retaining the new per-reservation comparison and separately rejecting growth that would push aggregate state.used above pool_size. Add a regression that exercises actual pool admission when a consumer registers after another has allocated, or when sibling reservations share a registration, while preserving the 10+6+10 success case. This can remain a narrow capacity fix without taking on the spillability-policy work in #5465.
| let used = state.used; | ||
| if limit < used + additional { | ||
| if let Some((used, limit)) = | ||
| fair_limit_exceeded(self.pool_size, num, reservation, additional) |
There was a problem hiding this comment.
[P2] Preserve the aggregate configured pool limit
This replaces the only aggregate admission bound with a per-reservation check. With spark.memory.offHeap.size=64m and spark.comet.exec.memoryPool.fraction=0.5, pool_size is 32 MiB. A can register and reserve 24 MiB; after B registers, B's request for 16 MiB now passes (0 + 16 <= 32 / 2), leaving 40 MiB reserved against the configured 32 MiB budget. The base rejects B's request. Similarly, new_empty() shares its consumer registration while starting a separate size counter, so sibling reservations can collectively exceed the budget; DataFusion's native sort merge uses this pattern.
Spark's JVM memory manager only applies its larger off-heap budget, not Comet's fraction, so it can grant these requests. This consumes the headroom that the tuning guide explicitly reserves to avoid OOM from untracked allocations. Please retain the corrected requester check and add a separate aggregate state.used versus pool_size bound; this does not require the spillability-policy change in #5465.
Both excess-allocation sequences were reproduced against unchanged head/base pool code in a harness with mocked JNI grants; Spark's larger backing limit was verified separately in source.
|
I'm currently experimenting with #4582 |
3 participants
Which issue does this PR close?
Addresses finding 1 and the fair-pool portion of finding 16 in #5212.
Follow-up #5465 tracks spillable versus unspillable consumer accounting separately.
Rationale for this change
CometFairMemoryPoolcalculates a per-consumer fair share but compares it with pool-wide usage. With a 32 MiB pool and two consumers, one reservation holding 10 MiB and the requester holding 6 MiB make total usage equal the 16 MiB fair share. Main then rejects any positive growth from the requester, even though that requester has 10 MiB of fair-share headroom.The fair pool also has a fixed configured capacity, so reporting an unknown memory limit is inaccurate.
What changes are included in this PR?
MemoryLimit::Finite(pool_size)forCometFairMemoryPool.How are these changes tested?
Focused native test:
cargo test --release -p datafusion-comet fair_share_uses_requesting_reservation_and_reports_pool_limit --libFormatting and diff checks:
A warmed local A/B used
local[1], a 32 MiB fair pool, one sort partition, three warmups, and seven measured runs:Spill counts and bytes were identical across the seven runs for each build. Runtime ranges overlap, so this PR makes no speed claim; the demonstrated benefit is restoring the requesting consumer's fair-share capacity and avoiding premature spills.