From 303875f28b6244755f7d030036556919e69d5ad8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 14 Sep 2026 13:34:30 -0600 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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/7] 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 9f20277123a40eda6def7bcf9b8965a2d3f3b133 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 08:24:02 -0600 Subject: [PATCH 6/7] build: enable the alloc-accounting feature by default Every build now installs the accounting wrapper and reports the native_allocated tracing metric. The feature guards stay so a --no-default-features build still compiles without the wrapper; the bench header and tracing guide document that opt-out invocation. --- docs/source/contributor-guide/tracing.md | 7 ++++--- native/core/Cargo.toml | 7 ++++--- native/core/benches/alloc_overhead.rs | 7 ++++--- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 4fb7a0e28de..e75b4c5553c 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -37,10 +37,11 @@ The `alloc-accounting` feature adds a second, allocator-independent measure of n 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: +allocator's own caching in the way. It is on by default, so the command above already includes it; +the two features are independent, and a build that drops the defaults can re-add either one: ```shell -make release COMET_FEATURES="jemalloc,alloc-accounting" +cd native && cargo build --release --no-default-features --features hdfs-opendal,jemalloc ``` Example output: @@ -121,6 +122,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 (`alloc-accounting` feature, on by default) | | 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 501592945d2..486126cd4ce 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -105,7 +105,7 @@ datafusion-functions-nested = { version = "55.1.0" } [features] backtrace = ["datafusion/backtrace"] -default = ["hdfs-opendal"] +default = ["hdfs-opendal", "alloc-accounting"] contrib-lance = ["dep:comet-contrib-lance"] hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] @@ -116,8 +116,9 @@ 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. +# against the memory pool's reservations, and gates the `greedy_unified_checked` memory pool. +# Never rejects an allocation. On by default; build with `--no-default-features` (re-adding the +# other defaults) to drop the wrapper and its per-allocation work. alloc-accounting = [] # exclude optional packages from cargo machete verifications diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index ba6b1663349..1254282b9b4 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -17,11 +17,12 @@ //! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation. //! -//! Run the same benchmark with and without the feature and compare: +//! The feature is on by default, so the "off" run has to drop the defaults and re-add the rest: //! //! ```shell -//! cargo bench --bench alloc_overhead -- --save-baseline off -//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off +//! cargo bench --bench alloc_overhead --no-default-features --features jemalloc,hdfs-opendal \ +//! -- --save-baseline off +//! cargo bench --bench alloc_overhead --features jemalloc -- --baseline off //! ``` //! //! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this From ed217bc0b26e98f927a188fbb9461d790fe4601d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 15 Sep 2026 08:24:08 -0600 Subject: [PATCH 7/7] feat: add greedy_unified_checked memory pool gated on real native usage CheckedMemoryPool refuses a reservation when the bytes the allocator has actually handed out, plus the request, would exceed Comet's off-heap budget, and otherwise defers to CometUnifiedMemoryPool. The new off-heap pool type greedy_unified_checked selects it. The balance and budget are process-wide, so once any task reaches the budget every task's next reservation is denied; allocations themselves are never refused. Selecting the pool without the alloc-accounting feature is a config error. --- docs/source/user-guide/latest/tuning.md | 12 +- native/core/src/alloc_accounting.rs | 11 +- .../execution/memory_pools/checked_pool.rs | 199 ++++++++++++++++++ .../core/src/execution/memory_pools/config.rs | 61 +++++- native/core/src/execution/memory_pools/mod.rs | 8 + .../scala/org/apache/comet/CometConf.scala | 3 +- .../apache/comet/exec/CometExecSuite.scala | 39 ++++ 7 files changed, 327 insertions(+), 6 deletions(-) create mode 100644 native/core/src/execution/memory_pools/checked_pool.rs diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 25b805d1ccb..ce9b8bfef3e 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -100,8 +100,9 @@ The valid pool types are: - `fair_unified` (default when `spark.memory.offHeap.enabled=true` is set) - `greedy_unified` +- `greedy_unified_checked` -Both pool types are shared across all native execution contexts within the same Spark task. When +All pool types are shared across all native execution contexts within the same Spark task. When Comet executes a shuffle, it runs two native execution contexts concurrently (e.g. one for pre-shuffle operators and one for the shuffle writer). The shared pool ensures that the combined memory usage stays within the per-task limit. @@ -114,6 +115,15 @@ when there is sufficient memory in order to leave enough memory for other operat The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not need to spill or have a single spillable operator. +The `greedy_unified_checked` pool type is `greedy_unified` with one extra check. Before asking Spark for memory, it +compares the bytes the native allocator has actually handed out, across the whole executor, against Comet's off-heap +budget (`spark.memory.offHeap.size` multiplied by `spark.comet.exec.memoryPool.fraction`). If the real usage plus the +request would exceed the budget, the reservation is denied and the operator spills, or fails if it cannot spill. +This catches native memory that operators never reserved, which the other pools cannot see. The check is process-wide, +so once any task pushes real usage to the budget every task's next reservation is denied, and it only gates +reservations: allocations themselves are never refused. It requires the native library to be built with the +`alloc-accounting` cargo feature, which is on by default. + [shuffle]: #shuffle [Advanced Memory Tuning]: #advanced-memory-tuning diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index f476c07271a..636d2a2a745 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -194,8 +194,7 @@ unsafe impl GlobalAlloc for AccountingAllocator { } #[cfg(test)] -mod tests { - use super::*; +pub(crate) mod test_support { use std::sync::{Mutex, MutexGuard}; /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it @@ -205,11 +204,17 @@ mod tests { /// allocates in that time. static SERIAL: Mutex<()> = Mutex::new(()); - fn serial() -> MutexGuard<'static, ()> { + pub(crate) fn serial() -> MutexGuard<'static, ()> { SERIAL .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } +} + +#[cfg(test)] +mod tests { + use super::test_support::serial; + use super::*; #[test] fn settle_accumulates_below_the_threshold() { diff --git a/native/core/src/execution/memory_pools/checked_pool.rs b/native/core/src/execution/memory_pools/checked_pool.rs new file mode 100644 index 00000000000..53b742931b3 --- /dev/null +++ b/native/core/src/execution/memory_pools/checked_pool.rs @@ -0,0 +1,199 @@ +// 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. + +use std::fmt::{Debug, Display, Formatter, Result as FmtResult}; + +use datafusion::{ + common::{resources_datafusion_err, DataFusionError}, + execution::memory_pool::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}, +}; + +use crate::alloc_accounting; + +/// A memory pool that refuses a reservation when the bytes the native allocator has actually +/// handed out, plus the request, would exceed a budget, and otherwise defers to an inner pool. +/// +/// Every other pool only counts what operators voluntarily reserve, so native memory that bypasses +/// the pool is invisible to it until the executor exceeds its container limit. This pool consults +/// the `alloc-accounting` balance before the inner pool's own check, which turns that overrun into +/// a `ResourcesExhausted` error at the next reservation instead. The operators that can spill do +/// so; the ones that cannot fail the task. +/// +/// Both the balance and the budget are process-wide. The balance counts every byte the Rust +/// allocator has served in this executor, not just this task's, and the budget is Comet's whole +/// off-heap allotment. So once any task pushes real usage to the budget, every task's next non-zero +/// reservation is denied. There is no per-task attribution, and allocations themselves are never +/// refused: this is a reservation gate, not a hard limit. +pub struct CheckedMemoryPool { + inner: P, + budget: usize, +} + +impl CheckedMemoryPool

{ + pub fn new(inner: P, budget: usize) -> Self { + Self { inner, budget } + } +} + +impl Debug for CheckedMemoryPool

{ + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + f.debug_struct("CheckedMemoryPool") + .field("budget", &self.budget) + .field("inner", &self.inner) + .finish() + } +} + +impl Display for CheckedMemoryPool

{ + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!( + f, + "CheckedMemoryPool(budget={}, inner={})", + self.budget, self.inner + ) + } +} + +impl MemoryPool for CheckedMemoryPool

{ + fn name(&self) -> &str { + "CheckedMemoryPool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer) + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer) + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.try_grow(reservation, additional).unwrap() + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<(), DataFusionError> { + if additional == 0 { + return Ok(()); + } + // Checked first because it is a single atomic load, whereas the inner unified pools cross + // JNI to ask Spark. A request denied here never reaches Spark's ledger. + let in_use = alloc_accounting::current_balance(); + if in_use.saturating_add(additional) > self.budget { + return Err(resources_datafusion_err!( + "Failed to reserve {additional} bytes for {}: native memory in use is {in_use} \ + bytes of a {} byte budget. Reserved: {}", + reservation.consumer().name(), + self.budget, + self.reserved() + )); + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + MemoryLimit::Finite(self.budget) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::UnboundedMemoryPool; + use std::sync::Arc; + + fn pool_with_budget(budget: usize) -> Arc { + Arc::new(CheckedMemoryPool::new( + UnboundedMemoryPool::default(), + budget, + )) + } + + #[test] + fn a_zero_byte_grow_never_fails() { + let pool = pool_with_budget(0); + let reservation = MemoryConsumer::new("zero").register(&pool); + reservation.try_grow(0).unwrap(); + } + + #[test] + fn reports_the_budget_as_its_limit() { + assert!(matches!( + pool_with_budget(4096).memory_limit(), + MemoryLimit::Finite(4096) + )); + } + + #[test] + fn successful_grows_and_shrinks_reach_the_inner_pool() { + let pool = pool_with_budget(usize::MAX); + let reservation = MemoryConsumer::new("delegate").register(&pool); + reservation.try_grow(1024).unwrap(); + assert_eq!(pool.reserved(), 1024); + reservation.shrink(1024); + assert_eq!(pool.reserved(), 0); + } + + /// The gate compares real bytes, not reservations: a block this pool never heard about is + /// enough to deny a one-byte request, and freeing it is enough to allow the same request. + /// + /// The block is touched so it is really allocated, and the margins are far wider than anything + /// the rest of the crate allocates in the microseconds between the checks. The serial lock + /// keeps the accounting tests that move the balance by tens of megabytes out of that window. + #[test] + #[cfg(feature = "alloc-accounting")] + fn denies_when_real_bytes_plus_request_exceed_the_budget() { + use std::hint::black_box; + + const HEADROOM: usize = 64 * 1024 * 1024; + const BLOCK: usize = 256 * 1024 * 1024; + + let _guard = alloc_accounting::test_support::serial(); + let budget = alloc_accounting::current_balance() + HEADROOM; + let pool = pool_with_budget(budget); + let reservation = MemoryConsumer::new("checked").register(&pool); + + let held: Vec = black_box(vec![1u8; BLOCK]); + let denied = reservation.try_grow(1).unwrap_err(); + black_box(&held); + assert!( + matches!(denied, DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted, got {denied:?}" + ); + let message = denied.to_string(); + assert!( + message.contains("native memory in use is") && message.contains("checked"), + "message should name the real bytes in use and the consumer: {message}" + ); + assert_eq!(pool.reserved(), 0, "a denied request must not be reserved"); + + drop(held); + reservation.try_grow(1).unwrap(); + assert_eq!(pool.reserved(), 1); + } +} diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index 312a3604383..1ee76dc77db 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -17,9 +17,11 @@ use crate::errors::{CometError, CometResult}; -#[derive(Copy, Clone, PartialEq, Eq)] +#[derive(Copy, Clone, Debug, PartialEq, Eq)] pub(crate) enum MemoryPoolType { GreedyUnified, + /// `GreedyUnified` behind a gate on the bytes the native allocator has actually handed out. + GreedyUnifiedChecked, FairUnified, Greedy, FairSpill, @@ -30,6 +32,7 @@ pub(crate) enum MemoryPoolType { Unbounded, } +#[derive(Debug)] pub(crate) struct MemoryPoolConfig { pub(crate) pool_type: MemoryPoolType, pub(crate) pool_size: usize, @@ -60,6 +63,20 @@ pub(crate) fn parse_memory_pool_config( // shared with Spark is set by `spark.memory.offHeap.size`. MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) } + // The checked pool gates on real native usage, so it needs the budget the balance is + // compared against: the same number `fair_unified` receives. + #[cfg(feature = "alloc-accounting")] + "greedy_unified_checked" => { + MemoryPoolConfig::new(MemoryPoolType::GreedyUnifiedChecked, pool_size) + } + #[cfg(not(feature = "alloc-accounting"))] + "greedy_unified_checked" => { + return Err(CometError::Config( + "Memory pool type greedy_unified_checked requires the native library to be \ + built with the alloc-accounting cargo feature" + .to_string(), + )) + } _ => { return Err(CometError::Config(format!( "Unsupported memory pool type for off-heap mode: {memory_pool_type}" @@ -92,3 +109,45 @@ pub(crate) fn parse_memory_pool_config( }; Ok(memory_pool_config) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[cfg(feature = "alloc-accounting")] + fn checked_pool_takes_the_off_heap_limit_as_its_budget() { + let config = + parse_memory_pool_config(true, "greedy_unified_checked".to_string(), 1 << 30, 1 << 20) + .unwrap(); + assert_eq!(config.pool_type, MemoryPoolType::GreedyUnifiedChecked); + assert_eq!(config.pool_size, 1 << 30); + } + + #[test] + #[cfg(not(feature = "alloc-accounting"))] + fn checked_pool_is_rejected_without_the_accounting_feature() { + let err = + parse_memory_pool_config(true, "greedy_unified_checked".to_string(), 1 << 30, 1 << 20) + .unwrap_err(); + assert!( + err.to_string().contains("alloc-accounting"), + "error should name the missing feature: {err}" + ); + } + + #[test] + fn checked_pool_is_off_heap_only() { + let err = parse_memory_pool_config( + false, + "greedy_unified_checked".to_string(), + 1 << 30, + 1 << 20, + ) + .unwrap_err(); + assert!( + err.to_string().contains("on-heap mode"), + "unexpected error: {err}" + ); + } +} diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index d7c2911f913..ba12bd6b8f0 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -15,12 +15,14 @@ // specific language governing permissions and limitations // under the License. +mod checked_pool; mod config; mod fair_pool; pub mod logging_pool; mod task_shared; mod unified_pool; +use checked_pool::CheckedMemoryPool; use datafusion::execution::memory_pool::{ FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool, }; @@ -62,6 +64,12 @@ pub(crate) fn create_memory_pool( task_attempt_id, )) }), + MemoryPoolType::GreedyUnifiedChecked => acquire_task_shared_pool(task_attempt_id, || { + tracked(CheckedMemoryPool::new( + CometUnifiedMemoryPool::new(comet_task_memory_manager, task_attempt_id), + pool_size, + )) + }), MemoryPoolType::FairUnified => acquire_task_shared_pool(task_attempt_id, || { tracked(CometFairMemoryPool::new( comet_task_memory_manager, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 43f030a7d9d..199e4ec0f0b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -824,7 +824,8 @@ object CometConf extends ShimCometConf { .category(CATEGORY_TUNING) .doc( "The type of memory pool to be used for Comet native execution when running Spark in " + - "off-heap mode. Available pool types are `greedy_unified` and `fair_unified`. " + + "off-heap mode. Available pool types are `greedy_unified`, `fair_unified`, and " + + "`greedy_unified_checked`. " + s"$TUNING_GUIDE.") .stringConf .createWithDefault("fair_unified") diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 5b5d43dbe8a..707dd5143df 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3053,6 +3053,45 @@ class CometExecSuite extends CometTestBase { } } + test("greedy_unified_checked pool completes a sort within budget") { + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key -> "greedy_unified_checked") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + makeRawTimeParquetFileColumns(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) + readParquetFile(path.toString) { df => + checkSparkAnswerAndOperator(df.sortWithinPartitions($"_0".desc_nulls_first)) + } + } + } + } + + test("greedy_unified_checked pool denies reservations once real native usage exceeds budget") { + // A fraction this small makes the budget a few kilobytes, which the executor's native code + // already exceeds before the query starts, so the sort's first reservation is denied by the + // real-bytes check rather than by Spark. Having reserved nothing, the sort cannot spill. + withSQLConf( + CometConf.COMET_OFFHEAP_MEMORY_POOL_TYPE.key -> "greedy_unified_checked", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.000001") { + withTempDir { dir => + val path = new Path(dir.toURI.toString, "part-r-0.parquet") + makeRawTimeParquetFileColumns(path, dictionaryEnabled = true, n = 1000, rowGroupSize = 10) + readParquetFile(path.toString) { df => + val thrown = intercept[Throwable] { + df.sortWithinPartitions($"_0".desc_nulls_first).collect() + } + val messages = Iterator + .iterate(thrown)(_.getCause) + .takeWhile(_ != null) + .map(e => Option(e.getMessage).getOrElse("")) + .toSeq + assert( + messages.exists(_.contains("native memory in use is")), + s"expected the checked pool's denial, got: ${messages.mkString(" <- ")}") + } + } + } + } + test("limit") { Seq("native", "jvm").foreach { columnarShuffleMode => withSQLConf(