From d529dc1d0eb06e381cb8728595c247d2679426cf Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 3 Sep 2026 12:30:46 -0700 Subject: [PATCH 1/2] feat: experiment with `RealUsagePool` --- native/core/Cargo.toml | 9 +- native/core/src/execution/jni_api.rs | 218 +++++++--- .../core/src/execution/memory_pools/config.rs | 39 ++ native/core/src/execution/memory_pools/mod.rs | 21 + .../src/execution/memory_pools/oom_guard.rs | 409 ++++++++++++++++++ .../execution/memory_pools/real_usage_pool.rs | 368 ++++++++++++++++ .../src/execution/memory_pools/task_shared.rs | 9 + native/core/src/execution/mod.rs | 2 +- native/core/src/execution/spark_config.rs | 4 + native/core/src/lib.rs | 35 +- .../scala/org/apache/comet/CometConf.scala | 28 +- 11 files changed, 1078 insertions(+), 64 deletions(-) create mode 100644 native/core/src/execution/memory_pools/oom_guard.rs create mode 100644 native/core/src/execution/memory_pools/real_usage_pool.rs diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 9ebb69b2f69..8b1ebf14bd6 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -103,9 +103,16 @@ datafusion-functions-nested = { version = "55.0.0" } [features] backtrace = ["datafusion/backtrace"] -default = ["hdfs-opendal"] +default = ["hdfs-opendal", "oom-guard"] hdfs-opendal = ["opendal", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] + +# Allocator-level OOM circuit breaker. Wraps the global allocator to track real +# allocated bytes and gate/abort over-budget query-worker threads. Enabled by default +# so `spark.comet.exec.memoryGuard.*` and the `real_usage` memory pool work without a +# special build; an idle guard is near-free (tracking stays off until a task arms it). +# Drop it from `default` for a bare allocator. +oom-guard = [] # Delta Lake integration. When enabled, links the `comet-contrib-delta` crate # into `libcomet` and activates the `OpStruct::DeltaScan` dispatcher arm. # Default builds carry zero Delta surface. diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 65a2d68ec18..9603eae4f60 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -109,11 +109,15 @@ use crate::execution::tracing::{ }; use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; +#[cfg(feature = "oom-guard")] +use crate::execution::memory_pools::{oom_guard, MemoryPoolType, RealUsagePool}; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; +#[cfg(feature = "oom-guard")] +use crate::execution::spark_config::{COMET_MEMORY_GUARD_ENABLED, COMET_MEMORY_GUARD_SIZE}; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::{info, warn}; @@ -224,6 +228,8 @@ fn parse_usize_env_var(name: &str) -> Option { fn build_runtime(default_worker_threads: Option) -> Runtime { let mut builder = tokio::runtime::Builder::new_multi_thread(); + #[cfg(feature = "oom-guard")] + builder.on_thread_start(oom_guard::stamp_current_thread); if let Some(n) = parse_usize_env_var("COMET_WORKER_THREADS") { info!("Comet tokio runtime: using COMET_WORKER_THREADS={n}"); builder.worker_threads(n); @@ -515,6 +521,31 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( memory_limit, memory_limit_per_task, )?; + + // Arm the hard breaker when the guard is enabled or the `real_usage` pool is + // selected (it carries the guard itself). It trips on *actual* over-budget + // usage; the cooperative gate below trips on *projected* usage and spills + // first. `spark.comet.exec.memoryGuard.size` gives the breaker headroom above + // the off-heap budget (e.g. up to the container RSS limit). + #[cfg(feature = "oom-guard")] + let (guard_enabled, is_real_usage) = ( + spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED), + memory_pool_config.pool_type == MemoryPoolType::RealUsage, + ); + #[cfg(feature = "oom-guard")] + if guard_enabled || is_real_usage { + let default_limit = memory_limit.max(0) as u64; + let limit = spark_config.get_u64(COMET_MEMORY_GUARD_SIZE, default_limit); + if limit == 0 { + warn!( + "Comet memory guard is active but the effective limit is 0 \ + (memory_limit={memory_limit}); the guard will not trip. Set \ + spark.comet.exec.memoryGuard.size explicitly." + ); + } + oom_guard::arm(limit as usize); + } + let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); @@ -525,6 +556,26 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( ThreadMemoryPoolRegistration::new(rust_thread_id, id, Arc::clone(&memory_pool)) }); + // Cooperative real-usage gate: reject growth (triggering a spill) once real + // allocator usage plus the request would exceed the off-heap budget. This is the + // first line of defense and fires before the hard breaker armed above, so + // over-budget work spills and retries rather than failing the task. The dedicated + // `real_usage` pool already gates internally, so it is not wrapped again. + #[cfg(feature = "oom-guard")] + let memory_pool = if guard_enabled && !is_real_usage { + let ceiling = memory_limit.max(0) as usize; + // Enable the fair-share guard for pools whose `reserved()` is per-task; + // `executor_cores` is the fallback divisor when no task count is known. + let fair_share = memory_pool_config + .pool_type + .has_per_task_budget() + .then_some(executor_cores); + Arc::new(RealUsagePool::new(memory_pool, ceiling, fair_share)) + as Arc + } else { + memory_pool + }; + let memory_pool = if logging_memory_pool { Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) } else { @@ -883,6 +934,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( schema_addrs: JLongArray, ) -> jlong { try_unwrap_or_throw(&e, |env| { + #[cfg(feature = "oom-guard")] + oom_guard::stamp_current_thread(); // Retrieve the query let exec_context = get_execution_context(exec_context); @@ -959,6 +1012,13 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( .await; if let Err(panic) = result { + #[cfg(feature = "oom-guard")] + if let Some(e) = oom_guard::map_panic_to_error(panic.as_ref()) { + // Runs on the tokio worker thread that panicked, so this clears + // that worker's UNWINDING flag (not the blocked JNI caller thread's). + let _ = tx.send(Err(e)).await; + return; + } let msg = match panic.downcast_ref::<&str>() { Some(s) => s.to_string(), None => match panic.downcast_ref::() { @@ -983,76 +1043,116 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( pull_input_batches(exec_context)?; } - if let Some(rx) = &mut exec_context.batch_receiver { - match rx.blocking_recv() { - Some(Ok(batch)) => { - update_metrics(env, exec_context)?; - return prepare_output( - env, - array_addrs, - schema_addrs, - batch, - exec_context.debug_native, - ); - } - Some(Err(e)) => { - return Err(e.into()); - } - None => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); + if exec_context.batch_receiver.is_some() { + let recv_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe( + || -> CometResult { + // Scope the rx borrow to just the blocking_recv call so that + // exec_context is free for update_metrics / prepare_output below. + let recv = exec_context + .batch_receiver + .as_mut() + .unwrap() + .blocking_recv(); + match recv { + Some(Ok(batch)) => { + update_metrics(env, exec_context)?; + prepare_output( + env, + array_addrs, + schema_addrs, + batch, + exec_context.debug_native, + ) + } + Some(Err(e)) => Err(e.into()), + None => { + log_plan_metrics(exec_context, stage_id, partition); + Ok(-1) + } + } + }, + )); + + match recv_result { + Ok(r) => return r, + Err(_panic) => { + // On a guard panic, drop the receiver so any re-entry re-initializes. + #[cfg(feature = "oom-guard")] + return Err(oom_guard::oom_error_or_resume(_panic, || { + exec_context.batch_receiver = None; + }) + .into()); + #[cfg(not(feature = "oom-guard"))] + std::panic::resume_unwind(_panic); } } } // ScanExec path: busy-poll to interleave JVM batch pulls with stream polling - get_runtime().block_on(async { - loop { - let next_item = exec_context.stream.as_mut().unwrap().next(); - let poll_output = poll!(next_item); - - // Only check time/tracing every 100 polls to reduce overhead - exec_context.poll_count_since_metrics_check += 1; - if exec_context.poll_count_since_metrics_check >= 100 { - exec_context.poll_count_since_metrics_check = 0; - if let Some(interval) = exec_context.metrics_update_interval { - let now = Instant::now(); - if now - exec_context.metrics_last_update_time >= interval { - update_metrics(env, exec_context)?; - exec_context.metrics_last_update_time = now; + let poll_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + get_runtime().block_on(async { + loop { + let next_item = exec_context.stream.as_mut().unwrap().next(); + let poll_output = poll!(next_item); + + // Only check time/tracing every 100 polls to reduce overhead + exec_context.poll_count_since_metrics_check += 1; + if exec_context.poll_count_since_metrics_check >= 100 { + exec_context.poll_count_since_metrics_check = 0; + if let Some(interval) = exec_context.metrics_update_interval { + let now = Instant::now(); + if now - exec_context.metrics_last_update_time >= interval { + update_metrics(env, exec_context)?; + exec_context.metrics_last_update_time = now; + } + } + if exec_context.tracing_enabled { + log_memory_usage( + &exec_context.tracing_memory_metric_name, + total_reserved_for_thread(exec_context.rust_thread_id) as u64, + ); } } - if exec_context.tracing_enabled { - log_memory_usage( - &exec_context.tracing_memory_metric_name, - total_reserved_for_thread(exec_context.rust_thread_id) as u64, - ); - } - } - match poll_output { - Poll::Ready(Some(output)) => { - return prepare_output( - env, - array_addrs, - schema_addrs, - output?, - exec_context.debug_native, - ); - } - Poll::Ready(None) => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); - } - Poll::Pending => { - // JNI call to pull batches from JVM into ScanExec operators. - // block_in_place lets tokio move other tasks off this worker - // while we wait for JVM data. - tokio::task::block_in_place(|| pull_input_batches(exec_context))?; + match poll_output { + Poll::Ready(Some(output)) => { + return prepare_output( + env, + array_addrs, + schema_addrs, + output?, + exec_context.debug_native, + ); + } + Poll::Ready(None) => { + log_plan_metrics(exec_context, stage_id, partition); + return Ok(-1); + } + Poll::Pending => { + // JNI call to pull batches from JVM into ScanExec operators. + // block_in_place lets tokio move other tasks off this worker + // while we wait for JVM data. + tokio::task::block_in_place(|| pull_input_batches(exec_context))?; + } } } + }) + })); + + match poll_result { + Ok(r) => r, + Err(_panic) => { + // The block_on future was dropped mid-poll; on a guard panic null the + // stream so any re-entry re-initializes rather than polling a half-consumed one. + #[cfg(feature = "oom-guard")] + return Err(oom_guard::oom_error_or_resume(_panic, || { + exec_context.stream = None; + }) + .into()); + #[cfg(not(feature = "oom-guard"))] + std::panic::resume_unwind(_panic); } - }) + } }); if exec_context.tracing_enabled { diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index 312a3604383..e5888609ec5 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -28,6 +28,30 @@ pub(crate) enum MemoryPoolType { GreedyGlobal, FairSpillGlobal, Unbounded, + #[cfg(feature = "oom-guard")] + RealUsage, +} + +impl MemoryPoolType { + /// True when this pool's `reserved()` reflects a single task's usage, so a per-task + /// fair-share comparison is meaningful (false for process-wide pools). The non-shared + /// per-task pools (`Greedy`/`FairSpill`) return true but keep no task registry, so the + /// fair-share divisor falls back to `executor_cores` rather than the active-task count. + #[cfg_attr(not(feature = "oom-guard"), allow(dead_code))] + pub(crate) fn has_per_task_budget(&self) -> bool { + // The dedicated `real_usage` pool gates on process-wide real usage + // (first-come), not a per-task reservation, so it has no per-task budget. + #[cfg(feature = "oom-guard")] + if matches!(self, MemoryPoolType::RealUsage) { + return false; + } + !matches!( + self, + MemoryPoolType::GreedyGlobal + | MemoryPoolType::FairSpillGlobal + | MemoryPoolType::Unbounded + ) + } } pub(crate) struct MemoryPoolConfig { @@ -60,6 +84,21 @@ pub(crate) fn parse_memory_pool_config( // shared with Spark is set by `spark.memory.offHeap.size`. MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) } + #[cfg(feature = "oom-guard")] + "real_usage" => { + // Gate growth on real allocator usage against the off-heap budget + // (`pool_size`) instead of delegating per-task accounting to Spark's + // TaskMemoryManager. See `RealUsagePool`. + MemoryPoolConfig::new(MemoryPoolType::RealUsage, pool_size) + } + #[cfg(not(feature = "oom-guard"))] + "real_usage" => { + return Err(CometError::Config( + "Memory pool type 'real_usage' requires a Comet build with the \ + 'oom-guard' native feature" + .to_string(), + )) + } _ => { return Err(CometError::Config(format!( "Unsupported memory pool type for off-heap mode: {memory_pool_type}" diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index d7c2911f913..3ee9c813068 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -18,6 +18,10 @@ mod config; mod fair_pool; pub mod logging_pool; +#[cfg(feature = "oom-guard")] +pub mod oom_guard; +#[cfg(feature = "oom-guard")] +mod real_usage_pool; mod task_shared; mod unified_pool; @@ -32,6 +36,8 @@ use std::sync::Arc; use unified_pool::CometUnifiedMemoryPool; pub(crate) use config::*; +#[cfg(feature = "oom-guard")] +pub(crate) use real_usage_pool::RealUsagePool; pub(crate) use task_shared::*; /// Creates the memory pool for a native plan. @@ -89,5 +95,20 @@ pub(crate) fn create_memory_pool( Arc::clone(memory_pool) } MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()), + #[cfg(feature = "oom-guard")] + MemoryPoolType::RealUsage => { + // Dedicated off-heap pool: `RealUsagePool` is the sole gate, comparing + // process-wide real usage against `pool_size` (first-come across tasks, so + // `fair_share` is `None`) instead of Spark's per-task TaskMemoryManager + // division. The inner `UnboundedMemoryPool` never rejects; `TrackConsumersPool` + // still reports top consumers on rejection. `enable_tracking()` because the + // gate reads the allocator balance even when the hard breaker is unarmed. + oom_guard::enable_tracking(); + tracked(RealUsagePool::new( + Arc::new(UnboundedMemoryPool::default()), + pool_size, + None, + )) + } } } diff --git a/native/core/src/execution/memory_pools/oom_guard.rs b/native/core/src/execution/memory_pools/oom_guard.rs new file mode 100644 index 00000000000..eec4776311a --- /dev/null +++ b/native/core/src/execution/memory_pools/oom_guard.rs @@ -0,0 +1,409 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::common::DataFusionError; +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering}; + +/// Per-thread drift is flushed into the shared balance once it crosses this. +const SETTLE_THRESHOLD: isize = 64 * 1024; + +/// Process-wide outstanding bytes (signed so transient under-settle is fine). +static BALANCE: AtomicIsize = AtomicIsize::new(0); +/// Enforcement limit in bytes; 0 means unset. +static LIMIT: AtomicUsize = AtomicUsize::new(0); +/// Master enforcement gate (single relaxed load on the hot path). +static ARMED: AtomicBool = AtomicBool::new(false); +/// Runtime tracking gate. The accounting allocator is always linked into the +/// default build, but balance tracking is skipped until a task turns it on (via +/// `arm` or `enable_tracking`), so an unused guard costs one relaxed load per +/// allocation. Stays on for the process lifetime once enabled. +static TRACKING_ENABLED: AtomicBool = AtomicBool::new(false); + +thread_local! { + /// Un-flushed per-thread delta. + static LOCAL_DRIFT: Cell = const { Cell::new(0) }; + /// Is this a query-worker thread eligible for enforcement? + static STAMPED: Cell = const { Cell::new(false) }; + /// Set while a guard panic is unwinding this thread, to avoid double-faults. + static UNWINDING: Cell = const { Cell::new(false) }; +} + +/// Payload of the panic raised when an armed, stamped thread exceeds the limit. +#[derive(Debug)] +pub struct OomGuardPanic { + pub balance: usize, + pub limit: usize, +} + +/// Arm the guard with a byte limit. Idempotent. +pub fn arm(limit_bytes: usize) { + LIMIT.store(limit_bytes, Ordering::Relaxed); + enable_tracking(); + ARMED.store(true, Ordering::Relaxed); +} + +/// Disarm the guard (enforcement off; tracking continues cheaply). +#[cfg(test)] +fn disarm() { + ARMED.store(false, Ordering::Relaxed); +} + +/// Turn on real-usage balance tracking. Called when the guard is armed or the +/// `real_usage` memory pool is created, so the process-wide balance is live for +/// the cooperative gate even when the hard breaker is not armed. Idempotent, and +/// tracking is never turned back off in production. +pub fn enable_tracking() { + TRACKING_ENABLED.store(true, Ordering::Relaxed); +} + +/// Mark the current thread as a query-worker thread eligible for enforcement. +pub fn stamp_current_thread() { + STAMPED.with(|s| s.set(true)); +} + +/// Reset the per-thread unwinding guard after a guard panic has been caught on +/// this thread. Safe to call when not unwinding. The JNI caller thread is +/// reused across tasks, so this must run after catching an OomGuardPanic. +fn clear_unwinding() { + UNWINDING.with(|u| u.set(false)); +} + +/// If `panic` is an `OomGuardPanic`, clear this thread's unwinding guard and +/// return the mapped retriable error. Returns `None` for any other panic. +/// Centralizes the downcast + unwinding-reset + error mapping for all catch sites. +pub fn map_panic_to_error(panic: &(dyn std::any::Any + Send)) -> Option { + let g = panic.downcast_ref::()?; + clear_unwinding(); + Some(DataFusionError::ResourcesExhausted(format!( + "Comet OomGuard: native allocation pushed usage to {} bytes, over the limit of {} \ + bytes; failing this task", + g.balance, g.limit + ))) +} + +/// Handle a panic caught by `catch_unwind` on a JNI caller thread. If it is an +/// `OomGuardPanic`, run `cleanup` (e.g. null a half-consumed stream/receiver so any +/// re-entry re-initializes) and return the mapped retriable error; otherwise re-raise +/// the original panic. Shared by the executePlan catch sites. +pub fn oom_error_or_resume( + panic: Box, + cleanup: impl FnOnce(), +) -> DataFusionError { + match map_panic_to_error(panic.as_ref()) { + Some(e) => { + cleanup(); + e + } + None => std::panic::resume_unwind(panic), + } +} + +/// Current process-wide balance in bytes (never reported negative). +pub fn current_balance() -> usize { + BALANCE.load(Ordering::Relaxed).max(0) as usize +} + +/// Record an allocation of `size` bytes; may trip the breaker. +#[inline] +fn record_alloc(size: usize) { + track(size as isize); +} + +/// Record a deallocation of `size` bytes; never trips (credit only). +#[inline] +fn record_dealloc(size: usize) { + track(-(size as isize)); +} + +/// Core tracking + enforcement. Flushes drift; on a debit flush that crosses the +/// limit on an armed, stamped, non-unwinding thread, panics with `OomGuardPanic`. +#[inline] +fn track(delta: isize) { + // Runtime gate: skip all balance bookkeeping until a task enables tracking. + // Keeps the always-linked accounting allocator near-free when the guard and + // the `real_usage` pool are both unused. + if !TRACKING_ENABLED.load(Ordering::Relaxed) { + return; + } + let new_balance = LOCAL_DRIFT.with(|d| { + let mut drift = d.get(); + let flushed = settle(&mut drift, delta, &BALANCE); + d.set(drift); + flushed + }); + + if delta <= 0 { + return; // credits never enforce + } + let Some(balance) = new_balance else { return }; + if !ARMED.load(Ordering::Relaxed) { + return; + } + if !STAMPED.with(|s| s.get()) { + return; + } + if UNWINDING.with(|u| u.get()) { + return; + } + let limit = LIMIT.load(Ordering::Relaxed); + if should_trip(balance, limit) { + // At most one thread may fire the guard panic per arm cycle. CAS the + // master gate true->false; threads that lose the race bail before + // panic_any. The relaxed load above (line ~121) is not a serialization + // point: several threads can all read ARMED=true and reach here in the + // same tight window. If each then dispatches a panic, Rust's unwind ABI + // can abort the process with "failed to initiate panic" instead of + // unwinding cleanly (observed on the 5-concurrent repro: ~4 threads + // firing within ~10 ms -> exit 133). The guard re-arms on the next + // createPlan. + if ARMED + .compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed) + .is_err() + { + return; + } + // panic_any boxes the payload, which re-enters this allocator and calls + // track() again. ARMED is now false so the re-entrant call short-circuits + // at the ARMED check above; setting UNWINDING adds defense in depth in + // case a concurrent createPlan re-arms mid-unwind. + UNWINDING.with(|u| u.set(true)); + std::panic::panic_any(OomGuardPanic { + balance: balance.max(0) as usize, + limit, + }); + } +} + +/// Pure helper: given the current shared balance and a limit, decide whether an +/// armed+stamped thread should trip the breaker. `limit == 0` means "unset". +fn should_trip(balance: isize, limit: usize) -> bool { + limit != 0 && balance > limit.try_into().unwrap_or(isize::MAX) +} + +/// Pure helper: add `delta` to `local_drift`; if it reaches or exceeds `SETTLE_THRESHOLD` +/// in magnitude, flush it into `shared` and return the new shared balance. +/// Otherwise return `None` (nothing flushed). +fn settle(local_drift: &mut isize, delta: isize, shared: &AtomicIsize) -> Option { + *local_drift = local_drift.wrapping_add(delta); + if local_drift.unsigned_abs() >= SETTLE_THRESHOLD as usize { + let flushed = *local_drift; + *local_drift = 0; + let prev = shared.fetch_add(flushed, Ordering::Relaxed); + Some(prev.wrapping_add(flushed)) + } else { + None + } +} + +/// Wraps an inner global allocator, tracking layout bytes for the OomGuard. +pub struct AccountingAllocator { + inner: A, +} + +impl AccountingAllocator { + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +unsafe impl GlobalAlloc for AccountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc(layout); + if !ptr.is_null() { + record_alloc(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.inner.dealloc(ptr, layout); + record_dealloc(layout.size()); + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = self.inner.alloc_zeroed(layout); + if !ptr.is_null() { + record_alloc(layout.size()); + } + ptr + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // Account for and enforce the size delta BEFORE delegating to the inner + // realloc. If this trips the breaker it panics here, while `ptr` is still + // valid, so the unwind frees it correctly. Panicking *after* inner.realloc + // would be unsound: realloc may have already freed/moved the old block, and + // the caller (which never received the new pointer) would free the dangling + // old pointer on unwind and segfault. Only growth can trip; over-counting on + // a (rare) realloc failure errs on the conservative side for an OOM guard. + // + // Casts and subtraction are safe in practice: a single allocation cannot + // exceed isize::MAX on any real platform, so no wrapping or overflow occurs. + let old = layout.size() as isize; + let new = new_size as isize; + track(new - old); + self.inner.realloc(ptr, layout, new_size) + } +} + +#[cfg(test)] +fn reset_for_test() { + BALANCE.store(0, Ordering::Relaxed); + LIMIT.store(0, Ordering::Relaxed); + ARMED.store(false, Ordering::Relaxed); + // Tests exercise the tracking path directly, so keep it on across resets. + TRACKING_ENABLED.store(true, Ordering::Relaxed); + LOCAL_DRIFT.with(|d| d.set(0)); + STAMPED.with(|s| s.set(false)); + UNWINDING.with(|u| u.set(false)); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + // Serializes tests that mutate the process-global guard state. + static GUARD: Mutex<()> = Mutex::new(()); + + #[test] + fn test_should_trip() { + assert!(!should_trip(100, 0)); // unset limit never trips + assert!(!should_trip(100, 200)); // under limit + assert!(!should_trip(200, 200)); // at limit (strictly greater required) + assert!(should_trip(201, 200)); // over limit + } + + #[test] + fn test_settle_accumulates_then_flushes() { + let shared = AtomicIsize::new(0); + let mut drift = 0isize; + // small allocs below threshold do not flush + assert_eq!(settle(&mut drift, 1024, &shared), None); + assert_eq!(shared.load(Ordering::Relaxed), 0); + // crossing the threshold flushes the accumulated drift + let new_balance = settle(&mut drift, SETTLE_THRESHOLD, &shared); + assert_eq!(new_balance, Some(1024 + SETTLE_THRESHOLD)); + assert_eq!(shared.load(Ordering::Relaxed), 1024 + SETTLE_THRESHOLD); + assert_eq!(drift, 0); // drift reset after flush + } + + #[test] + fn test_settle_flushes_negative_drift() { + let shared = AtomicIsize::new(1_000_000); + let mut drift = 0isize; + assert_eq!( + settle(&mut drift, -SETTLE_THRESHOLD, &shared), + Some(1_000_000 - SETTLE_THRESHOLD) + ); + assert_eq!(drift, 0); + } + + #[test] + fn test_settle_flushes_at_exact_threshold() { + let shared = AtomicIsize::new(0); + let mut drift = 0isize; + assert_eq!( + settle(&mut drift, SETTLE_THRESHOLD, &shared), + Some(SETTLE_THRESHOLD) + ); + assert_eq!(drift, 0); + } + + #[test] + fn test_disarmed_never_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + // not armed -> record_alloc must never panic regardless of size + record_alloc(usize::MAX / 2); + record_alloc(usize::MAX / 2); + } + + #[test] + fn test_unstamped_thread_never_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + // arm with a tiny limit relative to current balance, but DO NOT stamp + let limit = current_balance() + 1; + arm(limit); + record_alloc(SETTLE_THRESHOLD as usize * 4); // big enough to flush + disarm(); + } + + #[test] + fn test_stamped_over_budget_trips() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + let limit = current_balance() + SETTLE_THRESHOLD as usize; // headroom + arm(limit); + let result = std::panic::catch_unwind(|| { + // exceed the headroom in one flush + record_alloc(SETTLE_THRESHOLD as usize * 4); + }); + disarm(); + clear_unwinding(); + assert!(result.is_err(), "expected OomGuardPanic"); + let panic = result.unwrap_err(); + assert!( + panic.downcast_ref::().is_some(), + "panic payload should be OomGuardPanic" + ); + } + + // Drives a real heap allocation through the installed AccountingAllocator (only + // wrapped under the `oom-guard` feature) and confirms the guard trips. + #[test] + fn test_real_allocation_trips_guard() { + let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); + reset_for_test(); + stamp_current_thread(); + // 8 MiB headroom over the current (noisy) baseline. + let headroom = 8 * 1024 * 1024; + arm(current_balance() + headroom); + + let result = std::panic::catch_unwind(|| { + // Allocate well past the headroom in 1 MiB chunks so a flush crosses the limit. + let mut held: Vec> = Vec::new(); + for _ in 0..64 { + held.push(vec![0u8; 1024 * 1024]); + } + // Touch the data so the allocation cannot be optimized away. + held.iter().map(|v| v.len()).sum::() + }); + + // Disarm BEFORE clearing UNWINDING so no post-catch allocation on this still-armed, + // still-stamped thread can re-trip outside the catch. + disarm(); + clear_unwinding(); + + assert!( + result.is_err(), + "large allocation on a stamped, armed thread should trip the guard" + ); + assert!( + result + .unwrap_err() + .downcast_ref::() + .is_some(), + "panic payload should be OomGuardPanic" + ); + } +} diff --git a/native/core/src/execution/memory_pools/real_usage_pool.rs b/native/core/src/execution/memory_pools/real_usage_pool.rs new file mode 100644 index 00000000000..1e9f5c45fa0 --- /dev/null +++ b/native/core/src/execution/memory_pools/real_usage_pool.rs @@ -0,0 +1,368 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::execution::memory_pools::{active_task_count, oom_guard}; +use datafusion::common::{resources_datafusion_err, DataFusionError}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use std::sync::Arc; + +/// Source of the current process-wide real allocator usage in bytes. Production reads the +/// live `oom_guard` balance; tests inject a fixed value without touching global state. +enum BalanceSource { + Live, + #[cfg(test)] + Fixed(usize), +} + +impl BalanceSource { + #[inline] + fn current(&self) -> usize { + match self { + BalanceSource::Live => oom_guard::current_balance(), + #[cfg(test)] + BalanceSource::Fixed(bytes) => *bytes, + } + } +} + +/// A `MemoryPool` decorator that, on top of the inner pool's tracked-reservation +/// accounting, rejects growth when *real* allocator usage (untracked Arrow / join / +/// kernel bytes included) plus the requested amount would exceed a process-global +/// ceiling. Returning `ResourcesExhausted` lets DataFusion spill and retry. +pub(crate) struct RealUsagePool { + inner: Arc, + /// Process-global real-usage ceiling in bytes; 0 means unset (no gating). + ceiling: usize, + /// Fixed fallback divisor (concurrent-task count) used when the dynamic + /// active-task count is 0. `None` disables the fair-share guard (first-come), + /// used for pools whose `reserved()` is process-wide. + fair_share: Option, + balance_source: BalanceSource, +} + +impl std::fmt::Debug for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RealUsagePool") + .field("inner", &self.inner) + .field("ceiling", &self.ceiling) + .field("fair_share", &self.fair_share) + .finish_non_exhaustive() + } +} + +impl std::fmt::Display for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "RealUsagePool(ceiling={}, inner={})", + self.ceiling, self.inner + ) + } +} + +impl RealUsagePool { + /// Wrap `inner` with the real-usage gate using the live OomGuard balance. + pub(crate) fn new( + inner: Arc, + ceiling: usize, + fair_share: Option, + ) -> Self { + Self { + inner, + ceiling, + fair_share, + balance_source: BalanceSource::Live, + } + } + + /// Wrap `inner` with an explicit balance source (test seam). + #[cfg(test)] + fn with_balance_source( + inner: Arc, + ceiling: usize, + fair_share: Option, + balance_source: BalanceSource, + ) -> Self { + Self { + inner, + ceiling, + fair_share, + balance_source, + } + } +} + +/// Per-task fair share of `ceiling` given the number of concurrently active +/// tasks, or `cores_fallback` when the dynamic count is unavailable (0). The +/// divisor is floored at 1 so it is never zero. +fn fair_share_limit(ceiling: usize, active_tasks: usize, cores_fallback: usize) -> usize { + let n = if active_tasks > 0 { + active_tasks + } else { + cores_fallback + }; + ceiling / n.max(1) +} + +/// Given the process is already over the real-usage ceiling, decide whether to +/// reject this task's grow. `None` is first-come (reject whoever hit the ceiling); +/// `Some(s)` rejects only a task whose tracked reservation would exceed its fair +/// share `s`, sparing under-share tasks (the OomGuard breaker backstops runaway +/// cases). +fn should_reject_over_ceiling(reserved: usize, additional: usize, share: Option) -> bool { + match share { + None => true, + Some(s) => reserved.saturating_add(additional) > s, + } +} + +impl MemoryPool for RealUsagePool { + fn name(&self) -> &str { + "RealUsagePool" + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer) + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer) + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional) + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> Result<(), DataFusionError> { + // Check the real-usage ceiling before delegating, so an over-budget request is + // rejected without speculatively reserving the inner pool. When the process is + // over the ceiling, the fair-share guard rejects only a task whose own tracked + // reservation exceeds its fair share, sparing innocent small tasks; the OomGuard + // breaker backstops runaway cases. Returning `ResourcesExhausted` lets DataFusion + // spill and retry. + if self.ceiling != 0 && additional != 0 { + let real = self.balance_source.current(); + if real.saturating_add(additional) > self.ceiling { + let share = self + .fair_share + .map(|cores| fair_share_limit(self.ceiling, active_task_count(), cores)); + if should_reject_over_ceiling(self.inner.reserved(), additional, share) { + return Err(resources_datafusion_err!( + "Comet real-usage gate: native usage {real} bytes + requested \ + {additional} bytes exceeds the off-heap budget of {} bytes; \ + spilling/failing this consumer", + self.ceiling + )); + } + } + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::{GreedyMemoryPool, UnboundedMemoryPool}; + + #[test] + fn under_ceiling_succeeds_and_delegates() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + None, + BalanceSource::Fixed(100), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // real usage 100 + request 100 = 200 <= ceiling 1000 + assert!(pool.try_grow(&reservation, 100).is_ok()); + assert_eq!(inner.reserved(), 100); + } + + #[test] + fn over_ceiling_rejects_without_reserving_inner() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + None, + BalanceSource::Fixed(900), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // real usage 900 + request 200 = 1100 > ceiling 1000 -> reject + let result = pool.try_grow(&reservation, 200); + assert!(result.is_err(), "over-ceiling grow should be rejected"); + // inner pool is never touched on rejection, so there is nothing to roll back + assert_eq!(inner.reserved(), 0); + } + + #[test] + fn zero_ceiling_never_gates() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 0, + None, + BalanceSource::Fixed(usize::MAX / 2), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + assert!(pool.try_grow(&reservation, 1024).is_ok()); + assert_eq!(inner.reserved(), 1024); + } + + #[test] + fn shrink_delegates() { + let inner: Arc = Arc::new(UnboundedMemoryPool::default()); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1_000_000, + None, + BalanceSource::Fixed(0), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + pool.try_grow(&reservation, 500).unwrap(); + assert_eq!(pool.reserved(), 500); + pool.shrink(&reservation, 200); + assert_eq!(pool.reserved(), 300); + } + + // Drives a real heap allocation through the installed AccountingAllocator (only + // wrapped under the `oom-guard` feature) and confirms the real-usage gate rejects. + // Robust to parallel test noise: other allocations only raise the balance further, + // which can only make the over-ceiling assertion more true. + #[test] + fn real_allocation_trips_real_usage_gate() { + // The accounting allocator only updates the balance once tracking is on. + oom_guard::enable_tracking(); + let inner: Arc = Arc::new(UnboundedMemoryPool::default()); + let base = oom_guard::current_balance(); + // 4 MiB headroom over the (noisy) baseline. + let ceiling = base + 4 * 1024 * 1024; + let pool: Arc = + Arc::new(RealUsagePool::new(Arc::clone(&inner), ceiling, None)); + let reservation = MemoryConsumer::new("test").register(&pool); + + // Push real usage ~8 MiB above the baseline, held alive across the check so the + // balance stays elevated. 8 MiB > 64 KiB settle threshold, so it flushes to BALANCE. + let held: Vec = vec![0u8; 8 * 1024 * 1024]; + assert!( + oom_guard::current_balance() > ceiling, + "allocation should push balance over ceiling" + ); + + let result = pool.try_grow(&reservation, 1); + assert!( + result.is_err(), + "real usage over the ceiling should reject the grow" + ); + // Keep `held` alive until after the assertion above. + drop(held); + } + + #[test] + fn fair_share_limit_uses_active_count_when_positive() { + // active count wins over the fallback divisor + assert_eq!(fair_share_limit(1000, 4, 8), 250); + } + + #[test] + fn fair_share_limit_falls_back_when_no_active_tasks() { + assert_eq!(fair_share_limit(1000, 0, 5), 200); + } + + #[test] + fn fair_share_limit_floors_divisor_at_one() { + // active and fallback both zero -> divide by 1, no panic + assert_eq!(fair_share_limit(1000, 0, 0), 1000); + } + + #[test] + fn fair_share_limit_zero_when_ceiling_below_n() { + assert_eq!(fair_share_limit(3, 4, 8), 0); + } + + #[test] + fn should_reject_none_is_first_come() { + assert!(should_reject_over_ceiling(0, 1, None)); + assert!(should_reject_over_ceiling(1000, 0, None)); + } + + #[test] + fn should_reject_some_only_above_share() { + // strictly above share -> reject + assert!(should_reject_over_ceiling(400, 200, Some(500))); + // exactly at share -> allow + assert!(!should_reject_over_ceiling(300, 200, Some(500))); + // below share -> allow + assert!(!should_reject_over_ceiling(100, 100, Some(500))); + } + + #[test] + fn over_ceiling_rejects_task_over_fair_share() { + // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + Some(2), + BalanceSource::Fixed(900), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // Put this task above its 500-byte fair share. + inner.grow(&reservation, 600); + // Over ceiling (900 + 200 > 1000) AND over fair share (600 + 200 > 500) -> reject. + assert!(pool.try_grow(&reservation, 200).is_err()); + } + + #[test] + fn over_ceiling_spares_task_under_fair_share() { + // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool: Arc = Arc::new(RealUsagePool::with_balance_source( + Arc::clone(&inner), + 1000, + Some(2), + BalanceSource::Fixed(1000), + )); + let reservation = MemoryConsumer::new("test").register(&pool); + // This task holds only 100, under its 500 fair share. + inner.grow(&reservation, 100); + // Over ceiling (1000 + 50 > 1000) but under fair share (100 + 50 <= 500) -> allowed. + assert!(pool.try_grow(&reservation, 50).is_ok()); + // The grow was delegated to the inner pool. + assert_eq!(inner.reserved(), 150); + } +} diff --git a/native/core/src/execution/memory_pools/task_shared.rs b/native/core/src/execution/memory_pools/task_shared.rs index b5b4da61f9f..eaedfa180ef 100644 --- a/native/core/src/execution/memory_pools/task_shared.rs +++ b/native/core/src/execution/memory_pools/task_shared.rs @@ -30,6 +30,15 @@ use std::sync::{Arc, Weak}; static TASK_SHARED_MEMORY_POOLS: Lazy>>> = Lazy::new(|| Mutex::new(HashMap::new())); +/// Number of distinct task-attempt ids with a live task-shared memory pool, derived from +/// the registry so there is no separate counter to keep in sync. The real-usage fair-share +/// guard uses this as the divisor for each task's share of the budget; it returns 0 when no +/// task-shared pool is active, in which case the guard falls back to a fixed divisor. +#[cfg_attr(not(feature = "oom-guard"), allow(dead_code))] +pub(crate) fn active_task_count() -> usize { + TASK_SHARED_MEMORY_POOLS.lock().len() +} + /// A transparent `MemoryPool` wrapper whose lifetime also controls its registry entry. #[derive(Debug)] struct TaskSharedMemoryPool { diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 55da2c733aa..2067e53d6af 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -25,7 +25,7 @@ pub mod operators; pub(crate) mod planner; pub mod serde; pub use datafusion_comet_shuffle as shuffle; -mod memory_pools; +pub(crate) mod memory_pools; pub(crate) mod sort; pub(crate) mod spark_config; pub(crate) mod spark_plan; diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5de..fb27d48ef90 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -25,6 +25,10 @@ pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +#[cfg(feature = "oom-guard")] +pub(crate) const COMET_MEMORY_GUARD_ENABLED: &str = "spark.comet.exec.memoryGuard.enabled"; +#[cfg(feature = "oom-guard")] +pub(crate) const COMET_MEMORY_GUARD_SIZE: &str = "spark.comet.exec.memoryGuard.size"; pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index b5656dba102..02896d472c2 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -75,18 +75,49 @@ pub mod debug; #[cfg(all( not(target_env = "msvc"), feature = "jemalloc", - not(feature = "mimalloc") + not(feature = "mimalloc"), + not(feature = "oom-guard") ))] #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; #[cfg(all( feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")) + not(all(not(target_env = "msvc"), feature = "jemalloc")), + not(feature = "oom-guard") ))] #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +#[cfg(all( + not(target_env = "msvc"), + feature = "jemalloc", + not(feature = "mimalloc"), + feature = "oom-guard" +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(Jemalloc); + +#[cfg(all( + feature = "mimalloc", + not(all(not(target_env = "msvc"), feature = "jemalloc")), + feature = "oom-guard" +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(MiMalloc); + +// oom-guard enabled with system allocator (no mimalloc, and no jemalloc or on MSVC). +#[cfg(all( + feature = "oom-guard", + not(feature = "mimalloc"), + any(target_env = "msvc", not(feature = "jemalloc")) +))] +#[global_allocator] +static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = + crate::execution::memory_pools::oom_guard::AccountingAllocator::new(std::alloc::System); + #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init( e: EnvUnowned, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 43f030a7d9d..fd80aa0df35 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -824,7 +824,12 @@ object CometConf extends ShimCometConf { .category(CATEGORY_TUNING) .doc( "The type of memory pool to be used for Comet native execution when running Spark in " + - "off-heap mode. Available pool types are `greedy_unified` and `fair_unified`. " + + "off-heap mode. Available pool types are `greedy_unified`, `fair_unified`, and " + + "`real_usage`. The experimental `real_usage` pool gates growth on real allocator " + + "usage against the off-heap budget rather than delegating per-task accounting to " + + "Spark, and arms the last-resort OOM breaker on its own, so it needs no separate " + + "`spark.comet.exec.memoryGuard.enabled`. It relies on the `oom-guard` native " + + "feature, which is enabled by default. " + s"$TUNING_GUIDE.") .stringConf .createWithDefault("fair_unified") @@ -1001,6 +1006,27 @@ object CometConf extends ShimCometConf { .bytesConf(ByteUnit.BYTE) .createWithDefault(100L * 1024 * 1024 * 1024) // 100 GB + val COMET_EXEC_MEMORY_GUARD_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.memoryGuard.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental. When enabled, Comet tracks real native memory allocations and aborts " + + "an over-budget task with a retriable error instead of risking an executor-wide OOM " + + "kill. The `real_usage` memory pool arms this automatically, so this flag is only " + + "needed to add the guard on top of another pool type. Uses the 'oom-guard' native " + + "feature, which is enabled by default. Has no effect if that feature is compiled out.") + .booleanConf + .createWithDefault(false) + + val COMET_EXEC_MEMORY_GUARD_SIZE: OptionalConfigEntry[Long] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.memoryGuard.size") + .category(CATEGORY_EXEC) + .doc( + "Experimental. Memory budget for the Comet native OOM guard (accepts sizes like '4g'). " + + "Defaults to the executor off-heap memory size (spark.memory.offHeap.size) when unset.") + .bytesConf(ByteUnit.BYTE) + .createOptional + val COMET_RESPECT_DATAFUSION_CONFIGS: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.respectDataFusionConfigs") .category(CATEGORY_TESTING) From 1469019d8f97b855338595d690d12b33aeb1116e Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 9 Sep 2026 10:06:21 -0700 Subject: [PATCH 2/2] refactor: consolidate memory guard configs and simplify the OOM guard Collapse the two ways of enabling the real-usage gate into one. The `real_usage` memory pool type and `spark.comet.exec.memoryGuard.enabled` both armed the breaker and built a `RealUsageMemoryPool` ceilinged at the off-heap budget, differing only in the inner pool and in a `fair_share` value already derivable from the pool type. Drop the pool type, make `unbounded` selectable in off-heap mode, and have `memoryGuard.enabled` force it there. Layering the gate over a unified pool left both ceilings equal to the off-heap budget, so Spark's per-task accounting always rejected first and the gate never fired. Fix a hook that the rebase onto main silently clobbered. Tokio's `on_thread_start` is a setter rather than a list, so the `attach_thread_as_daemon` hook added upstream replaced the guard's `stamp_current_thread`. No worker thread was ever stamped, which left the breaker unable to fire on the threads doing the allocating. Both now run from a single `on_worker_thread_start`. Resolve `spark.comet.exec.memoryGuard.size` to a byte count before it crosses JNI. `bytesConf` only converts on read through the `ConfigEntry`, so the native side saw the raw `4g` string, failed to parse it, and silently fell back to the off-heap budget. Cleanups: five `#[global_allocator]` blocks collapse to two, the duplicate native reader for `spark.comet.exec.memoryPool` is gone, the gate skips `reserved()` and the task registry lock when no fair share is configured, `serializeCometSQLConfs` makes one pass with one `SQLConf.get`, `RealUsagePool` is renamed to `RealUsageMemoryPool` to match its siblings, and six repeated test fixtures fold into one. --- docs/source/user-guide/latest/tuning.md | 9 +- native/core/Cargo.toml | 4 +- native/core/src/execution/jni_api.rs | 93 ++++--- .../core/src/execution/memory_pools/config.rs | 36 +-- native/core/src/execution/memory_pools/mod.rs | 19 +- .../src/execution/memory_pools/oom_guard.rs | 49 ++-- .../execution/memory_pools/real_usage_pool.rs | 239 +++++------------- .../src/execution/memory_pools/task_shared.rs | 2 +- native/core/src/lib.rs | 71 ++---- .../scala/org/apache/comet/CometConf.scala | 20 +- .../org/apache/comet/CometExecIterator.scala | 28 +- 11 files changed, 225 insertions(+), 345 deletions(-) diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index ceec8b1a15e..3afff315e5e 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -60,8 +60,9 @@ The valid pool types are: - `fair_unified` (default when `spark.memory.offHeap.enabled=true` is set) - `greedy_unified` +- `unbounded` -Both pool types are shared across all native execution contexts within the same Spark task. When +The two `unified` pool types are shared across all native execution contexts within the same Spark task. When Comet executes a shuffle, it runs two native execution contexts concurrently (e.g. one for pre-shuffle operators and one for the shuffle writer). The shared pool ensures that the combined memory usage stays within the per-task limit. @@ -74,6 +75,12 @@ when there is sufficient memory in order to leave enough memory for other operat The `greedy_unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not need to spill or have a single spillable operator. +The `unbounded` pool does no accounting of its own and imposes no limit, so Comet's native memory is capped only by +the experimental `spark.comet.exec.memoryGuard.enabled`. Enabling that guard in off-heap mode selects `unbounded` +automatically and ignores this setting: growth is then gated on real allocator usage against the off-heap budget +rather than on Spark's per-task reservations. Note that Spark's `TaskMemoryManager` no longer sees Comet's off-heap +usage in that mode, so Spark cannot ask Comet to spill on behalf of its own operators. + [shuffle]: #shuffle [Advanced Memory Tuning]: #advanced-memory-tuning diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 8b1ebf14bd6..345bfb81e02 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -109,8 +109,8 @@ jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] # Allocator-level OOM circuit breaker. Wraps the global allocator to track real # allocated bytes and gate/abort over-budget query-worker threads. Enabled by default -# so `spark.comet.exec.memoryGuard.*` and the `real_usage` memory pool work without a -# special build; an idle guard is near-free (tracking stays off until a task arms it). +# so `spark.comet.exec.memoryGuard.*` works without a special build; an idle guard is +# near-free (tracking stays off until a task arms it). # Drop it from `default` for a bare allocator. oom-guard = [] # Delta Lake integration. When enabled, links the `comet-contrib-delta` crate diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 9603eae4f60..bbe64b022d5 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -110,7 +110,9 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; #[cfg(feature = "oom-guard")] -use crate::execution::memory_pools::{oom_guard, MemoryPoolType, RealUsagePool}; +use crate::execution::memory_pools::{ + oom_guard, MemoryPoolConfig, MemoryPoolType, RealUsageMemoryPool, +}; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, @@ -228,8 +230,6 @@ fn parse_usize_env_var(name: &str) -> Option { fn build_runtime(default_worker_threads: Option) -> Runtime { let mut builder = tokio::runtime::Builder::new_multi_thread(); - #[cfg(feature = "oom-guard")] - builder.on_thread_start(oom_guard::stamp_current_thread); if let Some(n) = parse_usize_env_var("COMET_WORKER_THREADS") { info!("Comet tokio runtime: using COMET_WORKER_THREADS={n}"); builder.worker_threads(n); @@ -244,12 +244,23 @@ fn build_runtime(default_worker_threads: Option) -> Runtime { } builder .enable_all() - .on_thread_start(attach_thread_as_daemon) + .on_thread_start(on_worker_thread_start) .on_thread_stop(detach_thread) .build() .expect("Failed to create Tokio runtime") } +/// Everything that has to run on a freshly spawned runtime thread. +/// +/// Tokio's `on_thread_start` is a setter, not a list: a second call silently replaces the +/// first. Register this one hook and add to it rather than calling `on_thread_start` again. +fn on_worker_thread_start() { + // Marks the thread as a query worker, which is what makes it eligible for the breaker. + #[cfg(feature = "oom-guard")] + oom_guard::stamp_current_thread(); + attach_thread_as_daemon(); +} + /// Attaches a runtime thread to the JVM as a daemon thread. /// /// jni-rs attaches threads lazily with `AttachCurrentThread`, which makes them non-daemon JVM @@ -517,34 +528,34 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( let memory_pool_type = memory_pool_type.try_to_string(env)?; let memory_pool_config = parse_memory_pool_config( off_heap_mode != JNI_FALSE, - memory_pool_type, + &memory_pool_type, memory_limit, memory_limit_per_task, )?; - // Arm the hard breaker when the guard is enabled or the `real_usage` pool is - // selected (it carries the guard itself). It trips on *actual* over-budget - // usage; the cooperative gate below trips on *projected* usage and spills - // first. `spark.comet.exec.memoryGuard.size` gives the breaker headroom above - // the off-heap budget (e.g. up to the container RSS limit). + // In off-heap mode the guard *replaces* Spark's per-task accounting rather than + // layering over it. Both the gate ceiling and the unified pools' budget are the + // off-heap size, so an inner unified pool would always reject first and the gate + // would never fire. #[cfg(feature = "oom-guard")] - let (guard_enabled, is_real_usage) = ( - spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED), - memory_pool_config.pool_type == MemoryPoolType::RealUsage, - ); + let guard_enabled = spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED); #[cfg(feature = "oom-guard")] - if guard_enabled || is_real_usage { - let default_limit = memory_limit.max(0) as u64; - let limit = spark_config.get_u64(COMET_MEMORY_GUARD_SIZE, default_limit); - if limit == 0 { - warn!( - "Comet memory guard is active but the effective limit is 0 \ - (memory_limit={memory_limit}); the guard will not trip. Set \ - spark.comet.exec.memoryGuard.size explicitly." - ); + let memory_pool_config = if guard_enabled && off_heap_mode != JNI_FALSE { + if memory_pool_type != "unbounded" { + // Once per executor: this runs on every plan creation. + static WARNED: std::sync::Once = std::sync::Once::new(); + WARNED.call_once(|| { + warn!( + "{COMET_MEMORY_GUARD_ENABLED}=true overrides \ + spark.comet.exec.memoryPool={memory_pool_type}; using `unbounded` \ + so that real allocator usage is the sole limit." + ) + }); } - oom_guard::arm(limit as usize); - } + MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0) + } else { + memory_pool_config + }; let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); @@ -556,22 +567,38 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( ThreadMemoryPoolRegistration::new(rust_thread_id, id, Arc::clone(&memory_pool)) }); - // Cooperative real-usage gate: reject growth (triggering a spill) once real - // allocator usage plus the request would exceed the off-heap budget. This is the - // first line of defense and fires before the hard breaker armed above, so - // over-budget work spills and retries rather than failing the task. The dedicated - // `real_usage` pool already gates internally, so it is not wrapped again. + // Two layers of defense against a native OOM kill, both driven by + // `spark.comet.exec.memoryGuard.enabled`: + // + // 1. The cooperative gate (`RealUsageMemoryPool`) rejects growth once *projected* + // real allocator usage would exceed the off-heap budget, so over-budget work + // spills and retries rather than failing the task. + // 2. The hard breaker (`oom_guard`) is the last resort and trips on *actual* + // over-budget usage. `spark.comet.exec.memoryGuard.size` gives it headroom + // above the off-heap budget (e.g. up to the container RSS limit). + // + // In off-heap mode the pool underneath is always `unbounded` (forced above), so + // the gate is the only thing rejecting growth. #[cfg(feature = "oom-guard")] - let memory_pool = if guard_enabled && !is_real_usage { + let memory_pool = if guard_enabled { let ceiling = memory_limit.max(0) as usize; + let limit = spark_config.get_u64(COMET_MEMORY_GUARD_SIZE, ceiling as u64); + if limit == 0 { + warn!( + "Comet memory guard is active but the effective limit is 0 \ + (memory_limit={memory_limit}); the guard will not trip. Set \ + spark.comet.exec.memoryGuard.size explicitly." + ); + } + oom_guard::arm(limit as usize); // Enable the fair-share guard for pools whose `reserved()` is per-task; // `executor_cores` is the fallback divisor when no task count is known. let fair_share = memory_pool_config .pool_type .has_per_task_budget() .then_some(executor_cores); - Arc::new(RealUsagePool::new(memory_pool, ceiling, fair_share)) - as Arc + Arc::new(RealUsageMemoryPool::new(memory_pool, ceiling, fair_share)) + as Arc } else { memory_pool }; diff --git a/native/core/src/execution/memory_pools/config.rs b/native/core/src/execution/memory_pools/config.rs index e5888609ec5..dcede933100 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -28,23 +28,15 @@ pub(crate) enum MemoryPoolType { GreedyGlobal, FairSpillGlobal, Unbounded, - #[cfg(feature = "oom-guard")] - RealUsage, } +#[cfg(feature = "oom-guard")] impl MemoryPoolType { /// True when this pool's `reserved()` reflects a single task's usage, so a per-task /// fair-share comparison is meaningful (false for process-wide pools). The non-shared /// per-task pools (`Greedy`/`FairSpill`) return true but keep no task registry, so the /// fair-share divisor falls back to `executor_cores` rather than the active-task count. - #[cfg_attr(not(feature = "oom-guard"), allow(dead_code))] pub(crate) fn has_per_task_budget(&self) -> bool { - // The dedicated `real_usage` pool gates on process-wide real usage - // (first-come), not a per-task reservation, so it has no per-task budget. - #[cfg(feature = "oom-guard")] - if matches!(self, MemoryPoolType::RealUsage) { - return false; - } !matches!( self, MemoryPoolType::GreedyGlobal @@ -70,13 +62,13 @@ impl MemoryPoolConfig { pub(crate) fn parse_memory_pool_config( off_heap_mode: bool, - memory_pool_type: String, + memory_pool_type: &str, memory_limit: i64, memory_limit_per_task: i64, ) -> CometResult { let pool_size = memory_limit as usize; let memory_pool_config = if off_heap_mode { - match memory_pool_type.as_str() { + match memory_pool_type { "fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size), "greedy_unified" => { // the `unified` memory pool interacts with Spark's memory pool to allocate @@ -84,20 +76,12 @@ pub(crate) fn parse_memory_pool_config( // shared with Spark is set by `spark.memory.offHeap.size`. MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) } - #[cfg(feature = "oom-guard")] - "real_usage" => { - // Gate growth on real allocator usage against the off-heap budget - // (`pool_size`) instead of delegating per-task accounting to Spark's - // TaskMemoryManager. See `RealUsagePool`. - MemoryPoolConfig::new(MemoryPoolType::RealUsage, pool_size) - } - #[cfg(not(feature = "oom-guard"))] - "real_usage" => { - return Err(CometError::Config( - "Memory pool type 'real_usage' requires a Comet build with the \ - 'oom-guard' native feature" - .to_string(), - )) + "unbounded" => { + // No accounting of its own. In off-heap mode this is what + // `spark.comet.exec.memoryGuard.enabled` forces, so the real-usage gate + // wrapped around it is the only thing rejecting growth, instead of + // delegating per-task accounting to Spark's TaskMemoryManager. + MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0) } _ => { return Err(CometError::Config(format!( @@ -108,7 +92,7 @@ pub(crate) fn parse_memory_pool_config( } else { // Use the memory pool from DF let pool_size_per_task = memory_limit_per_task as usize; - match memory_pool_type.as_str() { + match memory_pool_type { "fair_spill_task_shared" => { MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task) } diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 3ee9c813068..137787497a3 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -19,7 +19,7 @@ mod config; mod fair_pool; pub mod logging_pool; #[cfg(feature = "oom-guard")] -pub mod oom_guard; +pub(crate) mod oom_guard; #[cfg(feature = "oom-guard")] mod real_usage_pool; mod task_shared; @@ -37,7 +37,7 @@ use unified_pool::CometUnifiedMemoryPool; pub(crate) use config::*; #[cfg(feature = "oom-guard")] -pub(crate) use real_usage_pool::RealUsagePool; +pub(crate) use real_usage_pool::RealUsageMemoryPool; pub(crate) use task_shared::*; /// Creates the memory pool for a native plan. @@ -95,20 +95,5 @@ pub(crate) fn create_memory_pool( Arc::clone(memory_pool) } MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()), - #[cfg(feature = "oom-guard")] - MemoryPoolType::RealUsage => { - // Dedicated off-heap pool: `RealUsagePool` is the sole gate, comparing - // process-wide real usage against `pool_size` (first-come across tasks, so - // `fair_share` is `None`) instead of Spark's per-task TaskMemoryManager - // division. The inner `UnboundedMemoryPool` never rejects; `TrackConsumersPool` - // still reports top consumers on rejection. `enable_tracking()` because the - // gate reads the allocator balance even when the hard breaker is unarmed. - oom_guard::enable_tracking(); - tracked(RealUsagePool::new( - Arc::new(UnboundedMemoryPool::default()), - pool_size, - None, - )) - } } } diff --git a/native/core/src/execution/memory_pools/oom_guard.rs b/native/core/src/execution/memory_pools/oom_guard.rs index eec4776311a..a74f7a8c216 100644 --- a/native/core/src/execution/memory_pools/oom_guard.rs +++ b/native/core/src/execution/memory_pools/oom_guard.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use datafusion::common::DataFusionError; +use datafusion::common::{resources_datafusion_err, DataFusionError}; use std::alloc::{GlobalAlloc, Layout}; use std::cell::Cell; use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering}; @@ -64,10 +64,9 @@ fn disarm() { ARMED.store(false, Ordering::Relaxed); } -/// Turn on real-usage balance tracking. Called when the guard is armed or the -/// `real_usage` memory pool is created, so the process-wide balance is live for -/// the cooperative gate even when the hard breaker is not armed. Idempotent, and -/// tracking is never turned back off in production. +/// Turn on real-usage balance tracking. Called when the guard is armed, so the +/// process-wide balance is live for the cooperative gate. Idempotent, and tracking +/// is never turned back off in production. pub fn enable_tracking() { TRACKING_ENABLED.store(true, Ordering::Relaxed); } @@ -90,11 +89,12 @@ fn clear_unwinding() { pub fn map_panic_to_error(panic: &(dyn std::any::Any + Send)) -> Option { let g = panic.downcast_ref::()?; clear_unwinding(); - Some(DataFusionError::ResourcesExhausted(format!( + Some(resources_datafusion_err!( "Comet OomGuard: native allocation pushed usage to {} bytes, over the limit of {} \ bytes; failing this task", - g.balance, g.limit - ))) + g.balance, + g.limit + )) } /// Handle a panic caught by `catch_unwind` on a JNI caller thread. If it is an @@ -119,25 +119,12 @@ pub fn current_balance() -> usize { BALANCE.load(Ordering::Relaxed).max(0) as usize } -/// Record an allocation of `size` bytes; may trip the breaker. -#[inline] -fn record_alloc(size: usize) { - track(size as isize); -} - -/// Record a deallocation of `size` bytes; never trips (credit only). -#[inline] -fn record_dealloc(size: usize) { - track(-(size as isize)); -} - /// Core tracking + enforcement. Flushes drift; on a debit flush that crosses the /// limit on an armed, stamped, non-unwinding thread, panics with `OomGuardPanic`. #[inline] fn track(delta: isize) { // Runtime gate: skip all balance bookkeeping until a task enables tracking. - // Keeps the always-linked accounting allocator near-free when the guard and - // the `real_usage` pool are both unused. + // Keeps the always-linked accounting allocator near-free when the guard is unused. if !TRACKING_ENABLED.load(Ordering::Relaxed) { return; } @@ -165,7 +152,7 @@ fn track(delta: isize) { if should_trip(balance, limit) { // At most one thread may fire the guard panic per arm cycle. CAS the // master gate true->false; threads that lose the race bail before - // panic_any. The relaxed load above (line ~121) is not a serialization + // panic_any. The relaxed `ARMED` load above is not a serialization // point: several threads can all read ARMED=true and reach here in the // same tight window. If each then dispatches a panic, Rust's unwind ABI // can abort the process with "failed to initiate panic" instead of @@ -226,20 +213,20 @@ unsafe impl GlobalAlloc for AccountingAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { let ptr = self.inner.alloc(layout); if !ptr.is_null() { - record_alloc(layout.size()); + track(layout.size() as isize); } ptr } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { self.inner.dealloc(ptr, layout); - record_dealloc(layout.size()); + track(-(layout.size() as isize)); } unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { let ptr = self.inner.alloc_zeroed(layout); if !ptr.is_null() { - record_alloc(layout.size()); + track(layout.size() as isize); } ptr } @@ -331,9 +318,9 @@ mod tests { let _g = GUARD.lock().unwrap_or_else(|e| e.into_inner()); reset_for_test(); stamp_current_thread(); - // not armed -> record_alloc must never panic regardless of size - record_alloc(usize::MAX / 2); - record_alloc(usize::MAX / 2); + // not armed -> track must never panic regardless of size + track((usize::MAX / 2) as isize); + track((usize::MAX / 2) as isize); } #[test] @@ -343,7 +330,7 @@ mod tests { // arm with a tiny limit relative to current balance, but DO NOT stamp let limit = current_balance() + 1; arm(limit); - record_alloc(SETTLE_THRESHOLD as usize * 4); // big enough to flush + track(SETTLE_THRESHOLD * 4); // big enough to flush disarm(); } @@ -356,7 +343,7 @@ mod tests { arm(limit); let result = std::panic::catch_unwind(|| { // exceed the headroom in one flush - record_alloc(SETTLE_THRESHOLD as usize * 4); + track(SETTLE_THRESHOLD * 4); }); disarm(); clear_unwinding(); diff --git a/native/core/src/execution/memory_pools/real_usage_pool.rs b/native/core/src/execution/memory_pools/real_usage_pool.rs index 1e9f5c45fa0..f695cc75e58 100644 --- a/native/core/src/execution/memory_pools/real_usage_pool.rs +++ b/native/core/src/execution/memory_pools/real_usage_pool.rs @@ -24,6 +24,7 @@ use std::sync::Arc; /// Source of the current process-wide real allocator usage in bytes. Production reads the /// live `oom_guard` balance; tests inject a fixed value without touching global state. +#[derive(Debug)] enum BalanceSource { Live, #[cfg(test)] @@ -45,7 +46,8 @@ impl BalanceSource { /// accounting, rejects growth when *real* allocator usage (untracked Arrow / join / /// kernel bytes included) plus the requested amount would exceed a process-global /// ceiling. Returning `ResourcesExhausted` lets DataFusion spill and retry. -pub(crate) struct RealUsagePool { +#[derive(Debug)] +pub(crate) struct RealUsageMemoryPool { inner: Arc, /// Process-global real-usage ceiling in bytes; 0 means unset (no gating). ceiling: usize, @@ -56,27 +58,17 @@ pub(crate) struct RealUsagePool { balance_source: BalanceSource, } -impl std::fmt::Debug for RealUsagePool { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("RealUsagePool") - .field("inner", &self.inner) - .field("ceiling", &self.ceiling) - .field("fair_share", &self.fair_share) - .finish_non_exhaustive() - } -} - -impl std::fmt::Display for RealUsagePool { +impl std::fmt::Display for RealUsageMemoryPool { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "RealUsagePool(ceiling={}, inner={})", + "RealUsageMemoryPool(ceiling={}, inner={})", self.ceiling, self.inner ) } } -impl RealUsagePool { +impl RealUsageMemoryPool { /// Wrap `inner` with the real-usage gate using the live OomGuard balance. pub(crate) fn new( inner: Arc, @@ -90,22 +82,6 @@ impl RealUsagePool { balance_source: BalanceSource::Live, } } - - /// Wrap `inner` with an explicit balance source (test seam). - #[cfg(test)] - fn with_balance_source( - inner: Arc, - ceiling: usize, - fair_share: Option, - balance_source: BalanceSource, - ) -> Self { - Self { - inner, - ceiling, - fair_share, - balance_source, - } - } } /// Per-task fair share of `ceiling` given the number of concurrently active @@ -120,21 +96,9 @@ fn fair_share_limit(ceiling: usize, active_tasks: usize, cores_fallback: usize) ceiling / n.max(1) } -/// Given the process is already over the real-usage ceiling, decide whether to -/// reject this task's grow. `None` is first-come (reject whoever hit the ceiling); -/// `Some(s)` rejects only a task whose tracked reservation would exceed its fair -/// share `s`, sparing under-share tasks (the OomGuard breaker backstops runaway -/// cases). -fn should_reject_over_ceiling(reserved: usize, additional: usize, share: Option) -> bool { - match share { - None => true, - Some(s) => reserved.saturating_add(additional) > s, - } -} - -impl MemoryPool for RealUsagePool { +impl MemoryPool for RealUsageMemoryPool { fn name(&self) -> &str { - "RealUsagePool" + "RealUsageMemoryPool" } fn register(&self, consumer: &MemoryConsumer) { @@ -159,18 +123,20 @@ impl MemoryPool for RealUsagePool { additional: usize, ) -> Result<(), DataFusionError> { // Check the real-usage ceiling before delegating, so an over-budget request is - // rejected without speculatively reserving the inner pool. When the process is - // over the ceiling, the fair-share guard rejects only a task whose own tracked - // reservation exceeds its fair share, sparing innocent small tasks; the OomGuard - // breaker backstops runaway cases. Returning `ResourcesExhausted` lets DataFusion - // spill and retry. + // rejected without speculatively reserving the inner pool. if self.ceiling != 0 && additional != 0 { let real = self.balance_source.current(); if real.saturating_add(additional) > self.ceiling { - let share = self - .fair_share - .map(|cores| fair_share_limit(self.ceiling, active_task_count(), cores)); - if should_reject_over_ceiling(self.inner.reserved(), additional, share) { + // `None` is first-come. `Some` spares a task still under its fair share so + // one runaway task cannot starve small ones; the breaker backstops it. + let reject = match self.fair_share { + None => true, + Some(cores) => { + let share = fair_share_limit(self.ceiling, active_task_count(), cores); + self.inner.reserved().saturating_add(additional) > share + } + }; + if reject { return Err(resources_datafusion_err!( "Comet real-usage gate: native usage {real} bytes + requested \ {additional} bytes exceeds the off-heap budget of {} bytes; \ @@ -197,66 +163,69 @@ mod tests { use super::*; use datafusion::execution::memory_pool::{GreedyMemoryPool, UnboundedMemoryPool}; - #[test] - fn under_ceiling_succeeds_and_delegates() { + /// Pool with an injected real-usage balance, returned alongside its inner pool and a + /// registered reservation so tests can assert on what was delegated. + fn fixed_balance_pool( + ceiling: usize, + fair_share: Option, + real: usize, + ) -> (Arc, Arc, MemoryReservation) { let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 1000, - None, - BalanceSource::Fixed(100), - )); + let pool: Arc = Arc::new(RealUsageMemoryPool { + inner: Arc::clone(&inner), + ceiling, + fair_share, + balance_source: BalanceSource::Fixed(real), + }); let reservation = MemoryConsumer::new("test").register(&pool); + (inner, pool, reservation) + } + + #[test] + fn under_ceiling_delegates_grow_and_shrink() { // real usage 100 + request 100 = 200 <= ceiling 1000 + let (inner, pool, reservation) = fixed_balance_pool(1000, None, 100); assert!(pool.try_grow(&reservation, 100).is_ok()); assert_eq!(inner.reserved(), 100); + pool.shrink(&reservation, 40); + assert_eq!(inner.reserved(), 60); } #[test] fn over_ceiling_rejects_without_reserving_inner() { - let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 1000, - None, - BalanceSource::Fixed(900), - )); - let reservation = MemoryConsumer::new("test").register(&pool); - // real usage 900 + request 200 = 1100 > ceiling 1000 -> reject - let result = pool.try_grow(&reservation, 200); - assert!(result.is_err(), "over-ceiling grow should be rejected"); + // real usage 900 + request 200 = 1100 > ceiling 1000 + let (inner, pool, reservation) = fixed_balance_pool(1000, None, 900); + assert!(pool.try_grow(&reservation, 200).is_err()); // inner pool is never touched on rejection, so there is nothing to roll back assert_eq!(inner.reserved(), 0); } #[test] fn zero_ceiling_never_gates() { - let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 0, - None, - BalanceSource::Fixed(usize::MAX / 2), - )); - let reservation = MemoryConsumer::new("test").register(&pool); + let (inner, pool, reservation) = fixed_balance_pool(0, None, usize::MAX / 2); assert!(pool.try_grow(&reservation, 1024).is_ok()); assert_eq!(inner.reserved(), 1024); } + // The two fair-share cases below use ceiling 1000 and fallback divisor 2. No + // task-shared pool is registered in tests, so the active count is 0 and the divisor + // falls back to 2, giving a fair share of 500. + #[test] - fn shrink_delegates() { - let inner: Arc = Arc::new(UnboundedMemoryPool::default()); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 1_000_000, - None, - BalanceSource::Fixed(0), - )); - let reservation = MemoryConsumer::new("test").register(&pool); - pool.try_grow(&reservation, 500).unwrap(); - assert_eq!(pool.reserved(), 500); - pool.shrink(&reservation, 200); - assert_eq!(pool.reserved(), 300); + fn over_ceiling_rejects_task_over_fair_share() { + let (inner, pool, reservation) = fixed_balance_pool(1000, Some(2), 900); + inner.grow(&reservation, 600); + // over ceiling (900 + 200 > 1000) and over share (600 + 200 > 500) -> reject + assert!(pool.try_grow(&reservation, 200).is_err()); + } + + #[test] + fn over_ceiling_spares_task_at_or_under_fair_share() { + let (inner, pool, reservation) = fixed_balance_pool(1000, Some(2), 1000); + inner.grow(&reservation, 300); + // over ceiling, but exactly at the share boundary (300 + 200 == 500) -> allowed + assert!(pool.try_grow(&reservation, 200).is_ok()); + assert_eq!(inner.reserved(), 500); } // Drives a real heap allocation through the installed AccountingAllocator (only @@ -272,7 +241,7 @@ mod tests { // 4 MiB headroom over the (noisy) baseline. let ceiling = base + 4 * 1024 * 1024; let pool: Arc = - Arc::new(RealUsagePool::new(Arc::clone(&inner), ceiling, None)); + Arc::new(RealUsageMemoryPool::new(Arc::clone(&inner), ceiling, None)); let reservation = MemoryConsumer::new("test").register(&pool); // Push real usage ~8 MiB above the baseline, held alive across the check so the @@ -282,87 +251,19 @@ mod tests { oom_guard::current_balance() > ceiling, "allocation should push balance over ceiling" ); - - let result = pool.try_grow(&reservation, 1); assert!( - result.is_err(), + pool.try_grow(&reservation, 1).is_err(), "real usage over the ceiling should reject the grow" ); - // Keep `held` alive until after the assertion above. + // Keep `held` alive until after the assertions above. drop(held); } #[test] - fn fair_share_limit_uses_active_count_when_positive() { - // active count wins over the fallback divisor - assert_eq!(fair_share_limit(1000, 4, 8), 250); - } - - #[test] - fn fair_share_limit_falls_back_when_no_active_tasks() { - assert_eq!(fair_share_limit(1000, 0, 5), 200); - } - - #[test] - fn fair_share_limit_floors_divisor_at_one() { - // active and fallback both zero -> divide by 1, no panic - assert_eq!(fair_share_limit(1000, 0, 0), 1000); - } - - #[test] - fn fair_share_limit_zero_when_ceiling_below_n() { - assert_eq!(fair_share_limit(3, 4, 8), 0); - } - - #[test] - fn should_reject_none_is_first_come() { - assert!(should_reject_over_ceiling(0, 1, None)); - assert!(should_reject_over_ceiling(1000, 0, None)); - } - - #[test] - fn should_reject_some_only_above_share() { - // strictly above share -> reject - assert!(should_reject_over_ceiling(400, 200, Some(500))); - // exactly at share -> allow - assert!(!should_reject_over_ceiling(300, 200, Some(500))); - // below share -> allow - assert!(!should_reject_over_ceiling(100, 100, Some(500))); - } - - #[test] - fn over_ceiling_rejects_task_over_fair_share() { - // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 - let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 1000, - Some(2), - BalanceSource::Fixed(900), - )); - let reservation = MemoryConsumer::new("test").register(&pool); - // Put this task above its 500-byte fair share. - inner.grow(&reservation, 600); - // Over ceiling (900 + 200 > 1000) AND over fair share (600 + 200 > 500) -> reject. - assert!(pool.try_grow(&reservation, 200).is_err()); - } - - #[test] - fn over_ceiling_spares_task_under_fair_share() { - // ceiling 1000, fallback divisor 2, active count 0 in tests -> fair share 500 - let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let pool: Arc = Arc::new(RealUsagePool::with_balance_source( - Arc::clone(&inner), - 1000, - Some(2), - BalanceSource::Fixed(1000), - )); - let reservation = MemoryConsumer::new("test").register(&pool); - // This task holds only 100, under its 500 fair share. - inner.grow(&reservation, 100); - // Over ceiling (1000 + 50 > 1000) but under fair share (100 + 50 <= 500) -> allowed. - assert!(pool.try_grow(&reservation, 50).is_ok()); - // The grow was delegated to the inner pool. - assert_eq!(inner.reserved(), 150); + fn test_fair_share_limit() { + assert_eq!(fair_share_limit(1000, 4, 8), 250); // active count wins + assert_eq!(fair_share_limit(1000, 0, 5), 200); // falls back to cores + assert_eq!(fair_share_limit(1000, 0, 0), 1000); // divisor floored at 1 + assert_eq!(fair_share_limit(3, 4, 8), 0); // ceiling below the divisor } } diff --git a/native/core/src/execution/memory_pools/task_shared.rs b/native/core/src/execution/memory_pools/task_shared.rs index eaedfa180ef..47291658cf0 100644 --- a/native/core/src/execution/memory_pools/task_shared.rs +++ b/native/core/src/execution/memory_pools/task_shared.rs @@ -34,7 +34,7 @@ static TASK_SHARED_MEMORY_POOLS: Lazy usize { TASK_SHARED_MEMORY_POOLS.lock().len() } diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 02896d472c2..cee186e091e 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -41,18 +41,35 @@ use log4rs::{ Config, }; +// The allocator to install, named once so the `oom-guard` wrap below does not have to +// repeat the selection matrix. All three are unit structs, so the alias binds the value +// as well as the type. The final arm is `not(any(..))` of the other two, which also +// covers "both jemalloc and mimalloc enabled" -- neither is selected, as before. #[cfg(all( not(target_env = "msvc"), feature = "jemalloc", not(feature = "mimalloc") ))] -use tikv_jemallocator::Jemalloc; +use tikv_jemallocator::Jemalloc as InnerAllocator; #[cfg(all( feature = "mimalloc", not(all(not(target_env = "msvc"), feature = "jemalloc")) ))] -use mimalloc::MiMalloc; +use mimalloc::MiMalloc as InnerAllocator; + +#[cfg(not(any( + all( + not(target_env = "msvc"), + feature = "jemalloc", + not(feature = "mimalloc") + ), + all( + feature = "mimalloc", + not(all(not(target_env = "msvc"), feature = "jemalloc")) + ) +)))] +use std::alloc::System as InnerAllocator; // Re-export from jvm-bridge crate for internal use pub use datafusion_comet_jni_bridge::errors; @@ -65,6 +82,9 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; +#[cfg(feature = "oom-guard")] +use crate::execution::memory_pools::oom_guard; + pub mod cloud; pub mod execution; pub mod parquet; @@ -72,51 +92,14 @@ pub mod parquet; #[cfg(debug_assertions)] pub mod debug; -#[cfg(all( - not(target_env = "msvc"), - feature = "jemalloc", - not(feature = "mimalloc"), - not(feature = "oom-guard") -))] -#[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; - -#[cfg(all( - feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")), - not(feature = "oom-guard") -))] +#[cfg(feature = "oom-guard")] #[global_allocator] -static GLOBAL: MiMalloc = MiMalloc; +static GLOBAL: oom_guard::AccountingAllocator = + oom_guard::AccountingAllocator::new(InnerAllocator); -#[cfg(all( - not(target_env = "msvc"), - feature = "jemalloc", - not(feature = "mimalloc"), - feature = "oom-guard" -))] -#[global_allocator] -static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = - crate::execution::memory_pools::oom_guard::AccountingAllocator::new(Jemalloc); - -#[cfg(all( - feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")), - feature = "oom-guard" -))] -#[global_allocator] -static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = - crate::execution::memory_pools::oom_guard::AccountingAllocator::new(MiMalloc); - -// oom-guard enabled with system allocator (no mimalloc, and no jemalloc or on MSVC). -#[cfg(all( - feature = "oom-guard", - not(feature = "mimalloc"), - any(target_env = "msvc", not(feature = "jemalloc")) -))] +#[cfg(not(feature = "oom-guard"))] #[global_allocator] -static GLOBAL: crate::execution::memory_pools::oom_guard::AccountingAllocator = - crate::execution::memory_pools::oom_guard::AccountingAllocator::new(std::alloc::System); +static GLOBAL: InnerAllocator = InnerAllocator; #[no_mangle] pub extern "system" fn Java_org_apache_comet_NativeBase_init( diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index fd80aa0df35..9f611e4f55b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -825,11 +825,9 @@ object CometConf extends ShimCometConf { .doc( "The type of memory pool to be used for Comet native execution when running Spark in " + "off-heap mode. Available pool types are `greedy_unified`, `fair_unified`, and " + - "`real_usage`. The experimental `real_usage` pool gates growth on real allocator " + - "usage against the off-heap budget rather than delegating per-task accounting to " + - "Spark, and arms the last-resort OOM breaker on its own, so it needs no separate " + - "`spark.comet.exec.memoryGuard.enabled`. It relies on the `oom-guard` native " + - "feature, which is enabled by default. " + + "`unbounded`. `unbounded` does no accounting of its own, leaving Comet's native " + + "memory unlimited unless `spark.comet.exec.memoryGuard.enabled` is set. That " + + "setting overrides this one in off-heap mode. " + s"$TUNING_GUIDE.") .stringConf .createWithDefault("fair_unified") @@ -1010,11 +1008,13 @@ object CometConf extends ShimCometConf { conf(s"$COMET_EXEC_CONFIG_PREFIX.memoryGuard.enabled") .category(CATEGORY_EXEC) .doc( - "Experimental. When enabled, Comet tracks real native memory allocations and aborts " + - "an over-budget task with a retriable error instead of risking an executor-wide OOM " + - "kill. The `real_usage` memory pool arms this automatically, so this flag is only " + - "needed to add the guard on top of another pool type. Uses the 'oom-guard' native " + - "feature, which is enabled by default. Has no effect if that feature is compiled out.") + "Experimental. When enabled, Comet tracks real native memory allocations and gates " + + "growth against the off-heap budget, spilling rather than risking an executor-wide " + + "OOM kill, and aborts an over-budget task with a retriable error as a last resort. " + + "In off-heap mode this replaces Spark's per-task accounting rather than layering " + + "over it, so `spark.comet.exec.memoryPool` is ignored. Uses the 'oom-guard' " + + "native feature, which is enabled by default. Has no effect if that feature is " + + "compiled out.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e2c132904d5..dff7f5fce39 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -332,17 +332,16 @@ class CometExecIterator( object CometExecIterator extends Logging { - private def cometSqlConfs: Map[String, String] = - SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) - def serializeCometSQLConfs(): Array[Byte] = { val builder = ConfigMap.newBuilder() - cometSqlConfs.foreach { case (k, v) => - if (k.startsWith(s"${CometConf.COMET_PREFIX}.datafusion.")) { - if (CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.get(SQLConf.get)) { - builder.putEntries(k, v) - } - } else { + // Resolved once: on an executor `SQLConf.get` builds a fresh ReadOnlySQLConf when the + // thread-local is unset, so re-reading it per entry allocates for every config key. + val sqlConf = SQLConf.get + val datafusionPrefix = s"${CometConf.COMET_PREFIX}.datafusion." + val respectDatafusionConfs = CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.get(sqlConf) + sqlConf.getAllConfs.foreach { case (k, v) => + if (k.startsWith(CometConf.COMET_PREFIX) && + (respectDatafusionConfs || !k.startsWith(datafusionPrefix))) { builder.putEntries(k, v) } } @@ -352,11 +351,18 @@ object CometExecIterator extends Logging { builder.putEntries("spark.executor.cores", executorCores.toString) // Any Comet config that the native side reads must be added here manually. - // `cometSqlConfs` only carries values that were explicitly set, so defaults + // `getAllConfs` only carries values that were explicitly set, so defaults // from `createWithDefault(...)` would otherwise not cross JNI. builder.putEntries( CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, - CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(sqlConf).toString) + + // `getAllConfs` carries the raw string the user set, but the native side parses a + // plain integer, so resolve `4g`-style sizes to a byte count. Overwrites the raw entry. + CometConf.COMET_EXEC_MEMORY_GUARD_SIZE + .get(sqlConf) + .foreach(bytes => + builder.putEntries(CometConf.COMET_EXEC_MEMORY_GUARD_SIZE.key, bytes.toString)) builder.build().toByteArray }