feat: add native allocation accounting for memory observability

[andygrove]

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_store buffers, tokio's own machinery. Pool reservations are
therefore 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.fraction asks
operators 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-accounting cargo feature:

• native/core/src/alloc_accounting.rs: AccountingAllocator<A> wraps whichever
  global allocator the build selected and maintains one signed process-wide byte
  balance, exposed by current_balance().
• native/core/src/lib.rs: selects the backend (jemalloc, mimalloc or system)
  in one place and installs either it or the wrapper over it as the global
  allocator. The three backend cfgs partition every feature combination, so a
  combination matching none fails to compile rather than installing nothing.
• native/core/src/execution/jni_api.rs: reports the balance as the
  native_allocated tracing metric, logged in the same place as
  jemalloc_allocated and alongside the per-thread pool reservations it is meant
  to be compared against.
• native/common/src/bin/analyze_trace.rs: analyzes native_allocated when the
  trace has it, otherwise jemalloc_allocated, and names the counter it used.
• native/core/benches/alloc_overhead.rs: measures the wrapper's per-allocation
  cost with the feature off and on.
• tracing.md: documents the feature, the new metric, and the analyzer's
  counter selection.
• .github/actions/rust-test: lints, tests and checks the feature-on builds,
  which nothing else in CI compiles.

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, realloc accounts after delegating and
only on success. The prototype had to account before delegating, because
panicking after inner.realloc would leave a caller unwinding with a stale
pointer, a soundness constraint that simply does not exist here, and dropping it
also removes the over-count on a failed realloc. dealloc goes the other way and
settles before delegating, because jemalloc drops its own count at the start of
a large free and then spends milliseconds unmapping the pages; accounting
afterwards made native_allocated read above jemalloc_allocated during task
teardown.

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's
destructor 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-entrant
calls straight into the shared balance, and uses try_with so an allocation
during thread teardown cannot panic inside the allocator.

What this does not do

The balance counts Layout bytes, not resident pages, so it excludes allocator
fragmentation, jemalloc's retained pages, mmaped regions, and anything a C
dependency allocates through libc malloc. It is a lower bound on RSS, just a
much tighter one than pool reservations. It is also process-wide, not per-task,
and approximate to within 64 KiB per live thread. 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.rs cover the settle/flush threshold, the
dealloc ordering (the inner free must see the balance already reduced), the
realloc accounting (the balance moves by the size difference, after
delegating), and thread-exit settlement of remaining drift. The thread-exit test
injects a drift directly into the exiting thread's cell so the destructor is the
only path to the shared balance; it is confined to the default build because with
the wrapper installed, teardown's own allocations would flush the drift anyway.
With the feature on, a further test allocates 256 MiB through the real global
allocator and checks the balance moved, which is what catches a feature
combination that installs nothing.

CI runs the default build as before, and the Rust test job now also lints the
jemalloc,alloc-accounting build with --all-targets, runs the accounting tests
with the wrapper installed over jemalloc, and checks the alloc-accounting-only
build. Locally, clippy passes on all eight combinations of the three allocator
features.

The alloc_overhead benchmark, run with the feature off and on against jemalloc
(full numbers and setup in
#5934 (comment)):
the thread-local path costs about 2 ns per alloc/free pair, an uncontended flush
about 1.5 ns more, and the worst case for the shared counter (32 threads each
allocating and freeing exactly 64 KiB blocks with nothing in between) 25 ns per
pair, or 8% over jemalloc's own contended path. Filling a 64 KiB buffer and
growing a builder show no measurable difference. TPC-H SF100 with the feature on
ran 1.4% slower on the sum of per-query medians, with identical results
(#5934 (comment)).

[andygrove]

I ran TPC-H at SF100 against this PR to check the feature end to end and to get a first read on the per-allocation overhead that the description lists as not yet measured.

Setup

• Commit 303875f built twice in separate worktrees: make release COMET_FEATURES="jemalloc,alloc-accounting" and, as the baseline, make release COMET_FEATURES="jemalloc".
• Single node, Spark 4.1.1 standalone, 2 executors x 8 cores, 16 GB heap + 16 GB off-heap per executor, local Parquet at SF100. Driven by benchmarks/tpc/run.py.
• One traced run of each build (spark.comet.tracing.enabled=true), then three untraced iterations of each build for timing.

Correctness

Result hashes and row counts are identical between the two builds for all 22 queries, in both the traced and untraced runs. No errors, spills, or task retries in either.

Overhead (untraced, median of 3 iterations per query)

	jemalloc	jemalloc,alloc-accounting
sum of per-query medians	200.3 s	203.1 s (+1.4%)
per-iteration totals	199.7 / 206.1 / 206.8 s	202.9 / 208.3 / 209.1 s

The per-iteration ranges overlap. Only Q21 is slower in all three iterations (29.9 s vs 30.9 s, about +3%). Q10 is bimodal (7.5 s or 13 to 14 s) on both builds, so that is unrelated to this change.

What native_allocated shows

• Peaks around 2.1 GB per executor and tracks jemalloc_allocated at a median ratio of 0.90, which is consistent with jemalloc metadata and size-class rounding.
• It does expose the gap the PR is after. At the sample points below the per-thread pool reservations summed to far less than the bytes actually handed out:

query	peak native_allocated	pool total at that point
Q7	1627 MB	765 MB
Q8	2011 MB	428 MB
Q16	1192 MB	0 MB
Q17	1292 MB	34 MB

• One thing I could not explain: about 2% of sample points (mostly in Q10, Q17 and Q18) report native_allocated above jemalloc_allocated, by up to 160 MB. Both are logged a few microseconds apart on the same thread, and the un-flushed per-thread drift can only account for a few MB in total, so this is not the 64 KiB batching. It may be jemalloc's stats being approximate under concurrent frees on other threads, but it is worth understanding before treating the value as a strict lower bound relative to jemalloc.

Traces and result JSON are available if useful.

[andygrove]

Follow-up on the native_allocated > jemalloc_allocated samples from the comment above: root cause found, fixed in ce9b13a.

Cause

AccountingAllocator::dealloc called inner.dealloc first and subtracted afterwards, mirroring alloc and realloc. jemalloc 5.3 decrements stats.allocated at the very start of a large free (large_dalloc_prep_impl runs arena_large_dalloc_stats_update before large_dalloc_finish_impl releases the extent), and for blocks above its 8 MiB oversize threshold the release is eager: the pages go straight back to the kernel. Unmapping 100+ MB of resident memory takes milliseconds. For that whole window jemalloc had already forgotten the block but the balance still carried it, so native_allocated read high by exactly the size of the block being freed. The affected samples were all during task teardown, where many large buffers are dropped back to back.

Evidence

• It is not a sampling race. The two counters are logged 1 us apart on the same thread, and the anomalous samples have the same read gap as normal ones. The excess persists at a fixed value across consecutive samples from different threads, then the balance drops by that amount while jemalloc stays flat.
• A standalone repro (one thread allocating, touching and freeing a block in a loop, a sampler reading jemalloc then the balance) reproduces it deterministically and the episode length tracks the duration of the free:

ordering	block	samples with native > jemalloc	longest episode	slowest free
subtract after (previous)	128 MiB	16.5%	9.3 ms	9.3 ms
subtract after (previous)	4 MiB	4.4%	413 us	413 us
subtract before (fixed)	128 MiB	0.01%	15 us	11 ms
subtract before (fixed)	4 MiB	3.2%	25 us	355 us

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 alloc and realloc account after the fact does not apply to dealloc. With this the balance never includes memory the allocator has already handed back, so the "lower bound" relationship to jemalloc holds at every instant rather than modulo in-flight frees.

A new test, dealloc_settles_before_delegating, wraps a recording inner allocator and asserts the balance has already dropped by the time the inner free is called. Under the previous ordering it fails with inner dealloc saw balance 67182807, expected at most 33628375.

Re-run of the traced TPC-H SF100 suite with the fix

	before (303875f)	after (ce9b13a)
samples with native > jemalloc	156 of 9269	4 of 9272
worst excess	160 MB	9.4 MB
median native/jemalloc ratio	0.90	0.90

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 jemalloc-only baseline.

[sunchao]

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.

[sunchao]

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.

[andygrove]

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.

[andygrove]

Here is the alloc_overhead microbenchmark with accounting off and on, plus the parallel case around the flush threshold that was asked for (added in 2d78f87 as threshold_churn).

Before the numbers, one correction to the benchmark itself, fixed in ca2f8c8. In edition 2021 an --extern crate that nothing names is dropped from the crate graph. The accounting-off run of this benchmark named nothing in comet, so the rlib and its #[global_allocator] were never linked and the "jemalloc" baseline was in fact glibc malloc (the binary was 4 MB and contained no jemalloc symbols; with the feature on it names comet::alloc_accounting and is 28 MB with jemalloc linked). The bench now has an extern crate comet 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. The comparison below is jemalloc on both sides.

Setup: woody, Ryzen 9 7950X3D (16 cores / 32 threads, one socket), Linux, --features jemalloc vs --features jemalloc,alloc-accounting, Criterion with 2 s warm-up and 4 s measurement. Single-thread cases are pinned to one core with taskset because this CPU has two chiplets with different cache and unpinned single-thread numbers swing by 2x between runs; a pinned off-vs-off rerun agrees within 2%. Parallel cases use all 32 threads.

case	off (jemalloc)	on (jemalloc + accounting)	delta
alloc_free_16b	4.11 ns	6.05 ns	+1.9 ns
alloc_free_256b	4.26 ns	6.20 ns	+1.9 ns
alloc_free_4096b	6.10 ns	7.89 ns	+1.8 ns
alloc_free_32kb (never flushes)	23.0 ns	25.1 ns	+2.1 ns
alloc_free_64kb (flushes on every alloc and free)	200.9 ns	204.5 ns	+3.6 ns
alloc_fill_free_64kb	616 ns	611 ns	none (p = 0.56)
grow_vec_to_64kb	41.56 µs	41.38 µs	none
parallel_alloc_free_32kb_x32 (never flushes)	47.6 ns	49.0 ns	+1.4 ns (p = 0.24)
parallel_alloc_free_64kb_x32 (every thread flushes on every call)	295 ns	320 ns	+25 ns (+8%)

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:

• The thread-local path costs about 2 ns per alloc/free pair, or 1 ns per call, independent of size. That is the whole cost when a thread's drift stays under 64 KiB, which is the 32 KiB rows: the alloc and the free cancel in the thread-local cell and the shared counter is never touched, single-threaded or on 32 threads.
• A flush is two uncontended atomic read-modify-writes and adds another 1.5 ns on top of that single-threaded (64 KiB row). 64 KiB blocks are above jemalloc's thread-cache limit, so the allocator's own cost dominates at 200 ns.
• Contention is the last row: 32 threads each doing two atomic adds on the same cacheline per iteration, with nothing else in between, cost 25 ns per pair, or 8% over jemalloc's own contended large-allocation path. That is the upper bound for the shared counter, and it needs every core to do nothing but allocate and free exactly-threshold blocks. Anything below the threshold never reaches the counter, and anything doing real work between allocations amortizes it.
• The two rows closest to what Comet actually does, filling a batch-sized buffer and growing a builder, show no measurable difference.

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.

[sunchao]

Re-reviewed ca2f8c875318 against d1dd302d3fde. The previous P2 is fixed: the mixed-feature accounting case now installs the wrapper. The exact selector passed all eight feature combinations in a local harness using system-allocator aliases for the optional backend types. The allocator module passed six default and seven accounting-enabled tests, and the failure/realloc probe passed. Removing the exit flush makes the newly ungated thread-exit test fail as intended.

There is one new P2 in the benchmark liveness guard: enabling both allocator features selects System in the library but still requires jemalloc statistics in the benchmark. The inline comment identifies the condition to align. An extracted selector/guard harness reproduced that mismatch using a system-backed statistics test double. This was not a run with the actual jemalloc or mimalloc libraries.

The explicit crate linkage and added parallel threshold cases address the earlier benchmark requests. The revised report measures an 8% increase for the parallel 64 KiB case on the author's host. I treat that as a reported workload-specific result, not a general overhead bound, and did not reproduce the timings.

The current checks show 53 successes and 10 skips. The detailed CI log read stalled, so I am not claiming a verified CI checkout, accounting-enabled JNI execution, or a full local Comet build. No new issue was found in allocation ownership, failure accounting, or the observability-only design.

[sunchao]

[P2] Match the liveness guard to the selected allocator

Correctness

Could this guard use the same condition as the jemalloc backend in lib.rs? On non-MSVC targets, --features jemalloc,mimalloc deliberately selects the system allocator, but this condition still enables the jemalloc assertion. The held 8 MiB allocation therefore goes through System, while the guard checks jemalloc's separate allocation statistics. The benchmark aborts before collecting measurements for that accepted combination, both with and without alloc-accounting. Restricting this guard to all(not(target_env = "msvc"), feature = "jemalloc", not(feature = "mimalloc")), with the no-op using its complement, keeps the linkage check for actual jemalloc builds and permits the mixed-feature system baseline.

[andygrove]

Confirmed and fixed in ac51c5c. With jemalloc,mimalloc the library selects the system allocator but the bench still asserted on jemalloc's counters, so that combination aborted before measuring anything.

Rather than mirror the three-term predicate in the bench, lib.rs now exports the selection it made: each backend module declares a NAME and the crate re-exports it as comet::ALLOCATOR_BACKEND ("jemalloc", "mimalloc" or "system"). The bench announces that name at startup and runs the jemalloc liveness assertion only when the library reports jemalloc. The selection is still written once, so the guard cannot drift from it the way a copied cfg could, which is how the previous P2 crept in.

Verified with clippy on the bench target for all eight feature combinations, and a criterion --test run of alloc_overhead under jemalloc,mimalloc and jemalloc,mimalloc,alloc-accounting, both of which now report measuring the system allocator backend and complete; jemalloc alone still runs the assertion and passes.

[comphead]

checking this today

[sunchao]

Re-reviewed ac51c5c0 against d1dd302d, since review 5211754611. The benchmark P2 is fixed: the guard now reads the library’s selected backend, so mixed allocator features correctly use the system baseline. No new or remaining P1/P2 findings.

I reran the earlier extracted selector/guard reproduction. Both mixed-feature cases fail on the prior commit and pass on this head. All eight current feature combinations pass. Two negative controls still fail when jemalloc is selected but its statistics report zero. These use System-backed backend/statistics test doubles, not the actual jemalloc/mimalloc libraries or Criterion timings. The accounting implementation and timed benchmark loops are byte-identical to the prior review.

At September 15, 18:04 UTC, checks show 7 successes, 11 skips, 20 cancellations and 1 failure. Required Checks failed on upstream cancellations. Rust formatting passed on merge 95ccd751, whose six authored files match this head. The Rust job was cancelled during clippy. Full runtime and accounting-enabled JNI validation remain unverified.

[comphead]

Reviewed as an observability-only change, so the usual Spark-semantics questions do not apply here. Findings are inline. A few that do not anchor to a single line:

Comment volume. The module doc, the lib.rs backend header, the bench module doc, and several test doc comments restate the PR description at paragraph length. The parts genuinely worth keeping are the non-obvious invariants: why dealloc settles before delegating, why IN_TRACK has to be const-initialized and destructor-free, and why post-hoc realloc accounting is safe here but would not be under enforcement. The surrounding narrative could come out.

Scope. Appropriately scoped, and splitting measurement from enforcement is the right call. The benchmark is the one piece that could reasonably ship separately.

Tests. The allocator contract is well covered (settle/flush, clamp, dealloc ordering, thread-exit drift, real-allocation liveness), and mutation-checking the thread-exit test is a nice touch. Two gaps inline: nothing exercises realloc, and the thread-exit test is not mutation-proof in the feature-on build. Nothing here has a SQL surface, so no test belongs in an .slt file.

Performance. Default builds are unaffected. With the feature on, every allocation pays a thread-local read plus a compare, and current_balance() in executePlan is one relaxed load per JNI call. The missing piece is the overhead number itself.

No blockers. The analyze_trace gap is the one I would want closed before merge, since without it the feature's headline use case does not work end to end.

[comphead]

analyze_trace never sees this metric. native/common/src/bin/analyze_trace.rs:113 matches jemalloc_allocated by exact name and drops everything else in the trailing else { continue; }, so native_allocated is discarded.

That tool is what computes excess = allocated - pool_total, which is the comparison this PR exists to enable. In the combination the PR specifically motivates (alloc-accounting without jemalloc) it will report zero allocated and no excess.

Suggest making that arm accept either name, deciding which wins when both are logged, and updating the tracing.md prose at line 69 that currently names only jemalloc_allocated as the process-wide counter.

[andygrove]

Fixed in 4b71ad9. analyze_trace now recognises both counters. When a trace carries native_allocated it is analyzed against that alone, since it counts only what Rust code holds from the allocator; otherwise it falls back to jemalloc_allocated. The rule is order-independent: the first native_allocated event resets the peaks and violations so the report never mixes the two sources, and the output names the counter it used in every label. A trace with neither counter is rejected with a message naming the two features. tracing.md describes the selection and the sample output matches the new labels.

Checked with synthetic traces: 32 MiB native vs 8 MiB pool reports 24 MiB excess with or without jemalloc events present, in either event order, and a trace with only jvm_heap_used exits 1.

[andygrove]

The zero-reservation boundary is fixed in 5a4334a. The comparison now waits only until the trace has produced at least one pool sample, so an observed zero reservation is compared against like any other value: a zero pool sample with 32 MiB native_allocated reports 32 MiB excess, and reservations falling from 8 MiB to zero under a steady 32 MiB report a 32 MiB peak. When the trace never carries a pool sample the tool says so instead of reporting that allocation never exceeded reservations. The same commit reflows the tracing.md table that failed the Preflight prettier check.

[comphead]

Nothing in CI builds this feature. pr_benchmark_check.yml:55 runs cargo clippy --all-targets --workspace with default features, and no workflow passes --features. So the three backend arms, the wrapper, and the bench's jemalloc liveness check are compiled only on developer machines.

The mod backend partition is nicely self-checking (zero matches gives an unresolved backend, two gives a duplicate module), but only for combinations someone actually compiles. Adding cargo check --features alloc-accounting plus one jemalloc,alloc-accounting check to an existing job would keep that property honest.

[andygrove]

Fixed in 4b71ad9. The rust-test composite action now lints datafusion-comet with --all-targets --features jemalloc,alloc-accounting (which covers the bench guards for the jemalloc case), runs the alloc_accounting tests with the wrapper installed over jemalloc, and cargo checks the alloc-accounting-only build for the system-allocator arm. That is the Linux Rust test job; the extra steps reuse its cache.

[comphead]

The PR description still says "Not yet measured: the per-allocation overhead of the wrapper when the feature is on", but this benchmark is in the diff. 225 lines whose whole purpose is bounding that overhead, landing without the number, leaves a reviewer unable to judge the cost of turning the feature on.

Either paste the off versus alloc-accounting comparison this header describes, or land the bench separately once you have it.

[andygrove]

The numbers were posted in #5934 (comment) after the description was written. The description now summarises them and links there.

[comphead]

This stops being mutation-proof once the feature is on. The worker leaves LOCAL_DRIFT at 1 << 40, and thread teardown itself allocates, so the next track call sees drift.unsigned_abs() >= SETTLE_THRESHOLD and flushes immediately. BALANCE then moves by roughly INJECTED even with ThreadDrift::drop neutered.

The doc comment already says the default build is the one that catches it. Worth gating the test #[cfg(not(feature = "alloc-accounting"))] so it cannot quietly become a tautology in the build that actually ships the wrapper.

[andygrove]

Confirmed: with the wrapper installed, teardown's own allocations call track, see the oversized drift, and flush it before the destructor runs. Gated on not(feature = "alloc-accounting") in 4b71ad9, with the doc comment saying why.

[comphead]

This and the comment above it look unnecessary. assert_backend_is_live() names comet::ALLOCATOR_BACKEND unconditionally (no cfg), which in edition 2021 already resolves through the extern prelude and pulls the rlib, and with it the #[global_allocator], into the crate graph. The comment's premise, that nothing here names the crate, does not hold.

[andygrove]

Confirmed, it became redundant once the liveness check named comet::ALLOCATOR_BACKEND unconditionally. Removed in 4b71ad9; the note that naming the crate is what links the allocator moved onto that reference, since that is now the thing holding the link. Verified by running the bench under --features jemalloc with --test, where the jemalloc liveness assertion still passes.

[comphead]

clamp_balance exists only so a_transiently_negative_balance_reports_as_zero can assert on isize::max(0) as usize. Inlining it into current_balance and dropping that test loses no coverage. The clamp's real behavior is already implied by current_balance's documented contract.

[andygrove]

Inlined into current_balance and the test dropped in 4b71ad9.

[comphead]

Worth stating the accuracy bound here and in the tracing.md row: the returned value can lag the true balance by up to SETTLE_THRESHOLD times the number of live threads, since each thread holds un-flushed drift. Immaterial against GiB-scale footprints, but a reader comparing native_allocated against pool reservations byte-for-byte should know the number is approximate.

[andygrove]

Added to the current_balance doc and the tracing.md row in 4b71ad9: up to SETTLE_THRESHOLD per live thread.

[comphead]

realloc is the one method whose accounting order this PR deliberately changed from the prototype (after delegating rather than before), and nothing tests it. dealloc_settles_before_delegating already has the Recording inner-allocator harness. A grow and a shrink through it, asserting the balance moves by the delta rather than by the full new size, would lock the new ordering in.

[andygrove]

Added realloc_accounts_the_size_difference_after_delegating in 4b71ad9. It grows a 64 MiB block to 96 MiB and shrinks it to 32 MiB through the Recording harness, asserting the balance moves by +32 MiB and then -64 MiB (within a 16 MiB margin, so accounting the full new size would fail), and that the inner realloc sees the balance still carrying the old size, which pins the after-delegating order. The harness now records the balance at realloc as well as at dealloc.

[comphead]

Question: is it intentional that this is ungated? It makes current_balance() public API that always returns 0 in default builds, and a locally constructed AccountingAllocator mutates the process-wide BALANCE even when the wrapper is not installed, which dealloc_settles_before_delegating relies on. I assume the point is keeping the unit tests running in the default build, just want to confirm that is the reason rather than an oversight.

[andygrove]

Intentional. The module is ungated so the unit tests, including the dealloc-ordering test that drives a local AccountingAllocator against the shared balance, run in the default build that CI actually executes. current_balance documents that it returns 0 when the wrapper is not installed. Now that CI also runs the tests with the feature on, gating would be possible, but it would take the ordering tests out of the default build for no gain.

[comphead]

Question: this adds public crate API whose only consumer is a benchmark assertion. The reasoning in the doc comment is sound, since lib.rs owns the selection and re-deriving it from the feature set in the bench could disagree with it. Not arguing against it, just wondering whether #[doc(hidden)] is worth it to keep this out of the crate's documented surface.

[andygrove]

Added in 4b71ad9.

[sunchao]

Rechecked unchanged head ac51c5c0 against d1dd302d after the new review. One newly verified P2 remains, already covered by the analyzer thread.

Correctness

[P2] The analyzer ignores native_allocated. With accounting enabled and jemalloc disabled, its allocation side stays zero. In an isolated run of the exact post-parsing analyzer logic, native_allocated=32 MiB and pool reservations of 8 MiB produced zero excess and the “OK” report, in either event order. The equivalent jemalloc fixture reported 24 MiB excess. The raw trace remains usable, but the documented CLI misses this gap. Could the analyzer support the new counter with a defined source selection, or explicitly reject unsupported traces, and align its labels and documentation?

Other new review points

Two other claims need qualification:

• Off/on measurements are already posted, although the PR description is stale. They remain author-reported, workload-specific results.
• Filtering to fill/growth benchmarks still runs the guards: Criterion 0.7 calls every group target before filtering individual benchmark IDs.

The exact allocator module passed 6 default and 7 System-backed accounting tests. Independent grow/shrink and failed-realloc probes also passed. Removing the thread-exit settlement fails the default test but passes the accounting-enabled test, confirming that the latter alone does not prove destructor coverage. Feature-enabled CI coverage is still absent from the workflow.

At 2026-09-15 19:40 UTC, checks remain 7 successful, 11 skipped, 20 cancelled and 1 failed. Rust clippy was cancelled, and Required Checks failed on cancellations. No full Comet/JNI, actual jemalloc/mimalloc, or timing run was completed locally. This is a COMMENT follow-up. The existing approval is left in place.

[andygrove]

Third-round feedback addressed in 4b71ad9; each inline thread has a reply.

On the comment volume: trimmed in 4b71ad9. The module doc, the lib.rs header, the bench header and the test docs now keep only the invariants: why dealloc settles before delegating, why IN_TRACK is const-initialized and destructor-free, why post-hoc realloc accounting is safe here but not under enforcement, and why the thread-exit test is confined to the default build.

[sunchao]

Rechecked 4b71ad9c against d1dd302d. The counter-selection portion of the analyzer finding is fixed: the original 32 MiB native / 8 MiB pool case now reports 24 MiB excess. Native wins over jemalloc in either order, including a zero native value, labels identify the chosen source, and traces without an allocation counter are rejected.

Correctness

[P2] One analyzer boundary remains. The carried-over pool_total > 0 guard treats an explicitly observed zero reservation like missing data. This exclusion already exists in the base jemalloc analyzer, so this is not a new jemalloc regression. The newly supported native-counter path inherits it. With a zero pool sample and 32 MiB native_allocated, the exact post-parsing analyzer logic reports zero excess and “OK”. With reservations falling from 8 MiB to zero, it reports a 24 MiB peak instead of 32 MiB. Zero reservations are emitted by the tracing path. Could we compare against an observed zero and handle missing pool samples separately, rather than claim that allocation never exceeded reservations? This extends the existing analyzer thread, without another inline.

The updated allocator module passed 4 default and 4 System-backed feature tests. The default thread-exit test still detects a removed destructor. All 8 backend combinations passed isolated guard checks, and both negative controls failed as expected. These use test doubles and do not establish full Comet/JNI or real jemalloc/mimalloc execution.

At 2026-09-15 20:38 UTC, checks show 4 successful, 15 skipped and 2 failed. Preflight failed Markdown formatting in tracing.md, and Required Checks propagated that failure. The new feature-enabled CI commands have not run.

[sunchao]

Rechecked 5a4334a1 against d1dd302d. The remaining analyzer finding is fixed. Both an observed zero pool and reservations falling from 8 MiB to zero now report 32 MiB excess for 32 MiB allocated. Missing pool samples are reported separately, while missing allocation counters still fail and native/legacy counter precedence is preserved.

All 23 current-head typed-event fixtures passed, and the previous revision reproduces both failures. These execute the exact post-parsing logic, not the JSON parser or full Comet/JNI. This revision adds no analyzer regression tests. The allocator implementation and runtime hot path are unchanged.

No new or remaining P1/P2 findings. Preflight, including Markdown formatting, passed. The completed Rust CI log now confirms 1,455 default tests and four jemalloc,alloc-accounting tests passed, plus feature-enabled Clippy and the --features alloc-accounting compile. These ran on merge 96bfc50d, whose eight authored files match this head but whose base and full tree differ from the assigned pair.

At September 15, 21:42 UTC, CI shows 16 successful, 13 skipped and six running checks. Spark integration and TPC checks are still pending.

[comphead]

Thanks @andygrove lets try native alloc