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-6162-de2e70291460ed0310428493d8c3aa5223293d08 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit a8fc789b3cb8546592299c736e8a89a69adfb0b5 Author: Andy Grove <[email protected]> AuthorDate: Wed Sep 23 22:18:29 2026 +0000 feat: always count native allocations and log executor native memory usage (#6162) * feat: enable alloc-accounting by default and log executor native memory usage Build the alloc-accounting allocator wrapper by default (jemalloc stays opt-in) and have each executor periodically log its native memory usage while Comet native plans run: Comet native memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB (16 native plans, 8 memory pools) The first CometExecIterator on an executor starts one daemon thread that logs a single line per interval for the whole executor, plus one more after the last plan finishes. It samples on a timer rather than between batches because a plan rooted at a native shuffle writer, or fed by native scans, can spend its whole run inside one executePlan call. The interval is spark.comet.memory.logInterval (default 10s, 0 disables), read from the executor's SparkConf. The new Native.getMemoryUsage JNI call reads only the allocation counter and the pool registry, and reads pool reservations after releasing the registry lock: CometFairMemoryPool holds its own lock across the JNI acquire from Spark, which can wait for a finishing task whose releasePlan needs the registry lock. Also add a tuning guide section on using the log to size spark.executor.memoryOverhead, update the tracing and memory management guides, and switch the rust-test CI step to cover the jemalloc arm and the build without the feature. * fix: always install allocation accounting, keep pool reads off the registry lock Remove the alloc-accounting cargo feature: the accounting allocator now wraps whichever backend the build selects, unconditionally. The alloc_overhead bench measures the wrapper against the bare backend in one binary, since there is no longer a build without it. releasePlan read the remaining pools' reservations while holding the pool registry lock on every plan, not only when tracing. With fair_unified that can deadlock: the pool holds its lock across a Spark acquire that waits on another task, whose releasePlan needs the registry lock. Every reader now copies the pools out and reads them after releasing it, and releasePlan reads nothing unless tracing. A malformed spark.comet.memory.logInterval now disables the log with a warning instead of failing every Comet task, and a value set only in the session is reported as ignored. The log also warns when the executor's native footprint (untracked native memory plus Spark's off-heap memory in use) exceeds spark.memory.offHeap.size plus the memory overhead. * docs: stop recommending memoryPool.fraction in the native memory warning The fraction does not reserve room in Spark's off-heap pool: greedy_unified ignores it, and fair_unified applies it per task, where Spark's own per-task limit is tighter. Point the warning and the tuning guide at the memory overhead instead. --- .github/actions/rust-test/action.yaml | 15 +- docs/source/contributor-guide/memory_management.md | 54 ++-- docs/source/contributor-guide/tracing.md | 21 +- docs/source/user-guide/latest/tuning.md | 56 ++++ native/common/src/bin/analyze_trace.rs | 7 +- native/core/Cargo.toml | 6 - native/core/benches/alloc_overhead.rs | 206 +++++++------- native/core/src/alloc_accounting.rs | 41 ++- native/core/src/execution/jni_api.rs | 311 ++++++++++++++++----- native/core/src/lib.rs | 19 +- .../org/apache/spark/CometTaskMemoryManager.java | 10 + .../main/scala/org/apache/comet/CometConf.scala | 16 ++ .../scala/org/apache/comet/CometExecIterator.scala | 215 ++++++++++++++ spark/src/main/scala/org/apache/comet/Native.scala | 14 + .../spark/CometExecIteratorLifecycleSuite.scala | 118 ++++++++ 15 files changed, 842 insertions(+), 267 deletions(-) diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml index c1a0947926..85aa31979e 100644 --- a/.github/actions/rust-test/action.yaml +++ b/.github/actions/rust-test/action.yaml @@ -84,16 +84,13 @@ 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 + # The steps above lint and test the accounting allocator over the system allocator. Lint and + # test it over jemalloc too, since that is a different allocator backend and nothing else in CI + # builds it. + - name: Check and test the jemalloc allocator backend 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 - + cargo clippy --color=never -p datafusion-comet --all-targets --features jemalloc -- -D warnings + RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features jemalloc alloc_accounting diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 6296fefeee..93ba4e8754 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -363,14 +363,14 @@ diverge for several structural reasons: Freeing memory does not necessarily return pages to the OS. - **Non-Rust allocations.** Memory allocated by C dependencies through libc `malloc`, and anything `mmap`ed, never passes through Rust's `GlobalAlloc`, so neither the memory pool nor the - `jemalloc_allocated` metric sees it. In a default build the C dependencies are libzstd - (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`, pulled in by the default - `hdfs-opendal` feature), and the TLS stack used for cloud object stores (`aws-lc-sys`). Building - with the `jemalloc` or `mimalloc` feature adds the allocator itself (`tikv-jemalloc-sys`, - `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, because several names - suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy, - `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those - allocations do pass through `GlobalAlloc` and are counted. + allocation counters (`native_allocated`, `jemalloc_allocated`) see it. In a default build the C + dependencies are libzstd (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`, + pulled in by the default `hdfs-opendal` feature), and the TLS stack used for cloud object stores + (`aws-lc-sys`). Building with the `jemalloc` or `mimalloc` feature adds the allocator itself + (`tikv-jemalloc-sys`, `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, + because several names suggest otherwise: the other Parquet codecs are pure Rust in this build, + `snap` for Snappy, `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its + name, so those allocations do pass through `GlobalAlloc` and are counted. - **Batches in flight across the FFI boundary.** Reservations stop at the operator that made them. Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the @@ -386,9 +386,13 @@ The practical consequence is that `reserved()` is a lower bound on Comet's real gap is workload-dependent. `spark.comet.exec.memoryPool.fraction` exists purely so operators can hand-tune a margin that covers the gap for their workload. -To measure the gap on a real query, enable tracing with the `jemalloc` feature and compare -`jemalloc_allocated` against the summed `thread_NNN_comet_memory_reserved` values; see -[Tracing](tracing.md#analyzing-memory-usage). +To measure the gap on a real workload, read the executor's periodic memory usage log, which +reports the bytes Rust's allocator has handed out next to the pools' reservations; see +[Sizing the Overhead from the Memory Usage Log][memory-usage-log]. For a view per event rather +than per interval, enable tracing and compare `native_allocated` against +`comet_memory_reserved_total`; see [Tracing](tracing.md#analyzing-memory-usage). + +[memory-usage-log]: ../user-guide/latest/tuning.md#sizing-the-overhead-from-the-memory-usage-log ## What the container sees @@ -466,9 +470,9 @@ has bounds _declared reservations_, and the sections above describe several stru declared reservations are a lower bound on physical usage. The known gaps, roughly in order of how much they matter: -- **No signal for real native usage.** The only way to observe the gap today is to enable tracing - with the `jemalloc` feature and compare `jemalloc_allocated` against summed reservations after - the fact. There is no runtime value that an operator, a metric, or a policy could read. +- **Real native usage is observed but not acted on.** Comet counts the bytes Rust's allocator has + handed out, and each executor logs that count next to the pools' reservations. Nothing reads it at runtime, though: no operator, metric, or policy + responds to it, so an executor that outgrows its container is still stopped only by the kill. - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to guess a per-workload margin rather than measuring anything. - **`CometArrowAllocator` is unbounded** and participates in no budget. @@ -483,13 +487,14 @@ much they matter: ## Debugging memory issues -| Tool | What it gives you | -| ------------------------------------------------ | -------------------------------------------------------------------------- | -| `spark.comet.debug.memory=true` | `LoggingMemoryPool` logs every register/grow/shrink with the consumer name | -| `spark.comet.explain.native.enabled=true` | Native plan with per-operator metrics, including spill counts | -| [Tracing](tracing.md#analyzing-memory-usage) | `jemalloc_allocated` vs summed pool reservations; the accounting gap | -| `TrackConsumersPool` | Names the top 10 consumers in `ResourcesExhausted` messages (always on) | -| [`thresher`](https://github.com/cetra3/thresher) | Third-party crate that dumps a jemalloc heap profile at a threshold | +| Tool | What it gives you | +| ------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------- | +| `spark.comet.debug.memory=true` | `LoggingMemoryPool` logs every register/grow/shrink with the consumer name | +| `spark.comet.explain.native.enabled=true` | Native plan with per-operator metrics, including spill counts | +| [Memory usage log](../user-guide/latest/tuning.md#sizing-the-overhead-from-the-memory-usage-log) | Executor-wide native allocation vs pool reservations, logged every 10 seconds by default | +| [Tracing](tracing.md#analyzing-memory-usage) | `native_allocated` vs `comet_memory_reserved_total` per event; the accounting gap over time | +| `TrackConsumersPool` | Names the top 10 consumers in `ResourcesExhausted` messages (always on) | +| [`thresher`](https://github.com/cetra3/thresher) | Third-party crate that dumps a jemalloc heap profile at a threshold | A checklist for triaging an executor OOM kill: @@ -498,9 +503,10 @@ A checklist for triaging an executor OOM kill: treats it as fatal, so the executor is lost either way and the exit code is what distinguishes them. A failed task with `SparkOutOfMemoryError` and a surviving executor is Spark's managed memory pool, which is the only one of the three that is recoverable at task level. -2. Compare `jemalloc_allocated` against the summed pool reservations from a trace. A large excess - points at undeclared native allocations; a small excess points at the budget simply being too - small, or at the JVM side. +2. Compare `allocated` against `reserved` in the executor's `Comet native memory usage` log lines + leading up to the kill, or `native_allocated` against `comet_memory_reserved_total` in a trace. + A large excess points at undeclared native allocations; a small excess points at the budget + simply being too small, or at the JVM side. 3. Check `spark.comet.batchSize` against the schema width. Peak memory scales with `batch_size * columns`, and wide or deeply nested schemas amplify it. 4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index 4e7b82b914..507b58c8ce 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -27,20 +27,17 @@ directory with the filename `comet-event-trace.json`. [Trace Event Format]: https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview?tab=t.0#heading=h.yr4qxyxotyw -Additionally, enabling the `jemalloc` feature will enable tracing of native memory allocations. +Native memory is traced as `native_allocated`. Comet wraps whichever global allocator the build +selected and counts the bytes it has handed out. 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 same figure appears in +the executor's periodic memory usage log, which does not need tracing; see +[Sizing the Overhead from the Memory Usage Log](../user-guide/latest/tuning.md#sizing-the-overhead-from-the-memory-usage-log). -```shell -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: +Enabling the `jemalloc` feature adds a second measure, `jemalloc_allocated`, which also includes +jemalloc's own metadata and fragmentation: ```shell -make release COMET_FEATURES="jemalloc,alloc-accounting" +make release COMET_FEATURES="jemalloc" ``` Example output: @@ -165,7 +162,7 @@ samples: they are not an atomic per-query balance, and neither is a measure of R | jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | | jvm_arrow_allocated | Bytes charged to Comet's Arrow allocator tree on the JVM, including buffers imported from native over the Arrow C Data Interface | | jvm_arrow_imported | Bytes charged to the Arrow C Data Interface import allocator, a subset of `jvm_arrow_allocated`. An allocator charge, not a measure of where the bytes were allocated; see above. | -| 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. | +| native_allocated | Bytes handed out by the Rust global allocator, process-wide. Approximate to within 64 KiB of un-flushed delta per live thread. | | comet_memory_reserved_total | Total memory reserved across every live Comet memory pool, process-wide, whatever the configured pool type. Counts a pool shared between execution contexts once, so unlike the per-thread counters it can be compared directly against an allocation counter. | | 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. Do not sum these across threads: a shared pool reports its full reservation on every thread that references it. | | thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 84a3a16ede..368f46ece6 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -108,6 +108,7 @@ there is. That includes: Reserved memory is therefore a lower bound on what Comet really uses, and how far below it sits depends on the workload. This is why Comet can stay within the pool's limit and still push the executor past its container limit. +Each executor logs how far apart the two are while Comet runs; see [Sizing the Overhead from the Memory Usage Log]. To leave room for the part that is not counted, set `spark.comet.exec.memoryPool.fraction` to a value less than `1.0`, which restricts the amount of memory Comet is allowed to reserve. @@ -180,10 +181,65 @@ spark.executor.memoryOverheadFactor=0.2 Raise the value further if executors are killed by the cluster manager (on Kubernetes, `ExecutorLostFailure` with exit code 137) rather than failing with a task-level out-of-memory error. +To measure how much Comet needs rather than guessing, see [Sizing the Overhead from the Memory Usage Log]. Note that on Kubernetes and YARN the overhead is added to the container size, so raising it reduces how many executors fit on a node. +[Sizing the Overhead from the Memory Usage Log]: #sizing-the-overhead-from-the-memory-usage-log + +### Sizing the Overhead from the Memory Usage Log + +While Comet native plans are running, each executor logs its native memory usage at INFO level, +one line every 10 seconds for the whole executor: + +``` +Comet native memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB (16 native plans, 8 memory pools) +``` + +- `allocated` is the memory that Comet's native code has allocated and not yet freed, whether or not + a pool tracks it. +- `reserved` is the part that Comet's memory pools track. It is charged against + `spark.memory.offHeap.size`, so the container already has room for it. + +The difference between the two, `allocated - reserved`, is Comet's untracked native memory. It is +the part of Comet's footprint that has to fit in `spark.executor.memoryOverhead`, alongside the +JVM's own non-heap memory. To size the overhead from it: + +1. Run a representative workload and find the line with the largest difference in each executor's + log. Take both figures from the same line: they are sampled together, and figures from different + lines describe different moments. Setting `spark.comet.memory.logInterval=1s` for this run makes a + short-lived peak less likely to fall between samples. +2. Start from the overhead the executors had before Comet was enabled, which covers the JVM's own + non-heap memory, and add the largest difference seen on any executor. +3. Add a margin on top. The log can miss the true peak between samples, and neither figure includes + the allocator's fragmentation and retained pages, memory allocated by native C libraries such as + zstd, or Comet's Arrow buffers on the JVM side. + +For example, a 16 GiB executor derives an overhead of 1638 MiB. If the largest difference in its +log is the 1522.3 MiB in the line above, the overhead needs to be at least 1638 + 1523 = 3161 MiB +before any margin, so `spark.executor.memoryOverhead=4g` would be a reasonable setting. + +The executor also logs a warning when its native memory looks larger than its container allows: +when the difference, plus everything in use in Spark's off-heap memory pool (which includes Comet's +reservations), exceeds `spark.memory.offHeap.size` plus the memory overhead. This counts the part of +the off-heap pool that nothing has acquired at that moment, which untracked memory can occupy until +Spark hands it out, so a quiet log is not a sign that the overhead is large enough: size it from the +largest difference as described above. The overhead also has to hold the JVM's own non-heap memory, +so by the time the warning appears the executor has likely outgrown its container. It warns the first time this +happens, and again each time it happens after dropping back below. The overhead it uses is +`spark.executor.memoryOverhead` if set, otherwise `spark.executor.memoryOverheadFactor` of +`spark.executor.memory` with a minimum of `spark.executor.minMemoryOverhead`, as Spark sizes the +default container. There is no warning in local mode. + +Look more closely before raising the overhead if the difference keeps growing through a run rather +than levelling off: native memory that is not being released will exhaust any overhead eventually. +The executor logs one more line after its last native plan finishes, and an `allocated` figure there +that grows from one query to the next points the same way. + +`spark.comet.memory.logInterval` is read when an executor starts its first Comet native plan, so set +it when the application is submitted. Set it to `0` to turn the log off. + ### Determining How Much Memory to Allocate Generally, increasing the amount of memory allocated to Comet will improve query performance by reducing the diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs index 05da1d1ea7..b2cdb9eba9 100644 --- a/native/common/src/bin/analyze_trace.rs +++ b/native/common/src/bin/analyze_trace.rs @@ -31,8 +31,8 @@ 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. +/// `native_allocated` (emitted by every build) 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. @@ -250,7 +250,8 @@ fn main() { 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.", + Was the trace produced by a native library older than the `native_allocated` \ + counter?", ALLOCATED_COUNTERS.join(", ") ); std::process::exit(1); diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 501592945d..8ff9277a6d 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -114,12 +114,6 @@ 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"] diff --git a/native/core/benches/alloc_overhead.rs b/native/core/benches/alloc_overhead.rs index dbba3ed424..b29c45e38a 100644 --- a/native/core/benches/alloc_overhead.rs +++ b/native/core/benches/alloc_overhead.rs @@ -15,24 +15,32 @@ // specific language governing permissions and limitations // under the License. -//! Measures the cost the `alloc-accounting` global-allocator wrapper adds per allocation. +//! Measures the cost the accounting wrapper around the global allocator adds per allocation. //! -//! Run the same benchmark with and without the feature and compare: +//! The wrapper is always installed, so there is no build without it to compare against. Instead +//! each case runs twice in one binary, calling the backend allocator the build selected directly +//! (`direct`) and through an `AccountingAllocator` wrapping it (`accounted`). Both call the +//! allocator explicitly rather than through `#[global_allocator]`, and the accounted calls update +//! the same process-wide balance the installed wrapper does, so the difference between the two is +//! the wrapper's cost. //! //! ```shell -//! cargo bench --bench alloc_overhead -- --save-baseline off -//! cargo bench --bench alloc_overhead --features alloc-accounting -- --baseline off +//! cargo bench --bench alloc_overhead +//! cargo bench --bench alloc_overhead --features jemalloc //! ``` //! -//! `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 +//! `alloc_free` 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. +//! `alloc_fill_free` and `grow_to_64kb` are closer to what Comet does, where filling a batch-sized +//! buffer or growing a builder dwarfs the bookkeeping. +use comet::alloc_accounting::AccountingAllocator; +use comet::{AllocatorBackend, ALLOCATOR_BACKEND, BACKEND_ALLOCATOR}; use criterion::{criterion_group, criterion_main, Criterion, Throughput}; +use std::alloc::{GlobalAlloc, Layout}; use std::hint::black_box; use std::sync::Once; use std::thread; @@ -42,76 +50,30 @@ use std::time::Instant; /// 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(); - }); -} +static DIRECT: AllocatorBackend = BACKEND_ALLOCATOR; +static ACCOUNTED: AccountingAllocator<AllocatorBackend> = + AccountingAllocator::new(BACKEND_ALLOCATOR); -/// jemalloc keeps its own count of bytes it has served; if it is not the global allocator of this -/// binary that count stays at zero. -#[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); +/// The two allocators every case is measured against. +fn variants() -> [(&'static str, &'static (dyn GlobalAlloc + Sync)); 2] { + static ANNOUNCE: Once = Once::new(); + ANNOUNCE.call_once(|| { + eprintln!("alloc_overhead: measuring the `{ALLOCATOR_BACKEND}` allocator backend"); + }); + [("direct", &DIRECT), ("accounted", &ACCOUNTED)] } -#[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); +fn alloc_free(allocator: &dyn GlobalAlloc, size: usize) { + let layout = Layout::from_size_align(black_box(size), 8).unwrap(); + // SAFETY: the layout is valid and non-zero, and the block is freed through the allocator that + // produced it. + unsafe { + let ptr = allocator.alloc(layout); + assert!(!ptr.is_null()); + black_box(ptr); + allocator.dealloc(ptr, layout); + } } /// Alloc/free loops from well below the settle threshold up to exactly on it, single-threaded, and @@ -125,7 +87,6 @@ fn alloc_free(size: usize) { /// 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)); @@ -138,25 +99,30 @@ fn churn(c: &mut Criterion) { } 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); + for (variant, allocator) in variants() { + group.bench_function(format!("alloc_free_{label}/{variant}"), |b| { + b.iter(|| alloc_free(allocator, size)); + }); + if near_threshold { + group.bench_function( + format!("parallel_alloc_free_{label}_x{threads}/{variant}"), + |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(allocator, size); + } + }); } }); - } - }); - start.elapsed() - }); - }); + start.elapsed() + }); + }, + ); + } } } group.finish(); @@ -165,31 +131,53 @@ fn churn(c: &mut Criterion) { /// 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(); + const SIZE: usize = 64 * 1024; 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.throughput(Throughput::Bytes(SIZE as u64)); + for (variant, allocator) in variants() { + group.bench_function(format!("alloc_fill_free_64kb/{variant}"), |b| { + let layout = Layout::from_size_align(SIZE, 8).unwrap(); + b.iter(|| { + // SAFETY: the layout is valid and non-zero, the fill stays within the block, and + // the block is freed through the allocator that produced it. + unsafe { + let ptr = allocator.alloc(layout); + assert!(!ptr.is_null()); + ptr.write_bytes(1, black_box(SIZE)); + black_box(ptr); + allocator.dealloc(ptr, layout); + } + }); }); - }); + } group.finish(); } -/// Repeated growth, which is the `realloc` path: a builder doubling its buffer. +/// Repeated growth, which is the `realloc` path: a builder doubling its buffer from 64 bytes to +/// 64 KiB. 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()) + for (variant, allocator) in variants() { + group.bench_function(format!("grow_to_64kb/{variant}"), |b| { + b.iter(|| { + // SAFETY: each `realloc` passes the block's current layout, and the block is freed + // at its final size through the allocator that produced it. + unsafe { + let mut layout = Layout::from_size_align(64, 8).unwrap(); + let mut ptr = allocator.alloc(layout); + assert!(!ptr.is_null()); + while layout.size() < 64 * 1024 { + let new_size = layout.size() * 2; + ptr = allocator.realloc(ptr, layout, new_size); + assert!(!ptr.is_null()); + layout = Layout::from_size_align(new_size, 8).unwrap(); + black_box(ptr); + } + allocator.dealloc(ptr, layout); + } + }); }); - }); + } group.finish(); } diff --git a/native/core/src/alloc_accounting.rs b/native/core/src/alloc_accounting.rs index 394bbe774d..461e26739e 100644 --- a/native/core/src/alloc_accounting.rs +++ b/native/core/src/alloc_accounting.rs @@ -19,7 +19,7 @@ //! //! [`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 +//! memory pool's reservations in the executor's memory usage log and 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, @@ -71,7 +71,7 @@ impl Drop for ThreadDrift { /// Bytes currently handed out by the Rust global allocator, process-wide. /// -/// Returns 0 when the [`AccountingAllocator`] is not installed. Never reported negative: the +/// 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 @@ -263,10 +263,9 @@ mod tests { } /// 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. + /// wrapper is actually installed as the global allocator. 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; @@ -280,7 +279,7 @@ mod tests { 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?" + is the accounting wrapper installed as the global allocator?" ); drop(held); } @@ -359,35 +358,29 @@ mod tests { unsafe { allocator.dealloc(ptr, Layout::from_size_align(SHRUNK, 8).unwrap()) }; } - /// Threads must settle their remaining drift on exit. + /// A thread's remaining drift must reach the shared balance when the thread exits. /// - /// 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. + /// This drops a `ThreadDrift` holding a drift directly, rather than injecting one into a real + /// thread's `LOCAL_DRIFT` and letting the thread exit. With the wrapper installed, the thread's + /// teardown allocates, and those allocations flush an oversized drift through `track` before + /// the destructor runs, so a thread-exit test would pass without the destructor. That the + /// destructor runs when a thread exits is the `thread_local!` guarantee; what needs testing is + /// that it settles the drift. 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; - + fn dropping_a_thread_drift_settles_it() { 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(); + drop(ThreadDrift(Cell::new(INJECTED))); 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 {}", + "a dropped thread drift 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 7652428410..58011c4a65 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -139,7 +139,6 @@ fn log_jemalloc_usage() { /// /// 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", @@ -148,7 +147,14 @@ fn log_native_allocated() { } /// Registry of active memory pools per Rust thread ID. -/// Used to sum memory reservations across all contexts on the same thread for tracing. +/// Used to sum memory reservations across all contexts for the memory usage log and tracing. +/// +/// Never read a pool's reservation while holding this registry's lock; copy the pools out with +/// [`snapshot_registry`] and read them after it is released. `CometFairMemoryPool` holds its own +/// lock across the JNI call that acquires memory from Spark, and Spark can park that call until +/// another task frees memory. A finishing task frees its reservations only after `releasePlan` has +/// taken this lock to unregister, so a reservation read under this lock can wait on a pool that is +/// itself waiting on the lock. type ThreadPoolMap = HashMap<u64, HashMap<i64, Arc<dyn MemoryPool>>>; static THREAD_MEMORY_POOLS: OnceLock<Mutex<ThreadPoolMap>> = OnceLock::new(); @@ -165,10 +171,26 @@ fn register_memory_pool(thread_id: u64, context_id: i64, pool: Arc<dyn MemoryPoo .insert(context_id, pool); } +/// Removes a context's pool from the registry, without reading any reservation. +fn unregister_memory_pool(thread_id: u64, context_id: i64) { + let removed = { + let mut map = get_thread_memory_pools().lock(); + let Some(pools) = map.get_mut(&thread_id) else { + return; + }; + let removed = pools.remove(&context_id); + if pools.is_empty() { + map.remove(&thread_id); + } + removed + }; + // Dropped after the lock is released, in case it was the last reference to the pool. + drop(removed); +} + struct ThreadMemoryPoolRegistration { thread_id: u64, context_id: i64, - registered: bool, } impl ThreadMemoryPoolRegistration { @@ -177,63 +199,68 @@ impl ThreadMemoryPoolRegistration { Self { thread_id, context_id, - registered: true, } } - - fn unregister_and_total(mut self) -> usize { - self.registered = false; - unregister_and_total(self.thread_id, self.context_id) - } } impl Drop for ThreadMemoryPoolRegistration { fn drop(&mut self) { - if self.registered { - unregister_and_total(self.thread_id, self.context_id); - } + unregister_memory_pool(self.thread_id, self.context_id); } } -/// Sums `reserved()` over `pools`, skipping any pool whose identity is already in `seen`. +/// Pools copied out of the registry in one acquisition of its lock, so that their reservations can +/// be read after it is released; see [`ThreadPoolMap`]. /// -/// Execution contexts routinely share one pool — every context in a task under the task-shared -/// pool types, every context in the process under the global ones — and each of them registers -/// it, so a walk of the registry has to deduplicate by pool identity or it reports one -/// reservation several times. -fn sum_distinct_pools<'a>( - pools: impl IntoIterator<Item = &'a Arc<dyn MemoryPool>>, - seen: &mut HashSet<*const ()>, -) -> usize { - pools - .into_iter() - .filter(|pool| seen.insert(Arc::as_ptr(pool) as *const ())) - .map(|pool| pool.reserved()) - .sum() +/// Execution contexts routinely share one pool (every context in a task under the task-shared pool +/// types, every context in the process under the global ones) and each of them registers it, so +/// both lists are deduplicated by pool identity or a sum over them would report one reservation +/// several times. +struct RegistrySnapshot { + /// Distinct pools registered on the requested thread. Empty when no thread was requested. + thread_pools: Vec<Arc<dyn MemoryPool>>, + /// Distinct pools across every thread. Deduplicated across the whole registry, not within each + /// thread: a task-shared or global pool spans threads. + all_pools: Vec<Arc<dyn MemoryPool>>, + /// Registered contexts, which is one per native plan created and not yet released. + plans: usize, } -/// Unregister a context's pool and return the remaining total reserved for the thread. -fn unregister_and_total(thread_id: u64, context_id: i64) -> usize { - let mut map = get_thread_memory_pools().lock(); - if let Some(pools) = map.get_mut(&thread_id) { - pools.remove(&context_id); - if pools.is_empty() { - map.remove(&thread_id); - return 0; - } - return sum_distinct_pools(pools.values(), &mut HashSet::new()); +fn snapshot_registry(thread_id: Option<u64>) -> RegistrySnapshot { + fn distinct<'a>( + pools: impl IntoIterator<Item = &'a Arc<dyn MemoryPool>>, + seen: &mut HashSet<*const ()>, + ) -> Vec<Arc<dyn MemoryPool>> { + pools + .into_iter() + .filter(|pool| seen.insert(Arc::as_ptr(pool) as *const ())) + .cloned() + .collect() } - 0 + + let map = get_thread_memory_pools().lock(); + let thread_pools = thread_id + .and_then(|id| map.get(&id)) + .map(|pools| distinct(pools.values(), &mut HashSet::new())) + .unwrap_or_default(); + let all_pools = distinct(map.values().flat_map(HashMap::values), &mut HashSet::new()); + let plans = map.values().map(HashMap::len).sum(); + RegistrySnapshot { + thread_pools, + all_pools, + plans, + } +} + +fn sum_reserved(pools: &[Arc<dyn MemoryPool>]) -> usize { + pools.iter().map(|pool| pool.reserved()).sum() } fn total_reserved_for_thread(thread_id: u64) -> usize { - let map = get_thread_memory_pools().lock(); - map.get(&thread_id) - .map(|pools| sum_distinct_pools(pools.values(), &mut HashSet::new())) - .unwrap_or(0) + sum_reserved(&snapshot_registry(Some(thread_id)).thread_pools) } -/// Reservation totals read from the pool registry in one pass, under one lock. +/// Reservation totals read from one snapshot of the pool registry. struct ReservedTotals { /// Bytes reserved by the pools registered on the requested thread, deduplicated within it. /// Zero when no thread was requested. @@ -255,29 +282,47 @@ struct ReservedTotals { across_threads: usize, } -/// Reads both totals under a single lock. +/// Reads both totals from one snapshot of the registry. /// -/// They are emitted as a pair, and a pair that straddles two acquisitions describes two different -/// instants. Taking the lock once also halves the tracing traffic through a mutex the executor -/// needs in order to register and release pools. +/// They are emitted as a pair, and a pair taken from two snapshots can describe two different +/// sets of pools. Taking the lock once also halves the tracing traffic through a mutex the +/// executor needs in order to register and release pools. /// /// `thread_id` of `None` skips the per-thread figure; the caller is only after the process total. fn total_reserved(thread_id: Option<u64>) -> ReservedTotals { - let map = get_thread_memory_pools().lock(); - let for_thread = thread_id - .and_then(|id| map.get(&id)) - .map(|pools| sum_distinct_pools(pools.values(), &mut HashSet::new())) - .unwrap_or(0); - // Deduplicated across the whole map, not within each thread: the same pool is registered by - // every context that holds it, and a task-shared or global pool spans threads. - let mut seen = HashSet::new(); - let across_threads = map - .values() - .map(|pools| sum_distinct_pools(pools.values(), &mut seen)) - .sum(); + let snapshot = snapshot_registry(thread_id); ReservedTotals { - for_thread, - across_threads, + for_thread: sum_reserved(&snapshot.thread_pools), + across_threads: sum_reserved(&snapshot.all_pools), + } +} + +/// Executor-wide memory figures for one line of the periodic memory usage log. +#[derive(Debug, PartialEq)] +struct MemoryUsage { + /// Bytes handed out by the Rust global allocator, process-wide. + native_allocated: usize, + /// Bytes reserved across every live Comet memory pool, counting each pool once however many + /// plans share it. + pools_reserved: usize, + /// Live memory pools. With the task-shared pool types, which include both defaults, that is one + /// per task running native plans. + pools: usize, + /// Native plans that have been created and not yet released. + plans: usize, +} + +/// Reads the executor's memory usage for the periodic memory usage log. +/// +/// This runs on a timer thread, concurrently with every plan in the executor, so it reads only the +/// allocation counter and the pool registry, never an execution context. +fn memory_usage() -> MemoryUsage { + let snapshot = snapshot_registry(None); + MemoryUsage { + native_allocated: crate::alloc_accounting::current_balance(), + pools_reserved: sum_reserved(&snapshot.all_pools), + pools: snapshot.all_pools.len(), + plans: snapshot.plans, } } @@ -1166,7 +1211,6 @@ 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(); // Both totals come from one read of the registry, so the pair describes a single // instant, and both are emitted next to the allocation counter above so a trace can @@ -1208,15 +1252,14 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( Box::from_raw(exec_context as *mut ExecutionContext); // Unregister this context's pool and, when tracing, emit the remaining total for the - // thread. Every context registers, but only a traced one writes counters. - if let Some(memory_pool_registration) = execution_context.memory_pool_registration.take() { - let remaining = memory_pool_registration.unregister_and_total(); - if execution_context.tracing_enabled { - log_memory_usage( - &execution_context.tracing_memory_metric_name, - remaining as u64, - ); - } + // thread. Every context registers, but only a traced one writes counters, so the + // reservations are read only then. + drop(execution_context.memory_pool_registration.take()); + if execution_context.tracing_enabled { + log_memory_usage( + &execution_context.tracing_memory_metric_name, + total_reserved_for_thread(execution_context.rust_thread_id) as u64, + ); } // Flush metrics last, as it is the only fallible step here. @@ -1620,6 +1663,28 @@ pub extern "system" fn Java_org_apache_comet_Native_getRustThreadId( get_thread_id() as jlong } +#[no_mangle] +/// Returns the executor's memory usage for the periodic memory usage log, as +/// `[native_allocated, pools_reserved, pools, plans]`; see [`MemoryUsage`]. Safe to call from any +/// thread; see [`memory_usage`]. +pub extern "system" fn Java_org_apache_comet_Native_getMemoryUsage( + e: EnvUnowned, + _class: JClass, +) -> jlongArray { + try_unwrap_or_throw(&e, |env| { + let usage = memory_usage(); + let values = [ + usage.native_allocated as jlong, + usage.pools_reserved as jlong, + usage.pools as jlong, + usage.plans as jlong, + ]; + let long_array = env.new_long_array(values.len())?; + long_array.set_region(env, 0, &values)?; + Ok(long_array.into_raw()) + }) +} + // ============================================================================ // Native Columnar to Row Conversion // ============================================================================ @@ -2066,6 +2131,112 @@ mod tests { drop(untraced_reservation); } + /// The periodic memory usage log counts every plan, and every pool once. Two plans of one task + /// share a pool across threads, as a task-shared pool does, and a third plan has a pool of its + /// own. + #[test] + fn memory_usage_counts_every_plan_and_every_pool_once() { + let _guard = serial(); + let before = memory_usage(); + let (shared_pool, shared_reservation) = reserving(4096); + let (own_pool, own_reservation) = reserving(8192); + + let _first = ThreadMemoryPoolRegistration::new(18, -6001, Arc::clone(&shared_pool)); + let _second = ThreadMemoryPoolRegistration::new(19, -6002, Arc::clone(&shared_pool)); + let third = ThreadMemoryPoolRegistration::new(18, -6003, own_pool); + + let during = memory_usage(); + assert_eq!(during.plans - before.plans, 3); + assert_eq!( + during.pools - before.pools, + 2, + "a pool shared by two plans must be counted once" + ); + assert_eq!(during.pools_reserved - before.pools_reserved, 4096 + 8192); + + drop(third); + let after = memory_usage(); + assert_eq!(after.plans - before.plans, 2); + assert_eq!(after.pools - before.pools, 1); + assert_eq!(after.pools_reserved - before.pools_reserved, 4096); + + drop(shared_reservation); + drop(own_reservation); + } + + /// Stands in for a `CometFairMemoryPool` whose lock is held across a Spark acquire: it counts + /// its reservation reads, and notes whether the registry lock was held during any of them. + #[derive(Debug, Default)] + struct RegistryProbePool { + reads: std::sync::atomic::AtomicUsize, + read_under_registry_lock: std::sync::atomic::AtomicBool, + } + + impl std::fmt::Display for RegistryProbePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "RegistryProbePool") + } + } + + impl MemoryPool for RegistryProbePool { + fn name(&self) -> &str { + "RegistryProbePool" + } + + fn grow(&self, _: &MemoryReservation, _: usize) {} + + fn shrink(&self, _: &MemoryReservation, _: usize) {} + + fn try_grow(&self, _: &MemoryReservation, _: usize) -> DataFusionResult<()> { + Ok(()) + } + + fn reserved(&self) -> usize { + self.reads + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if get_thread_memory_pools().try_lock().is_none() { + self.read_under_registry_lock + .store(true, std::sync::atomic::Ordering::Relaxed); + } + 4096 + } + } + + /// No path may read a reservation while holding the registry lock; see `ThreadPoolMap`. The + /// memory usage log reads from a thread running no plan, tracing reads from a plan's thread, + /// and `releasePlan` unregisters on every plan, so all of them are covered. Unregistering must + /// not read a reservation at all, since it runs whether or not anything is traced. + #[test] + fn reservations_are_read_outside_the_registry_lock() { + use std::sync::atomic::Ordering::Relaxed; + + let _guard = serial(); + let before = memory_usage().pools_reserved; + let probe = Arc::new(RegistryProbePool::default()); + let registration = + ThreadMemoryPoolRegistration::new(20, -7001, Arc::clone(&probe) as Arc<dyn MemoryPool>); + let second = + ThreadMemoryPoolRegistration::new(20, -7002, Arc::clone(&probe) as Arc<dyn MemoryPool>); + + assert_eq!(memory_usage().pools_reserved - before, 4096); + assert_eq!(total_reserved(Some(20)).for_thread, 4096); + assert_eq!(total_reserved_for_thread(20), 4096); + let reads = probe.reads.load(Relaxed); + assert!(reads >= 3, "each reader should have read the probe"); + + drop(registration); + drop(second); + assert_eq!( + probe.reads.load(Relaxed), + reads, + "unregistering read a reservation" + ); + assert!( + !probe.read_under_registry_lock.load(Relaxed), + "a pool's reservation was read while the registry lock was held" + ); + } + #[test] fn length_resolves_to_spark_length_for_string_and_binary() { use datafusion::physical_expr::expressions::{CastExpr, Column}; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 0872a478f9..6ff3767dfb 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -63,7 +63,8 @@ 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. +// fail to compile. Whichever it is, it is installed wrapped in the `AccountingAllocator`, which +// counts the bytes Rust code holds for the memory usage log and the `native_allocated` metric. /// jemalloc, on targets where it builds, unless mimalloc was also requested. #[cfg(all( @@ -107,17 +108,15 @@ mod backend { 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. +/// The allocator backend this build selected, its name (`"jemalloc"`, `"mimalloc"` or +/// `"system"`), and an instance of it. The selection is decided here and nowhere else, so the +/// `alloc_overhead` benchmark takes the backend from here to measure it with and without the +/// accounting wrapper. #[doc(hidden)] -pub use backend::NAME as ALLOCATOR_BACKEND; - -#[cfg(not(feature = "alloc-accounting"))] -#[global_allocator] -static GLOBAL: backend::Backend = backend::BACKEND; +pub use backend::{ + Backend as AllocatorBackend, BACKEND as BACKEND_ALLOCATOR, NAME as ALLOCATOR_BACKEND, +}; -#[cfg(feature = "alloc-accounting")] #[global_allocator] static GLOBAL: alloc_accounting::AccountingAllocator<backend::Backend> = alloc_accounting::AccountingAllocator::new(backend::BACKEND); diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index a0131cb892..89af671d4a 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -53,6 +53,16 @@ public class CometTaskMemoryManager { this.nativeMemoryConsumer = new NativeMemoryConsumer(); } + /** + * Bytes of the executor's off-heap memory pool in use, for execution and storage. In off-heap + * mode this includes every reservation Comet's memory pools have acquired from Spark. Spark's + * memory manager is private to Spark, which is why this lives here. + */ + public static long sparkOffHeapUsed() { + org.apache.spark.memory.MemoryManager memoryManager = SparkEnv.get().memoryManager(); + return memoryManager.offHeapExecutionMemoryUsed() + memoryManager.offHeapStorageMemoryUsed(); + } + // Called by Comet native through JNI. // Returns the actual amount of memory (in bytes) granted. public long acquireMemory(long size) { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d371f1ba44..3ccf434142 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -317,6 +317,22 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_MEMORY_LOG_INTERVAL: ConfigEntry[Long] = conf("spark.comet.memory.logInterval") + .category(CATEGORY_TUNING) + .doc( + "How often each executor logs its native memory usage at INFO level while Comet native " + + "plans are running: the bytes the native allocator has handed out, and the bytes " + + "reserved in Comet's memory pools. The difference is native memory that the pools are " + + "not accounting for. The executor logs one line per interval however many tasks are " + + "running, and one more after the last plan finishes. It logs a warning when the " + + "native memory looks larger than the executor's container allows. This is an executor " + + "setting, read when an executor starts its first Comet native plan, so it must be set " + + "when the application is submitted. An invalid value disables the log with a warning. " + + s"Set to 0 to disable. $TUNING_GUIDE.") + .timeConf(TimeUnit.MILLISECONDS) + .checkValue(_ >= 0, "The memory usage log interval must not be negative") + .createWithDefault(TimeUnit.SECONDS.toMillis(10)) + val COMET_ONHEAP_MEMORY_OVERHEAD: ConfigEntry[Long] = conf("spark.comet.memoryOverhead") .category(CATEGORY_TESTING) .doc( diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 4de949b9de..6a51cd05ea 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -20,6 +20,11 @@ package org.apache.comet import java.lang.management.ManagementFactory +import java.util.Locale +import java.util.concurrent.{Executors, ThreadFactory, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.util.control.NonFatal import org.apache.arrow.c.ArrowArrayStream import org.apache.hadoop.conf.Configuration @@ -219,6 +224,8 @@ class CometExecIterator( this.close() } + CometExecIterator.startMemoryUsageLog() + private def getNextBatch: Option[ColumnarBatch] = { assert(partitionIndex >= 0 && partitionIndex < numParts) @@ -363,6 +370,214 @@ class CometExecIterator( object CometExecIterator extends Logging { + private val memoryUsageLogStarted = new AtomicBoolean(false) + + /** Native plans running at the previous memory usage log. Only the log's own thread uses it. */ + private var plansAtLastMemoryUsageLog = 0L + + /** + * Whether the native footprint exceeded the executor's native memory limit at the previous + * memory usage log. Only the log's own thread uses it. + */ + private var limitExceededAtLastLog = false + + /** + * Starts the executor's native memory usage log when the first native plan is created, unless + * `spark.comet.memory.logInterval` is 0. + * + * Both figures the log reports, the bytes the native allocator has handed out and the bytes + * reserved in Comet's memory pools, are executor-wide, so one daemon thread logs one line per + * interval for the whole executor, however many tasks are running. It runs on a timer rather + * than between batches, because a plan can spend its whole run inside one `executePlan` call: a + * plan rooted at a native shuffle writer consumes all of its input before it returns, and a + * plan fed directly by native scans parks the task thread until its next batch is ready. + */ + private def startMemoryUsageLog(): Unit = { + if (memoryUsageLogStarted.compareAndSet(false, true)) { + // Read from the executor's configuration rather than the session's, since the one log + // serves every session on the executor. + val conf = SparkEnv.get.conf + val intervalMs = memoryUsageLogInterval(conf.getOption(COMET_MEMORY_LOG_INTERVAL.key)) + // A value set only in the session would otherwise be ignored without a trace. Only the + // session of the plan that starts the log is checked, which covers the common case of an + // application with one session. + Option(SQLConf.get.getConfString(COMET_MEMORY_LOG_INTERVAL.key, null)) + .filterNot(conf.getOption(COMET_MEMORY_LOG_INTERVAL.key).contains) + .foreach { sessionValue => + logWarning( + s"Ignoring ${COMET_MEMORY_LOG_INTERVAL.key}=$sessionValue set in the session: the " + + "native memory usage log is executor-wide, so it is read from the executor's " + + "configuration. Set it when the application is submitted.") + } + if (intervalMs > 0) { + val nativeLib = new Native() + val limitBytes = nativeMemoryLimit(conf) + Executors + .newSingleThreadScheduledExecutor(new ThreadFactory { + override def newThread(runnable: Runnable): Thread = { + val thread = new Thread(runnable, "comet-memory-usage-log") + thread.setDaemon(true) + thread + } + }) + .scheduleWithFixedDelay( + new Runnable { + override def run(): Unit = logMemoryUsage(nativeLib, limitBytes) + }, + intervalMs, + intervalMs, + TimeUnit.MILLISECONDS) + } + } + } + + /** + * The memory usage log interval in milliseconds for the executor's configured value, if any. A + * value that does not parse, or is negative, disables the log with a warning: it is read when a + * native plan is created, and a malformed logging setting must not fail every Comet task. + * Disabling rather than falling back to the default respects an attempt to turn the log off + * with a value such as `false`. + */ + def memoryUsageLogInterval(configured: Option[String]): Long = + configured match { + case None => COMET_MEMORY_LOG_INTERVAL.defaultValue.get + case Some(value) => + try { + COMET_MEMORY_LOG_INTERVAL.valueConverter(value) + } catch { + case NonFatal(e) => + logWarning( + s"Disabling the native memory usage log: invalid value '$value' for " + + s"${COMET_MEMORY_LOG_INTERVAL.key}. Expected a non-negative duration such as " + + s"10s or 500ms, or 0 to disable. ${e.getMessage}") + 0L + } + } + + /** + * The executor's memory overhead in bytes, sized the way Spark sizes the default resource + * profile's container: `spark.executor.memoryOverhead` if set, otherwise + * `spark.executor.memoryOverheadFactor` of `spark.executor.memory`, but at least + * `spark.executor.minMemoryOverhead`. None in local mode, where there is no container, or if + * the settings do not parse. An executor running a non-default resource profile may have a + * different overhead. + */ + def executorMemoryOverhead(conf: SparkConf): Option[Long] = { + if (conf.get("spark.master", "").startsWith("local")) { + None + } else { + try { + val overheadMiB = conf.getOption("spark.executor.memoryOverhead") match { + case Some(_) => conf.getSizeAsMb("spark.executor.memoryOverhead") + case None => + val executorMiB = conf.getSizeAsMb("spark.executor.memory", "1g") + val factor = conf.getDouble("spark.executor.memoryOverheadFactor", 0.1) + val minimumMiB = conf.getSizeAsMb("spark.executor.minMemoryOverhead", "384m") + math.max((executorMiB * factor).toLong, minimumMiB) + } + Some(ByteUnit.MiB.toBytes(overheadMiB)) + } catch { + case NonFatal(_) => None + } + } + } + + /** + * The memory the executor's container has for native memory: `spark.memory.offHeap.size` plus + * the memory overhead; see [[executorMemoryOverhead]]. None, so that nothing is compared + * against it, in local mode, when off-heap memory is disabled (a testing-only mode in which + * Comet's reservations do not come from Spark's off-heap pool), or if the settings do not + * parse. + */ + def nativeMemoryLimit(conf: SparkConf): Option[Long] = { + if (!CometSparkSessionExtensions.isOffHeapEnabled(conf)) { + None + } else { + executorMemoryOverhead(conf).flatMap { overhead => + try { + Some(overhead + conf.getSizeAsBytes("spark.memory.offHeap.size", "0")) + } catch { + case NonFatal(_) => None + } + } + } + } + + private def logMemoryUsage(nativeLib: Native, limitBytes: Option[Long]): Unit = { + try { + val usage = nativeLib.getMemoryUsage() + memoryUsageMessage(usage, plansAtLastMemoryUsageLog).foreach(logInfo(_)) + plansAtLastMemoryUsageLog = usage(3) + val warning = limitBytes.flatMap( + nativeMemoryLimitWarning(usage, CometTaskMemoryManager.sparkOffHeapUsed(), _)) + // Warn when the footprint first exceeds the limit, not at every interval while it stays + // there: the INFO line above keeps reporting it. + if (!limitExceededAtLastLog) { + warning.foreach(logWarning(_)) + } + limitExceededAtLastLog = warning.isDefined + } catch { + case NonFatal(e) => + logWarning("Stopping the native memory usage log after a failure", e) + // Rethrown so that the scheduler stops running the log, rather than having it fail and + // warn again every interval. + throw e + } + } + + /** + * The memory usage log line for `usage`, as returned by [[Native.getMemoryUsage]], or None to + * stay quiet. The log reports while native plans are running, and once more after the last of + * them finishes, so that allocation that outlives them is visible, then waits for plans to run + * again. + */ + def memoryUsageMessage(usage: Array[Long], plansAtLastLog: Long): Option[String] = { + val (allocated, reserved, pools, plans) = (usage(0), usage(1), usage(2), usage(3)) + if (plans == 0 && plansAtLastLog == 0) { + None + } else { + Some( + s"Comet native memory usage: allocated ${toMiB(allocated)}, reserved " + + s"${toMiB(reserved)} ($plans native plans, $pools memory pools)") + } + } + + /** + * A warning if the executor's native footprint exceeds `limitBytes`, the container's memory + * outside the JVM heap; see [[nativeMemoryLimit]]. + * + * The footprint is the native memory Comet's pools do not track, `allocated - reserved`, plus + * `sparkOffHeapUsed`, everything in use in Spark's off-heap pool, which includes Comet's + * reservations as well as Spark's own off-heap execution and storage memory. Comparing the sum + * rather than the untracked part against the overhead alone counts the part of + * `spark.memory.offHeap.size` that nothing has acquired at that moment, which untracked memory + * can occupy until Spark hands it out. The limit also has to hold the JVM's own non-heap + * memory, so by the time the footprint exceeds it the executor has likely outgrown its + * container. + */ + def nativeMemoryLimitWarning( + usage: Array[Long], + sparkOffHeapUsed: Long, + limitBytes: Long): Option[String] = { + val untracked = math.max(usage(0) - usage(1), 0L) + val footprint = untracked + sparkOffHeapUsed + if (footprint > limitBytes) { + Some( + s"Comet native memory not tracked by any memory pool (${toMiB(untracked)}) plus " + + s"Spark's off-heap memory in use (${toMiB(sparkOffHeapUsed)}, including Comet's " + + s"reservations) is ${toMiB(footprint)}, more than the ${toMiB(limitBytes)} the " + + "executor's container has outside the JVM heap (spark.memory.offHeap.size plus the " + + "memory overhead), which also has to hold the JVM's own non-heap memory. The cluster " + + "manager may kill this executor for exceeding its container limit. Raise " + + s"spark.executor.memoryOverhead. ${CometConf.TUNING_GUIDE}.") + } else { + None + } + } + + private def toMiB(bytes: Long): String = + "%.1f MiB".formatLocal(Locale.ROOT, bytes / 1024.0 / 1024.0) + private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 93b396ce0f..d38efb12e5 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -265,6 +265,20 @@ class Native extends NativeBase { */ @native def getRustThreadId(): Long + /** + * Returns the executor's native memory usage, for the periodic memory usage log. Reads only + * process-wide counters, never a plan's execution context, so it is safe to call from any + * thread. + * + * @return + * `[nativeAllocated, poolsReserved, pools, plans]`. `nativeAllocated` is the bytes the native + * allocator has handed out. `poolsReserved` is the bytes reserved across every Comet memory + * pool, counting a pool shared by several plans once. `pools` is the number of live pools, + * which with the default task-shared pool types is one per task running native plans, and + * `plans` is the number of native plans created and not yet released. + */ + @native def getMemoryUsage(): Array[Long] + // Native Columnar to Row conversion methods /** diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index a5b98dd598..c1f29c7447 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -220,4 +220,122 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { } } } + + test("getMemoryUsage counts live plans and reports native allocation") { + val nativeLib = new Native() + // Other suites' plans can still be live, so the plan count is compared as a delta. + val plansBefore = nativeLib.getMemoryUsage()(3) + withTaskContext(4500000L) { + val limitOp = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val iter = new CometExecIterator( + id = 4500001L, + inputObjects = Array.empty[Object], + numOutputCols = 1, + protobufQueryPlan = limitOp.toByteArray, + nativeMetrics = CometMetricNode(Map.empty), + numParts = 1, + partitionIndex = 0) + try { + val usage = nativeLib.getMemoryUsage() + assert(usage(3) == plansBefore + 1, "a created plan must be counted until it is released") + assert(usage(2) >= 1, "a live plan must have a memory pool") + assert(usage(1) >= 0) + // The native library always installs the accounting allocator, so a zero allocation + // means it or its wiring was lost. + assert(usage(0) > 0, s"native allocation was reported as ${usage(0)}") + } finally { + iter.close() + } + } + assert(nativeLib.getMemoryUsage()(3) == plansBefore, "a released plan must not be counted") + } + + test("the memory usage log reports while plans run and once after the last one finishes") { + val mib = 1024L * 1024 + val busy = Array(300 * mib, 100 * mib, 2L, 3L) + assert( + CometExecIterator + .memoryUsageMessage(busy, plansAtLastLog = 0) + .contains("Comet native memory usage: allocated 300.0 MiB, reserved 100.0 MiB " + + "(3 native plans, 2 memory pools)")) + + // The line after the last plan finishes shows the allocation the plans left behind. + val idle = Array(20 * mib, 0L, 0L, 0L) + assert( + CometExecIterator + .memoryUsageMessage(idle, plansAtLastLog = 3) + .exists(_.contains("allocated 20.0 MiB, reserved 0.0 MiB (0 native plans"))) + assert(CometExecIterator.memoryUsageMessage(idle, plansAtLastLog = 0).isEmpty) + } + + test("the memory usage log warns when the native footprint exceeds the container") { + import CometExecIterator.nativeMemoryLimitWarning + val mib = 1024L * 1024 + // A 4 GiB off-heap pool with a 1 GiB overhead, and a pool running at 0.8 of the off-heap size. + val limit = 5120 * mib + val reserved = 3000 * mib + // 1500 MiB untracked is more than the overhead, but fits in what the pool left free, since + // Spark's off-heap pool holds only the reservation. + assert(nativeMemoryLimitWarning(Array(4500 * mib, reserved, 1L, 1L), reserved, limit).isEmpty) + // 2500 MiB untracked does not fit: 2500 + 3000 = 5500 MiB. + val warning = nativeMemoryLimitWarning(Array(5500 * mib, reserved, 1L, 1L), reserved, limit) + assert(warning.exists(_.contains("(2500.0 MiB) plus Spark's off-heap memory in use (3000.0"))) + assert(warning.exists(_.contains("is 5500.0 MiB, more than the 5120.0 MiB"))) + // Spark's own off-heap use counts against the same limit. + assert( + nativeMemoryLimitWarning(Array(4500 * mib, reserved, 1L, 1L), 4000 * mib, limit).isDefined) + // Reservations can exceed the allocation, since operators reserve before they allocate. + assert(nativeMemoryLimitWarning(Array(100 * mib, reserved, 1L, 1L), reserved, limit).isEmpty) + } + + test("the native memory limit is the off-heap size plus the memory overhead") { + import CometExecIterator.nativeMemoryLimit + val mib = 1024L * 1024 + val offHeap = new SparkConf(false) + .set("spark.master", "yarn") + .set("spark.executor.memory", "16g") + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "4g") + assert(nativeMemoryLimit(offHeap) == Some((4096 + 1638) * mib)) + assert(nativeMemoryLimit(offHeap.clone.set("spark.memory.offHeap.enabled", "false")).isEmpty) + assert(nativeMemoryLimit(offHeap.clone.set("spark.master", "local[*]")).isEmpty) + } + + test("the executor memory overhead is sized as Spark sizes the container") { + import CometExecIterator.executorMemoryOverhead + val mib = 1024L * 1024 + def conf(settings: (String, String)*): SparkConf = + new SparkConf(false).set("spark.master", "yarn").setAll(settings) + + assert( + executorMemoryOverhead(conf("spark.executor.memoryOverhead" -> "3g")) == Some(3072 * mib)) + // A bare number is in MiB, as Spark reads it. + assert( + executorMemoryOverhead(conf("spark.executor.memoryOverhead" -> "500")) == Some(500 * mib)) + assert(executorMemoryOverhead(conf("spark.executor.memory" -> "16g")) == Some(1638 * mib)) + // The 1g default executor would get 102 MiB from the factor, so the minimum applies. + assert(executorMemoryOverhead(conf()) == Some(384 * mib)) + assert( + executorMemoryOverhead( + conf( + "spark.executor.memory" -> "10g", + "spark.executor.memoryOverheadFactor" -> "0.25")) == Some(2560 * mib)) + assert(executorMemoryOverhead(conf("spark.executor.memoryOverhead" -> "lots")).isEmpty) + assert(executorMemoryOverhead(new SparkConf(false).set("spark.master", "local[4]")).isEmpty) + } + + test("the memory usage log interval disables the log on a value it cannot use") { + import CometExecIterator.memoryUsageLogInterval + assert(memoryUsageLogInterval(None) == 10000L) + assert(memoryUsageLogInterval(Some("1s")) == 1000L) + assert(memoryUsageLogInterval(Some("250ms")) == 250L) + // A bare number is in milliseconds, the unit the setting is declared with. + assert(memoryUsageLogInterval(Some("500")) == 500L) + assert(memoryUsageLogInterval(Some("0")) == 0L) + // Values that would otherwise throw from the plan that starts the log, failing its task. + assert(memoryUsageLogInterval(Some("-5s")) == 0L) + assert(memoryUsageLogInterval(Some("false")) == 0L) + assert(memoryUsageLogInterval(Some("10 seconds please")) == 0L) + } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
