Skip to content

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

Open
dwsmith1983 wants to merge 125 commits into
apache:mainfrom
dwsmith1983:perf/fair-pool-lock-free
Open

dwsmith1983 wants to merge 125 commits into
apache:mainfrom
dwsmith1983:perf/fair-pool-lock-free

Conversation

@dwsmith1983

@dwsmith1983 dwsmith1983 commented Sep 1, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

No dedicated issue.

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. An acquire parked inside Spark waits for memory to be released, and a release from another thread of the same task that has to take a lock the parked thread holds cannot land, so the task stalls until an unrelated task frees memory. The sibling unified pool already avoids this.

What changes are included in this PR?

The pool builds on main's SparkMemory wrapper and its SparkMemoryManager trait and adds two items to the wrapper. manager() exposes the raw Spark calls without the overcommit ledger, for a short grant the pool hands back itself. try_acquire_leaving_a_short_grant is the variant behind try_acquire that leaves a short grant with Spark instead of returning it, so the pool can keep those bytes charged until it has handed them back. SparkMemory also tracks which outstanding overcommit an in-flight try_grow has already put into its Spark request, described below. Apart from that, every operation makes the same Spark calls as on main.

The fair limit is checked and the bytes are charged in one short locked step, since the limit couples the used total and the consumer count. The lock is dropped before any Spark call and the bookkeeping is settled after it returns. grow charges under the lock the same way, skips the limit, and carries whatever Spark declines as overcommit, as on main. try_grow rolls its charge back if Spark errors or panics inside the JNI frame, and grow does the same on a panic. A short grant stays charged until Spark has it back, and with overcommit owed the request carries the debt, so a refused try_grow can receive more bytes than it asked for and hands all of them back. If handing a short grant back fails, the pool logs it, keeps the bytes charged, since Spark keeps them until the task ends, and refuses the request as usual, where main returned the JVM error from try_grow. shrink checks its bounds under the lock, releases to Spark with the lock dropped, and only then takes the bytes off the pool's total. Going fully lock-free like the unified pool was rejected, since separate atomics would let a register or unregister slip between reading the count and committing the charge.

Concurrent try_grow calls still cannot jointly take a consumer past its share or the pool past pool_size, because each is charged before Spark is asked. Settling after the call opens two windows, both on the conservative side, and the memory guide documents them. A grow's charge is visible before Spark answers, so a second try_grow in that window is checked against a total that includes the first and can be refused where waiting would have let it through. A release is visible to Spark before the pool's total drops, so a grow that only fits once those bytes are free is refused by the share or pool check rather than sent to Spark ahead of the release, while a grow that fits without them can still reach Spark first and come back short if the task is at its Spark share. Neither window admits a try_grow the limit would have refused, and a refusal is the ordinary error a spillable operator answers by spilling.

Spark's ExecutionMemoryPool removes a task's entry the moment its balance reaches zero, and an acquire waiting inside Spark reads that entry when it wakes, so a release that zeroes the balance under a waiting acquire makes it throw NoSuchElementException: key not found: <task attempt id>. With the lock no longer held across the call, one native thread's release can do that to another native thread's waiting acquire in the same task. #6310 makes CometTaskMemoryManager.acquireMemory ask Spark again in that case, which is why this PR should land after it. Its retry is bounded at three attempts and then returns a zero grant, which the pool treats as a refusal or overcommit like any short grant, so the bound is fine for this race too.

With the lock gone, two concurrent try_grow calls could also read the same overcommit and both ask Spark for it on top of their own request. If the first then repaid it, the second stayed parked in Spark on a request larger than what was free. SparkMemory now keeps the debt and the part of it that in-flight requests have claimed in one small ledger, behind a mutex held only for the arithmetic and never across a Spark call. Each try_grow claims only the debt no other one is asking for, a refusal, error or panic hands its claim back, and bytes a release repaid while Spark was granting them still go back to Spark. CometUnifiedMemoryPool called SparkMemory without a lock on main and on branch-1.1, so it had the same race and gets the same fix. One window is left: a release that repays the debt while a claimed request is parked in Spark leaves that request larger than needed, since Spark cannot shrink a parked request. Closing it would mean handing claimed debt to Spark on release, which gives Spark bytes it never granted if the claim is then refused.

releaseMemory now calls Spark before it moves its count, so a release Spark rejects moves nothing and used keeps matching what the pool still has charged. The memory guide now explains why a short grant is logged with getMemoryConsumptionForThisTask rather than showMemoryUsage, which main already does: showMemoryUsage would take the task's monitor while another acquire of the task may be waiting for those bytes. The jni_api.rs registry comment no longer describes the fair pool holding its lock across JNI. task_shared.rs, spill_replay.rs, mod.rs and the memory review skill are unchanged from main. The memory guide gains paragraphs on the lock rule, the two windows and the short-grant log, and its pool-stack diagram now shows PlanMemoryPool.

How are these changes tested?

Unit tests in fair_pool.rs run the pool against two stand-ins for Spark. Main's FakeSpark covers the overcommit paths, and a new test pins that a refused try_grow with overcommit outstanding hands the whole grant back with neither the pool's total nor the overcommit moving. A stub of Spark 4.1.3's ExecutionMemoryPool covers the rest: the per-task entry lifecycle, the 1/N and 1/(2N) shares, the wait loop and notifyAll on release, the bridge asking again when the entry is gone (standing in for #6310), a bounded wait that fails a test instead of hanging, and gates that pause a call so a second thread can be interleaved at a precise point.

Those tests cover the lock discipline: a release does not stall behind a parked acquire on another thread, a second acquire does not queue behind a parked one inside the pool, and an eight-thread stress test asserts no consumer passes its share and the pool never passes its size, and that the pool, Spark and the JVM count all end at zero. They cover the rollback paths: short grants, acquire failure, a panic inside the call for both grow and try_grow, and an over-shrink. They cover Spark's wait loop: an acquire parked below its minimum share woken by a full release, a late acquire surviving a release already on its way, and a sibling consumer freeing its last page under a parked acquire. They cover the exact fit: with another task holding 90 of 100 bytes, a first try_grow(10) or grow(10) is granted with one Spark call; with the other task at 70, try_grow(30) and grow(30) are granted in full; and the same holds after a full release. They cover the two windows: a grow racing a shrink is refused without a JVM call rather than admitted on bytes Spark still holds, a grow that fits without those bytes gets a short grant at the Spark share, a short-grant rollback keeps its bytes charged until Spark has them back, and an in-flight grow counts against the limit until it settles.

SparkMemory tests cover the claims: two concurrent requests ask for the debt once, a refusal, a Spark error and a panic each hand the claim back, a release during two claims hands Spark back only the excess, a claim made after a release while new debt arrives covers only what is unclaimed, and a multi-thread run keeps the ledger balanced. A fair pool test pins the case where one byte is owed and three are free: the first try_grow(1) pauses after its grant, and the second asks Spark for its own byte and completes before the first settles.

CometTaskMemoryManagerSuite gains two tests: a release Spark rejects leaves used unchanged, and a short grant is handed back while another acquire of the task waits in Spark, with debug logging on so the short-grant log runs. The suite passes on Spark 3.4, 3.5, 4.0 and 4.1.

Microbenchmarks against a real Spark 4.1.3 TaskMemoryManager over an off-heap UnifiedMemoryManager, base and head in one binary, are in the discussion. They predate the SparkMemory merge and the current Spark call pattern, and the contended rows are the claim: 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 (releases now take SparkMemory's ledger mutex, for the arithmetic only), throughput improves 1.5x to 2.5x, grow p99 does not improve since that is Spark's own work, and errors and final balances are zero on every row. TPC-H SF10 on c2d5a28 with 2g off-heap, base and head alternated, spilled identically on every arm with no task failures, and in the settled pass the head to base ratio on the sum of medians was 0.99. A rerun on 88d881d spilled the same bytes the same number of times. Both heads predate the current design, so they cover the lock discipline rather than this exact head.

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.
@dwsmith1983
dwsmith1983 force-pushed the perf/fair-pool-lock-free branch from 196d709 to 626bc39 Compare September 2, 2026 02:13

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines 176 to 177
state.used -= subtractive;
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The existing P2 missing-task-entry concern also remains reproducible after a declined anchor: a sibling releases 100 bytes, another task takes 90, and the real grow acquires 10 without an anchor. The next anchor retry parks. Releasing those 10 native bytes removes the entry and produces NoSuchElementException.

Confirmed. When Spark frees up between a declined anchor retry and the real request, the pool holds bytes from Spark without its anchor. A later retry can then park, and releasing those bytes took the task's balance to zero under it.

Now the first release that hands bytes back to Spark while the anchor is missing keeps one of them as the anchor. That covers a shrink and the rollback of a short grant. The new CometTaskMemoryManager.releaseKeepingAnchor releases all but one byte and moves that byte from used to the anchor count, so it is returned once at drop and close warnings stay quiet. The claim is made under the pool lock, so only one release keeps a byte, and a retry that lands afterwards hands its byte back as a duplicate, as before. The anchor stays outside the pool total and the overcommit ledger.

releasing the last bytes held without the anchor keeps a parked anchor retry's entry in CometTaskMemoryManagerSuite runs your sequence against a real UnifiedMemoryManager and TaskMemoryManager. With a plain releaseMemory(10) it fails with NoSuchElementException on Spark 3.5 and 4.1, and with the new call the parked retry is granted. release_while_unanchored_keeps_a_byte_under_a_parked_anchor_retry in fair_pool.rs drives the same sequence through the pool itself and fails with key not found without the change. The memory management guide now describes the release rule, and names the one window left, a JVM consumer of the task freeing its last bytes while the pool holds nothing from Spark. That window is tracked in #6224.

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.
@dwsmith1983
dwsmith1983 requested a review from sunchao September 2, 2026 10:37
@dwsmith1983

Copy link
Copy Markdown
Contributor Author

@sunchao any more feedback here?

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@dwsmith1983

Copy link
Copy Markdown
Contributor Author

@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
dwsmith1983 requested a review from sunchao September 4, 2026 04:23
@andygrove andygrove added enhancement New feature or request performance area:memory Memory pools, reservations, OOM handling labels Sep 6, 2026

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: fair_unified held its mutex across blocking Spark memory calls, delaying concurrent releases and memory-usage reads.
  • Design approach: Check and charge reservations under the mutex, call Spark unlocked, then settle accounting. A lazy anchor preserves Spark’s task entry, and replacement pools wait for completed teardown.
  • Correctness / compatibility analysis: The existing P2 exact-fit wait remains reproducible at native/core/src/execution/memory_pools/fair_pool.rs:159. With a 100-byte Spark pool and another task holding 90, the base’s unanchored 10-byte request succeeds immediately. Head takes its anchor first, leaving nine bytes. Spark parks the request because 1 + 9 < 25, its minimum share. The request cannot succeed or return a spillable refusal until another release supplies memory. A bounded probe using the head’s Java bridge and real Spark 4.1.3 reproduced WAITING, while the unanchored control completed immediately. Releasing one byte rescued the request. Relevant locking and task-entry semantics match Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Beyond this existing concern, no additional introduced P1/P2 issues found within this review.
  • Key design decisions: Admission checks and provisional charges remain atomic. Reusing SparkMemory avoids duplicating overcommit accounting. The teardown guard coordinates pool generations without holding the registry lock across JNI.
  • Implementation sketch: grow records provisional charges and overcommit. try_grow rolls back failed acquisitions and returns short grants. shrink settles after Spark accepts the release, and Java usage accounting follows successful release.
  • Behavioral changes worth calling out: Compared with branch-1.1, intended PR changes permit concurrent releases, retain an anchor byte, allow conservative refusals during settlement, and make replacement creation wait for teardown. Earlier-head contention benchmarks and spill comparisons were reviewed but not independently repeated. They do not resolve the exact-fit wait regression.
  • Suggested improvements: Resolve the existing P2 by preserving task-entry safety without parking an otherwise satisfiable allocation solely for the anchor byte. Documenting and testing the wait does not eliminate the regression.

Reviewed the complete eleven-file diff from 8c783aa88104616dcf0f7876b4f7a9ba71e2bf31 to cef941e424f68e7daf2a3c9bc03d856aa9ab3c91, including surrounding code and existing reviews, issue comments, inline comments and threads. Confirmed non-draft status. Read AGENTS.md and applied review-comet-pr, review-comet-memory-pr and review-comet-ffi-pr.

Exact-head CI: Comet CI and CodeQL report action_required. Labeling passed. No CI build/test verdict is available.

Validation: All 76 native memory-pool tests passed with cargo test -p datafusion-comet --lib memory_pools --offline --no-default-features. Recompiled the head’s Java bridge and Scala suite, then passed all eight memory-manager tests against Spark 4.1.3. The bounded exact-fit component probe confirmed the existing blocker. No full Maven build, JNI lifecycle/query suite, Spark SQL suite or benchmark ran. Other Spark versions were checked through source comparison. Project files remain unchanged.

Review state: Request changes for the unresolved existing P2.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

  • Prior state and problem: CometFairMemoryPool held its mutex across blocking Spark memory calls, preventing concurrent releases and memory-usage reads from progressing.
  • Design approach: Check limits and charge reservations under the mutex, call Spark unlocked, then settle accounting. A lazy anchor byte protects Spark’s task entry, and replacement pools wait for completed teardown.
  • Correctness: No additional introduced P1/P2 issues found within this review. The existing P2 exact-fit wait remains at native/core/src/execution/memory_pools/fair_pool.rs:159. With a 100-byte Spark pool and another task holding 90, the base can grant a first try_grow(10) immediately. Head takes its anchor first, leaving nine bytes. Spark parks the request because 1 + 9 < 25, its minimum share. The request cannot succeed or return a spillable refusal until another release supplies memory. The current native regression confirms this behavior. A bounded probe using the freshly compiled head’s Java bridge and real Spark 4.1.3 reproduced WAITING, while the unanchored control completed immediately. Earlier release, diagnostic-deadlock, and pool-replacement regressions pass their focused tests.
  • Compatibility analysis: Checked ExecutionMemoryPool and TaskMemoryManager sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Relevant waiting, locking, and task-entry semantics agree. Compared the affected paths with branch-1.1. Task-wide Java manager sharing and blocking-aware acquisition already exist in the base. No expression semantics, fallback rules, or configuration defaults change here.
  • Key design decisions: Admission checks and provisional charges remain atomic across consumer registration changes. The anchor stays outside reservation and overcommit totals but counts as an ordinary Java grant. Teardown releases it before removing the registry entry.
  • Implementation sketch: grow charges memory and carries Spark’s shortfall as overcommit. try_grow rolls back failed acquisitions and retains short-grant charges until their release succeeds. shrink settles after Spark accepts the release. Java usage accounting now follows successful release too.
  • Performance: Earlier-head JNI benchmarks report lower contended release p99 and 1.5–2.5× throughput. The reported off-heap TPC-H comparisons retained identical spill counts. These support the contention improvement but do not cover the exact-fit stall. The implementation adds short bookkeeping lock acquisitions, an anchor acquisition while missing, and an anchor release at teardown. Exact-head performance was not independently benchmarked.
  • Design: Keeping a short critical section around coupled limits and charges is straightforward and preserves fairness. Moving blocking calls outside it addresses the original contention problem. The anchor protocol still needs to preserve progress for otherwise satisfiable allocations.
  • Abstraction & complexity: Reusing SparkMemory avoids duplicating the overcommit ledger. TeardownComplete gives pool replacement an explicit lifetime boundary without holding the registry lock across JNI. No separate P1/P2 abstraction issue was identified.
  • Behavioral changes worth calling out: Concurrent releases and usage reads can proceed during acquisitions. Provisional charges can conservatively refuse competing requests. The anchor retains one byte and keeps the task active until pool destruction. Replacement-plan creation can wait for teardown. At an exact-fit boundary, the anchor can introduce a wait, short grant, or one byte of overcommit.
  • Suggested improvements: Resolve the existing P2 by preserving task-entry safety without parking an otherwise satisfiable allocation solely for the anchor byte. Documenting and testing that wait does not remove the regression. No additional P1/P2 findings are proposed.

Reviewed the complete eleven-file diff from b56349697b786ff2ad1c1bcf6ecf45b809af5f30 to b12b7ce883714455bba609e7ce3f8061cdeeb47d, including surrounding code and existing reviews, issue comments, inline comments, and threads. Confirmed non-draft status. Read AGENTS.md and applied review-comet-pr, review-comet-memory-pr, and review-comet-ffi-pr.

Exact-head CI: Comet CI and CodeQL report action_required. Labeling passed. No CI build/test verdict is available.

Validation: All 76 native memory-pool tests passed with cargo test -p datafusion-comet --lib memory_pools --offline --no-default-features. Freshly compiled the head’s Java bridge and Scala suite, then passed all eight CometTaskMemoryManagerSuite tests against Spark 4.1.3. The bounded exact-fit component probe confirmed the existing blocker. No full Maven build, JNI lifecycle/query suite, Spark SQL suite, or benchmark ran. Other Spark versions were checked through source comparison. Project files remain unchanged.

The anchor byte the fair pool kept with Spark made an exact-fit first
allocation park or come back one byte short, and holding a byte back
from a parked acquire of the same task can stall it. The pool now makes
the same Spark calls as main for every operation, so a request Spark can
grant is granted in one call. With nothing held at drop, the teardown
wait in the task-shared registry goes too, and task_shared.rs, mod.rs
and spill_replay.rs match main again.

A release that empties the task's balance under a parked acquire of the
same task is handled by the acquire retry in CometTaskMemoryManager
(apache#6310), which this change depends on.

Adds exact-fit tests, a test that a panicking grow rolls back its
charge, and runs the short-grant test with debug logging on so the log
path that must not take the task monitor is exercised.
The teardown wait it described is gone, and TaskSharedMemoryPool::drop is
back to comparing pointers, which main's wording already describes.
@dwsmith1983

Copy link
Copy Markdown
Contributor Author

Resolve the existing P2 by preserving task-entry safety without parking an otherwise satisfiable allocation solely for the anchor byte.

The anchor is gone. The pool now makes the same Spark calls as main for every operation, so with another task holding 90 of 100 bytes a first try_grow(10) or grow(10) is granted with one call, and with the other task at 70, try_grow(30) and grow(30) get all 30. Those cases are tests in fair_pool.rs now, with nothing pre-seeded.

Task-entry safety now comes from #6310, which you approved: CometTaskMemoryManager.acquireMemory asks Spark again when a release empties the task's balance under a waiting acquire. That covers the release coming from another native thread as well as from a JVM consumer of the task, so this PR should merge after #6310. With the anchor gone the teardown wait went too: task_shared.rs, spill_replay.rs and mod.rs match main again, and createPlan no longer waits for the previous pool's teardown.

@dwsmith1983

Copy link
Copy Markdown
Contributor Author

Spill pressure run on the current head (155e74060) against its merge base on main (f80042f78): TPC-H SF10, one 8-core executor, fair_unified with 2g off-heap so each task gets 256 MB, base, head, head, base at three iterations of the 22 queries in one Spark application per arm, then q18 and q21 at five iterations alternated.

workload metric base head
22 queries x3, per arm disk bytes spilled 4,215,250,560 4,215,250,560
22 queries x3, per arm spills 180 180
q18 x5, per arm spills / bytes 80 / 1,405,702,160 80 / 1,405,702,160
q21 x5, per arm spills 0 0
every arm failed tasks / executors lost 0 / 0 0 / 0

Every query result hashes the same across arms and matches the earlier runs. Executor stderr has no lines for fair limit, only got, overcommit, OutOfMemory, NoSuchElementException or key not found on either side.

This ran in Docker rather than directly on the host as the earlier runs did, with smaller heaps (driver 1g, executor 1536m) to fit the VM; the off-heap pool and cores are unchanged. Main at 8166586fa, run the same way, reproduces the earlier numbers exactly (360 spills and 22,361,276,112 bytes for the 22 queries), so the lower spill volume here comes from main itself, most likely the lz4 spill compression in #6698.

On wall time, the sum of per-query medians was 34.13 s on head against 34.32 s on base, q18 3.21 s against 3.30 s and q21 5.78 s against 5.90 s. The container swapped during the full-suite arms on both sides, so I read that as no slower rather than faster. The contention gain this PR targets shows in the microbenchmarks in the description, not at this level.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed commit 6a90a24309f65fe3d2158710a1c4abf5adee092d. New P1/P2 findings are detailed inline.

// its JNI frame; the optimistic reservation must not outlive it, or the leaked bytes
// poison the task-shared pool for every other consumer.
let refusal = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
self.spark.try_acquire_leaving_a_short_grant(additional)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Coordinate concurrent acquisitions’ overcommit coverage. Running this call unlocked lets two try_grow calls snapshot the same debt. With one byte of debt and three bytes free, the first try_grow(1) acquires two bytes. If it pauses after Spark returns but before repay, a second try_grow(1) also requests two bytes, although only one remains free. Below Spark’s minimum share, that request parks. The first then clears the debt, but the second remains blocked on its stale request and cannot reach the surplus-return logic. Expected behavior is successful acquisition of the available byte. Instead, a default fair_unified task can stall until another release supplies unnecessary capacity. Base serialization makes the second request one byte and complete. Please coordinate in-flight debt coverage while keeping releases able to progress.

Evidence: Bounded regression review_repaid_debt_does_not_park_an_exact_fit_sibling, using unchanged head production pool code: both local and Spark pools contain 100 bytes. Other tasks hold 83 and 1. grow(17) receives 16, creating debt 1. Release the third task’s byte and two bytes from the other task. Gate the first try_grow(1) after its two-byte Spark grant, start the second try_grow(1), then resume the first. Head reports requests=[17, 2, 2] debt=0 free=1 second_completed=false; pristine base reports requests=[17, 2, 1] second_completed=true. A freshly compiled head Java bridge against real Spark 4.1.3 confirms the two-byte request waits with task usage 18, other-task usage 81 and one byte free; the one-byte control completes immediately. Releasing one external byte rescues the waiter. Artifacts: /tmp/pr5613-review-final-6a90-8jdrim1x/native/head-probe.log, native/base-probe.log, and jvm/debt-probe.log.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please coordinate in-flight debt coverage while keeping releases able to progress.

Done in 8cf1739. SparkMemory now keeps the overcommit and the part of it that in-flight requests have already asked Spark for in one small ledger, behind a mutex held only for the arithmetic and never across a Spark call, so a release still repays and moves on. Each try_grow claims only the debt nobody else is asking for: in your case the second asks Spark for its own byte and completes while the first is still paused, and the requests come out as 17, 2, 1. A refusal, an error or a panic hands the claim back for the next try_grow, and bytes a release repaid while Spark was granting them still go back to Spark.

The unified pool reached the same code without a lock on main, so it gets the fix too. Your scenario is now a test in fair_pool.rs. One window remains, and the memory guide describes it: a release that repays the debt while a claimed request is parked in Spark leaves that request larger than needed, since Spark cannot shrink a parked request.

…grow calls

Two concurrent try_grow calls could both read the same overcommit and ask Spark for it on top of their own request. If the first repaid it, the second stayed parked in Spark on a request larger than what was free. SparkMemory now keeps the debt and the part of it in-flight requests have claimed in one ledger, behind a mutex held only for the arithmetic, so each try_grow asks only for debt no other one is covering. A refusal, error or panic hands the claim back, and debt a release repaid while Spark was granting goes back to Spark. The unified pool shared the race and gets the same fix.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ffi Arrow FFI / JNI boundary area:memory Memory pools, reservations, OOM handling enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants