feat: add greedy_unified_checked memory pool gated on real native usage [WIP]

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.                                      |

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
 

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

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<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 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<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) {
+    let mut group = c.benchmark_group("alloc_overhead");
+    group.bench_function("grow_vec_to_64kb", |b| {
+        b.iter(|| {
+            let mut v: Vec<u8> = 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<u8> = 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);