From 303875f28b6244755f7d030036556919e69d5ad8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 14 Sep 2026 13:34:30 -0600 Subject: [PATCH 1/8] feat: add native allocation accounting for memory observability Comet's memory pool counts declared reservations. Plenty of real allocation never goes through it -- Arrow builders, expression kernels, decompression buffers, Parquet metadata, object_store buffers, tokio itself -- so pool reservations are a lower bound on Comet's footprint and the size of the gap is currently unmeasurable at runtime. Diagnosing an OOM means guessing at it, and spark.comet.exec.memoryPool.fraction asks operators to hand-tune a haircut for a quantity nobody can see. AccountingAllocator wraps the selected global allocator (jemalloc, mimalloc, or system) and maintains one signed process-wide byte balance. executePlan reports it as the native_allocated tracing metric, next to the per-thread pool reservations it should be compared against. This is observability only: it never rejects an allocation, never panics, and does not touch the memory pool. Because it cannot fail an allocation, realloc can account after delegating rather than before, which avoids over-counting a failed realloc. Per-thread deltas are batched and flushed into the shared balance at 64 KiB, so the common path is a thread-local add-and-compare rather than an atomic RMW. ThreadDrift's destructor settles the remainder when a thread exits, which matters because the blocking pool churns on tokio's idle timeout and would otherwise bias the balance on a long-lived executor. Touching that destructor-bearing thread-local can itself allocate on first use, so track() keeps a destructor-free re-entrancy flag and settles re-entrant calls straight into the shared balance. Off by default; a build without the feature has no wrapper and no per-allocation work. Verified clippy -D warnings and the native test suite across default, alloc-accounting, jemalloc+alloc-accounting, and mimalloc+alloc-accounting. The thread-exit test was mutation-checked: neutering the destructor fails it. --- docs/source/contributor-guide/tracing.md | 11 + native/core/Cargo.toml | 6 + native/core/src/alloc_accounting.rs | 289 +++++++++++++++++++++++ native/core/src/execution/jni_api.rs | 14 ++ native/core/src/lib.rs | 42 +++- 5 files changed, 360 insertions(+), 2 deletions(-) create mode 100644 native/core/src/alloc_accounting.rs diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 88a291f421b..4fb7a0e28de 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -33,6 +33,16 @@ Additionally, enabling the `jemalloc` feature will enable tracing of native memo make release COMET_FEATURES="jemalloc" ``` +The `alloc-accounting` feature adds a second, allocator-independent measure of native memory. It +wraps whichever global allocator the build selected and reports the bytes it has handed out as +`native_allocated`. Unlike `jemalloc_allocated` it does not require jemalloc, and it counts only +what Rust code allocated, so it can be compared against the memory pool's reservations without the +allocator's own caching in the way. The two features are independent and can be combined: + +```shell +make release COMET_FEATURES="jemalloc,alloc-accounting" +``` + Example output: ```json @@ -111,5 +121,6 @@ Large or growing excess may indicate memory that is not being tracked by the poo | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | | jvm_heap_used | JVM heap memory usage of live objects for the executor process | | jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | +| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature) | | thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | | thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index f0c7735a503..763b1f39adf 100644 --- a/native/core/Cargo.toml +++ b/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 = [] + # exclude optional packages from cargo machete verifications [package.metadata.cargo-machete] ignored = ["hdfs-sys", "paste"] diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs new file mode 100644 index 00000000000..36e0d016165 --- /dev/null +++ b/native/core/src/alloc_accounting.rs @@ -0,0 +1,289 @@ +// 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 = 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); + +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, 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 { + inner: A, +} + +impl AccountingAllocator { + 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 GlobalAlloc for AccountingAllocator { + 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) { + self.inner.dealloc(ptr, layout); + track(-(layout.size() as isize)); + } + + 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::*; + + #[test] + fn settle_accumulates_below_the_threshold() { + let drift = Cell::new(0); + let before = BALANCE.load(Ordering::Relaxed); + settle(&drift, 1024); + assert_eq!(drift.get(), 1024, "small delta stays thread-local"); + assert_eq!(BALANCE.load(Ordering::Relaxed), before); + } + + #[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. Parallel test noise can only add to the + /// balance, so the assertion is one-sided. + #[test] + #[cfg(feature = "alloc-accounting")] + fn a_real_allocation_raises_the_balance() { + let before = current_balance(); + // Well above the settle threshold, so it is guaranteed to flush. + let held: Vec = vec![0u8; 8 * 1024 * 1024]; + let during = current_balance(); + assert!( + during >= before + 4 * 1024 * 1024, + "8 MiB allocation should raise the balance (before={before}, during={during})" + ); + drop(held); + } + + /// Threads must settle their remaining drift on exit. + /// + /// Each worker allocates a sub-threshold buffer — so the bytes are still sitting in its local + /// drift, never flushed — and hands ownership back to this thread before exiting. The matching + /// free therefore happens here, after the worker is gone, so the only way those bytes can ever + /// reach the shared balance is `ThreadDrift::drop`. Without the destructor the balance does not + /// move at all, and the later frees drive it *below* where it started. + #[test] + #[cfg(feature = "alloc-accounting")] + fn thread_exit_settles_remaining_drift() { + use std::sync::mpsc; + use std::thread; + + const THREADS: usize = 64; + const PER_THREAD: usize = 32 * 1024; + assert!( + (PER_THREAD as isize) < SETTLE_THRESHOLD, + "the per-thread buffer must stay in local drift for this test to mean anything" + ); + + let (tx, rx) = mpsc::channel(); + let before = current_balance() as isize; + + for _ in 0..THREADS { + let tx = tx.clone(); + thread::spawn(move || tx.send(vec![0u8; PER_THREAD]).unwrap()) + .join() + .unwrap(); + } + drop(tx); + + let held: Vec> = rx.iter().collect(); + assert_eq!(held.len(), THREADS); + + let moved = current_balance() as isize - before; + let allocated = (THREADS * PER_THREAD) as isize; + // Half the expected total is a wide margin against parallel test noise while still being + // far outside anything the mutation (a destructor that discards the drift) could produce. + assert!( + moved >= allocated / 2, + "drift from exited threads never reached the shared balance: \ + balance moved {moved} bytes, expected at least {}", + allocated / 2 + ); + + drop(held); + } +} diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 692b0d1bccf..0107847c9e0 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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() { + log_memory_usage( + "native_allocated", + crate::alloc_accounting::current_balance() as u64, + ); +} + /// Registry of active memory pools per Rust thread ID. /// Used to sum memory reservations across all contexts on the same thread for tracing. type ThreadPoolMap = HashMap>>; @@ -1089,6 +1101,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( if exec_context.tracing_enabled { #[cfg(feature = "jemalloc")] log_jemalloc_usage(); + #[cfg(feature = "alloc-accounting")] + log_native_allocated(); log_memory_usage( &exec_context.tracing_memory_metric_name, total_reserved_for_thread(exec_context.rust_thread_id) as u64, diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index b5656dba102..e6ba3158c6e 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -65,6 +65,7 @@ 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; @@ -72,21 +73,58 @@ pub mod parquet; #[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 = + 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 = + 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")) +))] +#[global_allocator] +static GLOBAL: alloc_accounting::AccountingAllocator = + alloc_accounting::AccountingAllocator::new(std::alloc::System); + #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init( e: EnvUnowned, From 93d0770c99b8c6c43857894ac308e9fa7e36918a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 14 Sep 2026 14:16:48 -0600 Subject: [PATCH 2/8] bench: measure the alloc-accounting wrapper's per-allocation cost Answers the question the feature has to answer before anyone proposes enabling it by default. The liveness assertion is the point of the harness as much as the timings are: without it a 'with the feature' run can silently be a second baseline, which is exactly what happened on the first attempt here. --- native/core/Cargo.toml | 4 + native/core/benches/alloc_overhead.rs | 112 ++++++++++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 native/core/benches/alloc_overhead.rs diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 763b1f39adf..501592945d2 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -129,6 +129,10 @@ name = "comet" # "rlib" is for benchmarking with criterion. crate-type = ["cdylib", "rlib"] +[[bench]] +name = "alloc_overhead" +harness = false + [[bench]] name = "array_element_append" harness = false diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs new file mode 100644 index 00000000000..247a1544cad --- /dev/null +++ b/native/core/benches/alloc_overhead.rs @@ -0,0 +1,112 @@ +// 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 is linked in +//! through the `rlib`. `assert_accounting_is_live` fails the run if the wrapper is somehow not in +//! effect — without it, a "with the feature" run could silently be a second baseline. +//! +//! `small_churn` is the worst case: allocations so small that the wrapper's bookkeeping is a +//! meaningful fraction of the allocator's own work. `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. + +use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput}; +use std::hint::black_box; + +/// 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 = 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_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 = 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 = 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) { + let mut group = c.benchmark_group("alloc_overhead"); + group.bench_function("grow_vec_to_64kb", |b| { + b.iter(|| { + let mut v: Vec = Vec::new(); + for _ in 0..(64 * 1024) { + v.push(black_box(1u8)); + } + black_box(v.len()) + }); + }); + group.finish(); +} + +criterion_group!(benches, small_churn, arrow_sized_churn, growth_churn); +criterion_main!(benches); From ce9b13a351e81e2972dc2777833e3f514609b3bd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 14 Sep 2026 14:58:12 -0600 Subject: [PATCH 3/8] fix: settle the balance before delegating a free in AccountingAllocator `dealloc` accounted after calling the inner allocator, mirroring `alloc` and `realloc`. A free cannot fail, so that ordering bought nothing, and it opened a window: jemalloc decrements `stats.allocated` at the start of a large free and then, for blocks above its 8 MiB oversize threshold, unmaps the pages eagerly, which takes milliseconds for a block of a few hundred megabytes. For that whole window the balance still carried a block the allocator had already given back, so `native_allocated` read above `jemalloc_allocated` by the size of the block in flight. On TPC-H SF100 about 2% of trace samples showed the excess, up to 160 MB, during task teardown in Q10, Q17 and Q18. Subtract before delegating. A new test wraps a recording inner allocator and asserts the balance has already dropped by the time the inner free is called; it fails with "inner dealloc saw balance 67182807, expected at most 33628375" under the previous ordering. --- native/core/src/alloc_accounting.rs | 61 ++++++++++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index 36e0d016165..dd2087fe1cf 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -168,8 +168,13 @@ unsafe impl GlobalAlloc for AccountingAllocator { } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - self.inner.dealloc(ptr, 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 { @@ -239,6 +244,60 @@ mod tests { 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 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. /// /// Each worker allocates a sub-threshold buffer — so the bytes are still sitting in its local From 2d78f87864a8d3c9182c353f01c845d8a15915a4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 01:35:05 -0600 Subject: [PATCH 4/8] fix: install the accounting wrapper for every allocator feature combination Select the allocator backend once, in a `backend` module whose three cfgs partition every feature combination, and let the single accounting `#[global_allocator]` refer to whatever that resolved to. Previously each backend predicate was repeated per arm and the accounting fallback to the system allocator excluded `mimalloc`, so `jemalloc,mimalloc,alloc-accounting` on a non-MSVC target matched no arm at all: the process ran on the unwrapped default allocator while `native_allocated` stayed enabled and read zero. A build without the feature is unchanged: the unwrapped allocator lives in the backend module that owns it, and no explicit allocator is installed when the selection is the system allocator. Make the accounting tests immune to parallel test noise. `BALANCE` is process-wide, so with the wrapper installed the crate's other tests move it concurrently and the margin-based assertions failed intermittently. The tests that observe the balance now serialize against each other, the real-allocation probe uses an untouched 256 MiB block that nothing else in the crate can mask, and the thread-exit test injects its drift directly so it no longer needs the feature and runs in the default CI build. Add `threshold_churn` to the `alloc_overhead` benchmark: alloc/free loops at 32 KiB (never flushes) and 64 KiB (flushes on every call), single threaded and from every core at once, so shared-counter contention can be quantified rather than inferred. --- native/core/benches/alloc_overhead.rs | 69 +++++++++++++++-- native/core/src/alloc_accounting.rs | 96 ++++++++++++----------- native/core/src/lib.rs | 107 ++++++++++++++------------ 3 files changed, 173 insertions(+), 99 deletions(-) diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index 247a1544cad..b60a32f3b30 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -28,13 +28,19 @@ //! through the `rlib`. `assert_accounting_is_live` fails the run if the wrapper is somehow not in //! effect — without it, a "with the feature" run could silently be a second baseline. //! -//! `small_churn` is the worst case: allocations so small that the wrapper's bookkeeping is a -//! meaningful fraction of the allocator's own work. `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. +//! `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. use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput}; use std::hint::black_box; +use std::thread; +use std::time::Instant; /// 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. @@ -108,5 +114,58 @@ fn growth_churn(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, small_churn, arrow_sized_churn, growth_churn); +/// Alloc/free of an untouched block, so the allocator call itself is most of the work and the +/// wrapper's share is largest. +fn alloc_free(size: usize) { + let v: Vec = Vec::with_capacity(black_box(size)); + black_box(&v); +} + +/// Alloc/free loops either side of the 64 KiB settle threshold, single-threaded and from every +/// core at once. +/// +/// A loop of one size never accumulates drift, so the threshold decides everything: at 32 KiB the +/// alloc and the free cancel inside the thread-local cell and the shared counter is never touched, +/// while at 64 KiB every alloc and every free flushes. The 64 KiB parallel case is therefore the +/// upper bound on shared-counter contention: `available_parallelism()` threads each doing two +/// atomic read-modify-writes per iteration on the same cacheline, with nothing else in between. +/// +/// Times are reported per alloc/free pair per thread, so a parallel number equal to its +/// single-threaded counterpart means the threads did not slow each other down at all. +fn threshold_churn(c: &mut Criterion) { + assert_accounting_is_live(); + let threads = thread::available_parallelism().map_or(4, |n| n.get()); + let mut group = c.benchmark_group("alloc_overhead"); + group.throughput(Throughput::Elements(1)); + for size in [32 * 1024usize, 64 * 1024] { + let kb = size / 1024; + group.bench_function(format!("alloc_free_{kb}kb"), |b| { + b.iter(|| alloc_free(size)); + }); + group.bench_function(format!("parallel_alloc_free_{kb}kb_x{threads}"), |b| { + b.iter_custom(|iters| { + let start = Instant::now(); + thread::scope(|scope| { + for _ in 0..threads { + scope.spawn(move || { + for _ in 0..iters { + alloc_free(size); + } + }); + } + }); + start.elapsed() + }); + }); + } + group.finish(); +} + +criterion_group!( + benches, + small_churn, + threshold_churn, + arrow_sized_churn, + growth_churn +); criterion_main!(benches); diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index dd2087fe1cf..f476c07271a 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -196,14 +196,28 @@ unsafe impl GlobalAlloc for AccountingAllocator { #[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); - let before = BALANCE.load(Ordering::Relaxed); 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"); - assert_eq!(BALANCE.load(Ordering::Relaxed), before); } #[test] @@ -228,18 +242,29 @@ mod tests { assert_eq!(clamp_balance(4096), 4096); } - /// A real allocation must move the reported balance. Parallel test noise can only add to the - /// balance, so the assertion is one-sided. + /// 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(); - // Well above the settle threshold, so it is guaranteed to flush. - let held: Vec = vec![0u8; 8 * 1024 * 1024]; + // `black_box` keeps the allocation observable so it cannot be elided. + let held: Vec = black_box(vec![0u8; SIZE]); let during = current_balance(); + black_box(&held); assert!( - during >= before + 4 * 1024 * 1024, - "8 MiB allocation should raise the balance (before={before}, during={during})" + 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); } @@ -275,6 +300,7 @@ mod tests { // 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), }); @@ -300,49 +326,33 @@ mod tests { /// Threads must settle their remaining drift on exit. /// - /// Each worker allocates a sub-threshold buffer — so the bytes are still sitting in its local - /// drift, never flushed — and hands ownership back to this thread before exiting. The matching - /// free therefore happens here, after the worker is gone, so the only way those bytes can ever - /// reach the shared balance is `ThreadDrift::drop`. Without the destructor the balance does not - /// move at all, and the later frees drive it *below* where it started. + /// 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] - #[cfg(feature = "alloc-accounting")] fn thread_exit_settles_remaining_drift() { - use std::sync::mpsc; use std::thread; - const THREADS: usize = 64; - const PER_THREAD: usize = 32 * 1024; - assert!( - (PER_THREAD as isize) < SETTLE_THRESHOLD, - "the per-thread buffer must stay in local drift for this test to mean anything" - ); + const INJECTED: isize = 1 << 40; + let _guard = serial(); - let (tx, rx) = mpsc::channel(); - let before = current_balance() as isize; - - for _ in 0..THREADS { - let tx = tx.clone(); - thread::spawn(move || tx.send(vec![0u8; PER_THREAD]).unwrap()) - .join() - .unwrap(); - } - drop(tx); - - let held: Vec> = rx.iter().collect(); - assert_eq!(held.len(), THREADS); + let before = BALANCE.load(Ordering::Relaxed); + thread::spawn(|| { + LOCAL_DRIFT.with(|drift| drift.0.set(drift.0.get() + INJECTED)); + }) + .join() + .unwrap(); + let moved = BALANCE.load(Ordering::Relaxed) - before; + BALANCE.fetch_sub(INJECTED, Ordering::Relaxed); - let moved = current_balance() as isize - before; - let allocated = (THREADS * PER_THREAD) as isize; - // Half the expected total is a wide margin against parallel test noise while still being - // far outside anything the mutation (a destructor that discards the drift) could produce. assert!( - moved >= allocated / 2, - "drift from exited threads never reached the shared balance: \ + moved >= INJECTED / 2, + "drift from an exited thread never reached the shared balance: \ balance moved {moved} bytes, expected at least {}", - allocated / 2 + INJECTED / 2 ); - - drop(held); } } diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index e6ba3158c6e..583ed908a73 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -41,19 +41,6 @@ use log4rs::{ Config, }; -#[cfg(all( - not(target_env = "msvc"), - feature = "jemalloc", - not(feature = "mimalloc") -))] -use tikv_jemallocator::Jemalloc; - -#[cfg(all( - feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")) -))] -use mimalloc::MiMalloc; - // Re-export from jvm-bridge crate for internal use pub use datafusion_comet_jni_bridge::errors; pub use datafusion_comet_jni_bridge::JAVA_VM; @@ -73,57 +60,75 @@ pub mod parquet; #[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. +// 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"), - not(feature = "alloc-accounting") + not(feature = "mimalloc") ))] -#[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; +mod backend { + pub type Backend = tikv_jemallocator::Jemalloc; + pub const BACKEND: Backend = tikv_jemallocator::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")), - not(feature = "alloc-accounting") + not(all(not(target_env = "msvc"), feature = "jemalloc")) ))] -#[global_allocator] -static GLOBAL: MiMalloc = MiMalloc; +mod backend { + pub type Backend = mimalloc::MiMalloc; + pub const BACKEND: Backend = mimalloc::MiMalloc; -#[cfg(all( - not(target_env = "msvc"), - feature = "jemalloc", - not(feature = "mimalloc"), - feature = "alloc-accounting" -))] -#[global_allocator] -static GLOBAL: alloc_accounting::AccountingAllocator = - alloc_accounting::AccountingAllocator::new(Jemalloc); + #[cfg(not(feature = "alloc-accounting"))] + #[global_allocator] + static GLOBAL: Backend = BACKEND; +} -#[cfg(all( - feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")), - feature = "alloc-accounting" -))] -#[global_allocator] -static GLOBAL: alloc_accounting::AccountingAllocator = - alloc_accounting::AccountingAllocator::new(MiMalloc); +/// 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; +} -// 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")) -))] +#[cfg(feature = "alloc-accounting")] #[global_allocator] -static GLOBAL: alloc_accounting::AccountingAllocator = - alloc_accounting::AccountingAllocator::new(std::alloc::System); +static GLOBAL: alloc_accounting::AccountingAllocator = + alloc_accounting::AccountingAllocator::new(backend::BACKEND); #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init( From ca2f8c87531866662261e2a07e29e0a1f445ceb6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 02:47:51 -0600 Subject: [PATCH 5/8] bench: link comet into alloc_overhead so the baseline uses the selected allocator An `--extern` crate that nothing names is dropped from the crate graph, and the accounting-off run of this benchmark named nothing in `comet`, so the `#[global_allocator]` in `lib.rs` never reached the binary: the "jemalloc" baseline was measuring glibc malloc. The accounting-on run names `comet::alloc_accounting`, so it did link the crate, and every off/on comparison so far was glibc against jemalloc plus the wrapper. Add `extern crate comet` so the crate is always linked, and a jemalloc liveness assertion next to the existing accounting one, so a run against the wrong allocator fails instead of producing a plausible number. --- native/core/benches/alloc_overhead.rs | 37 ++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index b60a32f3b30..ba6b1663349 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -24,9 +24,13 @@ //! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off //! ``` //! -//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which is linked in -//! through the `rlib`. `assert_accounting_is_live` fails the run if the wrapper is somehow not in -//! effect — without it, a "with the feature" run could silently be a second baseline. +//! 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 @@ -37,11 +41,36 @@ //! 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")] +fn assert_backend_is_live() { + use tikv_jemalloc_ctl::{epoch, stats}; + let held: Vec = 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 jemalloc feature is enabled but jemalloc is not the global allocator of this binary \ + (stats.allocated = {allocated}); the numbers below would be meaningless" + ); + drop(held); +} + +#[cfg(not(feature = "jemalloc"))] +fn assert_backend_is_live() {} + /// 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")] @@ -67,6 +96,7 @@ 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] { @@ -133,6 +163,7 @@ fn alloc_free(size: usize) { /// Times are reported per alloc/free pair per thread, so a parallel number equal to its /// single-threaded counterpart means the threads did not slow each other down at all. fn threshold_churn(c: &mut Criterion) { + assert_backend_is_live(); assert_accounting_is_live(); let threads = thread::available_parallelism().map_or(4, |n| n.get()); let mut group = c.benchmark_group("alloc_overhead"); From ac51c5c0f62ba360d455ed0abbb1af50bf2ebe3a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 09:03:55 -0600 Subject: [PATCH 6/8] bench: check jemalloc liveness only when the library selected jemalloc The alloc_overhead benchmark asserted on jemalloc's counters whenever the jemalloc feature was on, but lib.rs selects jemalloc only when mimalloc was not also requested and the target is not MSVC. With jemalloc,mimalloc the library falls back to the system allocator and the bench aborted before measuring anything. Export the selection lib.rs made as comet::ALLOCATOR_BACKEND, one NAME per backend module, and have the bench announce it and run the jemalloc assertion only when the library reports jemalloc. The predicate stays written once, so the guard cannot drift from the selection it verifies. --- native/core/benches/alloc_overhead.rs | 35 ++++++++++++++++++++++----- native/core/src/lib.rs | 9 +++++++ 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index ba6b1663349..75e119a76de 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -50,11 +50,31 @@ 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")] +/// 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 = black_box(vec![1u8; 8 * 1024 * 1024]); black_box(&held); @@ -62,14 +82,17 @@ fn assert_backend_is_live() { let allocated = stats::allocated::read().expect("jemalloc stats.allocated"); assert!( allocated >= 8 * 1024 * 1024, - "the jemalloc feature is enabled but jemalloc is not the global allocator of this binary \ + "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_backend_is_live() {} +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. diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 583ed908a73..6467711ab4b 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -83,6 +83,7 @@ pub mod debug; 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] @@ -97,6 +98,7 @@ mod backend { 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] @@ -123,8 +125,15 @@ mod backend { 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; + #[cfg(feature = "alloc-accounting")] #[global_allocator] static GLOBAL: alloc_accounting::AccountingAllocator = From 4b71ad9c69ad0a9c59577d5bb5f4daea9d2a7029 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 14:12:04 -0600 Subject: [PATCH 7/8] fix: teach analyze_trace about native_allocated and cover the accounting feature in CI Address the third review round on the alloc-accounting feature. - analyze_trace only matched jemalloc_allocated, so a trace from a build with alloc-accounting and no jemalloc reported zero allocated and no excess. It now analyzes native_allocated when present, falls back to jemalloc_allocated otherwise, names the counter in its output, and rejects a trace with neither. tracing.md describes the selection. - Nothing in CI compiled the feature. The rust-test action now lints the jemalloc,alloc-accounting build with --all-targets, runs the accounting tests with the wrapper installed, and checks the system-allocator arm. - lib.rs: the three backend modules only name the type; the two global_allocator statics sit together, so the default build now installs System explicitly. ALLOCATOR_BACKEND is doc(hidden). - alloc_accounting.rs: inline the clamp, table-drive the settle tests, add a realloc test that pins the size-difference accounting and the after-delegating order, and confine the thread-exit test to the default build, where a missing destructor is actually caught. Trim the prose. - alloc_overhead bench: drop the redundant extern crate, run both liveness guards from a single Once that every bench function calls, merge the two alloc/free loops, and use iter instead of iter_batched. Benchmark IDs are unchanged. --- .github/actions/rust-test/action.yaml | 13 ++ docs/source/contributor-guide/tracing.md | 44 +++-- native/common/src/bin/analyze_trace.rs | 91 ++++++--- native/core/benches/alloc_overhead.rs | 206 +++++++++----------- native/core/src/alloc_accounting.rs | 234 +++++++++++++---------- native/core/src/lib.rs | 43 ++--- 6 files changed, 339 insertions(+), 292 deletions(-) diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml index c39c2dcd4f9..8934fbb4dfb 100644 --- a/.github/actions/rust-test/action.yaml +++ b/.github/actions/rust-test/action.yaml @@ -70,3 +70,16 @@ runs: export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} RUST_BACKTRACE=1 cargo nextest run + # The `alloc-accounting` feature is off by default, so nothing else in CI compiles the + # allocator wrapper, its backend selection, or the benchmark's liveness guards. Lint them with + # the accounting wrapper over jemalloc, run the accounting tests with the wrapper installed, and + # check the system-allocator arm. + - name: Check and test the alloc-accounting feature + shell: bash + run: | + cd native + export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} + cargo clippy --color=never -p datafusion-comet --all-targets --features jemalloc,alloc-accounting -- -D warnings + RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features jemalloc,alloc-accounting alloc_accounting + cargo check -p datafusion-comet --features alloc-accounting + diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 4fb7a0e28de..3a6b9b5afe0 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -65,9 +65,9 @@ Example trace visualization: ## Analyzing Memory Usage -The `analyze_trace` tool parses a trace log and compares jemalloc usage against the sum of per-thread -Comet memory pool reservations. This is useful for detecting untracked native memory growth where jemalloc -allocations exceed what the memory pools account for. +The `analyze_trace` tool parses a trace log and compares the process-wide native allocation counter against +the sum of per-thread Comet memory pool reservations. This is useful for detecting untracked native memory +growth where native allocations exceed what the memory pools account for. Build and run: @@ -76,10 +76,12 @@ cd native cargo run --bin analyze_trace -- /path/to/comet-event-trace.json ``` -The tool reads counter events from the trace log. Because tracing logs metrics per thread, `jemalloc_allocated` -is a process-wide value (the same global allocation reported from whichever thread logs it), while -`thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to get the total -tracked memory. +The tool reads counter events from the trace log. Because tracing logs metrics per thread, `native_allocated` +and `jemalloc_allocated` are process-wide values (the same global allocation reported from whichever thread +logs it), while `thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to +get the total tracked memory. The tool analyzes `native_allocated` when the trace contains it, since that +counts only what Rust code holds from the allocator, and otherwise falls back to `jemalloc_allocated`. The +output names the counter it used. A trace with neither counter is rejected. Sample output: @@ -87,20 +89,21 @@ Sample output: === Comet Trace Memory Analysis === Counter events parsed: 193104 +Allocation counter: jemalloc_allocated Threads with memory pools: 8 -Peak jemalloc allocated: 3068.2 MB +Peak jemalloc_allocated: 3068.2 MB Peak pool total: 2864.6 MB -Peak excess (jemalloc - pool): 364.6 MB +Peak excess (jemalloc_allocated - pool): 364.6 MB -WARNING: jemalloc exceeded pool reservation at 138 sampled points: +WARNING: jemalloc_allocated exceeded pool reservation at 138 sampled points: - Time (us) jemalloc pool_total excess --------------------------------------------------------------- - 179578 210.8 MB 0.1 MB 210.7 MB - 429663 420.5 MB 145.1 MB 275.5 MB - 1304969 2122.5 MB 1797.2 MB 325.2 MB - 21974838 407.0 MB 42.3 MB 364.6 MB - 33543599 5.5 MB 0.1 MB 5.3 MB + Time (us) jemalloc_allocated pool_total excess +------------------------------------------------------------------ + 179578 210.8 MB 0.1 MB 210.7 MB + 429663 420.5 MB 145.1 MB 275.5 MB + 1304969 2122.5 MB 1797.2 MB 325.2 MB + 21974838 407.0 MB 42.3 MB 364.6 MB + 33543599 5.5 MB 0.1 MB 5.3 MB --- Final per-thread pool reservations --- @@ -112,8 +115,9 @@ WARNING: jemalloc exceeded pool reservation at 138 sampled points: Total: 0.0 MB ``` -Some excess is expected (jemalloc metadata, fragmentation, non-pool allocations like Arrow IPC buffers). -Large or growing excess may indicate memory that is not being tracked by the pool. +Some excess is expected (allocator metadata and fragmentation for `jemalloc_allocated`, and non-pool +allocations like Arrow IPC buffers for either counter). Large or growing excess may indicate memory that is +not being tracked by the pool. ## Definition of Labels @@ -121,6 +125,6 @@ Large or growing excess may indicate memory that is not being tracked by the poo | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | | jvm_heap_used | JVM heap memory usage of live objects for the executor process | | jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | -| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature) | +| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate: each live thread holds up to 64 KiB of un-flushed delta, so the value can lag the true balance by that much per thread. | | thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | | thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs index 8df83f7ceae..47039df9400 100644 --- a/native/common/src/bin/analyze_trace.rs +++ b/native/common/src/bin/analyze_trace.rs @@ -16,8 +16,9 @@ // under the License. //! Analyzes a Comet chrome trace event log (`comet-event-trace.json`) and -//! compares jemalloc usage against the sum of per-thread Comet memory pool -//! reservations. Reports any points where jemalloc exceeds the total pool size. +//! compares the process-wide native allocation counter against the sum of +//! per-thread Comet memory pool reservations. Reports any points where the +//! allocated bytes exceed the total pool size. //! //! Usage: //! cargo run --bin analyze_trace -- @@ -27,6 +28,15 @@ use std::collections::HashMap; use std::io::{BufRead, BufReader}; use std::{env, fs::File}; +/// The process-wide allocation counters the tool understands, most preferred first. +/// +/// `native_allocated` (the `alloc-accounting` feature) counts only the bytes Rust code holds from +/// the global allocator, so it is the tighter comparison against pool reservations. +/// `jemalloc_allocated` (the `jemalloc` feature) also includes jemalloc's own metadata. A trace +/// that carries both is analyzed against `native_allocated` alone; a trace with neither cannot be +/// analyzed. +const ALLOCATED_COUNTERS: [&str; 2] = ["native_allocated", "jemalloc_allocated"]; + /// A single Chrome trace event (only the fields we care about). #[derive(Deserialize)] struct TraceEvent { @@ -42,7 +52,7 @@ struct TraceEvent { /// Snapshot of memory state at a given timestamp. struct MemorySnapshot { ts: u64, - jemalloc: u64, + allocated: u64, pool_total: u64, } @@ -61,14 +71,16 @@ fn main() { let file = File::open(&args[1]).expect("Failed to open trace file"); let reader = BufReader::new(file); - // Latest jemalloc value (global, not per-thread) - let mut latest_jemalloc: u64 = 0; + // Index into ALLOCATED_COUNTERS of the counter being analyzed, once one has been seen + let mut source: Option = None; + // Latest allocated value (global, not per-thread) + let mut latest_allocated: u64 = 0; // Per-thread pool reservations: thread_NNN -> bytes let mut pool_by_thread: HashMap = HashMap::new(); - // Points where jemalloc exceeded pool total + // Points where allocated exceeded pool total let mut violations: Vec = Vec::new(); // Track peak values - let mut peak_jemalloc: u64 = 0; + let mut peak_allocated: u64 = 0; let mut peak_pool_total: u64 = 0; let mut peak_excess: u64 = 0; let mut counter_events: u64 = 0; @@ -110,11 +122,28 @@ fn main() { counter_events += 1; - if event.name == "jemalloc_allocated" { - if let Some(val) = event.args.get("jemalloc_allocated") { - latest_jemalloc = val.as_u64().unwrap_or(0); - if latest_jemalloc > peak_jemalloc { - peak_jemalloc = latest_jemalloc; + if let Some(rank) = ALLOCATED_COUNTERS + .iter() + .position(|name| *name == event.name) + { + match source { + // A preferred counter is present in this trace; ignore the other one. + Some(current) if current < rank => continue, + Some(current) if current == rank => {} + // First sighting of a more preferred counter. Start over so the peaks and + // violations reported all come from a single source. + _ => { + source = Some(rank); + latest_allocated = 0; + peak_allocated = 0; + peak_excess = 0; + violations.clear(); + } + } + if let Some(val) = event.args.get(&event.name) { + latest_allocated = val.as_u64().unwrap_or(0); + if latest_allocated > peak_allocated { + peak_allocated = latest_allocated; } } } else if event.name.contains("comet_memory_reserved") { @@ -129,14 +158,14 @@ fn main() { continue; } - // After each jemalloc or pool update, check the current state + // After each allocated or pool update, check the current state let pool_total: u64 = pool_by_thread.values().sum(); if pool_total > peak_pool_total { peak_pool_total = pool_total; } - if latest_jemalloc > 0 && pool_total > 0 && latest_jemalloc > pool_total { - let excess = latest_jemalloc - pool_total; + if latest_allocated > 0 && pool_total > 0 && latest_allocated > pool_total { + let excess = latest_allocated - pool_total; if excess > peak_excess { peak_excess = excess; } @@ -147,46 +176,56 @@ fn main() { { violations.push(MemorySnapshot { ts: event.ts, - jemalloc: latest_jemalloc, + allocated: latest_allocated, pool_total, }); } } } + let Some(source) = source.map(|rank| ALLOCATED_COUNTERS[rank]) else { + eprintln!( + "No process-wide allocation counter found in the trace: expected one of {}. \ + Build the native library with the `alloc-accounting` or `jemalloc` feature.", + ALLOCATED_COUNTERS.join(", ") + ); + std::process::exit(1); + }; + // Print summary println!("=== Comet Trace Memory Analysis ===\n"); println!("Counter events parsed: {counter_events}"); + println!("Allocation counter: {source}"); println!("Threads with memory pools: {}", pool_by_thread.len()); - println!("Peak jemalloc allocated: {}", format_bytes(peak_jemalloc)); + println!("Peak {source}: {}", format_bytes(peak_allocated)); println!( "Peak pool total: {}", format_bytes(peak_pool_total) ); println!( - "Peak excess (jemalloc - pool): {}", + "Peak excess ({source} - pool): {}", format_bytes(peak_excess) ); println!(); if violations.is_empty() { - println!("OK: jemalloc never exceeded the total pool reservation."); + println!("OK: {source} never exceeded the total pool reservation."); } else { println!( - "WARNING: jemalloc exceeded pool reservation at {} sampled points:\n", + "WARNING: {source} exceeded pool reservation at {} sampled points:\n", violations.len() ); println!( - "{:>14} {:>14} {:>14} {:>14}", - "Time (us)", "jemalloc", "pool_total", "excess" + "{:>14} {:>18} {:>14} {:>14}", + "Time (us)", source, "pool_total", "excess" ); - println!("{}", "-".repeat(62)); + println!("{}", "-".repeat(66)); for snap in &violations { - let excess = snap.jemalloc - snap.pool_total; + let excess = snap.allocated - snap.pool_total; println!( - "{:>14} {:>14} {:>14} {:>14}", + "{:>14} {:>18} {:>14} {:>14}", snap.ts, - format_bytes(snap.jemalloc), + format_bytes(snap.allocated), format_bytes(snap.pool_total), format_bytes(excess), ); diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index 75e119a76de..dbba3ed424c 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -24,55 +24,48 @@ //! 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; +//! `churn` allocates and frees untouched blocks, so the allocator call is most of the work and the +//! wrapper's share is largest. Sizes below the 64 KiB settle threshold only ever touch the +//! thread-local path; a loop of exactly 64 KiB blocks flushes to the shared atomic on every alloc +//! and every free, and its parallel variant does that from every core at once, so the gap between +//! the single-threaded and parallel 64 KiB numbers is the cost of contention on that cacheline. +//! `arrow_sized_churn` and `growth_churn` are closer to what Comet does, where filling a +//! batch-sized buffer or growing a builder dwarfs the bookkeeping. -use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput}; +use criterion::{criterion_group, criterion_main, Criterion, Throughput}; use std::hint::black_box; +use std::sync::Once; 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(|| { +/// Mirrors `alloc_accounting::SETTLE_THRESHOLD`: a thread flushes to the shared balance once its +/// un-flushed delta reaches this. +const SETTLE_THRESHOLD: usize = 64 * 1024; + +/// Fails the run if the allocator being measured is not the one the feature set asked for, since +/// a number measured against the wrong allocator would be worse than no number. Every benchmark +/// function calls this, so a filtered run cannot skip it. +fn assert_allocators_are_live() { + static ONCE: Once = Once::new(); + ONCE.call_once(|| { + // Which backend is in effect is `lib.rs`'s decision (with `jemalloc,mimalloc` together it + // falls back to the system allocator), so ask it rather than re-deriving the answer from + // the feature set. Naming `comet::ALLOCATOR_BACKEND` is also what links the `comet` rlib, + // and with it the `#[global_allocator]` it installs, into this binary: an `--extern` crate + // that nothing names is dropped from the crate graph along with its allocator. eprintln!( "alloc_overhead: measuring the `{}` allocator backend", comet::ALLOCATOR_BACKEND - ) + ); + if comet::ALLOCATOR_BACKEND == "jemalloc" { + assert_jemalloc_is_live(); + } + assert_accounting_is_live(); }); - 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. +/// binary that count stays at zero. #[cfg(feature = "jemalloc")] fn assert_jemalloc_is_live() { use tikv_jemalloc_ctl::{epoch, stats}; @@ -94,14 +87,13 @@ 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. +/// If the wrapper were not installed in this 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. + // `black_box` is load-bearing: in release mode LLVM elides an allocation whose contents are + // never observed, and the check would then fail against a wrapper that is in fact working. let held: Vec = black_box(vec![1u8; 8 * 1024 * 1024]); black_box(&held); let during = comet::alloc_accounting::current_balance(); @@ -116,20 +108,56 @@ fn assert_accounting_is_live() { #[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(); +/// Alloc/free of an untouched block. +fn alloc_free(size: usize) { + let v: Vec = Vec::with_capacity(black_box(size)); + black_box(&v); +} + +/// Alloc/free loops from well below the settle threshold up to exactly on it, single-threaded, and +/// either side of the threshold from every core at once. +/// +/// A loop of one size never accumulates drift, so the threshold decides everything: below it the +/// alloc and the free cancel inside the thread-local cell and the shared counter is never touched, +/// while at exactly 64 KiB every alloc and every free flushes. The 64 KiB parallel case is +/// therefore the upper bound on shared-counter contention: `available_parallelism()` threads each +/// doing two atomic read-modify-writes per iteration on the same cacheline, with nothing else in +/// between. Times are per alloc/free pair per thread, so a parallel number equal to its +/// single-threaded counterpart means the threads did not slow each other down at all. +fn churn(c: &mut Criterion) { + assert_allocators_are_live(); + let threads = thread::available_parallelism().map_or(4, |n| n.get()); 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 = Vec::with_capacity(black_box(size)); - black_box(&v); - }); + group.throughput(Throughput::Elements(1)); + for size in [16usize, 256, 4096, SETTLE_THRESHOLD / 2, SETTLE_THRESHOLD] { + // The two sizes either side of the threshold are the ones where the parallel variant + // says something. + let near_threshold = size >= SETTLE_THRESHOLD / 2; + let label = if near_threshold { + format!("{}kb", size / 1024) + } else { + format!("{size}b") + }; + group.bench_function(format!("alloc_free_{label}"), |b| { + b.iter(|| alloc_free(size)); }); + if near_threshold { + group.bench_function(format!("parallel_alloc_free_{label}_x{threads}"), |b| { + b.iter_custom(|iters| { + let start = Instant::now(); + thread::scope(|scope| { + for _ in 0..threads { + scope.spawn(move || { + for _ in 0..iters { + alloc_free(size); + } + }); + } + }); + start.elapsed() + }); + }); + } } group.finish(); } @@ -137,23 +165,21 @@ fn small_churn(c: &mut Criterion) { /// 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) { + assert_allocators_are_live(); 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 = vec![1u8; black_box(64 * 1024)]; - black_box(v.len()) - }, - BatchSize::SmallInput, - ); + b.iter(|| { + let v: Vec = vec![1u8; black_box(64 * 1024)]; + black_box(v.len()) + }); }); group.finish(); } /// Repeated growth, which is the `realloc` path: a builder doubling its buffer. fn growth_churn(c: &mut Criterion) { + assert_allocators_are_live(); let mut group = c.benchmark_group("alloc_overhead"); group.bench_function("grow_vec_to_64kb", |b| { b.iter(|| { @@ -167,59 +193,5 @@ fn growth_churn(c: &mut Criterion) { group.finish(); } -/// Alloc/free of an untouched block, so the allocator call itself is most of the work and the -/// wrapper's share is largest. -fn alloc_free(size: usize) { - let v: Vec = Vec::with_capacity(black_box(size)); - black_box(&v); -} - -/// Alloc/free loops either side of the 64 KiB settle threshold, single-threaded and from every -/// core at once. -/// -/// A loop of one size never accumulates drift, so the threshold decides everything: at 32 KiB the -/// alloc and the free cancel inside the thread-local cell and the shared counter is never touched, -/// while at 64 KiB every alloc and every free flushes. The 64 KiB parallel case is therefore the -/// upper bound on shared-counter contention: `available_parallelism()` threads each doing two -/// atomic read-modify-writes per iteration on the same cacheline, with nothing else in between. -/// -/// Times are reported per alloc/free pair per thread, so a parallel number equal to its -/// single-threaded counterpart means the threads did not slow each other down at all. -fn threshold_churn(c: &mut Criterion) { - assert_backend_is_live(); - assert_accounting_is_live(); - let threads = thread::available_parallelism().map_or(4, |n| n.get()); - let mut group = c.benchmark_group("alloc_overhead"); - group.throughput(Throughput::Elements(1)); - for size in [32 * 1024usize, 64 * 1024] { - let kb = size / 1024; - group.bench_function(format!("alloc_free_{kb}kb"), |b| { - b.iter(|| alloc_free(size)); - }); - group.bench_function(format!("parallel_alloc_free_{kb}kb_x{threads}"), |b| { - b.iter_custom(|iters| { - let start = Instant::now(); - thread::scope(|scope| { - for _ in 0..threads { - scope.spawn(move || { - for _ in 0..iters { - alloc_free(size); - } - }); - } - }); - start.elapsed() - }); - }); - } - group.finish(); -} - -criterion_group!( - benches, - small_churn, - threshold_churn, - arrow_sized_churn, - growth_churn -); +criterion_group!(benches, churn, arrow_sized_churn, growth_churn); criterion_main!(benches); diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index f476c07271a..394bbe774d3 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -17,23 +17,14 @@ //! 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. +//! process-wide byte balance, which [`current_balance`] exposes so it can be compared against the +//! memory pool's reservations in tracing output. This is observability only: it never rejects an +//! allocation, never panics, and never gates the memory pool. //! -//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation, +//! 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. +//! `malloc` rather than Rust's `GlobalAlloc`. See the [memory management contributor guide]. //! //! [memory management contributor guide]: //! https://datafusion.apache.org/comet/contributor-guide/memory_management.html @@ -56,7 +47,7 @@ thread_local! { /// 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 — + /// 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 = const { Cell::new(false) }; @@ -64,12 +55,9 @@ thread_local! { 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. +/// 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 discarded every time a +/// thread died, and tokio's blocking pool churns threads on its idle timeout. struct ThreadDrift(Cell); impl Drop for ThreadDrift { @@ -85,13 +73,12 @@ impl Drop for ThreadDrift { /// /// 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. +/// +/// The value is approximate. Each live thread holds up to [`SETTLE_THRESHOLD`] bytes of +/// un-flushed delta in either direction, so the reported balance can lag the true one by up to +/// that amount times the number of live threads. 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 + BALANCE.load(Ordering::Relaxed).max(0) as usize } /// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches @@ -196,6 +183,8 @@ unsafe impl GlobalAlloc for AccountingAllocator { #[cfg(test)] mod tests { use super::*; + use std::alloc::System; + use std::sync::atomic::AtomicUsize; use std::sync::{Mutex, MutexGuard}; /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it @@ -211,50 +200,77 @@ mod tests { .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"); + const MIB: usize = 1024 * 1024; + + /// Slack allowed between an observed balance and the expected one, to absorb whatever the + /// rest of the crate allocates during a test's window. It is half the smallest move any test + /// below expects, so a wrongly ordered or wrongly sized update still lands outside it. + const MARGIN: usize = 16 * MIB; + + fn about(actual: usize, expected: usize) -> bool { + actual.abs_diff(expected) <= MARGIN } - #[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"); + /// An inner allocator that records the reported balance at the moment each inner call is + /// made, which pins down whether the wrapper accounts before or after delegating. + struct Recording { + balance_at_dealloc: AtomicUsize, + balance_at_realloc: AtomicUsize, } - #[test] - fn settle_flushes_negative_drift() { - let drift = Cell::new(0); - settle(&drift, -SETTLE_THRESHOLD); - assert_eq!(drift.get(), 0); + impl Recording { + fn new() -> Self { + Self { + balance_at_dealloc: AtomicUsize::new(usize::MAX), + balance_at_realloc: AtomicUsize::new(usize::MAX), + } + } + } + + 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) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + self.balance_at_realloc + .store(current_balance(), Ordering::Relaxed); + System.realloc(ptr, layout, new_size) + } } #[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); + fn settle_flushes_only_at_the_threshold() { + for (delta, residue) in [ + (1024, 1024), + (-1024, -1024), + (SETTLE_THRESHOLD - 1, SETTLE_THRESHOLD - 1), + (SETTLE_THRESHOLD, 0), + (-SETTLE_THRESHOLD, 0), + ] { + let drift = Cell::new(0); + settle(&drift, delta); + // A flush resets the drift to zero, so the residue alone says whether the shared + // balance was touched. Reading `BALANCE` here would race with every other test. + assert_eq!(drift.get(), residue, "delta {delta}"); + } } - /// 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. + /// 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. The + /// block is zeroed and never touched, so it costs address space rather than resident memory. #[test] #[cfg(feature = "alloc-accounting")] fn a_real_allocation_raises_the_balance() { use std::hint::black_box; - const SIZE: usize = 256 * 1024 * 1024; + const SIZE: usize = 256 * MIB; let _guard = serial(); let before = current_balance(); // `black_box` keeps the allocation observable so it cannot be elided. @@ -269,41 +285,15 @@ mod tests { 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`. + /// The balance must drop before the inner allocator is asked to free the block, because + /// jemalloc drops its own count at the start of a large free and then spends milliseconds + /// unmapping the pages; see the comment on `dealloc`. #[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; + const SIZE: usize = 64 * MIB; let _guard = serial(); - let allocator = AccountingAllocator::new(Recording { - balance_at_dealloc: AtomicUsize::new(usize::MAX), - }); + let allocator = AccountingAllocator::new(Recording::new()); 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 @@ -314,25 +304,71 @@ mod tests { 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 \ + about(seen + SIZE, after_alloc), + "inner dealloc saw balance {seen}, expected about {} (balance after alloc was \ {after_alloc})", - after_alloc - SIZE / 2 + after_alloc.saturating_sub(SIZE) + ); + } + + /// `realloc` moves the balance by the size difference, not by the new size, and does so after + /// delegating: the inner allocator must see the balance still carrying the old size. + #[test] + fn realloc_accounts_the_size_difference_after_delegating() { + const OLD: usize = 64 * MIB; + const GROWN: usize = 96 * MIB; + const SHRUNK: usize = 32 * MIB; + let _guard = serial(); + let allocator = AccountingAllocator::new(Recording::new()); + let layout = Layout::from_size_align(OLD, 8).unwrap(); + + // SAFETY: each layout matches the block's current size, and the block is freed at the end + // through the same allocator that produced it. + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + let before_grow = current_balance(); + + let ptr = unsafe { allocator.realloc(ptr, layout, GROWN) }; + assert!(!ptr.is_null()); + let after_grow = current_balance(); + let seen = allocator.inner.balance_at_realloc.load(Ordering::Relaxed); + assert!( + about(seen, before_grow), + "inner realloc saw balance {seen}, expected about {before_grow}: the wrapper must \ + account after delegating" + ); + assert!( + about(after_grow, before_grow + (GROWN - OLD)), + "growing {OLD} -> {GROWN} moved the balance {before_grow} -> {after_grow}, expected \ + about +{}", + GROWN - OLD + ); + + let layout = Layout::from_size_align(GROWN, 8).unwrap(); + let ptr = unsafe { allocator.realloc(ptr, layout, SHRUNK) }; + assert!(!ptr.is_null()); + let after_shrink = current_balance(); + assert!( + about(after_shrink + (GROWN - SHRUNK), after_grow), + "shrinking {GROWN} -> {SHRUNK} moved the balance {after_grow} -> {after_shrink}, \ + expected about -{}", + GROWN - SHRUNK ); + + unsafe { allocator.dealloc(ptr, Layout::from_size_align(SHRUNK, 8).unwrap()) }; } /// 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. + /// The worker writes a drift straight into its `LOCAL_DRIFT` cell and exits, so the only path + /// by which that value can reach the shared balance is `ThreadDrift::drop`. That holds only + /// while the wrapper is not installed: with it, thread teardown's own allocations call `track` + /// and flush the oversized drift before the destructor runs, and the test would pass without + /// one. So the test is confined to the default build, which is the one CI runs. The injected + /// amount is far larger than any real allocation, and is taken back out afterwards. #[test] + #[cfg(not(feature = "alloc-accounting"))] fn thread_exit_settles_remaining_drift() { use std::thread; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 6467711ab4b..0872a478f93 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -60,19 +60,10 @@ pub mod parquet; #[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. +// Global allocator selection. `backend` names the allocator the feature set asks for: jemalloc +// where it builds, otherwise mimalloc, otherwise the system allocator. The three cfgs partition +// every feature combination, so exactly one `backend` exists and a combination matching none would +// fail to compile rather than install nothing and leave the accounting metric reading zero. /// jemalloc, on targets where it builds, unless mimalloc was also requested. #[cfg(all( @@ -84,10 +75,6 @@ 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. @@ -99,15 +86,10 @@ 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. +/// The system allocator: the complement of the two cases above, which covers neither feature, a +/// jemalloc request on MSVC, and both features together. #[cfg(not(any( all( not(target_env = "msvc"), @@ -119,9 +101,6 @@ mod backend { 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; @@ -129,11 +108,15 @@ mod backend { } /// 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. +/// `"system"`. The selection is decided here and nowhere else, so the `alloc_overhead` benchmark's +/// liveness check reads it from here rather than re-deriving it from the feature set. +#[doc(hidden)] pub use backend::NAME as ALLOCATOR_BACKEND; +#[cfg(not(feature = "alloc-accounting"))] +#[global_allocator] +static GLOBAL: backend::Backend = backend::BACKEND; + #[cfg(feature = "alloc-accounting")] #[global_allocator] static GLOBAL: alloc_accounting::AccountingAllocator = From 5a4334a1ade5a9fe6b524c25ed51360aa61eab08 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 15:06:40 -0600 Subject: [PATCH 8/8] fix: compare against observed zero reservations in analyze_trace and format tracing.md A zero pool reservation is a real sample that allocation can exceed, so the analyzer now defers the comparison only until it has seen at least one pool sample, and says so when it never does instead of reporting that allocation never exceeded reservations. Reflow the tracing.md label table with prettier, which the Preflight check enforces. --- docs/source/contributor-guide/tracing.md | 14 +++++++------- native/common/src/bin/analyze_trace.rs | 12 +++++++++--- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 3a6b9b5afe0..9757bd8a57a 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -121,10 +121,10 @@ not being tracked by the pool. ## Definition of Labels -| Label | Meaning | -| -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | -| jvm_heap_used | JVM heap memory usage of live objects for the executor process | -| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | -| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate: each live thread holds up to 64 KiB of un-flushed delta, so the value can lag the true balance by that much per thread. | -| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | -| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | +| Label | Meaning | +| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| jvm_heap_used | JVM heap memory usage of live objects for the executor process | +| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | +| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate to within 64 KiB of un-flushed delta per live thread. | +| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | +| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs index 47039df9400..30f55a2ae39 100644 --- a/native/common/src/bin/analyze_trace.rs +++ b/native/common/src/bin/analyze_trace.rs @@ -158,13 +158,15 @@ fn main() { continue; } - // After each allocated or pool update, check the current state + // After each allocated or pool update, check the current state. A comparison needs one + // sample of each side: an observed zero reservation is a real value that allocation can + // exceed, so only the absence of any pool sample defers the check. let pool_total: u64 = pool_by_thread.values().sum(); if pool_total > peak_pool_total { peak_pool_total = pool_total; } - if latest_allocated > 0 && pool_total > 0 && latest_allocated > pool_total { + if source.is_some() && !pool_by_thread.is_empty() && latest_allocated > pool_total { let excess = latest_allocated - pool_total; if excess > peak_excess { peak_excess = excess; @@ -208,7 +210,11 @@ fn main() { ); println!(); - if violations.is_empty() { + if pool_by_thread.is_empty() { + println!( + "No pool reservation samples in the trace, so there is nothing to compare against." + ); + } else if violations.is_empty() { println!("OK: {source} never exceeded the total pool reservation."); } else { println!(