You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/16 04:32:23 UTC
(datafusion-comet) 01/01: feat: add native allocation accounting for memory observability (#5934)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5934-41d6d4483890696937fae753f8e575d703d25f5b
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 6e243c21e93e9cbed3c2013710182772f92dd480
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Sep 15 23:53:54 2026 +0000
feat: add native allocation accounting for memory observability (#5934)
* feat: add native allocation accounting for memory observability
Comet's memory pool counts declared reservations. Plenty of real allocation
never goes through it -- Arrow builders, expression kernels, decompression
buffers, Parquet metadata, object_store buffers, tokio itself -- so pool
reservations are a lower bound on Comet's footprint and the size of the gap is
currently unmeasurable at runtime. Diagnosing an OOM means guessing at it, and
spark.comet.exec.memoryPool.fraction asks operators to hand-tune a haircut for
a quantity nobody can see.
AccountingAllocator wraps the selected global allocator (jemalloc, mimalloc, or
system) and maintains one signed process-wide byte balance. executePlan reports
it as the native_allocated tracing metric, next to the per-thread pool
reservations it should be compared against.
This is observability only: it never rejects an allocation, never panics, and
does not touch the memory pool. Because it cannot fail an allocation, realloc
can account after delegating rather than before, which avoids over-counting a
failed realloc.
Per-thread deltas are batched and flushed into the shared balance at 64 KiB, so
the common path is a thread-local add-and-compare rather than an atomic RMW.
ThreadDrift's destructor settles the remainder when a thread exits, which
matters because the blocking pool churns on tokio's idle timeout and would
otherwise bias the balance on a long-lived executor. Touching that
destructor-bearing thread-local can itself allocate on first use, so track()
keeps a destructor-free re-entrancy flag and settles re-entrant calls straight
into the shared balance.
Off by default; a build without the feature has no wrapper and no
per-allocation work.
Verified clippy -D warnings and the native test suite across default,
alloc-accounting, jemalloc+alloc-accounting, and mimalloc+alloc-accounting.
The thread-exit test was mutation-checked: neutering the destructor fails it.
* bench: measure the alloc-accounting wrapper's per-allocation cost
Answers the question the feature has to answer before anyone proposes enabling
it by default. The liveness assertion is the point of the harness as much as the
timings are: without it a 'with the feature' run can silently be a second
baseline, which is exactly what happened on the first attempt here.
* fix: settle the balance before delegating a free in AccountingAllocator
`dealloc` accounted after calling the inner allocator, mirroring `alloc`
and `realloc`. A free cannot fail, so that ordering bought nothing, and it
opened a window: jemalloc decrements `stats.allocated` at the start of a
large free and then, for blocks above its 8 MiB oversize threshold,
unmaps the pages eagerly, which takes milliseconds for a block of a few
hundred megabytes. For that whole window the balance still carried a
block the allocator had already given back, so `native_allocated` read
above `jemalloc_allocated` by the size of the block in flight. On TPC-H
SF100 about 2% of trace samples showed the excess, up to 160 MB, during
task teardown in Q10, Q17 and Q18.
Subtract before delegating. A new test wraps a recording inner allocator
and asserts the balance has already dropped by the time the inner free is
called; it fails with "inner dealloc saw balance 67182807, expected at
most 33628375" under the previous ordering.
* fix: install the accounting wrapper for every allocator feature combination
Select the allocator backend once, in a `backend` module whose three cfgs
partition every feature combination, and let the single accounting
`#[global_allocator]` refer to whatever that resolved to. Previously each
backend predicate was repeated per arm and the accounting fallback to the
system allocator excluded `mimalloc`, so `jemalloc,mimalloc,alloc-accounting`
on a non-MSVC target matched no arm at all: the process ran on the unwrapped
default allocator while `native_allocated` stayed enabled and read zero.
A build without the feature is unchanged: the unwrapped allocator lives in
the backend module that owns it, and no explicit allocator is installed
when the selection is the system allocator.
Make the accounting tests immune to parallel test noise. `BALANCE` is
process-wide, so with the wrapper installed the crate's other tests move it
concurrently and the margin-based assertions failed intermittently. The
tests that observe the balance now serialize against each other, the
real-allocation probe uses an untouched 256 MiB block that nothing else in
the crate can mask, and the thread-exit test injects its drift directly so
it no longer needs the feature and runs in the default CI build.
Add `threshold_churn` to the `alloc_overhead` benchmark: alloc/free loops
at 32 KiB (never flushes) and 64 KiB (flushes on every call), single
threaded and from every core at once, so shared-counter contention can be
quantified rather than inferred.
* bench: link comet into alloc_overhead so the baseline uses the selected allocator
An `--extern` crate that nothing names is dropped from the crate graph, and
the accounting-off run of this benchmark named nothing in `comet`, so the
`#[global_allocator]` in `lib.rs` never reached the binary: the "jemalloc"
baseline was measuring glibc malloc. The accounting-on run names
`comet::alloc_accounting`, so it did link the crate, and every off/on
comparison so far was glibc against jemalloc plus the wrapper.
Add `extern crate comet` so the crate is always linked, and a jemalloc
liveness assertion next to the existing accounting one, so a run against
the wrong allocator fails instead of producing a plausible number.
* bench: check jemalloc liveness only when the library selected jemalloc
The alloc_overhead benchmark asserted on jemalloc's counters whenever the
jemalloc feature was on, but lib.rs selects jemalloc only when mimalloc
was not also requested and the target is not MSVC. With jemalloc,mimalloc
the library falls back to the system allocator and the bench aborted
before measuring anything.
Export the selection lib.rs made as comet::ALLOCATOR_BACKEND, one NAME
per backend module, and have the bench announce it and run the jemalloc
assertion only when the library reports jemalloc. The predicate stays
written once, so the guard cannot drift from the selection it verifies.
* fix: teach analyze_trace about native_allocated and cover the accounting feature in CI
Address the third review round on the alloc-accounting feature.
- analyze_trace only matched jemalloc_allocated, so a trace from a build with
alloc-accounting and no jemalloc reported zero allocated and no excess.
It now analyzes native_allocated when present, falls back to
jemalloc_allocated otherwise, names the counter in its output, and rejects
a trace with neither. tracing.md describes the selection.
- Nothing in CI compiled the feature. The rust-test action now lints the
jemalloc,alloc-accounting build with --all-targets, runs the accounting
tests with the wrapper installed, and checks the system-allocator arm.
- lib.rs: the three backend modules only name the type; the two
global_allocator statics sit together, so the default build now installs
System explicitly. ALLOCATOR_BACKEND is doc(hidden).
- alloc_accounting.rs: inline the clamp, table-drive the settle tests, add a
realloc test that pins the size-difference accounting and the
after-delegating order, and confine the thread-exit test to the default
build, where a missing destructor is actually caught. Trim the prose.
- alloc_overhead bench: drop the redundant extern crate, run both liveness
guards from a single Once that every bench function calls, merge the two
alloc/free loops, and use iter instead of iter_batched. Benchmark IDs are
unchanged.
* fix: compare against observed zero reservations in analyze_trace and format tracing.md
A zero pool reservation is a real sample that allocation can exceed, so the
analyzer now defers the comparison only until it has seen at least one pool
sample, and says so when it never does instead of reporting that allocation
never exceeded reservations. Reflow the tracing.md label table with prettier,
which the Preflight check enforces.
---
.github/actions/rust-test/action.yaml | 13 +
docs/source/contributor-guide/tracing.md | 65 +++--
native/common/src/bin/analyze_trace.rs | 99 +++++---
native/core/Cargo.toml | 10 +
native/core/benches/alloc_overhead.rs | 197 ++++++++++++++++
native/core/src/alloc_accounting.rs | 394 +++++++++++++++++++++++++++++++
native/core/src/execution/jni_api.rs | 14 ++
native/core/src/lib.rs | 67 ++++--
8 files changed, 791 insertions(+), 68 deletions(-)
diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml
index c39c2dcd4f..8934fbb4df 100644
--- a/.github/actions/rust-test/action.yaml
+++ b/.github/actions/rust-test/action.yaml
@@ -70,3 +70,16 @@ runs:
export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
RUST_BACKTRACE=1 cargo nextest run
+ # The `alloc-accounting` feature is off by default, so nothing else in CI compiles the
+ # allocator wrapper, its backend selection, or the benchmark's liveness guards. Lint them with
+ # the accounting wrapper over jemalloc, run the accounting tests with the wrapper installed, and
+ # check the system-allocator arm.
+ - name: Check and test the alloc-accounting feature
+ shell: bash
+ run: |
+ cd native
+ export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH}
+ cargo clippy --color=never -p datafusion-comet --all-targets --features jemalloc,alloc-accounting -- -D warnings
+ RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features jemalloc,alloc-accounting alloc_accounting
+ cargo check -p datafusion-comet --features alloc-accounting
+
diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md
index 88a291f421..9757bd8a57 100644
--- a/docs/source/contributor-guide/tracing.md
+++ b/docs/source/contributor-guide/tracing.md
@@ -33,6 +33,16 @@ Additionally, enabling the `jemalloc` feature will enable tracing of native memo
make release COMET_FEATURES="jemalloc"
```
+The `alloc-accounting` feature adds a second, allocator-independent measure of native memory. It
+wraps whichever global allocator the build selected and reports the bytes it has handed out as
+`native_allocated`. Unlike `jemalloc_allocated` it does not require jemalloc, and it counts only
+what Rust code allocated, so it can be compared against the memory pool's reservations without the
+allocator's own caching in the way. The two features are independent and can be combined:
+
+```shell
+make release COMET_FEATURES="jemalloc,alloc-accounting"
+```
+
Example output:
```json
@@ -55,9 +65,9 @@ Example trace visualization:
## Analyzing Memory Usage
-The `analyze_trace` tool parses a trace log and compares jemalloc usage against the sum of per-thread
-Comet memory pool reservations. This is useful for detecting untracked native memory growth where jemalloc
-allocations exceed what the memory pools account for.
+The `analyze_trace` tool parses a trace log and compares the process-wide native allocation counter against
+the sum of per-thread Comet memory pool reservations. This is useful for detecting untracked native memory
+growth where native allocations exceed what the memory pools account for.
Build and run:
@@ -66,10 +76,12 @@ cd native
cargo run --bin analyze_trace -- /path/to/comet-event-trace.json
```
-The tool reads counter events from the trace log. Because tracing logs metrics per thread, `jemalloc_allocated`
-is a process-wide value (the same global allocation reported from whichever thread logs it), while
-`thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to get the total
-tracked memory.
+The tool reads counter events from the trace log. Because tracing logs metrics per thread, `native_allocated`
+and `jemalloc_allocated` are process-wide values (the same global allocation reported from whichever thread
+logs it), while `thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to
+get the total tracked memory. The tool analyzes `native_allocated` when the trace contains it, since that
+counts only what Rust code holds from the allocator, and otherwise falls back to `jemalloc_allocated`. The
+output names the counter it used. A trace with neither counter is rejected.
Sample output:
@@ -77,20 +89,21 @@ Sample output:
=== Comet Trace Memory Analysis ===
Counter events parsed: 193104
+Allocation counter: jemalloc_allocated
Threads with memory pools: 8
-Peak jemalloc allocated: 3068.2 MB
+Peak jemalloc_allocated: 3068.2 MB
Peak pool total: 2864.6 MB
-Peak excess (jemalloc - pool): 364.6 MB
+Peak excess (jemalloc_allocated - pool): 364.6 MB
-WARNING: jemalloc exceeded pool reservation at 138 sampled points:
+WARNING: jemalloc_allocated exceeded pool reservation at 138 sampled points:
- Time (us) jemalloc pool_total excess
---------------------------------------------------------------
- 179578 210.8 MB 0.1 MB 210.7 MB
- 429663 420.5 MB 145.1 MB 275.5 MB
- 1304969 2122.5 MB 1797.2 MB 325.2 MB
- 21974838 407.0 MB 42.3 MB 364.6 MB
- 33543599 5.5 MB 0.1 MB 5.3 MB
+ Time (us) jemalloc_allocated pool_total excess
+------------------------------------------------------------------
+ 179578 210.8 MB 0.1 MB 210.7 MB
+ 429663 420.5 MB 145.1 MB 275.5 MB
+ 1304969 2122.5 MB 1797.2 MB 325.2 MB
+ 21974838 407.0 MB 42.3 MB 364.6 MB
+ 33543599 5.5 MB 0.1 MB 5.3 MB
--- Final per-thread pool reservations ---
@@ -102,14 +115,16 @@ WARNING: jemalloc exceeded pool reservation at 138 sampled points:
Total: 0.0 MB
```
-Some excess is expected (jemalloc metadata, fragmentation, non-pool allocations like Arrow IPC buffers).
-Large or growing excess may indicate memory that is not being tracked by the pool.
+Some excess is expected (allocator metadata and fragmentation for `jemalloc_allocated`, and non-pool
+allocations like Arrow IPC buffers for either counter). Large or growing excess may indicate memory that is
+not being tracked by the pool.
## Definition of Labels
-| Label | Meaning |
-| -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ |
-| jvm_heap_used | JVM heap memory usage of live objects for the executor process |
-| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) |
-| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. |
-| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. |
+| Label | Meaning |
+| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| jvm_heap_used | JVM heap memory usage of live objects for the executor process |
+| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) |
+| native_allocated | Bytes handed out by the Rust global allocator, process-wide (requires `alloc-accounting` feature). Approximate to within 64 KiB of un-flushed delta per live thread. |
+| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. |
+| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. |
diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs
index 8df83f7cea..30f55a2ae3 100644
--- a/native/common/src/bin/analyze_trace.rs
+++ b/native/common/src/bin/analyze_trace.rs
@@ -16,8 +16,9 @@
// under the License.
//! Analyzes a Comet chrome trace event log (`comet-event-trace.json`) and
-//! compares jemalloc usage against the sum of per-thread Comet memory pool
-//! reservations. Reports any points where jemalloc exceeds the total pool size.
+//! compares the process-wide native allocation counter against the sum of
+//! per-thread Comet memory pool reservations. Reports any points where the
+//! allocated bytes exceed the total pool size.
//!
//! Usage:
//! cargo run --bin analyze_trace -- <path-to-comet-event-trace.json>
@@ -27,6 +28,15 @@ use std::collections::HashMap;
use std::io::{BufRead, BufReader};
use std::{env, fs::File};
+/// The process-wide allocation counters the tool understands, most preferred first.
+///
+/// `native_allocated` (the `alloc-accounting` feature) counts only the bytes Rust code holds from
+/// the global allocator, so it is the tighter comparison against pool reservations.
+/// `jemalloc_allocated` (the `jemalloc` feature) also includes jemalloc's own metadata. A trace
+/// that carries both is analyzed against `native_allocated` alone; a trace with neither cannot be
+/// analyzed.
+const ALLOCATED_COUNTERS: [&str; 2] = ["native_allocated", "jemalloc_allocated"];
+
/// A single Chrome trace event (only the fields we care about).
#[derive(Deserialize)]
struct TraceEvent {
@@ -42,7 +52,7 @@ struct TraceEvent {
/// Snapshot of memory state at a given timestamp.
struct MemorySnapshot {
ts: u64,
- jemalloc: u64,
+ allocated: u64,
pool_total: u64,
}
@@ -61,14 +71,16 @@ fn main() {
let file = File::open(&args[1]).expect("Failed to open trace file");
let reader = BufReader::new(file);
- // Latest jemalloc value (global, not per-thread)
- let mut latest_jemalloc: u64 = 0;
+ // Index into ALLOCATED_COUNTERS of the counter being analyzed, once one has been seen
+ let mut source: Option<usize> = None;
+ // Latest allocated value (global, not per-thread)
+ let mut latest_allocated: u64 = 0;
// Per-thread pool reservations: thread_NNN -> bytes
let mut pool_by_thread: HashMap<String, u64> = HashMap::new();
- // Points where jemalloc exceeded pool total
+ // Points where allocated exceeded pool total
let mut violations: Vec<MemorySnapshot> = Vec::new();
// Track peak values
- let mut peak_jemalloc: u64 = 0;
+ let mut peak_allocated: u64 = 0;
let mut peak_pool_total: u64 = 0;
let mut peak_excess: u64 = 0;
let mut counter_events: u64 = 0;
@@ -110,11 +122,28 @@ fn main() {
counter_events += 1;
- if event.name == "jemalloc_allocated" {
- if let Some(val) = event.args.get("jemalloc_allocated") {
- latest_jemalloc = val.as_u64().unwrap_or(0);
- if latest_jemalloc > peak_jemalloc {
- peak_jemalloc = latest_jemalloc;
+ if let Some(rank) = ALLOCATED_COUNTERS
+ .iter()
+ .position(|name| *name == event.name)
+ {
+ match source {
+ // A preferred counter is present in this trace; ignore the other one.
+ Some(current) if current < rank => continue,
+ Some(current) if current == rank => {}
+ // First sighting of a more preferred counter. Start over so the peaks and
+ // violations reported all come from a single source.
+ _ => {
+ source = Some(rank);
+ latest_allocated = 0;
+ peak_allocated = 0;
+ peak_excess = 0;
+ violations.clear();
+ }
+ }
+ if let Some(val) = event.args.get(&event.name) {
+ latest_allocated = val.as_u64().unwrap_or(0);
+ if latest_allocated > peak_allocated {
+ peak_allocated = latest_allocated;
}
}
} else if event.name.contains("comet_memory_reserved") {
@@ -129,14 +158,16 @@ fn main() {
continue;
}
- // After each jemalloc or pool update, check the current state
+ // After each allocated or pool update, check the current state. A comparison needs one
+ // sample of each side: an observed zero reservation is a real value that allocation can
+ // exceed, so only the absence of any pool sample defers the check.
let pool_total: u64 = pool_by_thread.values().sum();
if pool_total > peak_pool_total {
peak_pool_total = pool_total;
}
- if latest_jemalloc > 0 && pool_total > 0 && latest_jemalloc > pool_total {
- let excess = latest_jemalloc - pool_total;
+ if source.is_some() && !pool_by_thread.is_empty() && latest_allocated > pool_total {
+ let excess = latest_allocated - pool_total;
if excess > peak_excess {
peak_excess = excess;
}
@@ -147,46 +178,60 @@ fn main() {
{
violations.push(MemorySnapshot {
ts: event.ts,
- jemalloc: latest_jemalloc,
+ allocated: latest_allocated,
pool_total,
});
}
}
}
+ let Some(source) = source.map(|rank| ALLOCATED_COUNTERS[rank]) else {
+ eprintln!(
+ "No process-wide allocation counter found in the trace: expected one of {}. \
+ Build the native library with the `alloc-accounting` or `jemalloc` feature.",
+ ALLOCATED_COUNTERS.join(", ")
+ );
+ std::process::exit(1);
+ };
+
// Print summary
println!("=== Comet Trace Memory Analysis ===\n");
println!("Counter events parsed: {counter_events}");
+ println!("Allocation counter: {source}");
println!("Threads with memory pools: {}", pool_by_thread.len());
- println!("Peak jemalloc allocated: {}", format_bytes(peak_jemalloc));
+ println!("Peak {source}: {}", format_bytes(peak_allocated));
println!(
"Peak pool total: {}",
format_bytes(peak_pool_total)
);
println!(
- "Peak excess (jemalloc - pool): {}",
+ "Peak excess ({source} - pool): {}",
format_bytes(peak_excess)
);
println!();
- if violations.is_empty() {
- println!("OK: jemalloc never exceeded the total pool reservation.");
+ if pool_by_thread.is_empty() {
+ println!(
+ "No pool reservation samples in the trace, so there is nothing to compare against."
+ );
+ } else if violations.is_empty() {
+ println!("OK: {source} never exceeded the total pool reservation.");
} else {
println!(
- "WARNING: jemalloc exceeded pool reservation at {} sampled points:\n",
+ "WARNING: {source} exceeded pool reservation at {} sampled points:\n",
violations.len()
);
println!(
- "{:>14} {:>14} {:>14} {:>14}",
- "Time (us)", "jemalloc", "pool_total", "excess"
+ "{:>14} {:>18} {:>14} {:>14}",
+ "Time (us)", source, "pool_total", "excess"
);
- println!("{}", "-".repeat(62));
+ println!("{}", "-".repeat(66));
for snap in &violations {
- let excess = snap.jemalloc - snap.pool_total;
+ let excess = snap.allocated - snap.pool_total;
println!(
- "{:>14} {:>14} {:>14} {:>14}",
+ "{:>14} {:>18} {:>14} {:>14}",
snap.ts,
- format_bytes(snap.jemalloc),
+ format_bytes(snap.allocated),
format_bytes(snap.pool_total),
format_bytes(excess),
);
diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml
index f0c7735a50..501592945d 100644
--- a/native/core/Cargo.toml
+++ b/native/core/Cargo.toml
@@ -114,6 +114,12 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
# Default builds carry zero Delta surface.
contrib-delta = ["dep:comet-contrib-delta"]
+# Observability for real native memory usage. Wraps the global allocator to track the bytes it
+# hands out, and reports the total as the `native_allocated` tracing metric so it can be compared
+# against the memory pool's reservations. Never rejects an allocation. Off by default; a build
+# without it has no wrapper and no per-allocation work.
+alloc-accounting = []
+
# exclude optional packages from cargo machete verifications
[package.metadata.cargo-machete]
ignored = ["hdfs-sys", "paste"]
@@ -123,6 +129,10 @@ name = "comet"
# "rlib" is for benchmarking with criterion.
crate-type = ["cdylib", "rlib"]
+[[bench]]
+name = "alloc_overhead"
+harness = false
+
[[bench]]
name = "array_element_append"
harness = false
diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs
new file mode 100644
index 0000000000..dbba3ed424
--- /dev/null
+++ b/native/core/benches/alloc_overhead.rs
@@ -0,0 +1,197 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off
+//! ```
+//!
+//! `churn` allocates and frees untouched blocks, so the allocator call is most of the work and the
+//! wrapper's share is largest. Sizes below the 64 KiB settle threshold only ever touch the
+//! thread-local path; a loop of exactly 64 KiB blocks flushes to the shared atomic on every alloc
+//! and every free, and its parallel variant does that from every core at once, so the gap between
+//! the single-threaded and parallel 64 KiB numbers is the cost of contention on that cacheline.
+//! `arrow_sized_churn` and `growth_churn` are closer to what Comet does, where filling a
+//! batch-sized buffer or growing a builder dwarfs the bookkeeping.
+
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use std::hint::black_box;
+use std::sync::Once;
+use std::thread;
+use std::time::Instant;
+
+/// Mirrors `alloc_accounting::SETTLE_THRESHOLD`: a thread flushes to the shared balance once its
+/// un-flushed delta reaches this.
+const SETTLE_THRESHOLD: usize = 64 * 1024;
+
+/// Fails the run if the allocator being measured is not the one the feature set asked for, since
+/// a number measured against the wrong allocator would be worse than no number. Every benchmark
+/// function calls this, so a filtered run cannot skip it.
+fn assert_allocators_are_live() {
+ static ONCE: Once = Once::new();
+ ONCE.call_once(|| {
+ // Which backend is in effect is `lib.rs`'s decision (with `jemalloc,mimalloc` together it
+ // falls back to the system allocator), so ask it rather than re-deriving the answer from
+ // the feature set. Naming `comet::ALLOCATOR_BACKEND` is also what links the `comet` rlib,
+ // and with it the `#[global_allocator]` it installs, into this binary: an `--extern` crate
+ // that nothing names is dropped from the crate graph along with its allocator.
+ eprintln!(
+ "alloc_overhead: measuring the `{}` allocator backend",
+ comet::ALLOCATOR_BACKEND
+ );
+ if comet::ALLOCATOR_BACKEND == "jemalloc" {
+ assert_jemalloc_is_live();
+ }
+ assert_accounting_is_live();
+ });
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this
+/// binary that count stays at zero.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+ use tikv_jemalloc_ctl::{epoch, stats};
+ let held: Vec<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");
+}
+
+/// If the wrapper were not installed in this binary, every "with the feature" number would
+/// silently be a second baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+ let before = comet::alloc_accounting::current_balance();
+ // `black_box` is load-bearing: in release mode LLVM elides an allocation whose contents are
+ // never observed, and the check would then fail against a wrapper that is in fact working.
+ let held: Vec<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() {}
+
+/// Alloc/free of an untouched block.
+fn alloc_free(size: usize) {
+ let v: Vec<u8> = Vec::with_capacity(black_box(size));
+ black_box(&v);
+}
+
+/// Alloc/free loops from well below the settle threshold up to exactly on it, single-threaded, and
+/// either side of the threshold from every core at once.
+///
+/// A loop of one size never accumulates drift, so the threshold decides everything: below it the
+/// alloc and the free cancel inside the thread-local cell and the shared counter is never touched,
+/// while at exactly 64 KiB every alloc and every free flushes. The 64 KiB parallel case is
+/// therefore the upper bound on shared-counter contention: `available_parallelism()` threads each
+/// doing two atomic read-modify-writes per iteration on the same cacheline, with nothing else in
+/// between. Times are per alloc/free pair per thread, so a parallel number equal to its
+/// single-threaded counterpart means the threads did not slow each other down at all.
+fn churn(c: &mut Criterion) {
+ assert_allocators_are_live();
+ let threads = thread::available_parallelism().map_or(4, |n| n.get());
+ let mut group = c.benchmark_group("alloc_overhead");
+ group.throughput(Throughput::Elements(1));
+ for size in [16usize, 256, 4096, SETTLE_THRESHOLD / 2, SETTLE_THRESHOLD] {
+ // The two sizes either side of the threshold are the ones where the parallel variant
+ // says something.
+ let near_threshold = size >= SETTLE_THRESHOLD / 2;
+ let label = if near_threshold {
+ format!("{}kb", size / 1024)
+ } else {
+ format!("{size}b")
+ };
+ group.bench_function(format!("alloc_free_{label}"), |b| {
+ b.iter(|| alloc_free(size));
+ });
+ if near_threshold {
+ group.bench_function(format!("parallel_alloc_free_{label}_x{threads}"), |b| {
+ b.iter_custom(|iters| {
+ let start = Instant::now();
+ thread::scope(|scope| {
+ for _ in 0..threads {
+ scope.spawn(move || {
+ for _ in 0..iters {
+ alloc_free(size);
+ }
+ });
+ }
+ });
+ start.elapsed()
+ });
+ });
+ }
+ }
+ group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {
+ assert_allocators_are_live();
+ let mut group = c.benchmark_group("alloc_overhead");
+ group.throughput(Throughput::Bytes(64 * 1024));
+ group.bench_function("alloc_fill_free_64kb", |b| {
+ b.iter(|| {
+ let v: Vec<u8> = vec![1u8; black_box(64 * 1024)];
+ black_box(v.len())
+ });
+ });
+ group.finish();
+}
+
+/// Repeated growth, which is the `realloc` path: a builder doubling its buffer.
+fn growth_churn(c: &mut Criterion) {
+ assert_allocators_are_live();
+ let mut group = c.benchmark_group("alloc_overhead");
+ group.bench_function("grow_vec_to_64kb", |b| {
+ b.iter(|| {
+ let mut v: Vec<u8> = Vec::new();
+ for _ in 0..(64 * 1024) {
+ v.push(black_box(1u8));
+ }
+ black_box(v.len())
+ });
+ });
+ group.finish();
+}
+
+criterion_group!(benches, churn, arrow_sized_churn, growth_churn);
+criterion_main!(benches);
diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs
new file mode 100644
index 0000000000..394bbe774d
--- /dev/null
+++ b/native/core/src/alloc_accounting.rs
@@ -0,0 +1,394 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust global allocator.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains a single signed
+//! process-wide byte balance, which [`current_balance`] exposes so it can be compared against the
+//! memory pool's reservations in tracing output. This is observability only: it never rejects an
+//! allocation, never panics, and never gates the memory pool.
+//!
+//! The balance counts `Layout` bytes, not resident pages: it excludes allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc`. See the [memory management contributor guide].
+//!
+//! [memory management contributor guide]:
+//! https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by* `track` settles
+ /// directly instead of recursing. The only such allocation today is the one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates and never fails,
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<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 discarded every time a
+/// thread died, and tokio's blocking pool churns threads on its idle timeout.
+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.
+///
+/// The value is approximate. Each live thread holds up to [`SETTLE_THRESHOLD`] bytes of
+/// un-flushed delta in either direction, so the reported balance can lag the true one by up to
+/// that amount times the number of live threads.
+pub fn current_balance() -> usize {
+ BALANCE.load(Ordering::Relaxed).max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<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::alloc::System;
+ use std::sync::atomic::AtomicUsize;
+ use std::sync::{Mutex, MutexGuard};
+
+ /// `BALANCE` is process-wide and the crate's tests run in parallel, so a test that reads it
+ /// sees every other test's allocations. The tests that move it by tens of megabytes take this
+ /// lock so they cannot land inside each other's windows; the rest of the crate is kept out by
+ /// making each window microseconds wide and each expected move far larger than anything else
+ /// allocates in that time.
+ static SERIAL: Mutex<()> = Mutex::new(());
+
+ fn serial() -> MutexGuard<'static, ()> {
+ SERIAL
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ }
+
+ const MIB: usize = 1024 * 1024;
+
+ /// Slack allowed between an observed balance and the expected one, to absorb whatever the
+ /// rest of the crate allocates during a test's window. It is half the smallest move any test
+ /// below expects, so a wrongly ordered or wrongly sized update still lands outside it.
+ const MARGIN: usize = 16 * MIB;
+
+ fn about(actual: usize, expected: usize) -> bool {
+ actual.abs_diff(expected) <= MARGIN
+ }
+
+ /// An inner allocator that records the reported balance at the moment each inner call is
+ /// made, which pins down whether the wrapper accounts before or after delegating.
+ struct Recording {
+ balance_at_dealloc: AtomicUsize,
+ balance_at_realloc: AtomicUsize,
+ }
+
+ impl Recording {
+ fn new() -> Self {
+ Self {
+ balance_at_dealloc: AtomicUsize::new(usize::MAX),
+ balance_at_realloc: AtomicUsize::new(usize::MAX),
+ }
+ }
+ }
+
+ unsafe impl GlobalAlloc for Recording {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ System.alloc(layout)
+ }
+
+ unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+ self.balance_at_dealloc
+ .store(current_balance(), Ordering::Relaxed);
+ System.dealloc(ptr, layout)
+ }
+
+ unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+ self.balance_at_realloc
+ .store(current_balance(), Ordering::Relaxed);
+ System.realloc(ptr, layout, new_size)
+ }
+ }
+
+ #[test]
+ fn settle_flushes_only_at_the_threshold() {
+ for (delta, residue) in [
+ (1024, 1024),
+ (-1024, -1024),
+ (SETTLE_THRESHOLD - 1, SETTLE_THRESHOLD - 1),
+ (SETTLE_THRESHOLD, 0),
+ (-SETTLE_THRESHOLD, 0),
+ ] {
+ let drift = Cell::new(0);
+ settle(&drift, delta);
+ // A flush resets the drift to zero, so the residue alone says whether the shared
+ // balance was touched. Reading `BALANCE` here would race with every other test.
+ assert_eq!(drift.get(), residue, "delta {delta}");
+ }
+ }
+
+ /// A real allocation must move the reported balance. This is the one test that checks the
+ /// wrapper is actually installed as the global allocator for the current feature set. The
+ /// block is zeroed and never touched, so it costs address space rather than resident memory.
+ #[test]
+ #[cfg(feature = "alloc-accounting")]
+ fn a_real_allocation_raises_the_balance() {
+ use std::hint::black_box;
+
+ const SIZE: usize = 256 * MIB;
+ let _guard = serial();
+ let before = current_balance();
+ // `black_box` keeps the allocation observable so it cannot be elided.
+ let held: Vec<u8> = black_box(vec![0u8; SIZE]);
+ let during = current_balance();
+ black_box(&held);
+ assert!(
+ during >= before + SIZE / 2,
+ "a {SIZE} byte allocation should raise the balance (before={before}, during={during}); \
+ is the accounting wrapper installed for this feature set?"
+ );
+ drop(held);
+ }
+
+ /// The balance must drop before the inner allocator is asked to free the block, because
+ /// jemalloc drops its own count at the start of a large free and then spends milliseconds
+ /// unmapping the pages; see the comment on `dealloc`.
+ #[test]
+ fn dealloc_settles_before_delegating() {
+ // Well above the settle threshold, so both the allocation and the free flush immediately.
+ const SIZE: usize = 64 * MIB;
+ let _guard = serial();
+ let allocator = AccountingAllocator::new(Recording::new());
+ let layout = Layout::from_size_align(SIZE, 8).unwrap();
+
+ // SAFETY: the layout is valid and non-zero, and the block is freed below through the same
+ // allocator that produced it.
+ let ptr = unsafe { allocator.alloc(layout) };
+ assert!(!ptr.is_null());
+ let after_alloc = current_balance();
+ unsafe { allocator.dealloc(ptr, layout) };
+
+ let seen = allocator.inner.balance_at_dealloc.load(Ordering::Relaxed);
+ assert!(
+ about(seen + SIZE, after_alloc),
+ "inner dealloc saw balance {seen}, expected about {} (balance after alloc was \
+ {after_alloc})",
+ after_alloc.saturating_sub(SIZE)
+ );
+ }
+
+ /// `realloc` moves the balance by the size difference, not by the new size, and does so after
+ /// delegating: the inner allocator must see the balance still carrying the old size.
+ #[test]
+ fn realloc_accounts_the_size_difference_after_delegating() {
+ const OLD: usize = 64 * MIB;
+ const GROWN: usize = 96 * MIB;
+ const SHRUNK: usize = 32 * MIB;
+ let _guard = serial();
+ let allocator = AccountingAllocator::new(Recording::new());
+ let layout = Layout::from_size_align(OLD, 8).unwrap();
+
+ // SAFETY: each layout matches the block's current size, and the block is freed at the end
+ // through the same allocator that produced it.
+ let ptr = unsafe { allocator.alloc(layout) };
+ assert!(!ptr.is_null());
+ let before_grow = current_balance();
+
+ let ptr = unsafe { allocator.realloc(ptr, layout, GROWN) };
+ assert!(!ptr.is_null());
+ let after_grow = current_balance();
+ let seen = allocator.inner.balance_at_realloc.load(Ordering::Relaxed);
+ assert!(
+ about(seen, before_grow),
+ "inner realloc saw balance {seen}, expected about {before_grow}: the wrapper must \
+ account after delegating"
+ );
+ assert!(
+ about(after_grow, before_grow + (GROWN - OLD)),
+ "growing {OLD} -> {GROWN} moved the balance {before_grow} -> {after_grow}, expected \
+ about +{}",
+ GROWN - OLD
+ );
+
+ let layout = Layout::from_size_align(GROWN, 8).unwrap();
+ let ptr = unsafe { allocator.realloc(ptr, layout, SHRUNK) };
+ assert!(!ptr.is_null());
+ let after_shrink = current_balance();
+ assert!(
+ about(after_shrink + (GROWN - SHRUNK), after_grow),
+ "shrinking {GROWN} -> {SHRUNK} moved the balance {after_grow} -> {after_shrink}, \
+ expected about -{}",
+ GROWN - SHRUNK
+ );
+
+ unsafe { allocator.dealloc(ptr, Layout::from_size_align(SHRUNK, 8).unwrap()) };
+ }
+
+ /// Threads must settle their remaining drift on exit.
+ ///
+ /// The worker writes a drift straight into its `LOCAL_DRIFT` cell and exits, so the only path
+ /// by which that value can reach the shared balance is `ThreadDrift::drop`. That holds only
+ /// while the wrapper is not installed: with it, thread teardown's own allocations call `track`
+ /// and flush the oversized drift before the destructor runs, and the test would pass without
+ /// one. So the test is confined to the default build, which is the one CI runs. The injected
+ /// amount is far larger than any real allocation, and is taken back out afterwards.
+ #[test]
+ #[cfg(not(feature = "alloc-accounting"))]
+ fn thread_exit_settles_remaining_drift() {
+ use std::thread;
+
+ const INJECTED: isize = 1 << 40;
+ let _guard = serial();
+
+ let before = BALANCE.load(Ordering::Relaxed);
+ thread::spawn(|| {
+ LOCAL_DRIFT.with(|drift| drift.0.set(drift.0.get() + INJECTED));
+ })
+ .join()
+ .unwrap();
+ let moved = BALANCE.load(Ordering::Relaxed) - before;
+ BALANCE.fetch_sub(INJECTED, Ordering::Relaxed);
+
+ assert!(
+ moved >= INJECTED / 2,
+ "drift from an exited thread never reached the shared balance: \
+ balance moved {moved} bytes, expected at least {}",
+ INJECTED / 2
+ );
+ }
+}
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
index 692b0d1bcc..0107847c9e 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<u64, HashMap<i64, Arc<dyn MemoryPool>>>;
@@ -1089,6 +1101,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
if exec_context.tracing_enabled {
#[cfg(feature = "jemalloc")]
log_jemalloc_usage();
+ #[cfg(feature = "alloc-accounting")]
+ log_native_allocated();
log_memory_usage(
&exec_context.tracing_memory_metric_name,
total_reserved_for_thread(exec_context.rust_thread_id) as u64,
diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs
index b5656dba10..0872a478f9 100644
--- a/native/core/src/lib.rs
+++ b/native/core/src/lib.rs
@@ -41,19 +41,6 @@ use log4rs::{
Config,
};
-#[cfg(all(
- not(target_env = "msvc"),
- feature = "jemalloc",
- not(feature = "mimalloc")
-))]
-use tikv_jemallocator::Jemalloc;
-
-#[cfg(all(
- feature = "mimalloc",
- not(all(not(target_env = "msvc"), feature = "jemalloc"))
-))]
-use mimalloc::MiMalloc;
-
// Re-export from jvm-bridge crate for internal use
pub use datafusion_comet_jni_bridge::errors;
pub use datafusion_comet_jni_bridge::JAVA_VM;
@@ -65,6 +52,7 @@ pub mod jvm_bridge {
use errors::{try_unwrap_or_throw, CometError, CometResult};
+pub mod alloc_accounting;
pub mod cloud;
pub mod execution;
pub mod parquet;
@@ -72,20 +60,67 @@ pub mod parquet;
#[cfg(debug_assertions)]
pub mod debug;
+// Global allocator selection. `backend` names the allocator the feature set asks for: jemalloc
+// where it builds, otherwise mimalloc, otherwise the system allocator. The three cfgs partition
+// every feature combination, so exactly one `backend` exists and a combination matching none would
+// fail to compile rather than install nothing and leave the accounting metric reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
#[cfg(all(
not(target_env = "msvc"),
feature = "jemalloc",
not(feature = "mimalloc")
))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+ pub type Backend = tikv_jemallocator::Jemalloc;
+ pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+ pub const NAME: &str = "jemalloc";
+}
+/// mimalloc, unless a usable jemalloc was also requested.
#[cfg(all(
feature = "mimalloc",
not(all(not(target_env = "msvc"), feature = "jemalloc"))
))]
+mod backend {
+ pub type Backend = mimalloc::MiMalloc;
+ pub const BACKEND: Backend = mimalloc::MiMalloc;
+ pub const NAME: &str = "mimalloc";
+}
+
+/// The system allocator: the complement of the two cases above, which covers neither feature, a
+/// jemalloc request on MSVC, and both features together.
+#[cfg(not(any(
+ all(
+ not(target_env = "msvc"),
+ feature = "jemalloc",
+ not(feature = "mimalloc")
+ ),
+ all(
+ feature = "mimalloc",
+ not(all(not(target_env = "msvc"), feature = "jemalloc"))
+ )
+)))]
+mod backend {
+ pub type Backend = std::alloc::System;
+ pub const BACKEND: Backend = std::alloc::System;
+ pub const NAME: &str = "system";
+}
+
+/// The name of the allocator backend this build selected: `"jemalloc"`, `"mimalloc"` or
+/// `"system"`. The selection is decided here and nowhere else, so the `alloc_overhead` benchmark's
+/// liveness check reads it from here rather than re-deriving it from the feature set.
+#[doc(hidden)]
+pub use backend::NAME as ALLOCATOR_BACKEND;
+
+#[cfg(not(feature = "alloc-accounting"))]
+#[global_allocator]
+static GLOBAL: backend::Backend = backend::BACKEND;
+
+#[cfg(feature = "alloc-accounting")]
#[global_allocator]
-static GLOBAL: MiMalloc = MiMalloc;
+static GLOBAL: alloc_accounting::AccountingAllocator<backend::Backend> =
+ alloc_accounting::AccountingAllocator::new(backend::BACKEND);
#[no_mangle]
pub extern "system" fn Java_org_apache_comet_NativeBase_init(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org