andygrove commented on code in PR #5934: URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4019832858
########## native/core/benches/alloc_overhead.rs: ########## @@ -0,0 +1,225 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation. +//! +//! Run the same benchmark with and without the feature and compare: +//! +//! ```shell +//! cargo bench --bench alloc_overhead -- --save-baseline off +//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off +//! ``` +//! +//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs, which reaches this +//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern` +//! crate that nothing names is dropped from the crate graph along with its allocator, so the +//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet` +//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run +//! if either the selected backend or the wrapper is somehow not in effect, because a number +//! measured against the wrong allocator would be worse than no number. +//! +//! `small_churn` is the worst case for the thread-local path: allocations so small that the +//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is +//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold +//! flushes to the process-wide atomic on every call, and the parallel variant does that from every +//! core at once, so the gap between the single-threaded and parallel numbers is the cost of +//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where +//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below +//! `arrow_sized_churn`, because they do actual work between allocations. + +// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this +// binary even when the feature set leaves nothing here that names the crate. +extern crate comet; + +use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput}; +use std::hint::black_box; +use std::thread; +use std::time::Instant; + +/// Guards against measuring the wrong allocator, and says which one is being measured. +/// +/// Which backend is in effect is `lib.rs`'s decision, not this crate's feature flags': with +/// `jemalloc,mimalloc` together the library deliberately falls back to the system allocator, and +/// jemalloc on MSVC is not selected at all. So the check asks the library which backend it chose +/// rather than re-deriving that from the feature set, and can never disagree with the selection it +/// is meant to verify. +fn assert_backend_is_live() { + static ANNOUNCE: std::sync::Once = std::sync::Once::new(); + ANNOUNCE.call_once(|| { + eprintln!( + "alloc_overhead: measuring the `{}` allocator backend", + comet::ALLOCATOR_BACKEND + ) + }); + if comet::ALLOCATOR_BACKEND == "jemalloc" { + assert_jemalloc_is_live(); + } +} + +/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this +/// binary that count stays at zero, and a "jemalloc" baseline would in fact be the system +/// allocator. +#[cfg(feature = "jemalloc")] +fn assert_jemalloc_is_live() { + use tikv_jemalloc_ctl::{epoch, stats}; + let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]); + black_box(&held); + epoch::advance().expect("jemalloc epoch"); + let allocated = stats::allocated::read().expect("jemalloc stats.allocated"); + assert!( + allocated >= 8 * 1024 * 1024, + "the library selected jemalloc but jemalloc is not the global allocator of this binary \ + (stats.allocated = {allocated}); the numbers below would be meaningless" + ); + drop(held); +} + +/// Without the feature the library cannot have selected jemalloc, so this is never reached. +#[cfg(not(feature = "jemalloc"))] +fn assert_jemalloc_is_live() { + unreachable!("the library reports the jemalloc backend but the feature is not enabled"); +} + +/// Guards against measuring nothing. If the wrapper were not actually installed in the benchmark +/// binary, every "with the feature" number would silently be a second baseline run. +#[cfg(feature = "alloc-accounting")] +fn assert_accounting_is_live() { + let before = comet::alloc_accounting::current_balance(); + // `black_box` is load-bearing: benchmarks build in release mode, where LLVM will happily + // elide an allocation whose contents are never observed, and the check would then fail + // against a wrapper that is in fact working. + let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]); + black_box(&held); + let during = comet::alloc_accounting::current_balance(); + assert!( + during >= before + 4 * 1024 * 1024, + "alloc-accounting is enabled but the allocator is not installed in this binary \ + (balance {before} -> {during}); the numbers below would be meaningless" + ); + drop(held); +} + +#[cfg(not(feature = "alloc-accounting"))] +fn assert_accounting_is_live() {} + +/// Allocation sizes that stay under the 64 KiB settle threshold, so most iterations exercise only +/// the thread-local fast path rather than the atomic flush. +fn small_churn(c: &mut Criterion) { + assert_backend_is_live(); + assert_accounting_is_live(); + let mut group = c.benchmark_group("alloc_overhead"); + for size in [16usize, 256, 4096] { + group.throughput(Throughput::Elements(1)); + group.bench_function(format!("alloc_free_{size}b"), |b| { + b.iter(|| { + let v: Vec<u8> = Vec::with_capacity(black_box(size)); + black_box(&v); + }); + }); + } + group.finish(); +} + +/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation +/// Comet does in bulk. +fn arrow_sized_churn(c: &mut Criterion) { + let mut group = c.benchmark_group("alloc_overhead"); + group.throughput(Throughput::Bytes(64 * 1024)); + group.bench_function("alloc_fill_free_64kb", |b| { + b.iter_batched( + || (), + |()| { + let v: Vec<u8> = vec![1u8; black_box(64 * 1024)]; + black_box(v.len()) + }, + BatchSize::SmallInput, + ); + }); + group.finish(); +} + +/// Repeated growth, which is the `realloc` path: a builder doubling its buffer. +fn growth_churn(c: &mut Criterion) { Review Comment: Fixed in 4b71ad9: both asserts are behind a single `Once` in `assert_allocators_are_live`, which every benchmark function calls first. Checked with `cargo bench --bench alloc_overhead -- --test alloc_fill_free_64kb`, which now prints the backend announcement. ########## native/core/src/alloc_accounting.rs: ########## @@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide accounting of the bytes currently handed out by the Rust global allocator. +//! +//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared +//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes +//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata, +//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on +//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at +//! runtime. See the [memory management contributor guide] for the full picture. +//! +//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed +//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it +//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the +//! accounting gap can be seen in tracing output next to the pool reservations it should be +//! compared against. +//! +//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation, +//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc +//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much +//! tighter one than pool reservations. +//! +//! [memory management contributor guide]: +//! https://datafusion.apache.org/comet/contributor-guide/memory_management.html + +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicIsize, Ordering}; + +/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this. +/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic +/// read-modify-write per 64 KiB of churn touches the shared cacheline. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before +/// another flushes the matching positive one. +static BALANCE: AtomicIsize = AtomicIsize::new(0); + +thread_local! { + /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles + /// directly instead of recursing. The only such allocation today is the one some platforms + /// make when registering `LOCAL_DRIFT`'s destructor on first touch. + /// + /// Const-initialized and destructor-free, so reading it never allocates and never fails — + /// which is what makes it safe to consult before touching `LOCAL_DRIFT`. + static IN_TRACK: Cell<bool> = const { Cell::new(false) }; + + /// This thread's un-flushed delta. + static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) }; +} + +/// Owns a thread's un-flushed delta and settles the remainder when the thread exits. +/// +/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently +/// discarded every time a thread died. Worker threads live for the process lifetime, but the +/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a +/// slowly accumulating bias in the reported balance. +struct ThreadDrift(Cell<isize>); + +impl Drop for ThreadDrift { + fn drop(&mut self) { + let drift = self.0.replace(0); + if drift != 0 { + BALANCE.fetch_add(drift, Ordering::Relaxed); + } + } +} + +/// Bytes currently handed out by the Rust global allocator, process-wide. +/// +/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the +/// balance can dip below zero transiently while per-thread deltas settle out of order. +pub fn current_balance() -> usize { + clamp_balance(BALANCE.load(Ordering::Relaxed)) +} + +/// Clamps a signed balance to the unsigned value reported to callers. +fn clamp_balance(balance: isize) -> usize { + balance.max(0) as usize +} + +/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches +/// [`SETTLE_THRESHOLD`]. +fn settle(local_drift: &Cell<isize>, delta: isize) { + let drift = local_drift.get().wrapping_add(delta); + if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize { + local_drift.set(0); + BALANCE.fetch_add(drift, Ordering::Relaxed); + } else { + local_drift.set(drift); + } +} + +/// Records a signed byte delta against the process balance. +#[inline] +fn track(delta: isize) { + if delta == 0 { + return; + } + + // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will + // clear it, so this frame must only settle and return. + if IN_TRACK.with(|in_track| in_track.replace(true)) { + BALANCE.fetch_add(delta, Ordering::Relaxed); + return; + } + + // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already + // run, and any allocation after that point must not panic inside the allocator. + if LOCAL_DRIFT + .try_with(|thread_drift| settle(&thread_drift.0, delta)) + .is_err() + { + BALANCE.fetch_add(delta, Ordering::Relaxed); + } + + IN_TRACK.with(|in_track| in_track.set(false)); +} + +/// Wraps a global allocator, accounting the `Layout` bytes it hands out. +/// +/// Adapted from the `AccountingAllocator` in +/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in +/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly. +pub struct AccountingAllocator<A: GlobalAlloc> { + inner: A, +} + +impl<A: GlobalAlloc> AccountingAllocator<A> { + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The +// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not +// inspect, retain, or alter any pointer, and it cannot unwind. +unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc(layout); + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc_zeroed(layout); + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the + // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping + // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on + // the balance for that whole window, after the allocator's own statistics had already + // dropped it. + track(-(layout.size() as isize)); + self.inner.dealloc(ptr, layout); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = self.inner.realloc(ptr, layout, new_size); + if !new_ptr.is_null() { + // Accounting after the fact is only safe because this allocator cannot fail the + // allocation or unwind. A variant that enforced a limit would have to decide *before* + // delegating: `realloc` may free or move the old block, and a caller that never + // received the new pointer would free the stale one while unwinding. + // + // A single allocation cannot exceed `isize::MAX` on any real platform, so neither cast + // wraps. + track(new_size as isize - layout.size() as isize); + } + new_ptr + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Mutex, MutexGuard}; + + /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it + /// sees every other test's allocations. The tests that move it by tens of megabytes take this + /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by + /// making each window microseconds wide and each expected move far larger than anything else + /// allocates in that time. + static SERIAL: Mutex<()> = Mutex::new(()); + + fn serial() -> MutexGuard<'static, ()> { + SERIAL + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + #[test] + fn settle_accumulates_below_the_threshold() { Review Comment: Done in 4b71ad9, as one table over `[(1024, 1024), (-1024, -1024), (THRESHOLD - 1, THRESHOLD - 1), (THRESHOLD, 0), (-THRESHOLD, 0)]`. ########## native/core/src/alloc_accounting.rs: ########## @@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide accounting of the bytes currently handed out by the Rust global allocator. +//! +//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared +//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes +//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata, +//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on +//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at +//! runtime. See the [memory management contributor guide] for the full picture. +//! +//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed +//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it +//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the +//! accounting gap can be seen in tracing output next to the pool reservations it should be +//! compared against. +//! +//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation, +//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc +//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much +//! tighter one than pool reservations. +//! +//! [memory management contributor guide]: +//! https://datafusion.apache.org/comet/contributor-guide/memory_management.html + +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicIsize, Ordering}; + +/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this. +/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic +/// read-modify-write per 64 KiB of churn touches the shared cacheline. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before +/// another flushes the matching positive one. +static BALANCE: AtomicIsize = AtomicIsize::new(0); + +thread_local! { + /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles + /// directly instead of recursing. The only such allocation today is the one some platforms + /// make when registering `LOCAL_DRIFT`'s destructor on first touch. + /// + /// Const-initialized and destructor-free, so reading it never allocates and never fails — + /// which is what makes it safe to consult before touching `LOCAL_DRIFT`. + static IN_TRACK: Cell<bool> = const { Cell::new(false) }; + + /// This thread's un-flushed delta. + static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) }; +} + +/// Owns a thread's un-flushed delta and settles the remainder when the thread exits. +/// +/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently +/// discarded every time a thread died. Worker threads live for the process lifetime, but the +/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a +/// slowly accumulating bias in the reported balance. +struct ThreadDrift(Cell<isize>); + +impl Drop for ThreadDrift { + fn drop(&mut self) { + let drift = self.0.replace(0); + if drift != 0 { + BALANCE.fetch_add(drift, Ordering::Relaxed); + } + } +} + +/// Bytes currently handed out by the Rust global allocator, process-wide. +/// +/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the +/// balance can dip below zero transiently while per-thread deltas settle out of order. +pub fn current_balance() -> usize { + clamp_balance(BALANCE.load(Ordering::Relaxed)) +} + +/// Clamps a signed balance to the unsigned value reported to callers. +fn clamp_balance(balance: isize) -> usize { Review Comment: Inlined into `current_balance` and the test dropped in 4b71ad9. ########## native/core/src/alloc_accounting.rs: ########## @@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide accounting of the bytes currently handed out by the Rust global allocator. +//! +//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared +//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes +//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata, +//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on +//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at +//! runtime. See the [memory management contributor guide] for the full picture. +//! +//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed +//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it +//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the +//! accounting gap can be seen in tracing output next to the pool reservations it should be +//! compared against. +//! +//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation, +//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc +//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much +//! tighter one than pool reservations. +//! +//! [memory management contributor guide]: +//! https://datafusion.apache.org/comet/contributor-guide/memory_management.html + +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicIsize, Ordering}; + +/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this. +/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic +/// read-modify-write per 64 KiB of churn touches the shared cacheline. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before +/// another flushes the matching positive one. +static BALANCE: AtomicIsize = AtomicIsize::new(0); + +thread_local! { + /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles + /// directly instead of recursing. The only such allocation today is the one some platforms + /// make when registering `LOCAL_DRIFT`'s destructor on first touch. + /// + /// Const-initialized and destructor-free, so reading it never allocates and never fails — + /// which is what makes it safe to consult before touching `LOCAL_DRIFT`. + static IN_TRACK: Cell<bool> = const { Cell::new(false) }; + + /// This thread's un-flushed delta. + static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) }; +} + +/// Owns a thread's un-flushed delta and settles the remainder when the thread exits. +/// +/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently +/// discarded every time a thread died. Worker threads live for the process lifetime, but the +/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a +/// slowly accumulating bias in the reported balance. +struct ThreadDrift(Cell<isize>); + +impl Drop for ThreadDrift { + fn drop(&mut self) { + let drift = self.0.replace(0); + if drift != 0 { + BALANCE.fetch_add(drift, Ordering::Relaxed); + } + } +} + +/// Bytes currently handed out by the Rust global allocator, process-wide. +/// +/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the +/// balance can dip below zero transiently while per-thread deltas settle out of order. +pub fn current_balance() -> usize { Review Comment: Added to the `current_balance` doc and the `tracing.md` row in 4b71ad9: up to `SETTLE_THRESHOLD` per live thread. ########## native/core/src/alloc_accounting.rs: ########## @@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Process-wide accounting of the bytes currently handed out by the Rust global allocator. +//! +//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool) counts *declared +//! reservations*: bytes an operator explicitly asked for. Plenty of real allocation never goes +//! through it — Arrow builders, expression kernels, decompression buffers, Parquet metadata, +//! `object_store` buffers, tokio's own machinery — so pool reservations are a lower bound on +//! Comet's footprint, and the size of the gap is workload-dependent and currently unmeasurable at +//! runtime. See the [memory management contributor guide] for the full picture. +//! +//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed +//! process-wide byte balance, which [`current_balance`] exposes. This is **observability only**: it +//! never rejects an allocation, never panics, and never gates the memory pool. It exists so the +//! accounting gap can be seen in tracing output next to the pool reservations it should be +//! compared against. +//! +//! The balance counts `Layout` bytes, not resident pages. It excludes allocator fragmentation, +//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc +//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS as well, just a much +//! tighter one than pool reservations. +//! +//! [memory management contributor guide]: +//! https://datafusion.apache.org/comet/contributor-guide/memory_management.html + +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicIsize, Ordering}; + +/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this. +/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic +/// read-modify-write per 64 KiB of churn touches the shared cacheline. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before +/// another flushes the matching positive one. +static BALANCE: AtomicIsize = AtomicIsize::new(0); + +thread_local! { + /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles + /// directly instead of recursing. The only such allocation today is the one some platforms + /// make when registering `LOCAL_DRIFT`'s destructor on first touch. + /// + /// Const-initialized and destructor-free, so reading it never allocates and never fails — + /// which is what makes it safe to consult before touching `LOCAL_DRIFT`. + static IN_TRACK: Cell<bool> = const { Cell::new(false) }; + + /// This thread's un-flushed delta. + static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) }; +} + +/// Owns a thread's un-flushed delta and settles the remainder when the thread exits. +/// +/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting would be silently +/// discarded every time a thread died. Worker threads live for the process lifetime, but the +/// blocking pool churns on tokio's idle timeout, so on a long-lived executor that would be a +/// slowly accumulating bias in the reported balance. +struct ThreadDrift(Cell<isize>); + +impl Drop for ThreadDrift { + fn drop(&mut self) { + let drift = self.0.replace(0); + if drift != 0 { + BALANCE.fetch_add(drift, Ordering::Relaxed); + } + } +} + +/// Bytes currently handed out by the Rust global allocator, process-wide. +/// +/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the +/// balance can dip below zero transiently while per-thread deltas settle out of order. +pub fn current_balance() -> usize { + clamp_balance(BALANCE.load(Ordering::Relaxed)) +} + +/// Clamps a signed balance to the unsigned value reported to callers. +fn clamp_balance(balance: isize) -> usize { + balance.max(0) as usize +} + +/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches +/// [`SETTLE_THRESHOLD`]. +fn settle(local_drift: &Cell<isize>, delta: isize) { + let drift = local_drift.get().wrapping_add(delta); + if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize { + local_drift.set(0); + BALANCE.fetch_add(drift, Ordering::Relaxed); + } else { + local_drift.set(drift); + } +} + +/// Records a signed byte delta against the process balance. +#[inline] +fn track(delta: isize) { + if delta == 0 { + return; + } + + // A re-entrant call is one made by `track` itself; the outer frame owns the flag and will + // clear it, so this frame must only settle and return. + if IN_TRACK.with(|in_track| in_track.replace(true)) { + BALANCE.fetch_add(delta, Ordering::Relaxed); + return; + } + + // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s destructor has already + // run, and any allocation after that point must not panic inside the allocator. + if LOCAL_DRIFT + .try_with(|thread_drift| settle(&thread_drift.0, delta)) + .is_err() + { + BALANCE.fetch_add(delta, Ordering::Relaxed); + } + + IN_TRACK.with(|in_track| in_track.set(false)); +} + +/// Wraps a global allocator, accounting the `Layout` bytes it hands out. +/// +/// Adapted from the `AccountingAllocator` in +/// [apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626), which lives in +/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on directly. +pub struct AccountingAllocator<A: GlobalAlloc> { + inner: A, +} + +impl<A: GlobalAlloc> AccountingAllocator<A> { + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc` contract. The +// accounting is pure bookkeeping over an `AtomicIsize` and thread-local `Cell`s: it does not +// inspect, retain, or alter any pointer, and it cannot unwind. +unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc(layout); + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc_zeroed(layout); + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // Settle before delegating. A free cannot fail, so there is nothing to wait for, and the + // inner free can be slow: jemalloc returns oversize blocks to the OS eagerly, and unmapping + // a few hundred megabytes takes milliseconds. Accounting afterwards would keep the block on + // the balance for that whole window, after the allocator's own statistics had already + // dropped it. + track(-(layout.size() as isize)); + self.inner.dealloc(ptr, layout); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { Review Comment: Added `realloc_accounts_the_size_difference_after_delegating` in 4b71ad9. It grows a 64 MiB block to 96 MiB and shrinks it to 32 MiB through the `Recording` harness, asserting the balance moves by +32 MiB and then -64 MiB (within a 16 MiB margin, so accounting the full new size would fail), and that the inner `realloc` sees the balance still carrying the old size, which pins the after-delegating order. The harness now records the balance at `realloc` as well as at `dealloc`. ########## native/core/src/lib.rs: ########## @@ -65,27 +52,92 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; +pub mod alloc_accounting; Review Comment: Intentional. The module is ungated so the unit tests, including the dealloc-ordering test that drives a local `AccountingAllocator` against the shared balance, run in the default build that CI actually executes. `current_balance` documents that it returns 0 when the wrapper is not installed. Now that CI also runs the tests with the feature on, gating would be possible, but it would take the ordering tests out of the default build for no gain. ########## native/core/src/lib.rs: ########## @@ -65,27 +52,92 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; +pub mod alloc_accounting; pub mod cloud; pub mod execution; pub mod parquet; // this module is for non release only. Intended for debugging/profiling purposes #[cfg(debug_assertions)] pub mod debug; +// Global allocator selection. +// +// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise +// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature +// combination, so exactly one definition exists, and each backend predicate is written once. The +// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without +// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work, +// and no explicit allocator at all when the selection is the system allocator. +// +// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to +// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop +// silently: a feature combination with no backend would fail to compile rather than run with the +// metric enabled and reading zero. + +/// jemalloc, on targets where it builds, unless mimalloc was also requested. #[cfg(all( not(target_env = "msvc"), feature = "jemalloc", not(feature = "mimalloc") ))] -#[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; +mod backend { + pub type Backend = tikv_jemallocator::Jemalloc; + pub const BACKEND: Backend = tikv_jemallocator::Jemalloc; + pub const NAME: &str = "jemalloc"; + #[cfg(not(feature = "alloc-accounting"))] + #[global_allocator] + static GLOBAL: Backend = BACKEND; +} + +/// mimalloc, unless a usable jemalloc was also requested. #[cfg(all( feature = "mimalloc", not(all(not(target_env = "msvc"), feature = "jemalloc")) ))] +mod backend { + pub type Backend = mimalloc::MiMalloc; + pub const BACKEND: Backend = mimalloc::MiMalloc; + pub const NAME: &str = "mimalloc"; + + #[cfg(not(feature = "alloc-accounting"))] + #[global_allocator] + static GLOBAL: Backend = BACKEND; +} + +/// The system allocator: the complement of the two cases above. This covers neither feature, a +/// jemalloc request on MSVC, and both features together, which each backend cfg excludes in favour +/// of the other. +#[cfg(not(any( + all( + not(target_env = "msvc"), + feature = "jemalloc", + not(feature = "mimalloc") + ), + all( + feature = "mimalloc", + not(all(not(target_env = "msvc"), feature = "jemalloc")) + ) +)))] +// Without `alloc-accounting` nothing refers to this selection: the system allocator is the +// default, so no `#[global_allocator]` is installed. +#[cfg_attr(not(feature = "alloc-accounting"), allow(dead_code))] +mod backend { + pub type Backend = std::alloc::System; + pub const BACKEND: Backend = std::alloc::System; + pub const NAME: &str = "system"; +} + +/// The name of the allocator backend this build selected: `"jemalloc"`, `"mimalloc"` or +/// `"system"`. This is the one place the selection is decided, so anything that needs to know +/// which allocator is in effect (the `alloc_overhead` benchmark's liveness check, for instance) +/// reads it from here rather than re-deriving it from the feature set. +pub use backend::NAME as ALLOCATOR_BACKEND; Review Comment: Added in 4b71ad9. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
