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-5991-587e387681859121f12eb9bc6c8bdce9ffe5074e in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit a51ff5cd886328c91bfb8155bca12cb0192ec83c Author: Andy Grove <[email protected]> AuthorDate: Sun Sep 20 22:51:36 2026 +0000 fix: count each memory pool once in analyze_trace (#5991) * fix: count each memory pool once in analyze_trace analyze_trace derived the total pool reservation by summing the per-thread thread_NNN_comet_memory_reserved counters. Those come from total_reserved_for_thread, which deduplicates shared pools only within a thread, so a task-shared pool reports its full reservation on every thread that references it and summing multiplies it by the thread count. Measured on TPC-H SF100 with tracing enabled, the reported total reached 500x the real allocation, 27% of samples claimed more memory reserved than the allocator had handed out, and the median total was 0 because threads with no registered pools contribute nothing. The tool's peak gap then landed on samples where the total happened to be zero, reporting peak allocation rather than a gap. Add total_reserved_across_tasks(), which sums reserved() over the distinct entries of TASK_SHARED_MEMORY_POOLS, and emit it as a single comet_memory_reserved_total counter beside the existing allocation counter so both halves of a sample are true at the same instant. analyze_trace prefers it, and still reads older traces from the per-thread sum while saying plainly that those totals over-count. The new counter has to be matched before the existing name.contains("comet_memory_reserved") branch, which would otherwise fold the process-wide total into the per-thread map. The per-thread counters are unchanged: they remain useful for per-thread attribution in a trace viewer. * fix: cover every pool type in the total, and pair trace samples Addresses review feedback on the process-wide reservation counter. Source the total from the per-thread tracing registry instead of the task-shared one. The task-shared registry only holds pools created by `acquire_task_shared_pool`, so the `greedy`, `fair_spill`, the `_global` variants and `unbounded` modes reported zero while holding live reservations, and the analyzer then called those reservations untracked. Every plan registers whichever pool `create_memory_pool` gave it, so deduplicating that registry by pool identity covers all of them and still counts a shared pool once. That also removes the `Weak::upgrade` under the registry lock. An upgraded reference could become a pool's last owner if the owning task finished mid-scan, and `TaskSharedMemoryPool::drop` re-locks the registry, so the sampling thread would have deadlocked holding it. The thread registry holds strong references that the scan only borrows. In the analyzer, compare an allocation sample only against the pool total emitted next to it on the same thread. `executePlan` emits the two back to back, but a busy-poll reservation event in between meant the previous comparison paired a fresh allocation with a stale total and recorded excess that was never held. Traces without the total counter still take the legacy path, whose numbers are unchanged. A trace that carries the counter but never pairs a sample now says so rather than reporting OK. * fix: compare observed zero reservations in legacy traces Gate the legacy comparison on whether an allocation sample and a reservation sample have been observed, not on whether their values are positive. A reservation recorded as zero is a real value, and allocation standing above it is the signal the tool exists to find: memory still held after the pool released it. The previous gate dropped those samples, and a trace whose only reservation sample was zero was reported as having no pool samples at all. The legacy warning said the excess was understated. That is only half of it now: summing inflates the total when a shared pool is counted on every referencing thread, and deflates it when a thread has not reported yet, so the excess can err either way. Say that instead. The doc's sample output predated the counter and its numbers no longer match what the tool prints, so replace it with real output from a trace that carries the counter, elided the way the per-thread list already was. * fix: count untraced plans' pools in the process-wide total `createPlan` registered a plan's memory pool only when that plan had `spark.comet.tracing.enabled` set, but `tracing.enabled` is a session config and the allocation counter the total is compared against is process-wide. An executor running an untraced plan holding 100 MiB alongside a traced plan holding 20 MiB reported 20 MiB against 120 MiB of live reservations, and the analyzer called the difference untracked allocation. Register unconditionally; only the counter emission stays gated on tracing. Read both the per-thread and the process-wide total under one lock, so the two counters emitted as a pair describe the same instant, and so tracing takes the registry mutex once per sample rather than twice. Move the counter name to `datafusion_comet_common::tracing::POOL_TOTAL_METRIC` so the producer and `analyze_trace` cannot drift apart: renaming it on one side alone would silently put every new trace on the warned fallback path. * test: give each temp test file a name no other test can pick `get_temp_filename` drew from 65536 random values into one shared `target/debug/testdata` directory. `cargo nextest` runs each test in its own process with several in flight, so collisions happen, and `fs::File::create` truncates: one test empties the Parquet file another is mid-read of, and that test fails with object_store's "Requested range was invalid" from a start offset past a now-zero-length file. This is what failed `parquet_roundtrip_unsigned_int` on this PR's last CI run, unrelated to the change under review. Name files by process ID and a counter instead, which is unique across the processes alive at any one time. --- docs/source/contributor-guide/tracing.md | 85 +++++---- native/common/src/bin/analyze_trace.rs | 190 ++++++++++++++----- native/common/src/tracing.rs | 7 + native/core/src/execution/jni_api.rs | 249 ++++++++++++++++++++++--- native/spark-expr/src/test_common/file_util.rs | 18 +- 5 files changed, 434 insertions(+), 115 deletions(-) diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md index bc5690bc59..4e7b82b914 100644 --- a/docs/source/contributor-guide/tracing.md +++ b/docs/source/contributor-guide/tracing.md @@ -66,7 +66,7 @@ Example trace visualization: ## Analyzing Memory Usage The `analyze_trace` tool parses a trace log and compares the process-wide native allocation counter against -the sum of per-thread Comet memory pool reservations. This is useful for detecting untracked native memory +the total memory reserved by Comet's memory pools. This is useful for detecting untracked native memory growth where native allocations exceed what the memory pools account for. Build and run: @@ -78,46 +78,60 @@ cargo run --bin analyze_trace -- /path/to/comet-event-trace.json The tool reads counter events from the trace log. Because tracing logs metrics per thread, `native_allocated` and `jemalloc_allocated` are process-wide values (the same global allocation reported from whichever thread -logs it), while `thread_NNN_comet_memory_reserved` values are per-thread pool reservations that are summed to -get the total tracked memory. The tool analyzes `native_allocated` when the trace contains it, since that -counts only what Rust code holds from the allocator, and otherwise falls back to `jemalloc_allocated`. The -output names the counter it used. A trace with neither counter is rejected. +logs it). The total tracked memory comes from `comet_memory_reserved_total`, which covers every pool type and +counts each pool once process-wide. Every plan registers its pool for this total, including plans running +with tracing disabled, because the allocation counter it is compared against covers them too. The per-thread +`thread_NNN_comet_memory_reserved` values must not be summed to obtain it: a shared pool reports its full +reservation on every thread that references it, so summing multiplies it by the thread count. -Sample output: +The allocation counter and the total are emitted back to back on one thread when a traced plan finishes +executing, and the tool compares only samples paired that way, so a fresh allocation is never measured +against a stale reservation. Traces recorded before that counter existed are still analyzed from the +per-thread sum, and the tool warns that the sum is not the same measure. + +The tool analyzes `native_allocated` when the trace contains it, since that counts only what Rust code holds +from the allocator, and otherwise falls back to `jemalloc_allocated`. The output names the counter it used. A +trace with neither counter is rejected. + +Sample output, with the violation table and the per-thread list elided: ``` === Comet Trace Memory Analysis === -Counter events parsed: 193104 -Allocation counter: jemalloc_allocated -Threads with memory pools: 8 -Peak jemalloc_allocated: 3068.2 MB -Peak pool total: 2864.6 MB -Peak excess (jemalloc_allocated - pool): 364.6 MB +Counter events parsed: 2946 +Allocation counter: native_allocated +Pool total source: comet_memory_reserved_total (process-wide) +Peak native_allocated: 395.9 MB +Peak pool total: 250.9 MB +Peak excess (native_allocated - pool): 171.2 MB -WARNING: jemalloc_allocated exceeded pool reservation at 138 sampled points: +WARNING: native_allocated exceeded pool reservation at 87 sampled points: - Time (us) jemalloc_allocated pool_total excess + Time (us) native_allocated pool_total excess ------------------------------------------------------------------ - 179578 210.8 MB 0.1 MB 210.7 MB - 429663 420.5 MB 145.1 MB 275.5 MB - 1304969 2122.5 MB 1797.2 MB 325.2 MB - 21974838 407.0 MB 42.3 MB 364.6 MB - 33543599 5.5 MB 0.1 MB 5.3 MB + 14662 19.0 MB 0.7 MB 18.3 MB + 109895 48.2 MB 22.4 MB 25.8 MB + 1623833 94.0 MB 71.4 MB 22.6 MB + 3315951 193.8 MB 167.8 MB 26.0 MB + 4107671 233.7 MB 199.9 MB 33.8 MB + ... + 5441547 389.3 MB 218.1 MB 171.2 MB + 6445070 235.6 MB 209.8 MB 25.8 MB --- Final per-thread pool reservations --- - thread_60_comet_memory_reserved: 0.0 MB - thread_95_comet_memory_reserved: 0.0 MB - thread_96_comet_memory_reserved: 0.0 MB + thread_60_comet_memory_reserved: 39.2 MB + thread_61_comet_memory_reserved: 48.3 MB + thread_62_comet_memory_reserved: 27.2 MB ... - Total: 0.0 MB + Total: 245.8 MB ``` -Some excess is expected (allocator metadata and fragmentation for `jemalloc_allocated`, and non-pool -allocations like Arrow IPC buffers for either counter). Large or growing excess may indicate memory that is -not being tracked by the pool. +A steady excess is expected, since not every native allocation goes through a pool (Arrow IPC buffers, for +instance), and `jemalloc_allocated` additionally includes the allocator's own metadata and fragmentation. An +excess that grows over the run, as it does from 26 MB to 171 MB above, is the signal worth chasing: that is +memory the pool is not accounting for. Arrow memory on the JVM side is reported separately, because it is off-heap and so invisible to `jvm_heap_used`. Comet imports batches from native over the Arrow C Data Interface, and Arrow @@ -145,12 +159,13 @@ samples: they are not an atomic per-query balance, and neither is a measure of R ## Definition of Labels -| Label | Meaning | -| -------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| jvm_heap_used | JVM heap memory usage of live objects for the executor process | -| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | -| 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. | -| thread_NNN_comet_memory_reserved | Memory reserved by Comet's DataFusion memory pool (summed across all contexts on the thread). NNN is the Rust thread ID. | -| thread_NNN_comet_jvm_shuffle | Off-heap memory allocated by Comet for columnar shuffle. NNN is the Rust thread ID. | +| Label | Meaning | +| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| jvm_heap_used | JVM heap memory usage of live objects for the executor process | +| jemalloc_allocated | Native memory usage for the executor process (requires `jemalloc` feature) | +| 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. | +| 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/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs index 30f55a2ae3..05da1d1ea7 100644 --- a/native/common/src/bin/analyze_trace.rs +++ b/native/common/src/bin/analyze_trace.rs @@ -16,13 +16,14 @@ // under the License. //! Analyzes a Comet chrome trace event log (`comet-event-trace.json`) and -//! compares the process-wide native allocation counter against the sum of -//! per-thread Comet memory pool reservations. Reports any points where the -//! allocated bytes exceed the total pool size. +//! compares the process-wide native allocation counter against the total memory +//! reserved by Comet's memory pools. Reports any points where the allocated +//! bytes exceed the total pool size. //! //! Usage: //! cargo run --bin analyze_trace -- <path-to-comet-event-trace.json> +use datafusion_comet_common::tracing::POOL_TOTAL_METRIC as POOL_TOTAL_COUNTER; use serde::Deserialize; use std::collections::HashMap; use std::io::{BufRead, BufReader}; @@ -37,12 +38,19 @@ use std::{env, fs::File}; /// analyzed. const ALLOCATED_COUNTERS: [&str; 2] = ["native_allocated", "jemalloc_allocated"]; +// The process-wide total of Comet's memory pool reservations, imported from the producer so the +// two cannot drift apart. +// +// Preferred over summing the per-thread `thread_NNN_comet_memory_reserved` counters. Those report +// the full reservation of a task-shared pool once per thread that references it, so adding them +// across threads multiplies a shared pool by its thread count. A trace without this counter is +// still analyzed from the per-thread sum, with a warning, so older traces remain readable. + /// A single Chrome trace event (only the fields we care about). #[derive(Deserialize)] struct TraceEvent { name: String, ph: String, - #[allow(dead_code)] tid: u64, ts: u64, #[serde(default)] @@ -56,6 +64,50 @@ struct MemorySnapshot { pool_total: u64, } +/// What one pool-total source says about the trace. +/// +/// Both sources are accumulated in the same pass and one is reported, because whether the trace +/// carries [`POOL_TOTAL_COUNTER`] is only known once a line containing it has been read. +#[derive(Default)] +struct Analysis { + /// Highest pool total this source observed. + peak_pool_total: u64, + /// Largest `allocated - pool_total` over the comparisons this source could make. + peak_excess: u64, + /// A sample of the points where allocation exceeded the total. + violations: Vec<MemorySnapshot>, + /// How many comparisons this source made. Zero means it never produced a usable pair, which + /// is worth reporting rather than passing off as "allocation never exceeded the total". + comparisons: u64, +} + +impl Analysis { + fn observe_total(&mut self, pool_total: u64) { + self.peak_pool_total = self.peak_pool_total.max(pool_total); + } + + fn compare(&mut self, ts: u64, allocated: u64, pool_total: u64) { + self.comparisons += 1; + self.observe_total(pool_total); + + let Some(excess) = allocated.checked_sub(pool_total).filter(|e| *e > 0) else { + return; + }; + self.peak_excess = self.peak_excess.max(excess); + // Sample the violations rather than listing every one, but always keep the worst. + if self.violations.is_empty() + || ts.saturating_sub(self.violations.last().unwrap().ts) > 1_000_000 + || excess == self.peak_excess + { + self.violations.push(MemorySnapshot { + ts, + allocated, + pool_total, + }); + } + } +} + fn format_bytes(bytes: u64) -> String { const MB: f64 = 1024.0 * 1024.0; format!("{:.1} MB", bytes as f64 / MB) @@ -75,14 +127,22 @@ fn main() { let mut source: Option<usize> = None; // Latest allocated value (global, not per-thread) let mut latest_allocated: u64 = 0; - // Per-thread pool reservations: thread_NNN -> bytes + // Per-thread pool reservations: thread_NNN -> bytes. Reported at the end, and summed as the + // pool total only for traces that predate the process-wide counter. let mut pool_by_thread: HashMap<String, u64> = HashMap::new(); - // Points where allocated exceeded pool total - let mut violations: Vec<MemorySnapshot> = Vec::new(); + // Whether the trace carries the process-wide total, and its last value. + let mut pool_total_counter: Option<u64> = None; + // The allocation sample most recently emitted on each thread, waiting for that thread's + // matching pool total. `executePlan` emits the allocation counter and the process-wide total + // back to back on one thread, so pairing them by thread compares two values sampled at the + // same instant. Comparing the latest allocation against whatever total arrived last instead + // pairs a fresh allocation with a stale reservation and invents excess that was never held. + let mut allocated_awaiting_total: HashMap<u64, u64> = HashMap::new(); + // Comparisons against the process-wide counter, and against the legacy per-thread sum. + let mut paired = Analysis::default(); + let mut legacy = Analysis::default(); // Track peak values let mut peak_allocated: u64 = 0; - let mut peak_pool_total: u64 = 0; - let mut peak_excess: u64 = 0; let mut counter_events: u64 = 0; // Each line is one JSON event, possibly with a trailing comma. @@ -136,14 +196,28 @@ fn main() { source = Some(rank); latest_allocated = 0; peak_allocated = 0; - peak_excess = 0; - violations.clear(); + allocated_awaiting_total.clear(); + paired = Analysis::default(); + legacy = Analysis::default(); } } if let Some(val) = event.args.get(&event.name) { latest_allocated = val.as_u64().unwrap_or(0); - if latest_allocated > peak_allocated { - peak_allocated = latest_allocated; + peak_allocated = peak_allocated.max(latest_allocated); + allocated_awaiting_total.insert(event.tid, latest_allocated); + } + } else if event.name == POOL_TOTAL_COUNTER { + // Must be matched before the per-thread branch below, whose `contains` check would + // otherwise also match this name and fold the process-wide total into the map. + if let Some(val) = event.args.get(&event.name) { + let pool_total = val.as_u64().unwrap_or(0); + pool_total_counter = Some(pool_total); + paired.observe_total(pool_total); + // Only compare against the allocation sampled in this thread's own group. An + // observed zero reservation is a real value that allocation can exceed, so a + // paired zero is a genuine comparison, not a missing sample. + if let Some(allocated) = allocated_awaiting_total.remove(&event.tid) { + paired.compare(event.ts, allocated, pool_total); } } } else if event.name.contains("comet_memory_reserved") { @@ -158,30 +232,18 @@ fn main() { continue; } - // After each allocated or pool update, check the current state. A comparison needs one - // sample of each side: an observed zero reservation is a real value that allocation can - // exceed, so only the absence of any pool sample defers the check. - let pool_total: u64 = pool_by_thread.values().sum(); - if pool_total > peak_pool_total { - peak_pool_total = pool_total; - } - - if source.is_some() && !pool_by_thread.is_empty() && latest_allocated > pool_total { - let excess = latest_allocated - pool_total; - if excess > peak_excess { - peak_excess = excess; - } - // Record violation (sample - don't record every single one) - if violations.is_empty() - || event.ts.saturating_sub(violations.last().unwrap().ts) > 1_000_000 - || excess == peak_excess - { - violations.push(MemorySnapshot { - ts: event.ts, - allocated: latest_allocated, - pool_total, - }); - } + // Legacy association, for traces recorded before the process-wide total existed: compare + // the latest allocation against the running per-thread sum after every counter event. + // There is nothing to pair on in those traces, so this is the best they support. + // + // A comparison needs one sample of each side, not a positive one. A reservation that has + // been observed at zero is a real value, and allocation standing above it is exactly the + // signal worth finding: memory still held after the pool released it. Requiring a positive + // total instead would drop those samples silently. + let per_thread_sum: u64 = pool_by_thread.values().sum(); + legacy.observe_total(per_thread_sum); + if source.is_some() && !pool_by_thread.is_empty() { + legacy.compare(event.ts, latest_allocated, per_thread_sum); } } @@ -194,39 +256,68 @@ fn main() { std::process::exit(1); }; + let have_pool_total_counter = pool_total_counter.is_some(); + let analysis = if have_pool_total_counter { + &paired + } else { + &legacy + }; + // Print summary println!("=== Comet Trace Memory Analysis ===\n"); println!("Counter events parsed: {counter_events}"); println!("Allocation counter: {source}"); - println!("Threads with memory pools: {}", pool_by_thread.len()); + if have_pool_total_counter { + println!("Pool total source: {POOL_TOTAL_COUNTER} (process-wide)"); + } else { + println!( + "Pool total source: sum of {} per-thread counters", + pool_by_thread.len() + ); + println!( + "WARNING: this trace predates {POOL_TOTAL_COUNTER}, so the total below is the sum of\n\ + the per-thread counters. That is not the same measure: a shared pool reports its full\n\ + reservation on every thread referencing it, which inflates the total, while a thread\n\ + that has not reported yet contributes nothing, which deflates it. The excess below can\n\ + err in either direction." + ); + } println!("Peak {source}: {}", format_bytes(peak_allocated)); println!( "Peak pool total: {}", - format_bytes(peak_pool_total) + format_bytes(analysis.peak_pool_total) ); println!( "Peak excess ({source} - pool): {}", - format_bytes(peak_excess) + format_bytes(analysis.peak_excess) ); println!(); - if pool_by_thread.is_empty() { - println!( - "No pool reservation samples in the trace, so there is nothing to compare against." - ); - } else if violations.is_empty() { + if analysis.comparisons == 0 { + if have_pool_total_counter { + println!( + "No allocation sample was emitted next to a {POOL_TOTAL_COUNTER} sample on the\n\ + same thread, so there was nothing to compare. The two are emitted together when\n\ + a traced plan finishes executing." + ); + } else { + println!( + "No pool reservation samples in the trace, so there is nothing to compare against." + ); + } + } else if analysis.violations.is_empty() { println!("OK: {source} never exceeded the total pool reservation."); } else { println!( "WARNING: {source} exceeded pool reservation at {} sampled points:\n", - violations.len() + analysis.violations.len() ); println!( "{:>14} {:>18} {:>14} {:>14}", "Time (us)", source, "pool_total", "excess" ); println!("{}", "-".repeat(66)); - for snap in &violations { + for snap in &analysis.violations { let excess = snap.allocated - snap.pool_total; println!( "{:>14} {:>18} {:>14} {:>14}", @@ -245,5 +336,8 @@ fn main() { for (thread, bytes) in &threads { println!(" {thread}: {}", format_bytes(**bytes)); } - println!("\n Total: {}", format_bytes(pool_by_thread.values().sum())); + println!( + "\n Total: {}", + format_bytes(pool_total_counter.unwrap_or_else(|| pool_by_thread.values().sum())) + ); } diff --git a/native/common/src/tracing.rs b/native/common/src/tracing.rs index aad4e7269b..ce158e9f98 100644 --- a/native/common/src/tracing.rs +++ b/native/common/src/tracing.rs @@ -110,6 +110,13 @@ pub fn log_memory_usage(name: &str, value: u64) { RECORDER.log_memory_usage(name, value); } +/// Counter name for the process-wide total of Comet's memory pool reservations. +/// +/// Shared between the producer in `jni_api` and `analyze_trace`, which recognises this exact name +/// and falls back to a warned, over-counting estimate when a trace does not carry it. Renaming it +/// in one place only would silently put every new trace on the fallback path. +pub const POOL_TOTAL_METRIC: &str = "comet_memory_reserved_total"; + pub fn with_trace<T, F>(label: &str, tracing_enabled: bool, f: F) -> T where F: FnOnce() -> T, diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index faf4d5dfda..7652428410 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -108,7 +108,7 @@ use crate::execution::shuffle::{ use crate::execution::spark_plan::SparkPlan; use crate::execution::tracing::{ - get_thread_id, log_memory_usage, trace_begin, trace_end, with_trace, + get_thread_id, log_memory_usage, trace_begin, trace_end, with_trace, POOL_TOTAL_METRIC, }; use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; @@ -195,6 +195,23 @@ impl Drop for ThreadMemoryPoolRegistration { } } +/// Sums `reserved()` over `pools`, skipping any pool whose identity is already in `seen`. +/// +/// 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() +} + /// 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(); @@ -204,14 +221,7 @@ fn unregister_and_total(thread_id: u64, context_id: i64) -> usize { map.remove(&thread_id); return 0; } - let mut seen = HashSet::new(); - return pools - .values() - .filter_map(|p| { - let ptr = Arc::as_ptr(p) as *const (); - seen.insert(ptr).then(|| p.reserved()) - }) - .sum::<usize>(); + return sum_distinct_pools(pools.values(), &mut HashSet::new()); } 0 } @@ -219,21 +229,58 @@ fn unregister_and_total(thread_id: u64, context_id: i64) -> usize { fn total_reserved_for_thread(thread_id: u64) -> usize { let map = get_thread_memory_pools().lock(); map.get(&thread_id) - .map(|pools| { - // Deduplicate pools that share the same underlying allocation - // (e.g. task-shared pools registered by multiple execution contexts) - let mut seen = HashSet::new(); - pools - .values() - .filter_map(|p| { - let ptr = Arc::as_ptr(p) as *const (); - seen.insert(ptr).then(|| p.reserved()) - }) - .sum::<usize>() - }) + .map(|pools| sum_distinct_pools(pools.values(), &mut HashSet::new())) .unwrap_or(0) } +/// Reservation totals read from the pool registry in one pass, under one lock. +struct ReservedTotals { + /// Bytes reserved by the pools registered on the requested thread, deduplicated within it. + /// Zero when no thread was requested. + for_thread: usize, + /// Bytes reserved across every live Comet memory pool in the process. + /// + /// This is the figure to compare against a process-wide allocation counter, so it has to + /// account for every live pool rather than only the ones a traced plan created. Allocation is + /// process-wide and a plan running with tracing off still holds memory, which is why + /// `createPlan` registers its pool unconditionally: with two concurrent plans configured + /// differently, counting only the traced one would report the untraced plan's reservation as + /// untracked allocation. It covers every pool type, because each plan registers whichever pool + /// `create_memory_pool` gave it, and it counts a pool once however many contexts or threads + /// hold it. + /// + /// The per-thread `thread_NNN_comet_memory_reserved` counters must not be summed to obtain it: + /// a shared pool reports its full reservation on every thread that references it, so adding + /// them across threads multiplies that pool by its thread count. + across_threads: usize, +} + +/// Reads both totals under a single lock. +/// +/// 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. +/// +/// `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(); + ReservedTotals { + for_thread, + across_threads, + } +} + fn parse_usize_env_var(name: &str) -> Option<usize> { std::env::var_os(name).and_then(|n| n.to_str().and_then(|s| s.parse::<usize>().ok())) } @@ -536,10 +583,18 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( // Register the shared base pool before wrapping it for per-plan debug logging. The // guard removes the entry if any later plan setup step fails. + // + // Registration is not conditional on this plan's tracing setting. `tracing.enabled` is + // a session config, so an executor can run a traced and an untraced plan at once, + // while the allocation counter a trace compares against is process-wide. Registering + // only traced plans would leave the untraced plan's reservation out of the total and + // report it as allocation held outside any pool. let rust_thread_id = get_thread_id(); - let memory_pool_registration = tracing_enabled.then(|| { - ThreadMemoryPoolRegistration::new(rust_thread_id, id, Arc::clone(&memory_pool)) - }); + let memory_pool_registration = Some(ThreadMemoryPoolRegistration::new( + rust_thread_id, + id, + Arc::clone(&memory_pool), + )); let memory_pool = if logging_memory_pool { Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) @@ -1113,10 +1168,17 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( 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 + // compare them. The per-thread counter cannot be summed across threads to obtain the + // process-wide one: it reports a shared pool's full reservation once per referencing + // thread. + let totals = total_reserved(Some(exec_context.rust_thread_id)); log_memory_usage( &exec_context.tracing_memory_metric_name, - total_reserved_for_thread(exec_context.rust_thread_id) as u64, + totals.for_thread as u64, ); + log_memory_usage(POOL_TOTAL_METRIC, totals.across_threads as u64); } result @@ -1145,13 +1207,16 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( let mut execution_context: Box<ExecutionContext> = Box::from_raw(exec_context as *mut ExecutionContext); - // Unregister this context's pool and emit the remaining total for the thread + // 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(); - log_memory_usage( - &execution_context.tracing_memory_metric_name, - remaining as u64, - ); + if execution_context.tracing_enabled { + log_memory_usage( + &execution_context.tracing_memory_metric_name, + remaining as u64, + ); + } } // Flush metrics last, as it is the only fallible step here. @@ -1700,7 +1765,9 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( mod tests { use super::*; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion::execution::memory_pool::{MemoryConsumer, UnboundedMemoryPool}; + use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryReservation, UnboundedMemoryPool, + }; use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::ReturnFieldArgs; use datafusion_comet_proto::spark_expression; @@ -1829,6 +1896,7 @@ mod tests { #[test] fn thread_memory_pool_registration_is_scoped_and_deduplicates_base_pool() { + let _guard = serial(); const THREAD_ID: u64 = u64::MAX; let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); let reservation = MemoryConsumer::new("test").register(&pool); @@ -1879,6 +1947,125 @@ mod tests { assert!(weak.upgrade().is_none()); } + /// `THREAD_MEMORY_POOLS` is process-wide and the crate's tests run in parallel, so any test + /// that registers a pool perturbs another's view of the process-wide total. The tests below + /// take this lock so their deltas are exact; without it they observe each other's pools. + static SERIAL: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + fn serial() -> std::sync::MutexGuard<'static, ()> { + SERIAL + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn reserving(bytes: usize) -> (Arc<dyn MemoryPool>, MemoryReservation) { + let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + reservation.grow(bytes); + (pool, reservation) + } + + fn total_reserved_across_threads() -> usize { + total_reserved(None).across_threads + } + + /// The property that makes this total comparable against a process-wide allocation counter, + /// and the one summing the per-thread counters gets wrong: a pool shared by several contexts + /// on several threads contributes its reservation once. + #[test] + fn a_shared_pool_is_counted_once_however_many_contexts_hold_it() { + let _guard = serial(); + let before = total_reserved_across_threads(); + let (pool, reservation) = reserving(4096); + + let _first = ThreadMemoryPoolRegistration::new(11, 1, Arc::clone(&pool)); + assert_eq!( + total_reserved_across_threads() - before, + 4096, + "a registered pool's reservation must appear in the total" + ); + + // A second context in the same task, and a third on another thread, both register the + // same pool. Summing the per-thread counters would report 4096 three times over. + let _second = ThreadMemoryPoolRegistration::new(11, 2, Arc::clone(&pool)); + let _third = ThreadMemoryPoolRegistration::new(12, 3, Arc::clone(&pool)); + assert_eq!( + total_reserved_across_threads() - before, + 4096, + "a shared pool must be counted once, not once per holder" + ); + assert_eq!(total_reserved_for_thread(11), 4096); + assert_eq!(total_reserved_for_thread(12), 4096); + + drop(reservation); + } + + /// The total has to cover every pool type, not just the task-shared ones: `greedy`, + /// `fair_spill`, the `_global` variants and `unbounded` all bypass the task-shared registry, + /// and reporting zero for them would make the analyzer call live reservations untracked. + #[test] + fn independent_pools_each_contribute_to_the_total() { + let _guard = serial(); + let before = total_reserved_across_threads(); + let (first_pool, first_reservation) = reserving(4096); + let (second_pool, second_reservation) = reserving(8192); + + let _first = ThreadMemoryPoolRegistration::new(13, 1, first_pool); + let _second = ThreadMemoryPoolRegistration::new(13, 2, second_pool); + assert_eq!(total_reserved_across_threads() - before, 4096 + 8192); + + drop(first_reservation); + drop(second_reservation); + } + + #[test] + fn released_pools_leave_the_total() { + let _guard = serial(); + let before = total_reserved_across_threads(); + let (pool, reservation) = reserving(8192); + { + let _registration = ThreadMemoryPoolRegistration::new(14, 1, pool); + assert_eq!(total_reserved_across_threads() - before, 8192); + } + assert_eq!( + total_reserved_across_threads(), + before, + "unregistering the last context must remove the pool from the total" + ); + drop(reservation); + } + + /// What the traced thread sees when a plan on another thread holds memory it knows nothing + /// about, which is the case a tracing-gated registry got wrong: `tracing.enabled` is a session + /// config, so an untraced plan can run alongside a traced one, and its reservation is part of + /// the process-wide allocation the trace compares against. The per-thread figure stays local — + /// that counter is for attribution — while the process total has to include the other thread. + #[test] + fn the_process_total_includes_pools_registered_on_other_threads() { + let _guard = serial(); + let before = total_reserved_across_threads(); + let (traced_pool, traced_reservation) = reserving(20 * 1024 * 1024); + let (untraced_pool, untraced_reservation) = reserving(100 * 1024 * 1024); + + let _traced = ThreadMemoryPoolRegistration::new(15, 1, traced_pool); + let _untraced = ThreadMemoryPoolRegistration::new(16, 2, untraced_pool); + + let totals = total_reserved(Some(15)); + assert_eq!( + totals.for_thread, + 20 * 1024 * 1024, + "the per-thread counter attributes only this thread's pools" + ); + assert_eq!( + totals.across_threads - before, + 120 * 1024 * 1024, + "the process total must cover the pool held by the other thread" + ); + + drop(traced_reservation); + drop(untraced_reservation); + } + #[test] fn length_resolves_to_spark_length_for_string_and_binary() { use datafusion::physical_expr::expressions::{CastExpr, Column}; diff --git a/native/spark-expr/src/test_common/file_util.rs b/native/spark-expr/src/test_common/file_util.rs index 78e42d29e6..010863965f 100644 --- a/native/spark-expr/src/test_common/file_util.rs +++ b/native/spark-expr/src/test_common/file_util.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::sync::atomic::{AtomicU64, Ordering}; use std::{env, fs, io::Write, path::PathBuf}; /// Returns file handle for a temp file in 'target' directory with a provided content @@ -41,13 +42,28 @@ pub fn get_temp_file(file_name: &str, content: &[u8]) -> fs::File { file.unwrap() } +/// Returns a path, in the same `target/debug/testdata` directory, that no other caller will pick. +/// +/// The directory is shared by every test in the crate and `cargo nextest` runs each test in its +/// own process, several at a time, so the name has to be unique across processes as well as within +/// one. Drawing it from 65536 random values was not: `fs::File::create` truncates, so two tests +/// landing on the same name leave one of them reading a Parquet file that the other has just +/// emptied, and it fails with an out-of-range read that says nothing about the real cause. The +/// process ID plus a counter is unique among the processes alive at any one time, which is all +/// that is needed — a leftover file from an earlier run is simply overwritten. pub fn get_temp_filename() -> PathBuf { + static NEXT_ID: AtomicU64 = AtomicU64::new(0); + let mut path_buf = env::current_dir().unwrap(); path_buf.push("target"); path_buf.push("debug"); path_buf.push("testdata"); fs::create_dir_all(&path_buf).unwrap(); - path_buf.push(rand::random::<i16>().to_string()); + path_buf.push(format!( + "{}-{}", + std::process::id(), + NEXT_ID.fetch_add(1, Ordering::Relaxed) + )); path_buf } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
