comphead commented on code in PR #5934:
URL: https://github.com/apache/datafusion-comet/pull/5934#discussion_r4018853380
##########
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() {
Review Comment:
`analyze_trace` never sees this metric.
`native/common/src/bin/analyze_trace.rs:113` matches `jemalloc_allocated` by
exact name and drops everything else in the trailing `else { continue; }`, so
`native_allocated` is discarded.
That tool is what computes `excess = allocated - pool_total`, which is the
comparison this PR exists to enable. In the combination the PR specifically
motivates (`alloc-accounting` without `jemalloc`) it will report zero allocated
and no excess.
Suggest making that arm accept either name, deciding which wins when both
are logged, and updating the `tracing.md` prose at line 69 that currently names
only `jemalloc_allocated` as the process-wide counter.
##########
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 = []
Review Comment:
Nothing in CI builds this feature. `pr_benchmark_check.yml:55` runs `cargo
clippy --all-targets --workspace` with default features, and no workflow passes
`--features`. So the three backend arms, the wrapper, and the bench's jemalloc
liveness check are compiled only on developer machines.
The `mod backend` partition is nicely self-checking (zero matches gives an
unresolved `backend`, two gives a duplicate module), but only for combinations
someone actually compiles. Adding `cargo check --features alloc-accounting`
plus one `jemalloc,alloc-accounting` check to an existing job would keep that
property honest.
##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per
allocation.
Review Comment:
The PR description still says "Not yet measured: the per-allocation overhead
of the wrapper when the feature is on", but this benchmark is in the diff. 225
lines whose whole purpose is bounding that overhead, landing without the
number, leaves a reviewer unable to judge the cost of turning the feature on.
Either paste the `off` versus `alloc-accounting` comparison this header
describes, or land the bench separately once you have it.
##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per
allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting --
--baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs,
which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually
linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its
allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never
touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two
liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect,
because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so
small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own
work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the
64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant
does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers
is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet
actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at
or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its
feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
Review Comment:
This and the comment above it look unnecessary. `assert_backend_is_live()`
names `comet::ALLOCATOR_BACKEND` unconditionally (no cfg), which in edition
2021 already resolves through the extern prelude and pulls the rlib, and with
it the `#[global_allocator]`, into the crate graph. The comment's premise, that
nothing here names the crate, does not hold.
##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
use errors::{try_unwrap_or_throw, CometError, CometResult};
+pub mod alloc_accounting;
pub mod cloud;
pub mod execution;
pub mod parquet;
// this module is for non release only. Intended for debugging/profiling
purposes
#[cfg(debug_assertions)]
pub mod debug;
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it
builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs
partition every feature
+// combination, so exactly one definition exists, and each backend predicate
is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns
it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper,
no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below
refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper
impossible to drop
+// silently: a feature combination with no backend would fail to compile
rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
#[cfg(all(
not(target_env = "msvc"),
feature = "jemalloc",
not(feature = "mimalloc")
))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+ pub type Backend = tikv_jemallocator::Jemalloc;
+ pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+ pub const NAME: &str = "jemalloc";
+ #[cfg(not(feature = "alloc-accounting"))]
+ #[global_allocator]
Review Comment:
This `#[cfg(not(feature = "alloc-accounting"))] #[global_allocator]` is
repeated verbatim in the mimalloc module, and its absence from the system
module is what forces the `#[cfg_attr(..., allow(dead_code))]` there. Two
statics at a single site collapse all three:
```rust
#[cfg(not(feature = "alloc-accounting"))]
#[global_allocator]
static GLOBAL: backend::Backend = backend::BACKEND;
#[cfg(feature = "alloc-accounting")]
#[global_allocator]
static GLOBAL: alloc_accounting::AccountingAllocator<backend::Backend> =
alloc_accounting::AccountingAllocator::new(backend::BACKEND);
```
The one cost is that the default plus system build then installs
`std::alloc::System` explicitly rather than leaving the default in place. Same
allocator either way, but it does give up the "byte-for-byte the previous
arrangement" property the header comment claims, so your call whether that is
worth the deduplication.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust
global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool)
counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real
allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers,
Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a
lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and
currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full
picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains
a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is
**observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory
pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations
it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes
allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency
allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS
as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!
https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the
magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only
about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a
negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by*
`track` settles
+ /// directly instead of recursing. The only such allocation today is the
one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates
and never fails —
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+ /// This thread's un-flushed delta.
+ static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread
exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting
would be silently
+/// discarded every time a thread died. Worker threads live for the process
lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor
that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+ fn drop(&mut self) {
+ let drift = self.0.replace(0);
+ if drift != 0 {
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ }
+ }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never
reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out
of order.
+pub fn current_balance() -> usize {
+ clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
Review Comment:
`clamp_balance` exists only so
`a_transiently_negative_balance_reports_as_zero` can assert on `isize::max(0)
as usize`. Inlining it into `current_balance` and dropping that test loses no
coverage. The clamp's real behavior is already implied by `current_balance`'s
documented contract.
##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per
allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting --
--baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs,
which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually
linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its
allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never
touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two
liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect,
because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so
small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own
work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the
64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant
does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers
is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet
actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at
or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its
feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion,
Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being
measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's
feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the
system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library
which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree
with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+ static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+ ANNOUNCE.call_once(|| {
+ eprintln!(
+ "alloc_overhead: measuring the `{}` allocator backend",
+ comet::ALLOCATOR_BACKEND
+ )
+ });
+ if comet::ALLOCATOR_BACKEND == "jemalloc" {
+ assert_jemalloc_is_live();
+ }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the
global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact
be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+ use tikv_jemalloc_ctl::{epoch, stats};
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ epoch::advance().expect("jemalloc epoch");
+ let allocated = stats::allocated::read().expect("jemalloc
stats.allocated");
+ assert!(
+ allocated >= 8 * 1024 * 1024,
+ "the library selected jemalloc but jemalloc is not the global
allocator of this binary \
+ (stats.allocated = {allocated}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is
never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+ unreachable!("the library reports the jemalloc backend but the feature is
not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually
installed in the benchmark
+/// binary, every "with the feature" number would silently be a second
baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+ let before = comet::alloc_accounting::current_balance();
+ // `black_box` is load-bearing: benchmarks build in release mode, where
LLVM will happily
+ // elide an allocation whose contents are never observed, and the check
would then fail
+ // against a wrapper that is in fact working.
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ let during = comet::alloc_accounting::current_balance();
+ assert!(
+ during >= before + 4 * 1024 * 1024,
+ "alloc-accounting is enabled but the allocator is not installed in
this binary \
+ (balance {before} -> {during}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most
iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+ assert_backend_is_live();
+ assert_accounting_is_live();
+ let mut group = c.benchmark_group("alloc_overhead");
+ for size in [16usize, 256, 4096] {
+ group.throughput(Throughput::Elements(1));
+ group.bench_function(format!("alloc_free_{size}b"), |b| {
+ b.iter(|| {
+ let v: Vec<u8> = Vec::with_capacity(black_box(size));
+ black_box(&v);
+ });
+ });
+ }
+ group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is
the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {
+ let mut group = c.benchmark_group("alloc_overhead");
+ group.throughput(Throughput::Bytes(64 * 1024));
+ group.bench_function("alloc_fill_free_64kb", |b| {
+ b.iter_batched(
+ || (),
+ |()| {
+ let v: Vec<u8> = vec![1u8; black_box(64 * 1024)];
+ black_box(v.len())
+ },
+ BatchSize::SmallInput,
+ );
+ });
+ group.finish();
+}
+
+/// Repeated growth, which is the `realloc` path: a builder doubling its
buffer.
+fn growth_churn(c: &mut Criterion) {
Review Comment:
The liveness asserts only run from `small_churn` and `threshold_churn`, so
`cargo bench --bench alloc_overhead -- alloc_fill_free_64kb` (or
`grow_vec_to_64kb`) measures with no guard at all, which is exactly the
silent-wrong-allocator case the header says would be worse than no number.
Putting both asserts behind a single `Once` that every bench function calls
would close that.
##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
use errors::{try_unwrap_or_throw, CometError, CometResult};
+pub mod alloc_accounting;
Review Comment:
Question: is it intentional that this is ungated? It makes
`current_balance()` public API that always returns 0 in default builds, and a
locally constructed `AccountingAllocator` mutates the process-wide `BALANCE`
even when the wrapper is not installed, which
`dealloc_settles_before_delegating` relies on. I assume the point is keeping
the unit tests running in the default build, just want to confirm that is the
reason rather than an oversight.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust
global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool)
counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real
allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers,
Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a
lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and
currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full
picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains
a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is
**observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory
pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations
it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes
allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency
allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS
as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!
https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the
magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only
about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a
negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by*
`track` settles
+ /// directly instead of recursing. The only such allocation today is the
one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates
and never fails —
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+ /// This thread's un-flushed delta.
+ static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread
exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting
would be silently
+/// discarded every time a thread died. Worker threads live for the process
lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor
that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+ fn drop(&mut self) {
+ let drift = self.0.replace(0);
+ if drift != 0 {
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ }
+ }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never
reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out
of order.
+pub fn current_balance() -> usize {
+ clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+ balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the
magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+ let drift = local_drift.get().wrapping_add(delta);
+ if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+ local_drift.set(0);
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ } else {
+ local_drift.set(drift);
+ }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+ if delta == 0 {
+ return;
+ }
+
+ // A re-entrant call is one made by `track` itself; the outer frame owns
the flag and will
+ // clear it, so this frame must only settle and return.
+ if IN_TRACK.with(|in_track| in_track.replace(true)) {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ return;
+ }
+
+ // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s
destructor has already
+ // run, and any allocation after that point must not panic inside the
allocator.
+ if LOCAL_DRIFT
+ .try_with(|thread_drift| settle(&thread_drift.0, delta))
+ .is_err()
+ {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ }
+
+ IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+///
[apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626),
which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on
directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+ inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+ pub const fn new(inner: A) -> Self {
+ Self { inner }
+ }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc`
contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local
`Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc_zeroed(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+ // Settle before delegating. A free cannot fail, so there is nothing
to wait for, and the
+ // inner free can be slow: jemalloc returns oversize blocks to the OS
eagerly, and unmapping
+ // a few hundred megabytes takes milliseconds. Accounting afterwards
would keep the block on
+ // the balance for that whole window, after the allocator's own
statistics had already
+ // dropped it.
+ track(-(layout.size() as isize));
+ self.inner.dealloc(ptr, layout);
+ }
+
+ unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) ->
*mut u8 {
+ let new_ptr = self.inner.realloc(ptr, layout, new_size);
+ if !new_ptr.is_null() {
+ // Accounting after the fact is only safe because this allocator
cannot fail the
+ // allocation or unwind. A variant that enforced a limit would
have to decide *before*
+ // delegating: `realloc` may free or move the old block, and a
caller that never
+ // received the new pointer would free the stale one while
unwinding.
+ //
+ // A single allocation cannot exceed `isize::MAX` on any real
platform, so neither cast
+ // wraps.
+ track(new_size as isize - layout.size() as isize);
+ }
+ new_ptr
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::{Mutex, MutexGuard};
+
+ /// `BALANCE` is process-wide and the crate's tests run in parallel, so a
test that reads it
+ /// sees every other test's allocations. The tests that move it by tens of
megabytes take this
+ /// lock so they cannot land inside each other's windows; the rest of the
crate is kept out by
+ /// making each window microseconds wide and each expected move far larger
than anything else
+ /// allocates in that time.
+ static SERIAL: Mutex<()> = Mutex::new(());
+
+ fn serial() -> MutexGuard<'static, ()> {
+ SERIAL
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ }
+
+ #[test]
+ fn settle_accumulates_below_the_threshold() {
+ 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<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.
+ ///
+ /// 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() {
Review Comment:
This stops being mutation-proof once the feature is on. The worker leaves
`LOCAL_DRIFT` at `1 << 40`, and thread teardown itself allocates, so the next
`track` call sees `drift.unsigned_abs() >= SETTLE_THRESHOLD` and flushes
immediately. `BALANCE` then moves by roughly `INJECTED` even with
`ThreadDrift::drop` neutered.
The doc comment already says the default build is the one that catches it.
Worth gating the test `#[cfg(not(feature = "alloc-accounting"))]` so it cannot
quietly become a tautology in the build that actually ships the wrapper.
##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per
allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting --
--baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs,
which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually
linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its
allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never
touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two
liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect,
because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so
small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own
work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the
64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant
does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers
is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet
actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at
or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its
feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion,
Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being
measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's
feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the
system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library
which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree
with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+ static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+ ANNOUNCE.call_once(|| {
+ eprintln!(
+ "alloc_overhead: measuring the `{}` allocator backend",
+ comet::ALLOCATOR_BACKEND
+ )
+ });
+ if comet::ALLOCATOR_BACKEND == "jemalloc" {
+ assert_jemalloc_is_live();
+ }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the
global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact
be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+ use tikv_jemalloc_ctl::{epoch, stats};
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ epoch::advance().expect("jemalloc epoch");
+ let allocated = stats::allocated::read().expect("jemalloc
stats.allocated");
+ assert!(
+ allocated >= 8 * 1024 * 1024,
+ "the library selected jemalloc but jemalloc is not the global
allocator of this binary \
+ (stats.allocated = {allocated}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is
never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+ unreachable!("the library reports the jemalloc backend but the feature is
not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually
installed in the benchmark
+/// binary, every "with the feature" number would silently be a second
baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+ let before = comet::alloc_accounting::current_balance();
+ // `black_box` is load-bearing: benchmarks build in release mode, where
LLVM will happily
+ // elide an allocation whose contents are never observed, and the check
would then fail
+ // against a wrapper that is in fact working.
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ let during = comet::alloc_accounting::current_balance();
+ assert!(
+ during >= before + 4 * 1024 * 1024,
+ "alloc-accounting is enabled but the allocator is not installed in
this binary \
+ (balance {before} -> {during}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most
iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
+ assert_backend_is_live();
+ assert_accounting_is_live();
+ let mut group = c.benchmark_group("alloc_overhead");
+ for size in [16usize, 256, 4096] {
+ group.throughput(Throughput::Elements(1));
+ group.bench_function(format!("alloc_free_{size}b"), |b| {
+ b.iter(|| {
+ let v: Vec<u8> = Vec::with_capacity(black_box(size));
+ black_box(&v);
+ });
+ });
+ }
+ group.finish();
+}
+
+/// A batch-sized buffer, filled so the pages are actually touched. This is
the shape of allocation
+/// Comet does in bulk.
+fn arrow_sized_churn(c: &mut Criterion) {
Review Comment:
`iter_batched(|| (), |()| ..., BatchSize::SmallInput)` with a unit setup is
`b.iter(...)` with extra machinery. The batching only earns its keep when setup
allocates something you do not want in the measurement.
##########
native/core/benches/alloc_overhead.rs:
##########
@@ -0,0 +1,225 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per
allocation.
+//!
+//! Run the same benchmark with and without the feature and compare:
+//!
+//! ```shell
+//! cargo bench --bench alloc_overhead -- --save-baseline off
+//! cargo bench --bench alloc_overhead --features alloc-accounting --
--baseline off
+//! ```
+//!
+//! The benchmark relies on the `#[global_allocator]` that `lib.rs` installs,
which reaches this
+//! binary through the `rlib`. That only happens if the crate is actually
linked, and an `--extern`
+//! crate that nothing names is dropped from the crate graph along with its
allocator, so the
+//! `extern crate` below is load-bearing: without it a baseline run that never
touches `comet`
+//! silently measures the system allocator instead of jemalloc. The two
liveness checks fail the run
+//! if either the selected backend or the wrapper is somehow not in effect,
because a number
+//! measured against the wrong allocator would be worse than no number.
+//!
+//! `small_churn` is the worst case for the thread-local path: allocations so
small that the
+//! wrapper's bookkeeping is a meaningful fraction of the allocator's own
work. `threshold_churn` is
+//! the worst case for the shared counter: an alloc/free loop at exactly the
64 KiB settle threshold
+//! flushes to the process-wide atomic on every call, and the parallel variant
does that from every
+//! core at once, so the gap between the single-threaded and parallel numbers
is the cost of
+//! contention on that cacheline. `arrow_sized_churn` is closer to what Comet
actually does, where
+//! a batch-sized buffer dwarfs the bookkeeping. Real query workloads sit at
or below
+//! `arrow_sized_churn`, because they do actual work between allocations.
+
+// Pulls `comet`, and with it the `#[global_allocator]` selected by its
feature set, into this
+// binary even when the feature set leaves nothing here that names the crate.
+extern crate comet;
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion,
Throughput};
+use std::hint::black_box;
+use std::thread;
+use std::time::Instant;
+
+/// Guards against measuring the wrong allocator, and says which one is being
measured.
+///
+/// Which backend is in effect is `lib.rs`'s decision, not this crate's
feature flags': with
+/// `jemalloc,mimalloc` together the library deliberately falls back to the
system allocator, and
+/// jemalloc on MSVC is not selected at all. So the check asks the library
which backend it chose
+/// rather than re-deriving that from the feature set, and can never disagree
with the selection it
+/// is meant to verify.
+fn assert_backend_is_live() {
+ static ANNOUNCE: std::sync::Once = std::sync::Once::new();
+ ANNOUNCE.call_once(|| {
+ eprintln!(
+ "alloc_overhead: measuring the `{}` allocator backend",
+ comet::ALLOCATOR_BACKEND
+ )
+ });
+ if comet::ALLOCATOR_BACKEND == "jemalloc" {
+ assert_jemalloc_is_live();
+ }
+}
+
+/// jemalloc keeps its own count of bytes it has served; if it is not the
global allocator of this
+/// binary that count stays at zero, and a "jemalloc" baseline would in fact
be the system
+/// allocator.
+#[cfg(feature = "jemalloc")]
+fn assert_jemalloc_is_live() {
+ use tikv_jemalloc_ctl::{epoch, stats};
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ epoch::advance().expect("jemalloc epoch");
+ let allocated = stats::allocated::read().expect("jemalloc
stats.allocated");
+ assert!(
+ allocated >= 8 * 1024 * 1024,
+ "the library selected jemalloc but jemalloc is not the global
allocator of this binary \
+ (stats.allocated = {allocated}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+/// Without the feature the library cannot have selected jemalloc, so this is
never reached.
+#[cfg(not(feature = "jemalloc"))]
+fn assert_jemalloc_is_live() {
+ unreachable!("the library reports the jemalloc backend but the feature is
not enabled");
+}
+
+/// Guards against measuring nothing. If the wrapper were not actually
installed in the benchmark
+/// binary, every "with the feature" number would silently be a second
baseline run.
+#[cfg(feature = "alloc-accounting")]
+fn assert_accounting_is_live() {
+ let before = comet::alloc_accounting::current_balance();
+ // `black_box` is load-bearing: benchmarks build in release mode, where
LLVM will happily
+ // elide an allocation whose contents are never observed, and the check
would then fail
+ // against a wrapper that is in fact working.
+ let held: Vec<u8> = black_box(vec![1u8; 8 * 1024 * 1024]);
+ black_box(&held);
+ let during = comet::alloc_accounting::current_balance();
+ assert!(
+ during >= before + 4 * 1024 * 1024,
+ "alloc-accounting is enabled but the allocator is not installed in
this binary \
+ (balance {before} -> {during}); the numbers below would be
meaningless"
+ );
+ drop(held);
+}
+
+#[cfg(not(feature = "alloc-accounting"))]
+fn assert_accounting_is_live() {}
+
+/// Allocation sizes that stay under the 64 KiB settle threshold, so most
iterations exercise only
+/// the thread-local fast path rather than the atomic flush.
+fn small_churn(c: &mut Criterion) {
Review Comment:
`small_churn` inlines the body of the `alloc_free` helper that
`threshold_churn` calls a few lines down. Both are alloc/free loops over a list
of sizes in the same group, so one function over `[16, 256, 4096, 32 * 1024, 64
* 1024]` covers both and keeps the threshold story in one place.
Minor: `group.throughput(...)` is re-set on every loop iteration. It only
needs to be set once before the loop.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust
global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool)
counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real
allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers,
Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a
lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and
currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full
picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains
a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is
**observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory
pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations
it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes
allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency
allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS
as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!
https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the
magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only
about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a
negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by*
`track` settles
+ /// directly instead of recursing. The only such allocation today is the
one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates
and never fails —
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+ /// This thread's un-flushed delta.
+ static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread
exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting
would be silently
+/// discarded every time a thread died. Worker threads live for the process
lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor
that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+ fn drop(&mut self) {
+ let drift = self.0.replace(0);
+ if drift != 0 {
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ }
+ }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never
reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out
of order.
+pub fn current_balance() -> usize {
+ clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+ balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the
magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+ let drift = local_drift.get().wrapping_add(delta);
+ if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+ local_drift.set(0);
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ } else {
+ local_drift.set(drift);
+ }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+ if delta == 0 {
+ return;
+ }
+
+ // A re-entrant call is one made by `track` itself; the outer frame owns
the flag and will
+ // clear it, so this frame must only settle and return.
+ if IN_TRACK.with(|in_track| in_track.replace(true)) {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ return;
+ }
+
+ // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s
destructor has already
+ // run, and any allocation after that point must not panic inside the
allocator.
+ if LOCAL_DRIFT
+ .try_with(|thread_drift| settle(&thread_drift.0, delta))
+ .is_err()
+ {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ }
+
+ IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+///
[apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626),
which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on
directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+ inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+ pub const fn new(inner: A) -> Self {
+ Self { inner }
+ }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc`
contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local
`Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc_zeroed(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+ // Settle before delegating. A free cannot fail, so there is nothing
to wait for, and the
+ // inner free can be slow: jemalloc returns oversize blocks to the OS
eagerly, and unmapping
+ // a few hundred megabytes takes milliseconds. Accounting afterwards
would keep the block on
+ // the balance for that whole window, after the allocator's own
statistics had already
+ // dropped it.
+ track(-(layout.size() as isize));
+ self.inner.dealloc(ptr, layout);
+ }
+
+ unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) ->
*mut u8 {
+ let new_ptr = self.inner.realloc(ptr, layout, new_size);
+ if !new_ptr.is_null() {
+ // Accounting after the fact is only safe because this allocator
cannot fail the
+ // allocation or unwind. A variant that enforced a limit would
have to decide *before*
+ // delegating: `realloc` may free or move the old block, and a
caller that never
+ // received the new pointer would free the stale one while
unwinding.
+ //
+ // A single allocation cannot exceed `isize::MAX` on any real
platform, so neither cast
+ // wraps.
+ track(new_size as isize - layout.size() as isize);
+ }
+ new_ptr
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::{Mutex, MutexGuard};
+
+ /// `BALANCE` is process-wide and the crate's tests run in parallel, so a
test that reads it
+ /// sees every other test's allocations. The tests that move it by tens of
megabytes take this
+ /// lock so they cannot land inside each other's windows; the rest of the
crate is kept out by
+ /// making each window microseconds wide and each expected move far larger
than anything else
+ /// allocates in that time.
+ static SERIAL: Mutex<()> = Mutex::new(());
+
+ fn serial() -> MutexGuard<'static, ()> {
+ SERIAL
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ }
+
+ #[test]
+ fn settle_accumulates_below_the_threshold() {
Review Comment:
Three tests over a six-line function.
`settle_accumulates_below_the_threshold`, `settle_flushes_at_the_threshold`,
and `settle_flushes_negative_drift` differ only in the delta and the expected
residue, so one table-driven case over `[(1024, 1024), (SETTLE_THRESHOLD, 0),
(-SETTLE_THRESHOLD, 0)]` says the same thing.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust
global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool)
counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real
allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers,
Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a
lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and
currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full
picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains
a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is
**observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory
pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations
it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes
allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency
allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS
as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!
https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the
magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only
about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a
negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by*
`track` settles
+ /// directly instead of recursing. The only such allocation today is the
one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates
and never fails —
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+ /// This thread's un-flushed delta.
+ static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread
exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting
would be silently
+/// discarded every time a thread died. Worker threads live for the process
lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor
that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+ fn drop(&mut self) {
+ let drift = self.0.replace(0);
+ if drift != 0 {
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ }
+ }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never
reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out
of order.
+pub fn current_balance() -> usize {
Review Comment:
Worth stating the accuracy bound here and in the `tracing.md` row: the
returned value can lag the true balance by up to `SETTLE_THRESHOLD` times the
number of live threads, since each thread holds un-flushed drift. Immaterial
against GiB-scale footprints, but a reader comparing `native_allocated` against
pool reservations byte-for-byte should know the number is approximate.
##########
native/core/src/alloc_accounting.rs:
##########
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Process-wide accounting of the bytes currently handed out by the Rust
global allocator.
+//!
+//! Comet's [`MemoryPool`](datafusion::execution::memory_pool::MemoryPool)
counts *declared
+//! reservations*: bytes an operator explicitly asked for. Plenty of real
allocation never goes
+//! through it — Arrow builders, expression kernels, decompression buffers,
Parquet metadata,
+//! `object_store` buffers, tokio's own machinery — so pool reservations are a
lower bound on
+//! Comet's footprint, and the size of the gap is workload-dependent and
currently unmeasurable at
+//! runtime. See the [memory management contributor guide] for the full
picture.
+//!
+//! [`AccountingAllocator`] wraps the selected global allocator and maintains
a single signed
+//! process-wide byte balance, which [`current_balance`] exposes. This is
**observability only**: it
+//! never rejects an allocation, never panics, and never gates the memory
pool. It exists so the
+//! accounting gap can be seen in tracing output next to the pool reservations
it should be
+//! compared against.
+//!
+//! The balance counts `Layout` bytes, not resident pages. It excludes
allocator fragmentation,
+//! jemalloc's retained pages, `mmap`ed regions, and anything a C dependency
allocates through libc
+//! `malloc` rather than Rust's `GlobalAlloc` — so it is a lower bound on RSS
as well, just a much
+//! tighter one than pool reservations.
+//!
+//! [memory management contributor guide]:
+//!
https://datafusion.apache.org/comet/contributor-guide/memory_management.html
+
+use std::alloc::{GlobalAlloc, Layout};
+use std::cell::Cell;
+use std::sync::atomic::{AtomicIsize, Ordering};
+
+/// A thread flushes its accumulated delta into the shared balance once the
magnitude reaches this.
+/// Batching keeps the common path to a thread-local add-and-compare, so only
about one atomic
+/// read-modify-write per 64 KiB of churn touches the shared cacheline.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+/// Outstanding bytes, process-wide. Signed because a thread can flush a
negative delta before
+/// another flushes the matching positive one.
+static BALANCE: AtomicIsize = AtomicIsize::new(0);
+
+thread_local! {
+ /// Set while this thread is inside [`track`], so an allocation made *by*
`track` settles
+ /// directly instead of recursing. The only such allocation today is the
one some platforms
+ /// make when registering `LOCAL_DRIFT`'s destructor on first touch.
+ ///
+ /// Const-initialized and destructor-free, so reading it never allocates
and never fails —
+ /// which is what makes it safe to consult before touching `LOCAL_DRIFT`.
+ static IN_TRACK: Cell<bool> = const { Cell::new(false) };
+
+ /// This thread's un-flushed delta.
+ static LOCAL_DRIFT: ThreadDrift = const { ThreadDrift(Cell::new(0)) };
+}
+
+/// Owns a thread's un-flushed delta and settles the remainder when the thread
exits.
+///
+/// Without the destructor, up to [`SETTLE_THRESHOLD`] bytes of accounting
would be silently
+/// discarded every time a thread died. Worker threads live for the process
lifetime, but the
+/// blocking pool churns on tokio's idle timeout, so on a long-lived executor
that would be a
+/// slowly accumulating bias in the reported balance.
+struct ThreadDrift(Cell<isize>);
+
+impl Drop for ThreadDrift {
+ fn drop(&mut self) {
+ let drift = self.0.replace(0);
+ if drift != 0 {
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ }
+ }
+}
+
+/// Bytes currently handed out by the Rust global allocator, process-wide.
+///
+/// Returns 0 when the [`AccountingAllocator`] is not installed. Never
reported negative: the
+/// balance can dip below zero transiently while per-thread deltas settle out
of order.
+pub fn current_balance() -> usize {
+ clamp_balance(BALANCE.load(Ordering::Relaxed))
+}
+
+/// Clamps a signed balance to the unsigned value reported to callers.
+fn clamp_balance(balance: isize) -> usize {
+ balance.max(0) as usize
+}
+
+/// Adds `delta` to `local_drift`, flushing into the shared balance once the
magnitude reaches
+/// [`SETTLE_THRESHOLD`].
+fn settle(local_drift: &Cell<isize>, delta: isize) {
+ let drift = local_drift.get().wrapping_add(delta);
+ if drift.unsigned_abs() >= SETTLE_THRESHOLD as usize {
+ local_drift.set(0);
+ BALANCE.fetch_add(drift, Ordering::Relaxed);
+ } else {
+ local_drift.set(drift);
+ }
+}
+
+/// Records a signed byte delta against the process balance.
+#[inline]
+fn track(delta: isize) {
+ if delta == 0 {
+ return;
+ }
+
+ // A re-entrant call is one made by `track` itself; the outer frame owns
the flag and will
+ // clear it, so this frame must only settle and return.
+ if IN_TRACK.with(|in_track| in_track.replace(true)) {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ return;
+ }
+
+ // `try_with` rather than `with`: during thread teardown `LOCAL_DRIFT`'s
destructor has already
+ // run, and any allocation after that point must not panic inside the
allocator.
+ if LOCAL_DRIFT
+ .try_with(|thread_drift| settle(&thread_drift.0, delta))
+ .is_err()
+ {
+ BALANCE.fetch_add(delta, Ordering::Relaxed);
+ }
+
+ IN_TRACK.with(|in_track| in_track.set(false));
+}
+
+/// Wraps a global allocator, accounting the `Layout` bytes it hands out.
+///
+/// Adapted from the `AccountingAllocator` in
+///
[apache/datafusion#22626](https://github.com/apache/datafusion/pull/22626),
which lives in
+/// DataFusion's test-only `sqllogictest` crate and so cannot be depended on
directly.
+pub struct AccountingAllocator<A: GlobalAlloc> {
+ inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+ pub const fn new(inner: A) -> Self {
+ Self { inner }
+ }
+}
+
+// SAFETY: every method delegates to `inner`, which upholds the `GlobalAlloc`
contract. The
+// accounting is pure bookkeeping over an `AtomicIsize` and thread-local
`Cell`s: it does not
+// inspect, retain, or alter any pointer, and it cannot unwind.
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+ let ptr = self.inner.alloc_zeroed(layout);
+ if !ptr.is_null() {
+ track(layout.size() as isize);
+ }
+ ptr
+ }
+
+ unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+ // Settle before delegating. A free cannot fail, so there is nothing
to wait for, and the
+ // inner free can be slow: jemalloc returns oversize blocks to the OS
eagerly, and unmapping
+ // a few hundred megabytes takes milliseconds. Accounting afterwards
would keep the block on
+ // the balance for that whole window, after the allocator's own
statistics had already
+ // dropped it.
+ track(-(layout.size() as isize));
+ self.inner.dealloc(ptr, layout);
+ }
+
+ unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) ->
*mut u8 {
Review Comment:
`realloc` is the one method whose accounting order this PR deliberately
changed from the prototype (after delegating rather than before), and nothing
tests it. `dealloc_settles_before_delegating` already has the `Recording`
inner-allocator harness. A grow and a shrink through it, asserting the balance
moves by the delta rather than by the full new size, would lock the new
ordering in.
##########
native/core/src/lib.rs:
##########
@@ -65,27 +52,92 @@ pub mod jvm_bridge {
use errors::{try_unwrap_or_throw, CometError, CometResult};
+pub mod alloc_accounting;
pub mod cloud;
pub mod execution;
pub mod parquet;
// this module is for non release only. Intended for debugging/profiling
purposes
#[cfg(debug_assertions)]
pub mod debug;
+// Global allocator selection.
+//
+// `backend` names the allocator the feature set asks for: jemalloc where it
builds, otherwise
+// mimalloc, otherwise the system allocator. The three `backend` cfgs
partition every feature
+// combination, so exactly one definition exists, and each backend predicate
is written once. The
+// unwrapped `#[global_allocator]` lives inside the backend module that owns
it, so a build without
+// `alloc-accounting` is byte-for-byte the previous arrangement: no wrapper,
no per-allocation work,
+// and no explicit allocator at all when the selection is the system allocator.
+//
+// With `alloc-accounting`, the single wrapped `#[global_allocator]` below
refers to
+// `backend::Backend` whatever it resolved to. That is what makes the wrapper
impossible to drop
+// silently: a feature combination with no backend would fail to compile
rather than run with the
+// metric enabled and reading zero.
+
+/// jemalloc, on targets where it builds, unless mimalloc was also requested.
#[cfg(all(
not(target_env = "msvc"),
feature = "jemalloc",
not(feature = "mimalloc")
))]
-#[global_allocator]
-static GLOBAL: Jemalloc = Jemalloc;
+mod backend {
+ pub type Backend = tikv_jemallocator::Jemalloc;
+ pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
+ pub const NAME: &str = "jemalloc";
+ #[cfg(not(feature = "alloc-accounting"))]
+ #[global_allocator]
+ static GLOBAL: Backend = BACKEND;
+}
+
+/// mimalloc, unless a usable jemalloc was also requested.
#[cfg(all(
feature = "mimalloc",
not(all(not(target_env = "msvc"), feature = "jemalloc"))
))]
+mod backend {
+ pub type Backend = mimalloc::MiMalloc;
+ pub const BACKEND: Backend = mimalloc::MiMalloc;
+ pub const NAME: &str = "mimalloc";
+
+ #[cfg(not(feature = "alloc-accounting"))]
+ #[global_allocator]
+ static GLOBAL: Backend = BACKEND;
+}
+
+/// The system allocator: the complement of the two cases above. This covers
neither feature, a
+/// jemalloc request on MSVC, and both features together, which each backend
cfg excludes in favour
+/// of the other.
+#[cfg(not(any(
+ all(
+ not(target_env = "msvc"),
+ feature = "jemalloc",
+ not(feature = "mimalloc")
+ ),
+ all(
+ feature = "mimalloc",
+ not(all(not(target_env = "msvc"), feature = "jemalloc"))
+ )
+)))]
+// Without `alloc-accounting` nothing refers to this selection: the system
allocator is the
+// default, so no `#[global_allocator]` is installed.
+#[cfg_attr(not(feature = "alloc-accounting"), allow(dead_code))]
+mod backend {
+ pub type Backend = std::alloc::System;
+ pub const BACKEND: Backend = std::alloc::System;
+ pub const NAME: &str = "system";
+}
+
+/// The name of the allocator backend this build selected: `"jemalloc"`,
`"mimalloc"` or
+/// `"system"`. This is the one place the selection is decided, so anything
that needs to know
+/// which allocator is in effect (the `alloc_overhead` benchmark's liveness
check, for instance)
+/// reads it from here rather than re-deriving it from the feature set.
+pub use backend::NAME as ALLOCATOR_BACKEND;
Review Comment:
Question: this adds public crate API whose only consumer is a benchmark
assertion. The reasoning in the doc comment is sound, since `lib.rs` owns the
selection and re-deriving it from the feature set in the bench could disagree
with it. Not arguing against it, just wondering whether `#[doc(hidden)]` is
worth it to keep this out of the crate's documented surface.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
