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.
|
Follow-up on the Cause
Evidence
The residual in the fixed rows is the sampler's own 1 us gap catching an allocation on the other thread, not a lag. Fix Subtract before delegating. A free cannot fail, so the reason A new test, Re-run of the traced TPC-H SF100 suite with the fix
The four remaining samples have a 1 to 3 us read gap, consistent with the residual above. Result hashes are unchanged and the traced total is 208.2 s against 207.7 s for the |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
This adds visibility into Rust allocations that are absent from the memory pool's declared reservations. AccountingAllocator wraps the selected backend, accumulates signed thread-local deltas, and flushes them to a process-wide balance at 64 KiB or thread exit. native_allocated exposes that balance in tracing. The feature measures allocation sizes without enforcing a limit or changing the memory pool.
I found one P2 issue at ce9b13a351e8 against d1dd302d3fde: on non-MSVC targets, enabling both jemalloc and mimalloc together with alloc-accounting selects none of the allocator definitions. The metric remains enabled but reports zero. The inline comment covers the exact branch and reproduction.
For the individual backend selections, allocation and zeroed allocation update the balance only after success. A failed realloc leaves both the old pointer and its accounting intact. Successful realloc records the size difference. Freeing subtracts before calling the inner allocator, addressing the lag described in the author's follow-up. This matches the ownership and failure rules in Rust's GlobalAlloc contract.
The destructor-free re-entry flag protects initialization of the drift state. try_with falls back to the atomic balance after that state has been destroyed, and ThreadDrift::drop settles its remainder. The balance is an approximate allocation-size sample, with unflushed deltas on live threads. It should not be read as exact RSS, per-task usage, or a memory limit.
No Spark expression, operator result, error mode, or fallback decision changes. The wrapper delegates allocation requests and adds bookkeeping. No new equivalence claim is made for Spark's types, nulls, overflow, or ANSI behavior. Maintained Spark 3.4 and 4.1 sources were unavailable and are not claimed as reviewed.
Validation
The exact allocator module passed five tests without the feature and seven with an installed system-allocator wrapper in an isolated local harness. A separate probe passed allocation failure, zeroed allocation failure, failed realloc ownership, successful zeroed allocation, grow, shrink, and free cases. These are component checks, not a full Comet/JNI build. The configuration probe used the exact allocator-selection block with system-allocator aliases for the two optional backend types.
The public checks show 53 successes and 10 skips. Rust CI ran the five ungated accounting tests, but its default feature set excludes alloc-accounting. The two tests requiring the installed wrapper therefore have no CI result here. The Spark 4.1 execution suite passed 867 tests and consumed the native artifact whose ID and digest match the producer.
Those jobs checked out d97e84153938, with this head and base 5cff668396e7, one commit beyond the assigned base. All six changed files match that tested merge, but the complete trees differ. I credit that as source-matched CI coverage, not execution of the assigned pair or of the accounting-enabled JNI path.
Performance
Without the feature, allocator selection retains the previous behavior and does not add per-allocation bookkeeping. With it enabled, small changes use thread-local state. Large allocations and frees flush to the shared atomic counter, so contention is still relevant for allocation-heavy parallel workloads.
The new benchmark covers small allocation/free cycles, a filled 64 KiB buffer, and vector growth. I checked the optimized code for the extracted filled-buffer body and found that allocation, fill, and free were retained, so I am not raising an allocation-elision finding. This was an assembly check, not a timing measurement or a run of the Criterion target.
The author's TPC-H report gives a 1.4% increase in the sum of per-query medians over three untraced iterations at the earlier 303875f revision. The current-head follow-up is traced. These are useful initial observations, but they do not establish a general overhead bound or isolate atomic contention. I did not reproduce those timings. Could you report the new alloc_overhead microbenchmark with accounting off/on and add a parallel allocation/free case around the 64 KiB flush threshold to quantify shared-counter contention?
Design
Keeping measurement separate from enforcement is a clear boundary. It avoids introducing new allocation failures, pool limits, or retry behavior while collecting evidence about the gap between reservations and allocated bytes. Process-wide scope and exclusions such as C-library allocations and memory mappings are explicit.
The backend selection needs to cover every accepted feature combination. Defining the system fallback as the complement of the selected backend cases, or explicitly rejecting the conflicting combination, would resolve the reported zero-metric failure. The probe should accompany that adjustment.
Abstraction & complexity
The generic wrapper is a useful small abstraction because all three backends share the same bookkeeping. ThreadDrift gives the exit flush a clear owner, and the signed balance handles cross-thread allocation/free ordering without underflow. The new code does not introduce task attribution or enforcement machinery. The duplicated cfg predicates are the one place where complexity has produced an observable gap.
| #[cfg(all( | ||
| feature = "alloc-accounting", | ||
| not(feature = "mimalloc"), | ||
| any(target_env = "msvc", not(feature = "jemalloc")) |
There was a problem hiding this comment.
Correctness
[P2] Could the system fallback also cover the case where both allocator features are enabled, or reject that combination explicitly? On non-MSVC targets, jemalloc,mimalloc,alloc-accounting makes every GLOBAL definition false: each backend excludes the other, and this fallback excludes mimalloc. The process therefore uses an unwrapped default allocator while log_native_allocated is still enabled, so the new metric silently stays zero. I reproduced the selection with this exact cfg block and the exact accounting module: holding an 8 MiB buffer moved the balance with either individual backend configuration, but left it at zero with both features. The probe substitutes System for the backend types, so it tests allocator selection rather than jemalloc/mimalloc behavior. A combination check would prevent publishing a plausible zero metric when accounting was requested.
There was a problem hiding this comment.
Confirmed and fixed in 2d78f87. On main, jemalloc,mimalloc on a non-MSVC target already selects neither arm and falls through to the system allocator; the accounting fallback here excluded mimalloc, so that combination with alloc-accounting matched nothing.
The selection is now a backend module chosen by three cfgs that partition every feature combination: jemalloc where it builds and mimalloc was not requested, mimalloc otherwise, and the system allocator as the exact complement of those two (which is where jemalloc,mimalloc still lands, as on main). The unwrapped #[global_allocator] lives inside the jemalloc and mimalloc modules, so a build without the feature is unchanged and still installs nothing for the system case. The single accounting #[global_allocator] refers to backend::Backend, so a combination with no backend is a compile error rather than a silent zero.
Verified with cargo check on all eight combinations of the three features, plus cargo test --features jemalloc,mimalloc,alloc-accounting alloc_accounting, where a_real_allocation_raises_the_balance (the test that checks the wrapper is really installed for the current feature set) passes.
While doing that I found the two feature-gated tests were flaky under the feature (2 of 20 runs): BALANCE is process-wide and the rest of the crate's tests, plus the 64 MiB block in dealloc_settles_before_delegating, move it concurrently. The same commit makes them noise-proof, and the thread-exit test no longer needs the feature, so it now runs in the default CI build.
…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.
|
Here is the Before the numbers, one correction to the benchmark itself, fixed in ca2f8c8. In edition 2021 an Setup: woody, Ryzen 9 7950X3D (16 cores / 32 threads, one socket), Linux, Times are per alloc/free pair, and for the parallel rows per pair per thread, so a parallel number equal to its single-thread counterpart would mean no interference at all. Reading it:
This is consistent with the 1.4% on the TPC-H SF100 sum of medians reported above, where the executors spend a small fraction of their time in the allocator. |
2 participants
Which issue does this PR close?
Relates to #4576. This is the first of the pieces extracted from the closed
prototype in #4582, and it deliberately stops short of enforcement, so it does
not close that issue.
Rationale for this change
Comet's memory pool counts declared reservations: bytes an operator explicitly
asked for. A lot of real allocation never goes through it — Arrow builders,
expression kernels producing intermediates, decompression buffers, Parquet
metadata,
object_storebuffers, tokio's own machinery. Pool reservations aretherefore a lower bound on Comet's footprint, and the size of the gap is
workload-dependent.
Today that gap is invisible at runtime. Diagnosing an out-of-memory report means
reasoning about it indirectly, and
spark.comet.exec.memoryPool.fractionasksoperators to hand-tune a haircut for a quantity nobody can measure.
#4582 tried to both measure the gap and enforce on it. Reviewing it convinced me
the enforcement half needs a question answered first — how well does a tracked
byte balance actually track RSS on real workloads? — and that question is much
easier to answer if the measurement lands on its own. So this PR is the
measurement only.
What changes are included in this PR?
Behind a new, off-by-default
alloc-accountingcargo feature:native/core/src/alloc_accounting.rs:AccountingAllocator<A>wraps whicheverglobal allocator the build selected and maintains one signed process-wide byte
balance, exposed by
current_balance().native/core/src/lib.rs: installs the wrapper over jemalloc / mimalloc /system. The cfgs are mutually exclusive, so a build without the feature is
exactly the previous arrangement — no wrapper, no per-allocation work.
native/core/src/execution/jni_api.rs: reports the balance as thenative_allocatedtracing metric, logged in the same place asjemalloc_allocatedand alongside the per-thread pool reservations it is meantto be compared against.
tracing.md: documents the feature and the new metric.It is observability only. It never rejects an allocation, never panics, and does
not touch the memory pool. Two things follow from that which are worth calling
out:
Because it cannot fail an allocation,
reallocaccounts after delegating andonly on success. The prototype had to account before delegating, because
panicking after
inner.reallocwould leave a caller unwinding with a stalepointer — a soundness constraint that simply does not exist here, and dropping it
also removes the over-count on a failed realloc.
Per-thread deltas are batched and flushed at 64 KiB, so the common path is a
thread-local add-and-compare rather than an atomic read-modify-write. The
prototype leaked up to 64 KiB of accounting every time a thread died, which
matters because the blocking pool churns on tokio's idle timeout.
ThreadDrift'sdestructor settles the remainder on exit. That is slightly more delicate than it
looks: touching a destructor-bearing thread-local can itself allocate on first
use, so
track()keeps a destructor-free re-entrancy flag and settles re-entrantcalls straight into the shared balance, and uses
try_withso an allocationduring thread teardown cannot panic inside the allocator.
What this does not do
The balance counts
Layoutbytes, not resident pages, so it excludes allocatorfragmentation, jemalloc's retained pages,
mmaped regions, and anything a Cdependency allocates through libc
malloc. It is a lower bound on RSS — just amuch tighter one than pool reservations. It is also process-wide, not per-task.
Whether it is a good enough proxy to enforce on is exactly what I would like to
learn from it before proposing that.
How are these changes tested?
Unit tests in
alloc_accounting.rscover the settle/flush helper and thenegative-balance clamp, and two tests drive real allocations through the
installed allocator: one asserts an 8 MiB allocation moves the reported balance,
and one asserts that drift from exited threads reaches the shared balance.
The thread-exit test is mutation-checked. Each worker allocates a sub-threshold
buffer and hands ownership back before exiting, so the matching free happens on
the main thread after the worker is gone and the destructor is the only path by
which those bytes can be counted. Neutering the destructor fails it with
"balance moved 0 bytes, expected at least 1048576". An earlier version of the
test, where each thread freed what it allocated, passed under the same mutation
and was replaced.
cargo clippy --all-targets -- -D warningsand the full native test suite passunder
default,alloc-accounting,jemalloc,alloc-accounting, andmimalloc,alloc-accounting.Not yet measured: the per-allocation overhead of the wrapper when the feature is
on. That is worth a benchmark before anyone considers enabling it by default,
and I have not done it.