From 0986e1c909146240e303a007e717c0ede7f33dc8 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 3 Sep 2026 12:30:46 -0700 Subject: [PATCH] feat: experiment with `RealUsagePool` --- native/core/Cargo.toml | 9 +- native/core/src/execution/jni_api.rs | 220 +++++++--- .../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, 1080 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 8dc8d73273f..d2ca5b5238c 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -98,9 +98,16 @@ datafusion-functions-nested = { version = "54.1.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 4d4094b1472..d2ae3f566f1 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -109,14 +109,20 @@ 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; +#[cfg(feature = "oom-guard")] +use log::warn; use std::sync::OnceLock; #[cfg(feature = "jemalloc")] use tikv_jemalloc_ctl::{epoch, stats}; @@ -224,6 +230,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); @@ -475,6 +483,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); @@ -485,6 +518,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 { @@ -841,6 +894,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); @@ -917,6 +972,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::() { @@ -941,76 +1003,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 6cfe33223f1..806a1ba91e0 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 a7d7b80db1f..a248916d9a6 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -780,7 +780,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") @@ -939,6 +944,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)