perf: stop holding the fair pool lock across blocking memory calls - #5613
perf: stop holding the fair pool lock across blocking memory calls#5613dwsmith1983 wants to merge 10 commits into
Conversation
CometFairMemoryPool held its mutex across the JNI calls into Spark task memory manager, which can block for seconds while Spark spills, so every native thread sharing a task pool serialized behind whichever thread was acquiring. The fairness check and the reservation are now one short locked step, the blocking call runs unlocked, and the reservation rolls back if the JVM fails to back it or the call panics. Fairness semantics are unchanged: concurrent grows still cannot jointly exceed pool_size divided by the consumer count. The JNI boundary moved behind a small trait so the pool finally has tests: fairness rejection, limit tightening on register, partial-grant rollback, panic rollback, a blocking test that took ten seconds on the old code and 30ms now, and an eight-thread stress test.
196d709 to
626bc39
Compare
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
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.
| state.used -= subtractive; | ||
| } |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
A release that would zero the JVM-side balance while other acquires are still in flight frees all but one byte immediately and holds the last byte until the in-flight acquires complete. Spark drops the task's accounting entry when its balance hits zero, so a parked acquire waking after that point indexes a missing entry and fails. Blocking the release instead can deadlock because the parked acquire may be waiting for the very memory the release frees. The stub task memory now models the entry lifecycle so the regression is covered.
|
@sunchao any more feedback here? |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
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.
| .jvm_held | ||
| .checked_sub(bytes) | ||
| .expect("released more bytes than the JVM side holds"); | ||
| if bytes > 0 && state.jvm_held == 0 && state.pending_acquires > 0 { |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
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.
Spark's ExecutionMemoryPool removes a task's entry when its balance hits zero, and an acquire parked inside Spark indexes that entry on wake. The previous fix held one byte back from a zeroing release while acquires were in flight and paid it back later. That withheld byte starved Spark's minimum share check, which grants a parked request only when the freed bytes cover it in full, so the waiter slept forever; and a release planned before a late acquire parked could still remove the entry. The pool now asks for one extra byte with every acquire that starts before the anchor lands, keeps the first one until the pool drops, and sends every release whole. The task's balance never returns to zero mid task, so both failure modes are impossible by construction. The test double now models ExecutionMemoryPool's entry lifecycle, share rules and wait loop, and covers the minimum share scenario, the late acquire race and the in flight anchor race.
…e task entry Until the anchor lands, a short grant is handed back to Spark whole. In the gap between that grant returning and its rollback release landing, a sibling acquire of the same task can enter Spark and park if another task has left the pool meanwhile, and the release then removes the task's entry under it. Carriers now run one at a time from the bridge acquire through the rollback release. Releases never take that lock, so a parked carrier cannot hold up the release it waits for, and Spark serializes a task's acquires anyway.
2 participants
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
CometFairMemoryPoolheld its internal mutex across the JNI acquire and release calls into Spark'sTaskMemoryManager. 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.
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?
Ten tests, all new: fairness rejection without reaching Spark, limit tightening on register, partial-grant rollback with the excess returned to the JVM, acquire-failure accounting, zero-size no-op, over-shrink panic, a panic inside the bridge rolling back the reservation, a blocking test where a parked acquire must not stall a concurrent release (it hung for its full ten-second timeout on the old code and completes in 30ms now), and an eight-thread by 500-iteration stress test asserting the accounting never exceeds the fair limit and nets to zero after quiesce. The stress and blocking tests were looped 50 times in both debug and release with no failures. Full core crate suite passes in both profiles, clippy with warnings denied and fmt are clean.