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 88a291f421b..9757bd8a57a 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 @@ -55,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: @@ -66,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: @@ -77,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 --- @@ -102,14 +115,16 @@ 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 -| 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) | -| 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 8df83f7ceae..30f55a2ae39 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,16 @@ fn main() { continue; } - // After each jemalloc 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_jemalloc > 0 && pool_total > 0 && latest_jemalloc > pool_total { - let excess = latest_jemalloc - 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; } @@ -147,46 +178,60 @@ 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."); + 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!( - "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/Cargo.toml b/native/core/Cargo.toml index f0c7735a503..501592945d2 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"] @@ -123,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..dbba3ed424c --- /dev/null +++ b/native/core/benches/alloc_overhead.rs @@ -0,0 +1,197 @@ +// 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 +//! ``` +//! +//! `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, Criterion, Throughput}; +use std::hint::black_box; +use std::sync::Once; +use std::thread; +use std::time::Instant; + +/// 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(); + }); +} + +/// 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. +#[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); + epoch::advance().expect("jemalloc epoch"); + let allocated = stats::allocated::read().expect("jemalloc stats.allocated"); + assert!( + allocated >= 8 * 1024 * 1024, + "the library selected jemalloc but jemalloc is not the global allocator of this binary \ + (stats.allocated = {allocated}); the numbers below would be meaningless" + ); + drop(held); +} + +/// Without the feature the library cannot have selected jemalloc, so this is never reached. +#[cfg(not(feature = "jemalloc"))] +fn assert_jemalloc_is_live() { + unreachable!("the library reports the jemalloc backend but the feature is not enabled"); +} + +/// 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: 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(); + 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() {} + +/// 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"); + 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(); +} + +/// 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(|| { + 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(|| { + 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, 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 new file mode 100644 index 00000000000..394bbe774d3 --- /dev/null +++ b/native/core/src/alloc_accounting.rs @@ -0,0 +1,394 @@ +// 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. +//! +//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed +//! 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, +//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc +//! `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 + +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 discarded every time a +/// thread died, and tokio's blocking pool churns threads on its idle timeout. +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. +/// +/// 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 { + BALANCE.load(Ordering::Relaxed).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) { + // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the + // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping + // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on + // the balance for that whole window, after the allocator's own statistics had already + // dropped it. + track(-(layout.size() as isize)); + self.inner.dealloc(ptr, layout); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = self.inner.realloc(ptr, layout, new_size); + if !new_ptr.is_null() { + // Accounting after the fact is only safe because this allocator cannot fail the + // allocation or unwind. A variant that enforced a limit would have to decide *before* + // delegating: `realloc` may free or move the old block, and a caller that never + // received the new pointer would free the stale one while unwinding. + // + // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast + // wraps. + track(new_size as isize - layout.size() as isize); + } + new_ptr + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::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 + /// 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()) + } + + 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 + } + + /// 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, + } + + 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 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. 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 * MIB; + let _guard = serial(); + let before = current_balance(); + // `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 + SIZE / 2, + "a {SIZE} byte allocation should raise the balance (before={before}, during={during}); \ + is the accounting wrapper installed for this feature set?" + ); + drop(held); + } + + /// The balance must drop before the inner allocator is asked to free the block, 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() { + // Well above the settle threshold, so both the allocation and the free flush immediately. + const SIZE: usize = 64 * MIB; + let _guard = serial(); + 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 + // 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); + assert!( + about(seen + SIZE, after_alloc), + "inner dealloc saw balance {seen}, expected about {} (balance after alloc was \ + {after_alloc})", + 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, 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; + + const INJECTED: isize = 1 << 40; + let _guard = serial(); + + 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); + + assert!( + moved >= INJECTED / 2, + "drift from an exited thread never reached the shared balance: \ + balance moved {moved} bytes, expected at least {}", + INJECTED / 2 + ); + } +} 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..0872a478f93 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; @@ -65,6 +52,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,20 +60,67 @@ 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 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( not(target_env = "msvc"), feature = "jemalloc", not(feature = "mimalloc") ))] -#[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; +mod backend { + pub type Backend = tikv_jemallocator::Jemalloc; + pub const BACKEND: Backend = tikv_jemallocator::Jemalloc; + pub const NAME: &str = "jemalloc"; +} +/// mimalloc, unless a usable jemalloc was also requested. #[cfg(all( feature = "mimalloc", not(all(not(target_env = "msvc"), feature = "jemalloc")) ))] +mod backend { + pub type Backend = mimalloc::MiMalloc; + pub const BACKEND: Backend = mimalloc::MiMalloc; + pub const NAME: &str = "mimalloc"; +} + +/// 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"), + feature = "jemalloc", + not(feature = "mimalloc") + ), + all( + feature = "mimalloc", + not(all(not(target_env = "msvc"), feature = "jemalloc")) + ) +)))] +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"`. 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: MiMalloc = MiMalloc; +static GLOBAL: alloc_accounting::AccountingAllocator = + alloc_accounting::AccountingAllocator::new(backend::BACKEND); #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init(