[PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove opened a new pull request, #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934

   ## 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`: 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 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.
   - `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, `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.
   
   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, `mmap`ed 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.
   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 helper and the
   negative-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 warnings` and the full native test suite pass
   under `default`, `alloc-accounting`, `jemalloc,alloc-accounting`, and
   `mimalloc,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.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove merged PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5687513783

   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4016507834


##########
native/core/src/lib.rs:
##########
@@ -65,28 +65,66 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// The global allocator is the selected backend (jemalloc, mimalloc, or the system allocator),
+// optionally wrapped in `AccountingAllocator` when the `alloc-accounting` feature is on. The cfgs
+// below are mutually exclusive so exactly one `#[global_allocator]` is defined; a build without
+// the feature is byte-for-byte the previous arrangement, with no wrapper and no per-allocation
+// work.
+
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
-    not(feature = "mimalloc")
+    not(feature = "mimalloc"),
+    not(feature = "alloc-accounting")
 ))]
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
 
 #[cfg(all(
     feature = "mimalloc",
-    not(all(not(target_env = "msvc"), feature = "jemalloc"))
+    not(all(not(target_env = "msvc"), feature = "jemalloc")),
+    not(feature = "alloc-accounting")
 ))]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
+#[cfg(all(
+    not(target_env = "msvc"),
+    feature = "jemalloc",
+    not(feature = "mimalloc"),
+    feature = "alloc-accounting"
+))]
+#[global_allocator]
+static GLOBAL: alloc_accounting::AccountingAllocator<Jemalloc> =
+    alloc_accounting::AccountingAllocator::new(Jemalloc);
+
+#[cfg(all(
+    feature = "mimalloc",
+    not(all(not(target_env = "msvc"), feature = "jemalloc")),
+    feature = "alloc-accounting"
+))]
+#[global_allocator]
+static GLOBAL: alloc_accounting::AccountingAllocator<MiMalloc> =
+    alloc_accounting::AccountingAllocator::new(MiMalloc);
+
+// Accounting over the system allocator: neither mimalloc nor a usable jemalloc.
+#[cfg(all(
+    feature = "alloc-accounting",
+    not(feature = "mimalloc"),
+    any(target_env = "msvc", not(feature = "jemalloc"))

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



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4017097065


##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,202 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator. jemalloc keeps its own count of bytes it has
+/// served; if it is not the global allocator of this binary that count stays at zero, and a
+/// "jemalloc" baseline would in fact be the system allocator.
+#[cfg(feature = "jemalloc")]

Review Comment:
   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.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4020221925


##########
native/core/src/execution/jni_api.rs:
##########
@@ -132,6 +132,18 @@ fn log_jemalloc_usage() {
     log_memory_usage("jemalloc_allocated", allocated.read().unwrap() as u64);
 }
 
+/// Reports the bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Logged alongside the per-thread pool reservations so the two can be compared directly: a large
+/// and growing excess is native memory the pool is not accounting for.
+#[cfg(feature = "alloc-accounting")]
+fn log_native_allocated() {

Review Comment:
   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.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5683774411

   checking this today


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["sunchao (via GitHub)" <gi...@apache.org>]

sunchao commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4012178025


##########
native/core/src/lib.rs:
##########
@@ -65,28 +65,66 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// The global allocator is the selected backend (jemalloc, mimalloc, or the system allocator),
+// optionally wrapped in `AccountingAllocator` when the `alloc-accounting` feature is on. The cfgs
+// below are mutually exclusive so exactly one `#[global_allocator]` is defined; a build without
+// the feature is byte-for-byte the previous arrangement, with no wrapper and no per-allocation
+// work.
+
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
-    not(feature = "mimalloc")
+    not(feature = "mimalloc"),
+    not(feature = "alloc-accounting")
 ))]
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
 
 #[cfg(all(
     feature = "mimalloc",
-    not(all(not(target_env = "msvc"), feature = "jemalloc"))
+    not(all(not(target_env = "msvc"), feature = "jemalloc")),
+    not(feature = "alloc-accounting")
 ))]
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
+#[cfg(all(
+    not(target_env = "msvc"),
+    feature = "jemalloc",
+    not(feature = "mimalloc"),
+    feature = "alloc-accounting"
+))]
+#[global_allocator]
+static GLOBAL: alloc_accounting::AccountingAllocator<Jemalloc> =
+    alloc_accounting::AccountingAllocator::new(Jemalloc);
+
+#[cfg(all(
+    feature = "mimalloc",
+    not(all(not(target_env = "msvc"), feature = "jemalloc")),
+    feature = "alloc-accounting"
+))]
+#[global_allocator]
+static GLOBAL: alloc_accounting::AccountingAllocator<MiMalloc> =
+    alloc_accounting::AccountingAllocator::new(MiMalloc);
+
+// Accounting over the system allocator: neither mimalloc nor a usable jemalloc.
+#[cfg(all(
+    feature = "alloc-accounting",
+    not(feature = "mimalloc"),
+    any(target_env = "msvc", not(feature = "jemalloc"))

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



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["comphead (via GitHub)" <gi...@apache.org>]

comphead commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4018853380


##########
native/core/src/execution/jni_api.rs:
##########
@@ -132,6 +132,18 @@ fn log_jemalloc_usage() {
     log_memory_usage("jemalloc_allocated", allocated.read().unwrap() as u64);
 }
 
+/// Reports the bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Logged alongside the per-thread pool reservations so the two can be compared directly: a large
+/// and growing excess is native memory the pool is not accounting for.
+#[cfg(feature = "alloc-accounting")]
+fn log_native_allocated() {

Review Comment:
   `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.



##########
native/core/Cargo.toml:
##########
@@ -114,6 +114,12 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
 # Default builds carry zero Delta surface.
 contrib-delta = ["dep:comet-contrib-delta"]
 
+# Observability for real native memory usage. Wraps the global allocator to track the bytes it
+# hands out, and reports the total as the `native_allocated` tracing metric so it can be compared
+# against the memory pool's reservations. Never rejects an allocation. Off by default; a build
+# without it has no wrapper and no per-allocation work.
+alloc-accounting = []

Review Comment:
   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.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.

Review Comment:
   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.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;

Review Comment:
   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.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature
+// combination, so exactly one definition exists, and each backend predicate is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop
+// silently: a feature combination with no backend would fail to compile rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
     not(feature = "mimalloc")
 ))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+    pub type Backend = tikv_jemallocator::Jemalloc;
+    pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+    pub const NAME: &str = "jemalloc";
 
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]

Review Comment:
   This `#[cfg(not(feature = "alloc-accounting"))] #[global_allocator]` is repeated verbatim in the mimalloc module, and its absence from the system module is what forces the `#[cfg_attr(..., allow(dead_code))]` there. Two statics at a single site collapse all three:
   
   ```rust
   #[cfg(not(feature = "alloc-accounting"))]
   #[global_allocator]
   static GLOBAL: backend::Backend = backend::BACKEND;
   
   #[cfg(feature = "alloc-accounting")]
   #[global_allocator]
   static GLOBAL: alloc_accounting::AccountingAllocator<backend::Backend> =
       alloc_accounting::AccountingAllocator::new(backend::BACKEND);
   ```
   
   The one cost is that the default plus system build then installs `std::alloc::System` explicitly rather than leaving the default in place. Same allocator either way, but it does give up the "byte-for-byte the previous arrangement" property the header comment claims, so your call whether that is worth the deduplication.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {

Review Comment:
   `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.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+    assert_backend_is_live();
+    assert_accounting_is_live();
+    let mut group = c.benchmark_group("alloc_overhead");
+    for size in [16usize, 256, 4096] {
+        group.throughput(Throughput::Elements(1));
+        group.bench_function(format!("alloc_free_{size}b"), |b| {
+            b.iter(|| {
+                let v: Vec<u8> = Vec::with_capacity(black_box(size));
+                black_box(&v);
+            });
+        });
+    }
+    group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {
+    let mut group = c.benchmark_group("alloc_overhead");
+    group.throughput(Throughput::Bytes(64 * 1024));
+    group.bench_function("alloc_fill_free_64kb", |b| {
+        b.iter_batched(
+            || (),
+            |()| {
+                let v: Vec<u8> = vec![1u8; black_box(64 * 1024)];
+                black_box(v.len())
+            },
+            BatchSize::SmallInput,
+        );
+    });
+    group.finish();
+}
+
+/// Repeated growth, which is the `realloc` path: a builder doubling its buffer.
+fn growth_churn(c: &mut Criterion) {

Review Comment:
   The liveness asserts only run from `small_churn` and `threshold_churn`, so `cargo bench --bench alloc_overhead -- alloc_fill_free_64kb` (or `grow_vec_to_64kb`) measures with no guard at all, which is exactly the silent-wrong-allocator case the header says would be worse than no number. Putting both asserts behind a single `Once` that every bench function calls would close that.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;

Review Comment:
   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.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let new_ptr = self.inner.realloc(ptr, layout, new_size);
+        if !new_ptr.is_null() {
+            // Accounting after the fact is only safe because this allocator cannot fail the
+            // allocation or unwind. A variant that enforced a limit would have to decide *before*
+            // delegating: `realloc` may free or move the old block, and a caller that never
+            // received the new pointer would free the stale one while unwinding.
+            //
+            // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast
+            // wraps.
+            track(new_size as isize - layout.size() as isize);
+        }
+        new_ptr
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{Mutex, MutexGuard};
+
+    /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
+    /// sees every other test's allocations. The tests that move it by tens of megabytes take this
+    /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
+    /// making each window microseconds wide and each expected move far larger than anything else
+    /// allocates in that time.
+    static SERIAL: Mutex<()> = Mutex::new(());
+
+    fn serial() -> MutexGuard<'static, ()> {
+        SERIAL
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    #[test]
+    fn settle_accumulates_below_the_threshold() {
+        let drift = Cell::new(0);
+        settle(&drift, 1024);
+        // A flush would have reset the drift to zero, so this alone shows the shared balance was
+        // not touched. Reading `BALANCE` here would race with every other test's allocations.
+        assert_eq!(drift.get(), 1024, "small delta stays thread-local");
+    }
+
+    #[test]
+    fn settle_flushes_at_the_threshold() {
+        let drift = Cell::new(0);
+        settle(&drift, SETTLE_THRESHOLD);
+        assert_eq!(drift.get(), 0, "drift resets once flushed");
+    }
+
+    #[test]
+    fn settle_flushes_negative_drift() {
+        let drift = Cell::new(0);
+        settle(&drift, -SETTLE_THRESHOLD);
+        assert_eq!(drift.get(), 0);
+    }
+
+    #[test]
+    fn a_transiently_negative_balance_reports_as_zero() {
+        assert_eq!(clamp_balance(-1), 0);
+        assert_eq!(clamp_balance(isize::MIN), 0);
+        assert_eq!(clamp_balance(0), 0);
+        assert_eq!(clamp_balance(4096), 4096);
+    }
+
+    /// A real allocation must move the reported balance: this is the one test that checks the
+    /// wrapper is actually installed as the global allocator for the current feature set, rather
+    /// than exercising it through a local instance.
+    ///
+    /// The block is zeroed and never touched, so it costs address space rather than resident
+    /// memory, and it is large enough that nothing else in the crate can free half of it inside the
+    /// microseconds between the two reads.
+    #[test]
+    #[cfg(feature = "alloc-accounting")]
+    fn a_real_allocation_raises_the_balance() {
+        use std::hint::black_box;
+
+        const SIZE: usize = 256 * 1024 * 1024;
+        let _guard = serial();
+        let before = current_balance();
+        // `black_box` keeps the allocation observable so it cannot be elided.
+        let held: Vec<u8> = black_box(vec![0u8; SIZE]);
+        let during = current_balance();
+        black_box(&held);
+        assert!(
+            during >= before + SIZE / 2,
+            "a {SIZE} byte allocation should raise the balance (before={before}, during={during}); \
+             is the accounting wrapper installed for this feature set?"
+        );
+        drop(held);
+    }
+
+    /// The balance must drop before the inner allocator is asked to free the block.
+    ///
+    /// jemalloc decrements its own `stats.allocated` at the start of a large free and then, for
+    /// blocks above its oversize threshold, unmaps the pages eagerly, which takes milliseconds for
+    /// a block of a few hundred megabytes. If the subtraction happened after delegating, the balance
+    /// would keep reporting a block the allocator had already given back for that whole window,
+    /// and `native_allocated` would read above `jemalloc_allocated`.
+    #[test]
+    fn dealloc_settles_before_delegating() {
+        use std::alloc::System;
+        use std::sync::atomic::AtomicUsize;
+
+        /// Records the reported balance at the moment the inner free is called.
+        struct Recording {
+            balance_at_dealloc: AtomicUsize,
+        }
+
+        unsafe impl GlobalAlloc for Recording {
+            unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+                System.alloc(layout)
+            }
+
+            unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+                self.balance_at_dealloc
+                    .store(current_balance(), Ordering::Relaxed);
+                System.dealloc(ptr, layout)
+            }
+        }
+
+        // Well above the settle threshold, so both the allocation and the free flush immediately.
+        const SIZE: usize = 64 * 1024 * 1024;
+        let _guard = serial();
+        let allocator = AccountingAllocator::new(Recording {
+            balance_at_dealloc: AtomicUsize::new(usize::MAX),
+        });
+        let layout = Layout::from_size_align(SIZE, 8).unwrap();
+
+        // SAFETY: the layout is valid and non-zero, and the block is freed below through the same
+        // allocator that produced it.
+        let ptr = unsafe { allocator.alloc(layout) };
+        assert!(!ptr.is_null());
+        let after_alloc = current_balance();
+        unsafe { allocator.dealloc(ptr, layout) };
+
+        let seen = allocator.inner.balance_at_dealloc.load(Ordering::Relaxed);
+        // Half the block is a wide margin against parallel test noise while still being far
+        // outside anything the mutation (subtracting after delegating) could produce.
+        assert!(
+            seen + SIZE / 2 <= after_alloc,
+            "inner dealloc saw balance {seen}, expected at most {} (balance after alloc was \
+             {after_alloc})",
+            after_alloc - SIZE / 2
+        );
+    }
+
+    /// Threads must settle their remaining drift on exit.
+    ///
+    /// The worker writes a drift straight into its `LOCAL_DRIFT` cell and exits. Without the
+    /// wrapper installed nothing else ever calls `track`, so the only path by which that value can
+    /// reach the shared balance is `ThreadDrift::drop`; that is the build CI runs, and the one in
+    /// which a missing destructor is caught. The value is far larger than any real allocation,
+    /// which makes the check immune to whatever the rest of the crate is allocating meanwhile.
+    /// The injected amount is taken back out afterwards so later tests see an unchanged balance.
+    #[test]
+    fn thread_exit_settles_remaining_drift() {

Review Comment:
   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.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+    assert_backend_is_live();
+    assert_accounting_is_live();
+    let mut group = c.benchmark_group("alloc_overhead");
+    for size in [16usize, 256, 4096] {
+        group.throughput(Throughput::Elements(1));
+        group.bench_function(format!("alloc_free_{size}b"), |b| {
+            b.iter(|| {
+                let v: Vec<u8> = Vec::with_capacity(black_box(size));
+                black_box(&v);
+            });
+        });
+    }
+    group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {

Review Comment:
   `iter_batched(|| (), |()| ..., BatchSize::SmallInput)` with a unit setup is `b.iter(...)` with extra machinery. The batching only earns its keep when setup allocates something you do not want in the measurement.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {

Review Comment:
   `small_churn` inlines the body of the `alloc_free` helper that `threshold_churn` calls a few lines down. Both are alloc/free loops over a list of sizes in the same group, so one function over `[16, 256, 4096, 32 * 1024, 64 * 1024]` covers both and keeps the threshold story in one place.
   
   Minor: `group.throughput(...)` is re-set on every loop iteration. It only needs to be set once before the loop.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let new_ptr = self.inner.realloc(ptr, layout, new_size);
+        if !new_ptr.is_null() {
+            // Accounting after the fact is only safe because this allocator cannot fail the
+            // allocation or unwind. A variant that enforced a limit would have to decide *before*
+            // delegating: `realloc` may free or move the old block, and a caller that never
+            // received the new pointer would free the stale one while unwinding.
+            //
+            // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast
+            // wraps.
+            track(new_size as isize - layout.size() as isize);
+        }
+        new_ptr
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{Mutex, MutexGuard};
+
+    /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
+    /// sees every other test's allocations. The tests that move it by tens of megabytes take this
+    /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
+    /// making each window microseconds wide and each expected move far larger than anything else
+    /// allocates in that time.
+    static SERIAL: Mutex<()> = Mutex::new(());
+
+    fn serial() -> MutexGuard<'static, ()> {
+        SERIAL
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    #[test]
+    fn settle_accumulates_below_the_threshold() {

Review Comment:
   Three tests over a six-line function. `settle_accumulates_below_the_threshold`, `settle_flushes_at_the_threshold`, and `settle_flushes_negative_drift` differ only in the delta and the expected residue, so one table-driven case over `[(1024, 1024), (SETTLE_THRESHOLD, 0), (-SETTLE_THRESHOLD, 0)]` says the same thing.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {

Review Comment:
   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.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {

Review Comment:
   `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.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature
+// combination, so exactly one definition exists, and each backend predicate is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop
+// silently: a feature combination with no backend would fail to compile rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
     not(feature = "mimalloc")
 ))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+    pub type Backend = tikv_jemallocator::Jemalloc;
+    pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+    pub const NAME: &str = "jemalloc";
 
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]
+    static GLOBAL: Backend = BACKEND;
+}
+
+/// mimalloc, unless a usable jemalloc was also requested.
 #[cfg(all(
     feature = "mimalloc",
     not(all(not(target_env = "msvc"), feature = "jemalloc"))
 ))]
+mod backend {
+    pub type Backend = mimalloc::MiMalloc;
+    pub const BACKEND: Backend = mimalloc::MiMalloc;
+    pub const NAME: &str = "mimalloc";
+
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]
+    static GLOBAL: Backend = BACKEND;
+}
+
+/// The system allocator: the complement of the two cases above. This covers neither feature, a
+/// jemalloc request on MSVC, and both features together, which each backend cfg excludes in favour
+/// of the other.
+#[cfg(not(any(
+    all(
+        not(target_env = "msvc"),
+        feature = "jemalloc",
+        not(feature = "mimalloc")
+    ),
+    all(
+        feature = "mimalloc",
+        not(all(not(target_env = "msvc"), feature = "jemalloc"))
+    )
+)))]
+// Without `alloc-accounting` nothing refers to this selection: the system allocator is the
+// default, so no `#[global_allocator]` is installed.
+#[cfg_attr(not(feature = "alloc-accounting"), allow(dead_code))]
+mod backend {
+    pub type Backend = std::alloc::System;
+    pub const BACKEND: Backend = std::alloc::System;
+    pub const NAME: &str = "system";
+}
+
+/// The name of the allocator backend this build selected: `"jemalloc"`, `"mimalloc"` or
+/// `"system"`. This is the one place the selection is decided, so anything that needs to know
+/// which allocator is in effect (the `alloc_overhead` benchmark's liveness check, for instance)
+/// reads it from here rather than re-deriving it from the feature set.
+pub use backend::NAME as ALLOCATOR_BACKEND;

Review Comment:
   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.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4019832858


##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+    assert_backend_is_live();
+    assert_accounting_is_live();
+    let mut group = c.benchmark_group("alloc_overhead");
+    for size in [16usize, 256, 4096] {
+        group.throughput(Throughput::Elements(1));
+        group.bench_function(format!("alloc_free_{size}b"), |b| {
+            b.iter(|| {
+                let v: Vec<u8> = Vec::with_capacity(black_box(size));
+                black_box(&v);
+            });
+        });
+    }
+    group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {
+    let mut group = c.benchmark_group("alloc_overhead");
+    group.throughput(Throughput::Bytes(64 * 1024));
+    group.bench_function("alloc_fill_free_64kb", |b| {
+        b.iter_batched(
+            || (),
+            |()| {
+                let v: Vec<u8> = vec![1u8; black_box(64 * 1024)];
+                black_box(v.len())
+            },
+            BatchSize::SmallInput,
+        );
+    });
+    group.finish();
+}
+
+/// Repeated growth, which is the `realloc` path: a builder doubling its buffer.
+fn growth_churn(c: &mut Criterion) {

Review Comment:
   Fixed in 4b71ad9: both asserts are behind a single `Once` in `assert_allocators_are_live`, which every benchmark function calls first. Checked with `cargo bench --bench alloc_overhead -- --test alloc_fill_free_64kb`, which now prints the backend announcement.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let new_ptr = self.inner.realloc(ptr, layout, new_size);
+        if !new_ptr.is_null() {
+            // Accounting after the fact is only safe because this allocator cannot fail the
+            // allocation or unwind. A variant that enforced a limit would have to decide *before*
+            // delegating: `realloc` may free or move the old block, and a caller that never
+            // received the new pointer would free the stale one while unwinding.
+            //
+            // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast
+            // wraps.
+            track(new_size as isize - layout.size() as isize);
+        }
+        new_ptr
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{Mutex, MutexGuard};
+
+    /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
+    /// sees every other test's allocations. The tests that move it by tens of megabytes take this
+    /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
+    /// making each window microseconds wide and each expected move far larger than anything else
+    /// allocates in that time.
+    static SERIAL: Mutex<()> = Mutex::new(());
+
+    fn serial() -> MutexGuard<'static, ()> {
+        SERIAL
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    #[test]
+    fn settle_accumulates_below_the_threshold() {

Review Comment:
   Done in 4b71ad9, as one table over `[(1024, 1024), (-1024, -1024), (THRESHOLD - 1, THRESHOLD - 1), (THRESHOLD, 0), (-THRESHOLD, 0)]`.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {

Review Comment:
   Inlined into `current_balance` and the test dropped in 4b71ad9.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {

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



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {

Review Comment:
   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`.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;

Review Comment:
   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.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature
+// combination, so exactly one definition exists, and each backend predicate is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop
+// silently: a feature combination with no backend would fail to compile rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
     not(feature = "mimalloc")
 ))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+    pub type Backend = tikv_jemallocator::Jemalloc;
+    pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+    pub const NAME: &str = "jemalloc";
 
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]
+    static GLOBAL: Backend = BACKEND;
+}
+
+/// mimalloc, unless a usable jemalloc was also requested.
 #[cfg(all(
     feature = "mimalloc",
     not(all(not(target_env = "msvc"), feature = "jemalloc"))
 ))]
+mod backend {
+    pub type Backend = mimalloc::MiMalloc;
+    pub const BACKEND: Backend = mimalloc::MiMalloc;
+    pub const NAME: &str = "mimalloc";
+
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]
+    static GLOBAL: Backend = BACKEND;
+}
+
+/// The system allocator: the complement of the two cases above. This covers neither feature, a
+/// jemalloc request on MSVC, and both features together, which each backend cfg excludes in favour
+/// of the other.
+#[cfg(not(any(
+    all(
+        not(target_env = "msvc"),
+        feature = "jemalloc",
+        not(feature = "mimalloc")
+    ),
+    all(
+        feature = "mimalloc",
+        not(all(not(target_env = "msvc"), feature = "jemalloc"))
+    )
+)))]
+// Without `alloc-accounting` nothing refers to this selection: the system allocator is the
+// default, so no `#[global_allocator]` is installed.
+#[cfg_attr(not(feature = "alloc-accounting"), allow(dead_code))]
+mod backend {
+    pub type Backend = std::alloc::System;
+    pub const BACKEND: Backend = std::alloc::System;
+    pub const NAME: &str = "system";
+}
+
+/// The name of the allocator backend this build selected: `"jemalloc"`, `"mimalloc"` or
+/// `"system"`. This is the one place the selection is decided, so anything that needs to know
+/// which allocator is in effect (the `alloc_overhead` benchmark's liveness check, for instance)
+/// reads it from here rather than re-deriving it from the feature set.
+pub use backend::NAME as ALLOCATOR_BACKEND;

Review Comment:
   Added in 4b71ad9.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5681578940

   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.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["sunchao (via GitHub)" <gi...@apache.org>]

sunchao commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4017037738


##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,202 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator. jemalloc keeps its own count of bytes it has
+/// served; if it is not the global allocator of this binary that count stays at zero, and a
+/// "jemalloc" baseline would in fact be the system allocator.
+#[cfg(feature = "jemalloc")]

Review Comment:
   [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.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4019830971


##########
native/core/src/execution/jni_api.rs:
##########
@@ -132,6 +132,18 @@ fn log_jemalloc_usage() {
     log_memory_usage("jemalloc_allocated", allocated.read().unwrap() as u64);
 }
 
+/// Reports the bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Logged alongside the per-thread pool reservations so the two can be compared directly: a large
+/// and growing excess is native memory the pool is not accounting for.
+#[cfg(feature = "alloc-accounting")]
+fn log_native_allocated() {

Review Comment:
   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.



##########
native/core/Cargo.toml:
##########
@@ -114,6 +114,12 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
 # Default builds carry zero Delta surface.
 contrib-delta = ["dep:comet-contrib-delta"]
 
+# Observability for real native memory usage. Wraps the global allocator to track the bytes it
+# hands out, and reports the total as the `native_allocated` tracing metric so it can be compared
+# against the memory pool's reservations. Never rejects an allocation. Off by default; a build
+# without it has no wrapper and no per-allocation work.
+alloc-accounting = []

Review Comment:
   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 check`s the `alloc-accounting`-only build for the system-allocator arm. That is the Linux Rust test job; the extra steps reuse its cache.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.

Review Comment:
   The numbers were posted in https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5681578940 after the description was written. The description now summarises them and links there.



##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!     https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+    /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+    /// directly instead of recursing. The only such allocation today is the one some platforms
+    /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+    ///
+    /// Const-initialized and destructor-free, so reading it never allocates and never fails —
+    /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+    static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+    /// This thread's un-flushed delta.
+    static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently
+/// discarded every time a thread died. Worker threads live for the process lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+    fn drop(&mut self) {
+        let drift = self.0.replace(0);
+        if drift != 0 {
+            BALANCE.fetch_add(drift, Ordering::Relaxed);
+        }
+    }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out of order.
+pub fn current_balance() -> usize {
+    clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+    balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+    let drift = local_drift.get().wrapping_add(delta);
+    if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+        local_drift.set(0);
+        BALANCE.fetch_add(drift, Ordering::Relaxed);
+    } else {
+        local_drift.set(drift);
+    }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+    if delta == 0 {
+        return;
+    }
+
+    // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will
+    // clear it, so this frame must only settle and return.
+    if IN_TRACK.with(|in_track| in_track.replace(true)) {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+        return;
+    }
+
+    // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already
+    // run, and any allocation after that point must not panic inside the allocator.
+    if LOCAL_DRIFT
+        .try_with(|thread_drift| settle(&thread_drift.0, delta))
+        .is_err()
+    {
+        BALANCE.fetch_add(delta, Ordering::Relaxed);
+    }
+
+    IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = self.inner.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            track(layout.size() as isize);
+        }
+        ptr
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the
+        // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping
+        // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on
+        // the balance for that whole window, after the allocator's own statistics had already
+        // dropped it.
+        track(-(layout.size() as isize));
+        self.inner.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let new_ptr = self.inner.realloc(ptr, layout, new_size);
+        if !new_ptr.is_null() {
+            // Accounting after the fact is only safe because this allocator cannot fail the
+            // allocation or unwind. A variant that enforced a limit would have to decide *before*
+            // delegating: `realloc` may free or move the old block, and a caller that never
+            // received the new pointer would free the stale one while unwinding.
+            //
+            // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast
+            // wraps.
+            track(new_size as isize - layout.size() as isize);
+        }
+        new_ptr
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::{Mutex, MutexGuard};
+
+    /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
+    /// sees every other test's allocations. The tests that move it by tens of megabytes take this
+    /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
+    /// making each window microseconds wide and each expected move far larger than anything else
+    /// allocates in that time.
+    static SERIAL: Mutex<()> = Mutex::new(());
+
+    fn serial() -> MutexGuard<'static, ()> {
+        SERIAL
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    #[test]
+    fn settle_accumulates_below_the_threshold() {
+        let drift = Cell::new(0);
+        settle(&drift, 1024);
+        // A flush would have reset the drift to zero, so this alone shows the shared balance was
+        // not touched. Reading `BALANCE` here would race with every other test's allocations.
+        assert_eq!(drift.get(), 1024, "small delta stays thread-local");
+    }
+
+    #[test]
+    fn settle_flushes_at_the_threshold() {
+        let drift = Cell::new(0);
+        settle(&drift, SETTLE_THRESHOLD);
+        assert_eq!(drift.get(), 0, "drift resets once flushed");
+    }
+
+    #[test]
+    fn settle_flushes_negative_drift() {
+        let drift = Cell::new(0);
+        settle(&drift, -SETTLE_THRESHOLD);
+        assert_eq!(drift.get(), 0);
+    }
+
+    #[test]
+    fn a_transiently_negative_balance_reports_as_zero() {
+        assert_eq!(clamp_balance(-1), 0);
+        assert_eq!(clamp_balance(isize::MIN), 0);
+        assert_eq!(clamp_balance(0), 0);
+        assert_eq!(clamp_balance(4096), 4096);
+    }
+
+    /// A real allocation must move the reported balance: this is the one test that checks the
+    /// wrapper is actually installed as the global allocator for the current feature set, rather
+    /// than exercising it through a local instance.
+    ///
+    /// The block is zeroed and never touched, so it costs address space rather than resident
+    /// memory, and it is large enough that nothing else in the crate can free half of it inside the
+    /// microseconds between the two reads.
+    #[test]
+    #[cfg(feature = "alloc-accounting")]
+    fn a_real_allocation_raises_the_balance() {
+        use std::hint::black_box;
+
+        const SIZE: usize = 256 * 1024 * 1024;
+        let _guard = serial();
+        let before = current_balance();
+        // `black_box` keeps the allocation observable so it cannot be elided.
+        let held: Vec<u8> = black_box(vec![0u8; SIZE]);
+        let during = current_balance();
+        black_box(&held);
+        assert!(
+            during >= before + SIZE / 2,
+            "a {SIZE} byte allocation should raise the balance (before={before}, during={during}); \
+             is the accounting wrapper installed for this feature set?"
+        );
+        drop(held);
+    }
+
+    /// The balance must drop before the inner allocator is asked to free the block.
+    ///
+    /// jemalloc decrements its own `stats.allocated` at the start of a large free and then, for
+    /// blocks above its oversize threshold, unmaps the pages eagerly, which takes milliseconds for
+    /// a block of a few hundred megabytes. If the subtraction happened after delegating, the balance
+    /// would keep reporting a block the allocator had already given back for that whole window,
+    /// and `native_allocated` would read above `jemalloc_allocated`.
+    #[test]
+    fn dealloc_settles_before_delegating() {
+        use std::alloc::System;
+        use std::sync::atomic::AtomicUsize;
+
+        /// Records the reported balance at the moment the inner free is called.
+        struct Recording {
+            balance_at_dealloc: AtomicUsize,
+        }
+
+        unsafe impl GlobalAlloc for Recording {
+            unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+                System.alloc(layout)
+            }
+
+            unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+                self.balance_at_dealloc
+                    .store(current_balance(), Ordering::Relaxed);
+                System.dealloc(ptr, layout)
+            }
+        }
+
+        // Well above the settle threshold, so both the allocation and the free flush immediately.
+        const SIZE: usize = 64 * 1024 * 1024;
+        let _guard = serial();
+        let allocator = AccountingAllocator::new(Recording {
+            balance_at_dealloc: AtomicUsize::new(usize::MAX),
+        });
+        let layout = Layout::from_size_align(SIZE, 8).unwrap();
+
+        // SAFETY: the layout is valid and non-zero, and the block is freed below through the same
+        // allocator that produced it.
+        let ptr = unsafe { allocator.alloc(layout) };
+        assert!(!ptr.is_null());
+        let after_alloc = current_balance();
+        unsafe { allocator.dealloc(ptr, layout) };
+
+        let seen = allocator.inner.balance_at_dealloc.load(Ordering::Relaxed);
+        // Half the block is a wide margin against parallel test noise while still being far
+        // outside anything the mutation (subtracting after delegating) could produce.
+        assert!(
+            seen + SIZE / 2 <= after_alloc,
+            "inner dealloc saw balance {seen}, expected at most {} (balance after alloc was \
+             {after_alloc})",
+            after_alloc - SIZE / 2
+        );
+    }
+
+    /// Threads must settle their remaining drift on exit.
+    ///
+    /// The worker writes a drift straight into its `LOCAL_DRIFT` cell and exits. Without the
+    /// wrapper installed nothing else ever calls `track`, so the only path by which that value can
+    /// reach the shared balance is `ThreadDrift::drop`; that is the build CI runs, and the one in
+    /// which a missing destructor is caught. The value is far larger than any real allocation,
+    /// which makes the check immune to whatever the rest of the crate is allocating meanwhile.
+    /// The injected amount is taken back out afterwards so later tests see an unchanged balance.
+    #[test]
+    fn thread_exit_settles_remaining_drift() {

Review Comment:
   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.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;

Review Comment:
   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.



##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
 
 use errors::{try_unwrap_or_throw, CometError, CometResult};
 
+pub mod alloc_accounting;
 pub mod cloud;
 pub mod execution;
 pub mod parquet;
 // this module is for non release only. Intended for debugging/profiling purposes
 #[cfg(debug_assertions)]
 pub mod debug;
 
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature
+// combination, so exactly one definition exists, and each backend predicate is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop
+// silently: a feature combination with no backend would fail to compile rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
 #[cfg(all(
     not(target_env = "msvc"),
     feature = "jemalloc",
     not(feature = "mimalloc")
 ))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+    pub type Backend = tikv_jemallocator::Jemalloc;
+    pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+    pub const NAME: &str = "jemalloc";
 
+    #[cfg(not(feature = "alloc-accounting"))]
+    #[global_allocator]

Review Comment:
   Taken in 4b71ad9. The three backend modules now only name the type, and the two `#[global_allocator]` statics sit together at the end. The default build installs `System` explicitly instead of leaving the default in place, which is the same allocator, so I dropped the "byte-for-byte" claim from the header.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {

Review Comment:
   Merged into a single `churn` in 4b71ad9: one loop over `[16, 256, 4096, 32 KiB, 64 KiB]` with the parallel variant for the two threshold sizes, throughput set once. Benchmark IDs are unchanged so the posted numbers still line up.



##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect, because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+    static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+    ANNOUNCE.call_once(|| {
+        eprintln!(
+            "alloc_overhead: measuring the `{}` allocator backend",
+            comet::ALLOCATOR_BACKEND
+        )
+    });
+    if comet::ALLOCATOR_BACKEND == "jemalloc" {
+        assert_jemalloc_is_live();
+    }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+    use tikv_jemalloc_ctl::{epoch, stats};
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    epoch::advance().expect("jemalloc epoch");
+    let allocated = stats::allocated::read().expect("jemalloc stats.allocated");
+    assert!(
+        allocated >= 8 * 1024 * 1024,
+        "the library selected jemalloc but jemalloc is not the global allocator of this binary \
+         (stats.allocated = {allocated}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+    unreachable!("the library reports the jemalloc backend but the feature is not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark
+/// binary, every "with the feature" number would silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+    let before = comet::alloc_accounting::current_balance();
+    // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily
+    // elide an allocation whose contents are never observed, and the check would then fail
+    // against a wrapper that is in fact working.
+    let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+    black_box(&held);
+    let during = comet::alloc_accounting::current_balance();
+    assert!(
+        during >= before + 4 * 1024 * 1024,
+        "alloc-accounting is enabled but the allocator is not installed in this binary \
+         (balance {before} -> {during}); the numbers below would be meaningless"
+    );
+    drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+    assert_backend_is_live();
+    assert_accounting_is_live();
+    let mut group = c.benchmark_group("alloc_overhead");
+    for size in [16usize, 256, 4096] {
+        group.throughput(Throughput::Elements(1));
+        group.bench_function(format!("alloc_free_{size}b"), |b| {
+            b.iter(|| {
+                let v: Vec<u8> = Vec::with_capacity(black_box(size));
+                black_box(&v);
+            });
+        });
+    }
+    group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {

Review Comment:
   Replaced with `b.iter` in 4b71ad9.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5670790447

   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.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] feat: add native allocation accounting for memory observability [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#issuecomment-5670362800

   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.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org