diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 88a291f421b..e75b4c5553c 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -33,6 +33,17 @@ 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. 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 +cd native && cargo build --release --no-default-features --features hdfs-opendal,jemalloc +``` + Example output: ```json @@ -111,5 +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 (`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/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/Cargo.toml b/native/core/Cargo.toml index f0c7735a503..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"] @@ -114,6 +114,13 @@ 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, 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 [package.metadata.cargo-machete] ignored = ["hdfs-sys", "paste"] @@ -123,6 +130,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..1254282b9b4 --- /dev/null +++ b/native/core/benches/alloc_overhead.rs @@ -0,0 +1,203 @@ +// 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. +//! +//! 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 --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 +//! binary through the `rlib`. That only happens if the crate is actually linked, and an `--extern` +//! crate that nothing names is dropped from the crate graph along with its allocator, so the +//! `extern crate` below is load-bearing: without it a baseline run that never touches `comet` +//! silently measures the system allocator instead of jemalloc. The two liveness checks fail the run +//! if either the selected backend or the wrapper is somehow not in effect, because a number +//! measured against the wrong allocator would be worse than no number. +//! +//! `small_churn` is the worst case for the thread-local path: allocations so small that the +//! wrapper's bookkeeping is a meaningful fraction of the allocator's own work. `threshold_churn` is +//! the worst case for the shared counter: an alloc/free loop at exactly the 64 KiB settle threshold +//! flushes to the process-wide atomic on every call, and the parallel variant does that from every +//! core at once, so the gap between the single-threaded and parallel numbers is the cost of +//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet actually does, where +//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at or below +//! `arrow_sized_churn`, because they do actual work between allocations. + +// Pulls `comet`, and with it the `#[global_allocator]` selected by its feature set, into this +// binary even when the feature set leaves nothing here that names the crate. +extern crate comet; + +use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput}; +use std::hint::black_box; +use std::thread; +use std::time::Instant; + +/// Guards against measuring the wrong allocator. jemalloc keeps its own count of bytes it has +/// served; if it is not the global allocator of this binary that count stays at zero, and a +/// "jemalloc" baseline would in fact be the system allocator. +#[cfg(feature = "jemalloc")] +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")] +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_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 = 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(); +} + +/// Alloc/free of an untouched block, so the allocator call itself is most of the work and the +/// wrapper's share is largest. +fn alloc_free(size: usize) { + let v: Vec = Vec::with_capacity(black_box(size)); + black_box(&v); +} + +/// Alloc/free loops either side of the 64 KiB settle threshold, single-threaded and from every +/// core at once. +/// +/// A loop of one size never accumulates drift, so the threshold decides everything: at 32 KiB the +/// alloc and the free cancel inside the thread-local cell and the shared counter is never touched, +/// while at 64 KiB every alloc and every free flushes. The 64 KiB parallel case is therefore the +/// upper bound on shared-counter contention: `available_parallelism()` threads each doing two +/// atomic read-modify-writes per iteration on the same cacheline, with nothing else in between. +/// +/// Times are reported per alloc/free pair per thread, so a parallel number equal to its +/// single-threaded counterpart means the threads did not slow each other down at all. +fn threshold_churn(c: &mut Criterion) { + assert_backend_is_live(); + assert_accounting_is_live(); + let threads = thread::available_parallelism().map_or(4, |n| n.get()); + let mut group = c.benchmark_group("alloc_overhead"); + group.throughput(Throughput::Elements(1)); + for size in [32 * 1024usize, 64 * 1024] { + let kb = size / 1024; + group.bench_function(format!("alloc_free_{kb}kb"), |b| { + b.iter(|| alloc_free(size)); + }); + group.bench_function(format!("parallel_alloc_free_{kb}kb_x{threads}"), |b| { + b.iter_custom(|iters| { + let start = Instant::now(); + thread::scope(|scope| { + for _ in 0..threads { + scope.spawn(move || { + for _ in 0..iters { + alloc_free(size); + } + }); + } + }); + start.elapsed() + }); + }); + } + group.finish(); +} + +criterion_group!( + benches, + small_churn, + threshold_churn, + arrow_sized_churn, + growth_churn +); +criterion_main!(benches); diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs new file mode 100644 index 00000000000..636d2a2a745 --- /dev/null +++ b/native/core/src/alloc_accounting.rs @@ -0,0 +1,363 @@ +// 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) { + // 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)] +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 + /// 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(()); + + 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() { + let drift = Cell::new(0); + settle(&drift, 1024); + // A flush would have reset the drift to zero, so this alone shows the shared balance was + // not touched. Reading `BALANCE` here would race with every other test's allocations. + assert_eq!(drift.get(), 1024, "small delta stays thread-local"); + } + + #[test] + fn settle_flushes_at_the_threshold() { + let drift = Cell::new(0); + settle(&drift, SETTLE_THRESHOLD); + assert_eq!(drift.get(), 0, "drift resets once flushed"); + } + + #[test] + fn settle_flushes_negative_drift() { + let drift = Cell::new(0); + settle(&drift, -SETTLE_THRESHOLD); + assert_eq!(drift.get(), 0); + } + + #[test] + fn a_transiently_negative_balance_reports_as_zero() { + assert_eq!(clamp_balance(-1), 0); + assert_eq!(clamp_balance(isize::MIN), 0); + assert_eq!(clamp_balance(0), 0); + assert_eq!(clamp_balance(4096), 4096); + } + + /// A real allocation must move the reported balance: this is the one test that checks the + /// wrapper is actually installed as the global allocator for the current feature set, rather + /// than exercising it through a local instance. + /// + /// The block is zeroed and never touched, so it costs address space rather than resident + /// memory, and it is large enough that nothing else in the crate can free half of it inside the + /// microseconds between the two reads. + #[test] + #[cfg(feature = "alloc-accounting")] + fn a_real_allocation_raises_the_balance() { + use std::hint::black_box; + + const SIZE: usize = 256 * 1024 * 1024; + let _guard = serial(); + let before = current_balance(); + // `black_box` keeps the allocation observable so it cannot be elided. + let held: Vec = black_box(vec![0u8; SIZE]); + let during = current_balance(); + black_box(&held); + assert!( + during >= before + SIZE / 2, + "a {SIZE} byte allocation should raise the balance (before={before}, during={during}); \ + is the accounting wrapper installed for this feature set?" + ); + drop(held); + } + + /// The balance must drop before the inner allocator is asked to free the block. + /// + /// jemalloc decrements its own `stats.allocated` at the start of a large free and then, for + /// blocks above its oversize threshold, unmaps the pages eagerly, which takes milliseconds for + /// a block of a few hundred megabytes. If the subtraction happened after delegating, the balance + /// would keep reporting a block the allocator had already given back for that whole window, + /// and `native_allocated` would read above `jemalloc_allocated`. + #[test] + fn dealloc_settles_before_delegating() { + use std::alloc::System; + use std::sync::atomic::AtomicUsize; + + /// Records the reported balance at the moment the inner free is called. + struct Recording { + balance_at_dealloc: AtomicUsize, + } + + unsafe impl GlobalAlloc for Recording { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + System.alloc(layout) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.balance_at_dealloc + .store(current_balance(), Ordering::Relaxed); + System.dealloc(ptr, layout) + } + } + + // Well above the settle threshold, so both the allocation and the free flush immediately. + const SIZE: usize = 64 * 1024 * 1024; + let _guard = serial(); + let allocator = AccountingAllocator::new(Recording { + balance_at_dealloc: AtomicUsize::new(usize::MAX), + }); + let layout = Layout::from_size_align(SIZE, 8).unwrap(); + + // SAFETY: the layout is valid and non-zero, and the block is freed below through the same + // allocator that produced it. + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + let after_alloc = current_balance(); + unsafe { allocator.dealloc(ptr, layout) }; + + let seen = allocator.inner.balance_at_dealloc.load(Ordering::Relaxed); + // Half the block is a wide margin against parallel test noise while still being far + // outside anything the mutation (subtracting after delegating) could produce. + assert!( + seen + SIZE / 2 <= after_alloc, + "inner dealloc saw balance {seen}, expected at most {} (balance after alloc was \ + {after_alloc})", + after_alloc - SIZE / 2 + ); + } + + /// Threads must settle their remaining drift on exit. + /// + /// The worker writes a drift straight into its `LOCAL_DRIFT` cell and exits. Without the + /// wrapper installed nothing else ever calls `track`, so the only path by which that value can + /// reach the shared balance is `ThreadDrift::drop`; that is the build CI runs, and the one in + /// which a missing destructor is caught. The value is far larger than any real allocation, + /// which makes the check immune to whatever the rest of the crate is allocating meanwhile. + /// The injected amount is taken back out afterwards so later tests see an unchanged balance. + #[test] + fn thread_exit_settles_remaining_drift() { + 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/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/native/core/src/lib.rs b/native/core/src/lib.rs index b5656dba102..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; @@ -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,75 @@ pub mod parquet; #[cfg(debug_assertions)] pub mod debug; +// Global allocator selection. +// +// `backend` names the allocator the feature set asks for: jemalloc where it builds, otherwise +// mimalloc, otherwise the system allocator. The three `backend` cfgs partition every feature +// combination, so exactly one definition exists, and each backend predicate is written once. The +// unwrapped `#[global_allocator]` lives inside the backend module that owns it, so a build without +// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper, no per-allocation work, +// and no explicit allocator at all when the selection is the system allocator. +// +// With `alloc-accounting`, the single wrapped `#[global_allocator]` below refers to +// `backend::Backend` whatever it resolved to. That is what makes the wrapper impossible to drop +// silently: a feature combination with no backend would fail to compile rather than run with the +// metric enabled and reading zero. + +/// 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; + + #[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; + + #[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; +} + +#[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( 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(