perf: stop holding the fair pool lock across blocking memory calls

[dwsmith1983]

Which issue does this PR close?

No dedicated issue. Adjacent to #5494, which made task-shared pools the norm and so widened the blast radius of this lock.

Rationale for this change

CometFairMemoryPool held its internal mutex across the JNI acquire and release calls into Spark's TaskMemoryManager. That call can block for a long time while Spark spills other consumers, and while it blocked, every other native thread sharing the task pool sat behind the lock, including plain releases that needed nothing from the JVM. The sibling unified pool already avoids this.

What changes are included in this PR?

The admission check and the reservation now happen as one short locked step (the fair limit couples used bytes and the consumer count, so this part genuinely needs mutual exclusion), then the blocking JVM call runs with no lock held, and the reservation rolls back if the JVM declines, grants partially, or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count, and registering a new consumer still only blocks further growth rather than clawing back existing reservations. Going fully lock-free like the unified pool was considered and rejected, since separate atomics would let a register or unregister slip between reading the count and committing the reservation.

Spark's ExecutionMemoryPool removes a task's accounting entry the moment its balance reaches zero, and an acquire parked inside Spark indexes that entry when it wakes, so a release that zeroes the balance while another acquire is parked crashes the waiter. With the lock no longer held across the call that window is real, so the pool takes a one byte anchor once at setup, before it exists, and returns it when it drops; no request ever carries it and the balance never reaches zero mid task. A parked one-byte acquire at setup blocks plan creation until memory frees, which is Spark's normal behavior for a starved task. Spark declines the byte outright, with a zero grant, for a task already at its share, which a task holding JVM shuffle pages when its second native plan is created can reach; the pool then exists without its anchor and each grow retries the byte as a request of its own until it is held, so the task spills on the short grant instead of failing at plan creation. That costs one extra JNI call per grow only while the anchor is missing, and the window without an anchor is bounded and strictly shorter than the whole first acquire on main. Because construction can now park, the task-shared pool registry creates the pool outside its process-wide lock and re-checks on insert; a concurrent loser drops its own pool and returns its anchor after the lock is released.

Two side effects worth naming. The JNI boundary moved behind a small internal trait so the pool can be tested without a live JVM (neither pool had any tests before). And the old code could deadlock if a blocked acquire ever re-entered the pool on the same thread via a spill callback, since the lock was held across the call; that hazard is gone by construction.

How are these changes tested?

Twenty-five pool tests and nine registry tests, all new. The stub models Spark 4.1.3's ExecutionMemoryPool: the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and notifyAll on release, with a bounded wait that fails a test instead of hanging. Covered: fairness rejection without reaching Spark, limit tightening on register, partial-grant rollback with the excess returned, acquire failure and panic rollback, zero-size no-op, over-shrink panic, a parked acquire not stalling a concurrent release, an acquire below its minimum share being granted after the holder frees memory in full, a release already on its way when a late acquire parks, a sibling consumer of the same task freeing its last page while an acquire is parked, a second acquire not queuing behind a parked first one, an exact-fit first grant being accepted, the anchor taken at setup and returned once at drop, construction failing cleanly when the bridge errors on the anchor, a declined anchor leaving a usable pool that reports short grants and takes the byte on the first grow after the share frees with no second request once held, two grows retrying the missing anchor at once keeping exactly one byte, the registry creating outside its lock and keeping exactly one anchor when two creates race, and an eight-thread stress test asserting the accounting never exceeds the fair limit and nets to zero. All looped 30 times in debug and release with no failures. Full core crate suite passes, clippy with warnings denied and fmt are clean.

Microbenchmarks against a real Spark 4.1.3 TaskMemoryManager over an off-heap UnifiedMemoryManager, driven through the JNI bridge with base and head in one binary (details and the full table in the discussion): at 4 and 8 native threads the release p99 falls from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire, mean grow latency and throughput improve 1.5x to 2.5x under contention, grow p99 is unchanged since that is Spark's own work, and errors and final balances are zero on every row.

[sunchao]

Reviewed 626bc395091f313ae93b8867c1cb1c846dc7fc53 against 8729f6e6adf7091e18a48670e790d4ba8fd41e51. I found one P2 in the new concurrent release path, detailed inline.

The focused Spark memory-pool component probe reproduced the missing-task exception. No full Comet native/JNI query suite ran. CI, CodeQL, and Delta Contrib Build Gate currently report action_required, with no test checks recorded.

For this performance change, please include matched BASE/HEAD microbenchmarks with one and multiple native threads, full and partial grants, and consumer registration changes. Report completed operations, grow/release latency, error counts, and final Rust/Spark balances. The parked stub test establishes lock behavior but does not measure production JNI throughput.

[sunchao]

[P2] Keep waiting Spark acquisitions registered during a full release

Could you handle Spark's same-task wait/release contract before allowing this release to overlap an acquire? With fair_unified, two native consumers can now enter the same CometTaskMemoryManager concurrently. In a 100-unit executor pool, let another task hold 90 and this task hold 10. This task's next 10-unit grow waits below Spark's 1/(2N) minimum. Freeing its last 10 units on another native thread removes its entry from ExecutionMemoryPool.memoryForTask and wakes the grower. The grower then indexes the removed entry and throws NoSuchElementException: key not found. Spark's release bypasses the task monitor held by the waiting acquire, so that monitor does not prevent this interleaving. The Rust provisional reservation does not keep Spark's entry alive.

I reproduced the failure using unchanged Spark 3.5.9 pool source with only logging/annotation/memory-mode scaffolding. A scheduling control modeling the previous serialization completed after the other task freed its memory. The relevant map lifecycle is also present in 4.0.4 source. This was a component probe plus JNI source tracing, not a full Comet query reproduction. Please make full releases safe while grows are pending and add a regression that exercises Spark's memory manager.

[dwsmith1983]

Confirmed against the pool source, thanks for the repro. Blocking the release until the in-flight acquires drain turned out to deadlock in testing, since the parked acquire can be waiting on exactly the memory that release frees. So the fix defers instead: a release that would zero the task's balance while acquires are in flight frees n-1 bytes right away (that is what wakes the waiter) and holds the last byte, which the final completing acquire pays off. At most one byte is ever deferred and it always settles once the acquires finish. The test stub now models the entry lifecycle (created on acquire, removed at zero, a woken waiter fails if the entry is gone) and reproduced this crash before the fix. It also turned out one of our existing tests was exercising the same broken pattern.

[sunchao]

The new deferral fixes the already-pending case, but I can still reproduce the same missing-task failure through a late-arriving acquire on current head eb410e51.

At current lines 252-259, plan_release can see pending_acquires == 0, schedule the whole balance, and drop the state lock before release reaches Spark. A new try_grow can then increment pending_acquires and park in Spark while the old balance is still present. The already-planned release removes the task entry, and the waiter resumes with NoSuchElementException.

I reproduced this with the exact head's production state machine and a gated bridge: hold 10, plan and pause the full release, start and park a grow of 10, then resume the release. The waiter panics with key not found: task entry removed while acquire waited. paying_deferred fences only deferred payments, so could you also coordinate ordinary zeroing releases with newly starting acquires?

[dwsmith1983]

Also covered by db1f1bc. With the anchor there is no zeroing release left to coordinate: once any reservation exists the anchor is held, so a full release leaves the task at one byte and the entry survives.

Two related windows came out of review and are closed in the same commit. A grant that covers the request but not the extra byte is handed back as a short grant rather than running unanchored, and every acquire that starts while the anchor request is still parked carries its own extra byte, with the first full grant keeping it and later ones returning theirs. Your gated repro is pinned as late_acquire_survives_a_release_already_on_its_way, alongside a test for the concurrent in-flight case; both failed on eb410e5 with key not found and pass now. 30 loops each in debug and release are clean.

The earlier build-gate failure compared dylib sizes on a change confined to fair_pool.rs, so it looks like the size check rather than this branch.

[dwsmith1983]

One more window closed in 353541d, found while probing the anchor bootstrap with a third task. Before the anchor lands, a short grant is handed back whole; if another task's entry disappears in that gap (raising this task's minimum share) and a sibling acquire of this task then parks, the rollback release zeroes the entry under it. A small bootstrap mutex now serializes anchor carriers from the bridge acquire through the rollback release. Releases never take it, so a parked carrier cannot starve anyone, and non-carriers only exist once the anchor is held. Pinned as short_grant_rollback_cannot_land_under_a_sibling_parked_in_spark, which failed with key not found before the change.

[dwsmith1983]

@sunchao any more feedback here?

[sunchao]

Reviewed eb410e5175417256f0223deaee1b122d760ae077 against ef62b46306e925bc51e7d7f29922c1870eb729e7. I found two actionable issues in the revised release/acquire protocol: the held-back byte can leave a valid Spark acquire waiting indefinitely, and an ordinary full release can still be overtaken by a newly starting acquire. The first is inline below. I added the second as a follow-up on the existing unresolved thread because it is the same missing-task failure under a different ordering.

Local validation: the focused native fair-pool suite passed 12/12 tests with 241 filtered out. An exact-head component probe reproduced the ordinary-release ordering failure. A probe using unmodified Spark 4.1.3 ExecutionMemoryPool reproduced both the missing-key failure and the one-byte wait. These are component probes, not a full Comet native/JNI query run.

Current CI has 63 successful checks, 9 skipped checks, and one failed Delta build gate. Its log fails the contrib-enabled-versus-default libcomet size invariant; this PR changes only fair_pool.rs.

[sunchao]

[P1] Avoid retaining the byte a waiting acquire may need

Could you avoid withholding one byte until finish_acquire? This can deadlock the default off-heap pool. In a 1 GiB Spark execution pool, let another task hold 900 MiB, let this task hold 100 MiB, and let a second native consumer for this task request 100 MiB. Spark parks that request because this task is below its 1/(2N) minimum share. When the holder frees 100 MiB, this branch sends only 100 MiB - 1. On wake, Spark computes toGrant = 100 MiB - 1; because that is short of the request and curMem + toGrant = 100 MiB is still below 256 MiB, it waits again. The deferred byte is paid only by finish_acquire, which cannot run while this acquire is waiting.

I reproduced this against unmodified Spark 4.1.3 ExecutionMemoryPool: retaining one byte left the grower in WAITING, and freeing one additional byte from the other task let it complete. The new stub test misses this because it grants the full request after any release without reapplying Spark's free-memory and minimum-share checks. Could you preserve the task entry without withholding capacity needed by the waiter?

[dwsmith1983]

Fixed in db1f1bc. The deferral is gone: every release now goes to the JVM whole, so a parked acquire sees the full freed amount and Spark's minimum-share check passes. The task entry is kept alive by a permanent anchor instead. The first acquire asks for one extra byte and the pool holds it until it drops, so the balance never reaches zero mid task. That means the task stays in Spark's active-task set for the pool's lifetime and retains one byte, which the header comment now states.

Your scenario is pinned as min_share_wait_is_granted_after_the_holder_frees_its_memory_in_full. The stub now models ExecutionMemoryPool.acquireMemory from 4.1.3 with the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop, and notifyAll on release, with a bounded wait that fails the test instead of hanging. It failed on the previous head with the waiter timing out and passes now.

[dwsmith1983]

@sunchao both of your findings are addressed on the branch head (353541d), with replies on each thread. Ready for another look whenever you have time.

[dwsmith1983]

@sunchao the anchor redesign replacing the deferral is pushed, with the Spark-faithful stub and your P1 and late-arrival cases as tests. Ready for another look.

[andygrove]

The anchor invariant this design rests on does not hold during the window in which the anchor is being established, and bootstrap only covers half of that window. In native/core/src/execution/memory_pools/fair_pool.rs around line 345 every carrier takes bootstrap and holds it across the bridge acquire, so until the anchor lands the pool is back to one blocking JNI call at a time, which is the serialization this PR set out to remove. That lock also serializes only this pool's own carriers. CometTaskMemoryManager.internal is the task-wide TaskMemoryManager, and CometUnifiedShuffleMemoryAllocator is another off-heap consumer on the same task, so if it frees its last page while a pre-anchor carrier is parked inside Spark, memoryForTask loses the entry and the waiter hits the same NoSuchElementException the anchor exists to prevent. Could the anchor instead be taken once as a standalone one byte acquire when the pool is set up, so that no ordinary request ever carries it and Anchor::Requested, bootstrap, and the surplus path all go away?

The extra byte on the request path also turns grants Spark would have satisfied into failures. exact_fit_grant_without_the_anchor_is_rejected_as_short pins the case where Spark backs additional but not additional + 1 and the whole grant is handed back as short. grow at line 272 is try_grow(...).unwrap(), so a caller on that path gets a panic rather than a spill. It is a knife edge, but it is a knife edge that only shows up under memory pressure. Is that acceptable, or is it another argument for keeping the anchor off the request path entirely?

#5466 rewrites the same admission check in this file and moves the fair-share basis from the pool-wide state.used to reservation.size(). DataFusion does not update reservation.size() until try_grow returns, so that change would undercut the reasoning in the comment at line 307 that reserving under the lock is what makes the unlocked window safe. Whichever of the two lands second needs a deliberate rebase rather than a textual merge.

Worth noting that CI has not run on b7007ea9 at all, since the workflows are still sitting at action_required. The microbenchmarks sunchao asked for in the first round are also still open, so the only evidence behind the performance claim right now is the stub timing.

[dwsmith1983]

@andygrove the anchor is now a standalone one byte acquire at pool setup (d004f04). try_new takes it before the pool exists and fails construction if Spark grants less than a byte or errors, so every live pool holds the anchor from birth to drop and no request carries it: Requested, the bootstrap mutex, the surplus release, and the exact-fit rejection are gone, and try_grow asks for exactly additional again with the reviewed reserve-under-lock, call-unlocked, roll-back shape. Your sibling-consumer case is a test now: the stub models another consumer of the same task freeing its last page while an acquire is parked, and the entry survives because the anchor byte is held. So is the serialization point: a second acquire no longer queues behind a parked first one.

One consequence worth calling out: construction can now park inside Spark, and acquire_task_shared_pool used to run the create step under the process-wide registry lock, which would have stalled plan creation for every task on the executor. It now creates outside the lock and re-checks on insert; a concurrent loser drops its own pool and its anchor goes back over JNI after the lock is released, with a test through the real registry.

On grow, the unwrap predates this branch and is DataFusion's contract for the infallible path; with the anchor off the request path the only panic left there is a short grant a plain request would also have hit.

On #5466: it moves the admission basis from state.used to reservation.size(), which DataFusion updates only after try_grow returns, so two concurrent grows of one reservation would both pass the check before either lands. Whichever lands second needs a deliberate rebase: keep the reserve-under-lock step and apply the per-reservation limit to it, and rewrite the two comments in try_grow and shrink that explain the state.used basis. Nothing in this commit changes that shape.

[dwsmith1983]

@sunchao the microbenchmarks from your first pass, against a real Spark task memory manager rather than the stub: the harness launches a JVM in-process through the jni crate's invocation feature, builds a Spark 4.1.3 TaskMemoryManager over an off-heap UnifiedMemoryManager with a 1 GiB pool, wraps it in CometTaskMemoryManager, and drives base (main's fair_pool.rs compiled as a sibling module) and head through the real JNI bridge in one binary. Three recorded rounds after warmup, 20,000 operations per thread. Latencies are mean / p50 / p99 in microseconds. Fair-limit rejects and other errors were zero on every row; Rust used and the Spark balance netted to zero after quiesce on every row, with head holding its one anchor byte until drop and zero after.

scenario	threads	ops/s base to head	grow base to head	release base to head	short grants base / head
full grants	1	1.52M to 1.92M (1.26x)	0.7/0.7/0.9 to 0.5/0.5/0.7	0.4/0.4/0.5 to 0.4/0.3/0.5	0 / 0
partial grants	1	1.15M to 1.57M (1.36x)	0.8/0.7/7.0 to 0.6/0.5/3.0	0.4/0.4/1.2 to 0.4/0.4/0.5	1941 / 1941
registration churn	1	1.89M to 1.94M (1.03x)	0.5/0.5/1.5 to 0.5/0.5/1.2	0.4/0.4/1.1 to 0.4/0.4/0.6	0 / 0
full grants	4	0.92M to 1.37M (1.50x)	4.4/1.3/32.4 to 3.6/1.2/26.5	4.2/1.0/31.8 to 0.8/0.6/3.7	0 / 0
partial grants	4	0.72M to 1.34M (1.86x)	5.2/1.5/43.9 to 3.1/1.1/40.3	7.5/1.4/67.4 to 0.8/0.7/5.8	29314 / 23544
registration churn	4	0.97M to 1.60M (1.65x)	4.1/0.7/37.3 to 3.0/1.1/26.2	4.0/0.5/37.3 to 0.7/0.6/1.5	0 / 0
full grants	8	0.61M to 1.40M (2.27x)	12.9/3.7/83.4 to 7.2/1.2/93.1	12.9/4.0/83.3 to 0.8/0.6/2.5	0 / 0
partial grants	8	0.49M to 1.23M (2.49x)	16.5/12.4/93.5 to 7.1/1.2/129.5	13.4/1.8/89.2 to 0.7/0.7/1.8	110853 / 78614
registration churn	8	0.62M to 1.40M (2.27x)	12.9/2.6/92.1 to 7.2/1.1/91.9	12.8/2.5/92.0 to 0.8/0.6/3.0	0 / 0

The release path is where the change shows: at 4 and 8 threads its p99 drops from tens of microseconds to single digits because a release no longer waits behind an in-flight acquire. Grow p99 is unchanged, since that is Spark's own work under the call, while mean grow latency and throughput improve 1.5x to 2.5x under contention. Head sees fewer short grants in the partial scenarios because it cycles past the release faster. Registration churn counts only depend on run length. The harness needs a JDK and a full Spark classpath so it is not a criterion bench in tree; the source and raw output are available if wanted. macOS arm64, single machine.

[andygrove]

This is the shape I was hoping for, and it is a lot simpler than what it replaces. Anchor::Requested, the bootstrap mutex, the surplus release and the exact-fit rejection are all gone, and try_grow is back to reserve-under-lock, call-unlocked, roll-back-on-shortfall with no extra byte riding along. Moving the create step out from under the process-wide registry lock was the right call and I would not have thought to check it until it bit someone: with the anchor acquire able to park, holding TASK_SHARED_MEMORY_POOLS across it would have stalled plan creation for every task on the executor. The std::ptr::eq guard in TaskSharedMemoryPool::drop already covers the loser, so a concurrent create hands its anchor back without deregistering the winner. Thanks also for running the benchmarks against a real TaskMemoryManager rather than the stub; the release p99 collapsing from tens of microseconds to single digits at 4 and 8 threads is the number that actually demonstrates the claim, and it is the one the stub could not show.

One thing came out of the move that I do not think we should ship as is.

A declined anchor now fails the task instead of spilling

with_bridge returns CometError::Internal when Spark grants less than the byte, and createPlan propagates it, so the task dies before any work is attempted. Walking Spark's ExecutionMemoryPool.acquireMemory with numBytes = 1: maxToGrant = min(1, max(0, maxMemoryPerTask - curMem)), so a task whose curMem has already reached maxMemoryPerTask gets maxToGrant = 0, and the park guard curMem + toGrant < minMemoryPerTask is false for it because minMemoryPerTask is half of maxMemoryPerTask. It does not wait. It returns 0, and the pool refuses to exist.

That is reachable without anything exotic. The task only needs to be holding its share of execution memory at the moment the first native plan is created, and CometTaskMemoryManager.internal is the task-wide manager, so CometUnifiedShuffleMemoryAllocator and any Spark-side operator in a mixed plan count toward curMem. The concrete sequence I would worry about is a task that runs one native plan, closes it so the shared pool drops, allocates JVM shuffle pages, and then creates a second native plan: lookup misses, try_new runs again, and Spark declines. On main that situation produces a short grant on the first real try_grow and DataFusion spills. On this branch it is a task failure with no spill available as an escape, and the message reads like an internal error rather than memory pressure.

Could a declined anchor be non-fatal instead? Keep a flag for "anchor not held" and re-attempt a standalone one byte acquire at the top of try_grow while it is unset. That keeps everything this commit won: the anchor never rides on a real request, so there is no exact-fit hazard, no surplus to hand back and no bootstrap mutex, and the request path stays plain. The cost is one extra JNI call per try_grow only while the anchor is missing, which is the rare path by construction. The window where the pool holds reservation bytes without the anchor is then real but bounded, and it is strictly better than the window on main, which is the whole first acquire. A test for the declined-at-setup case would pin it either way; right now the stub only covers the grant.

The parking case reads fine to me and I convinced myself it cannot deadlock: a task parked in try_new holds nothing, so a pool fully consumed by tasks that are all parked in construction is not a reachable state, and maybeGrowPool covers the storage-borrow side.

Two smaller things. with_bridge's short-grant branch returns without releasing granted, which is dead today because ANCHOR_BYTES is 1 so a short grant is necessarily 0, but it would leak silently if that constant ever moved; a debug_assert or a one-line comment tying the two together would be cheap. And bridge.acquire at construction is the one bridge call not wrapped in catch_unwind. That is correct as written, since there is no optimistic reservation to roll back and try_unwrap_or_throw catches the panic at the JNI boundary, but the comment above the catch_unwind in try_grow explains why it is needed there without saying why it is not needed here.

Your read on #5466 matches mine, including that the per-reservation basis has to be applied to the reserve-under-lock step rather than replacing it, and that the two comments explaining the state.used basis need rewriting rather than merging. Worth agreeing an order with @peterxcli so whoever is second plans for it.

[dwsmith1983]

Could a declined anchor be non-fatal instead?

Done that way. with_bridge no longer errors on a zero grant: the pool is created with anchor_held unset and jvm_held at zero, and try_grow calls take_missing_anchor before the reserve step, which issues a standalone one byte acquire while the flag is unset and never touches the real request. Once the byte lands the flag is set, so the retry is one extra JNI call per grow only while Spark keeps declining. Two grows retrying at once both get a byte; the second to lock hands its own back and books nothing. Drop returns the anchor only when it is held. A bridge error at construction still fails construction, since that is a JNI failure rather than memory pressure.

The short-grant branch in with_bridge is gone with it, and const _: () = assert!(ANCHOR_BYTES == 1) next to the constant pins the reasoning that a declined anchor is a zero grant with nothing to hand back. The comment above the construction acquire now says why it needs no catch_unwind (nothing reserved yet to roll back, and try_unwrap_or_throw catches the panic at the JNI boundary), and the retry carries the same note.

Tests: the declined-at-setup case now asserts the pool exists holding nothing, reports a short grant while the sibling still holds the share, takes the byte on the first grow after the share frees, makes exactly one JVM call per grow from then on, and returns the byte at drop; a second test pins that a pool that never held the anchor releases nothing on drop; a third drives two retries into the stub's acquire gate before either wins and checks one anchor is kept and one byte returned. Thirty loops in debug and release, no failures.

On #5466: left a note there proposing that whoever lands second rebases, with the per-reservation basis applied to the reserve-under-lock step rather than replacing it.