Conversation
Comet's memory pool counts declared reservations. Plenty of real allocation never goes through it -- Arrow builders, expression kernels, decompression buffers, Parquet metadata, object_store buffers, tokio itself -- so pool reservations are a lower bound on Comet's footprint and the size of the gap is currently unmeasurable at runtime. Diagnosing an OOM means guessing at it, and spark.comet.exec.memoryPool.fraction asks operators to hand-tune a haircut for a quantity nobody can see. AccountingAllocator wraps the selected global allocator (jemalloc, mimalloc, or system) and maintains one signed process-wide byte balance. executePlan reports it as the native_allocated tracing metric, next to the per-thread pool reservations it should be compared against. This is observability only: it never rejects an allocation, never panics, and does not touch the memory pool. Because it cannot fail an allocation, realloc can account after delegating rather than before, which avoids over-counting a failed realloc. Per-thread deltas are batched and flushed into the shared balance at 64 KiB, so the common path is a thread-local add-and-compare rather than an atomic RMW. ThreadDrift's destructor settles the remainder when a thread exits, which matters because the blocking pool churns on tokio's idle timeout and would otherwise bias the balance on a long-lived executor. Touching that destructor-bearing thread-local can itself allocate on first use, so track() keeps a destructor-free re-entrancy flag and settles re-entrant calls straight into the shared balance. Off by default; a build without the feature has no wrapper and no per-allocation work. Verified clippy -D warnings and the native test suite across default, alloc-accounting, jemalloc+alloc-accounting, and mimalloc+alloc-accounting. The thread-exit test was mutation-checked: neutering the destructor fails it.
Answers the question the feature has to answer before anyone proposes enabling it by default. The liveness assertion is the point of the harness as much as the timings are: without it a 'with the feature' run can silently be a second baseline, which is exactly what happened on the first attempt here.
`dealloc` accounted after calling the inner allocator, mirroring `alloc` and `realloc`. A free cannot fail, so that ordering bought nothing, and it opened a window: jemalloc decrements `stats.allocated` at the start of a large free and then, for blocks above its 8 MiB oversize threshold, unmaps the pages eagerly, which takes milliseconds for a block of a few hundred megabytes. For that whole window the balance still carried a block the allocator had already given back, so `native_allocated` read above `jemalloc_allocated` by the size of the block in flight. On TPC-H SF100 about 2% of trace samples showed the excess, up to 160 MB, during task teardown in Q10, Q17 and Q18. Subtract before delegating. A new test wraps a recording inner allocator and asserts the balance has already dropped by the time the inner free is called; it fails with "inner dealloc saw balance 67182807, expected at most 33628375" under the previous ordering.
…nation Select the allocator backend once, in a `backend` module whose three cfgs partition every feature combination, and let the single accounting `#[global_allocator]` refer to whatever that resolved to. Previously each backend predicate was repeated per arm and the accounting fallback to the system allocator excluded `mimalloc`, so `jemalloc,mimalloc,alloc-accounting` on a non-MSVC target matched no arm at all: the process ran on the unwrapped default allocator while `native_allocated` stayed enabled and read zero. A build without the feature is unchanged: the unwrapped allocator lives in the backend module that owns it, and no explicit allocator is installed when the selection is the system allocator. Make the accounting tests immune to parallel test noise. `BALANCE` is process-wide, so with the wrapper installed the crate's other tests move it concurrently and the margin-based assertions failed intermittently. The tests that observe the balance now serialize against each other, the real-allocation probe uses an untouched 256 MiB block that nothing else in the crate can mask, and the thread-exit test injects its drift directly so it no longer needs the feature and runs in the default CI build. Add `threshold_churn` to the `alloc_overhead` benchmark: alloc/free loops at 32 KiB (never flushes) and 64 KiB (flushes on every call), single threaded and from every core at once, so shared-counter contention can be quantified rather than inferred.
…ed allocator An `--extern` crate that nothing names is dropped from the crate graph, and the accounting-off run of this benchmark named nothing in `comet`, so the `#[global_allocator]` in `lib.rs` never reached the binary: the "jemalloc" baseline was measuring glibc malloc. The accounting-on run names `comet::alloc_accounting`, so it did link the crate, and every off/on comparison so far was glibc against jemalloc plus the wrapper. Add `extern crate comet` so the crate is always linked, and a jemalloc liveness assertion next to the existing accounting one, so a run against the wrong allocator fails instead of producing a plausible number.
Every build now installs the accounting wrapper and reports the native_allocated tracing metric. The feature guards stay so a --no-default-features build still compiles without the wrapper; the bench header and tracing guide document that opt-out invocation.
CheckedMemoryPool refuses a reservation when the bytes the allocator has actually handed out, plus the request, would exceed Comet's off-heap budget, and otherwise defers to CometUnifiedMemoryPool. The new off-heap pool type greedy_unified_checked selects it. The balance and budget are process-wide, so once any task reaches the budget every task's next reservation is denied; allocations themselves are never refused. Selecting the pool without the alloc-accounting feature is a config error.
github-actions
Bot
added
enhancement
andygrove
marked this pull request as draft
September 15, 2026 14:25
andygrove
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Relates to #4576. Stacks on #5934 and should merge after it; until then the diff includes that PR's commits.
Rationale for this change
Every Comet memory pool only counts what operators voluntarily reserve. Native memory that bypasses the pool (scratch buffers in kernels, intermediate Arrow arrays, allocations inside DataFusion operators that are not reserved) is invisible to it until the executor exceeds its container limit and is killed. #5934 added a process-wide count of the bytes the Rust allocator has actually handed out, but only as a tracing metric.
This PR makes that count available in every build and adds the smallest pool that acts on it: before asking Spark for memory, it refuses a reservation when real native usage plus the request would exceed Comet's off-heap budget. The operators that can spill then spill; the ones that cannot fail the task instead of the executor. The pool is opt-in and the default stays
fair_unified.What changes are included in this PR?
alloc-accountingis now a default cargo feature. All#[cfg(feature)]guards stay, so a build with--no-default-featuresstill compiles without the wrapper. The bench and the tracing guide document the opt-out invocation.CheckedMemoryPool<P: MemoryPool>(native/core/src/execution/memory_pools/checked_pool.rs).try_growreturnsResourcesExhaustedwhennative_allocated + additional > budget, naming the request, the bytes in use, the budget and the reserved total, and otherwise delegates to the inner pool.grow,shrink,register,unregisterandreserveddelegate;memory_limitreports the budget. The inner pool is generic so the gate is unit-testable without a JVM.greedy_unified_checkedforspark.comet.exec.memoryPool:CheckedMemoryPoolaroundCometUnifiedMemoryPool, budget =spark.memory.offHeap.sizexspark.comet.exec.memoryPool.fraction(the valuefair_unifiedalready receives), inside the same task-shared and consumer-tracking wrappers asgreedy_unified. If the native library was built without the feature, selecting this pool is a config error rather than a silently unchecked pool.spark.comet.exec.memoryPooldescription and atuning.mdentry describing what the check is and is not.Semantics worth stating up front: the balance and the budget are both process-wide. When any task pushes real usage to the budget, every task's next non-zero reservation is denied. There is no per-task attribution, and allocations themselves are never refused; this is a reservation gate, not a hard limit. Per-task attribution and changing the default pool are follow-ups.
How are these changes tested?
Rust unit tests:
denies_when_real_bytes_plus_request_exceed_the_budget: pool overUnboundedMemoryPoolwith budget = current balance + 64 MiB; holding a 256 MiB block the pool never heard about makes a 1-bytetry_growfail withResourcesExhaustedand nothing reserved; dropping the block makes the same call succeed.memory_limitreports the budget, successful grows and shrinks reach the inner pool.parse_memory_pool_configacceptsgreedy_unified_checkedin off-heap mode with the memory limit as budget, rejects it in on-heap mode, and (in a--no-default-featuresbuild) rejects it with a message naming the feature.JVM tests in
CometExecSuite:greedy_unified_checked pool completes a sort within budget: a dictionary-heavy sort under the new pool matches Spark and runs natively.greedy_unified_checked pool denies reservations once real native usage exceeds budget: withmemoryPool.fraction=0.000001the budget is a few KiB, below what the executor's native code already holds, so the sort's first reservation is denied with the checked pool's message.Both feature sets were run locally: default (
alloc-accountingon) and--no-default-features --features hdfs-opendal.