fix: CometFairMemoryPool compared pool-wide usage against a per-consumer limit

[andygrove]

Which issue does this PR close?

Part of #5212 (position 1).

Rationale for this change

CometFairMemoryPool::try_grow computed the per-consumer fair share pool_size / num but compared it against the pool-wide state.used:

let limit = self.pool_size.checked_div(num).expect("overflow in checked_div");
let used = state.used;              // pool-wide total
if limit < used + additional {
    return resources_err!(...);
}

So a task could never use more than a single consumer's share in total, and usable memory shrank linearly as consumers registered — a plan with 10 memory consumers could use 10% of its off-heap share before erroring or spilling. This is the default off-heap pool type (fair_unified, CometConf.scala:686), so it affects the default configuration whenever spark.memory.offHeap.enabled=true.

The comment justifying the comparison was incorrect:

We use state.used instead of reservation.size() because DataFusion 53+ calls pool.try_grow() before incrementing the reservation's atomic size, so reservation.size() would not include prior grows.

MemoryReservation::try_grow (datafusion-execution 54.1.0, memory_pool/mod.rs:471-475) is:

pub fn try_grow(&self, capacity: usize) -> Result<()> {
    self.registration.pool.try_grow(self, capacity)?;
    self.size.fetch_add(capacity, atomic::Ordering::Relaxed);
    Ok(())
}

reservation.size() therefore does include every prior grow. It excludes only the current additional, which is exactly why upstream FairSpillPool compares reservation.size() + additional > available (memory_pool/pool.rs:249). The corresponding comment on shrink is correct — MemoryReservation::shrink decrements the atomic before calling the pool — and the two cases look to have been conflated during the DataFusion 53 upgrade in #3629.

What changes are included in this PR?

• Compare the fair share against reservation.size() rather than the pool-wide state.used, matching upstream FairSpillPool semantics.
• Extract the decision into a pure check_fair_share function so the policy is unit testable without a JVM, and document that reserved must be the requesting consumer's own usage.
• Replace expect("overflow in checked_div") with a fallback to the whole pool. num cannot be zero while a reservation exists (a consumer registers before it can grow), but the previous code would have panicked with a message describing the wrong failure.
• Include the consumer name and the pool-wide total in the error message, which should make future reports of this class of failure easier to read.

Behaviour is otherwise unchanged: the pool still tracks state.used for reserved() and the shrink guard, and the hard bound is still whatever Spark's TaskMemoryManager actually grants (the acquired < additional path below the check).

How are these changes tested?

Six new unit tests in fair_pool.rs covering the fair-share policy: a sole consumer may use the whole pool, the share divides evenly, each consumer may independently reach its own share (the regression guard), a consumer may not exceed its share, zero registered consumers does not divide by zero, and an oversized request does not overflow.

running 6 tests
test execution::memory_pools::fair_pool::tests::each_consumer_may_reach_its_own_share ... ok
test execution::memory_pools::fair_pool::tests::consumer_may_not_exceed_its_own_share ... ok
test execution::memory_pools::fair_pool::tests::oversized_request_does_not_overflow ... ok
test execution::memory_pools::fair_pool::tests::share_is_divided_evenly_between_consumers ... ok
test execution::memory_pools::fair_pool::tests::sole_consumer_may_use_whole_pool ... ok
test execution::memory_pools::fair_pool::tests::no_registered_consumers_does_not_panic ... ok

Full crate suite passes (141 passed, 4 ignored), and cargo clippy --all-targets is clean.

Testing gap worth flagging: these tests cover the policy function, not the try_grow call site — that reservation.size() rather than state.used is passed in is verified by inspection only. Testing the pool end to end needs a JNI harness, since CometFairMemoryPool::new requires a live CometTaskMemoryManager global ref. That is tracked as the test-coverage item in #5212; I did not want to grow this PR into building that harness.

I have not measured the performance effect on a real workload. The change strictly widens what the pool admits, so I would expect fewer spurious spills and ResourcesExhausted errors on off-heap plans with several memory consumers, but that is reasoning rather than a measurement.

Are there any user-facing changes?

Yes, on the default off-heap configuration: plans that previously spilled or failed with "Failed to acquire N bytes … the fair limit is M bytes" will now be able to use their intended fair share of the pool. No configuration changes.

Follow-ups not in this PR

register/unregister count every consumer, whereas upstream FairSpillPool divides by num_spill (spillable consumers only) and gives unspillable consumers the remainder. Comet's limit is therefore still tighter than upstream's. Aligning that means tracking spillable and unspillable usage separately, which is a larger behavioural change than this fix, so I left it out — noted as a secondary divergence under position 1 in #5212.

[andygrove]

I dug into this against the values the JVM actually passes in. The core defect is real, but I think two parts of the description need correcting, and there is one consequence of the fix that isn't called out.

The defect and the reasoning behind it check out. limit = pool_size / num is a per-consumer share and state.used is the sum over every consumer in the pool, so two consumers each staying inside their own share can still trip the check. And the comment being removed is indeed wrong: MemoryReservation::try_grow calls pool.try_grow() before size.fetch_add in both datafusion-execution 53.1.0 (memory_pool/mod.rs:454-458) and 54.1.0 (:471-475), so reservation.size() does include prior grows. The shrink comment above it is correct (:404-414 decrements before calling the pool), which supports the "conflated during the DF 53 upgrade" reading — blame agrees, #1369 originally passed reservation.size(). docs/source/user-guide/latest/tuning.md:69 also documents the intended semantics as per-operator ("prevents operators from using more than an even fraction of the available memory (i.e. pool_size / num_reservations)"), which is what this restores.

Correction 1: pool_size is the executor-wide off-heap pool, not the task's share. For fair_unified, pool_size = memory_limit = spark.memory.offHeap.size × spark.comet.exec.memoryPool.fraction (CometExecIterator.scala:293, fraction defaults to 1.0 at CometConf.scala:707), passed straight through at config.rs:65,68. memoryLimitPerTask (CometExecIterator.scala:294) is computed, logged, and shipped over JNI, then ignored by both off-heap pool types.

Since the pool instance is per task attempt (TASK_SHARED_MEMORY_POOLS keyed by task_attempt_id), the old effective cap on a task was offHeap.size / num_consumers_in_this_task. On an 8-core executor with ~8 consumers that lands close to the task's fair 1/N share, so the old code accidentally approximated per-task fairness; it only bites when a task has many consumers, and it is looser than fair when a task has few. So "a plan with 10 memory consumers could use 10% of its off-heap share" reads as 10% of the task's share, when it is 10% of the executor-wide pool. The symptom is real; I'd soften the stated severity.

Correction 2: the fix removes the only aggregate per-task cap in the native pool. After this change, num consumers × pool_size/num = pool_size, so as far as this check goes a single task may reserve the entire executor off-heap pool. That is defensible — the real per-task bound is Spark's ExecutionMemoryPool.acquireMemory via CometTaskMemoryManager.acquireMemory, which grants each active task between 1/2N and 1/N and returns a partial grant, hitting the acquired < additional path and triggering a spill; it's the same enforcement greedy_unified relies on with pool_size = 0. But it's worth saying explicitly, because it reads against tuning.md:66-68 ("the shared pool ensures that the combined memory usage stays within the per-task limit") — that guarantee comes from TaskMemoryManager, not from this check, and the buggy code was supplying a second, accidental aggregate bound that this deletes. Please check whether that doc sentence needs rewording.

It also makes the num vs upstream's num_spill divergence noted under follow-ups more load-bearing: with the aggregate cap gone, dividing by all consumers rather than spillable ones is the only fairness knob left.

Smaller notes

• reservation.consumer().name() is public in 54.1.0 (mod.rs:333) and TrackConsumersPool::try_grow forwards the caller's reservation unchanged (pool.rs:570-572), so reservation.size() at the call site is the right consumer's usage. Note that TrackConsumersPool appends the top-consumers report on top of the now-longer message, so the combined error text gets fairly long.
• checked_div(num).unwrap_or(pool_size) silently admits everything at num == 0. Unreachable in practice, so either behaviour is fine, but the fallback is the permissive one rather than the safe one.
• The flagged test gap is the right one to flag: the tests cover check_fair_share, not that reservation.size() is what gets passed to it — which is precisely the line the DF 53 upgrade got wrong.

[andygrove]

Follow-up to my review above: rather than reason about how tight the old cap was, I measured it. I temporarily instrumented CometFairMemoryPool::register to print the consumer count per task pool and ran two Comet queries with spark.memory.offHeap.enabled=true, spark.memory.offHeap.size=512m, fair_unified, on local[5] (instrumentation reverted afterwards).

Consumers registering per task pool:

Plan	Peak num	Consumers
broadcast join + group-by + shuffle	3	HashJoinInput[0], GroupedHashAggregateStream[0], ShuffleRepartitioner[0]
final sort stage	2	ExternalSorter[0], ExternalSorterMerge[0]
2x SMJ + window + countDistinct	7	two join inputs, the multi-phase distinct-agg chain, ShuffleRepartitioner[0]

So num is 1-3 for ordinary stages and reaches ~7 on a wide one. Worth noting a single sort costs two consumers (sorts/sort.rs:283,288).

That lets us say where the old cap actually binds. The old code caps a task's total native usage at offHeap.size / num, while Spark's ExecutionMemoryPool independently caps each active task at roughly pool / N_active_tasks. Whichever is smaller wins:

Scenario	old Comet cap	Spark's grant	binding constraint
8 cores packed, num=3	33% of pool	~12.5%	Spark - old cap invisible
8 cores packed, num=7	14%	~12.5%	Spark, marginally
16 cores packed, num=7	14%	~6%	Spark - old cap invisible
1-2 active tasks, num=7	14%	up to ~100%	Comet, ~7x tighter
2 cores, num=7	14%	50%	Comet, ~3.5x tighter

The conclusion I'd draw: the defect is real, but on a busy executor Spark's own per-task fair share is tighter than the buggy cap, so it is masked. It bites when consumers-per-task exceeds active-tasks-per-executor - stage tails, AQE-coalesced partitions, low partition counts, small executors - where Spark would hand a lone task most of the pool and this check clamps it to 1/num. The expected symptom is unnecessary spilling on long-tail tasks rather than a uniform throughput loss, which I think is a more defensible framing for the PR than "the default config can only use 10% of its off-heap share".

Caveat on my own numbers: the two probe queries were small enough that the check never actually fired (zero fair limit errors in the run), so this measures num and the resulting arithmetic, not a performance delta.

If you want a reproducer that demonstrates the fix rather than arguing for it, the low-concurrency case is the one to target: local[1] with spark.sql.shuffle.partitions=1, a modest offHeap.size, and a wide join + distinct query should hit "fair limit is M bytes, 7 registered" on main and get further after the fix. Happy to build that out if it would help the review.