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 9ebb69b2f69..345bfb81e02 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.*` 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 # 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..bbe64b022d5 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -109,11 +109,17 @@ use crate::execution::tracing::{ }; use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; +#[cfg(feature = "oom-guard")] +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, 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}; @@ -238,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 @@ -511,10 +528,35 @@ 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, )?; + + // 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 = spark_config.get_bool(COMET_MEMORY_GUARD_ENABLED); + #[cfg(feature = "oom-guard")] + 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." + ) + }); + } + MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0) + } else { + memory_pool_config + }; + let memory_pool = create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); @@ -525,6 +567,42 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( ThreadMemoryPoolRegistration::new(rust_thread_id, id, Arc::clone(&memory_pool)) }); + // 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 { + 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(RealUsageMemoryPool::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 +961,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 +1039,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 +1070,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..dcede933100 100644 --- a/native/core/src/execution/memory_pools/config.rs +++ b/native/core/src/execution/memory_pools/config.rs @@ -30,6 +30,22 @@ pub(crate) enum MemoryPoolType { Unbounded, } +#[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. + pub(crate) fn has_per_task_budget(&self) -> bool { + !matches!( + self, + MemoryPoolType::GreedyGlobal + | MemoryPoolType::FairSpillGlobal + | MemoryPoolType::Unbounded + ) + } +} + pub(crate) struct MemoryPoolConfig { pub(crate) pool_type: MemoryPoolType, pub(crate) pool_size: usize, @@ -46,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 @@ -60,6 +76,13 @@ pub(crate) fn parse_memory_pool_config( // shared with Spark is set by `spark.memory.offHeap.size`. MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0) } + "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!( "Unsupported memory pool type for off-heap mode: {memory_pool_type}" @@ -69,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 d7c2911f913..137787497a3 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(crate) 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::RealUsageMemoryPool; pub(crate) use task_shared::*; /// Creates the memory pool for a native plan. 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..a74f7a8c216 --- /dev/null +++ b/native/core/src/execution/memory_pools/oom_guard.rs @@ -0,0 +1,396 @@ +// 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::{resources_datafusion_err, 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, 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); +} + +/// 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(resources_datafusion_err!( + "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 +} + +/// 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 is 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 `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 + // 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() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.inner.dealloc(ptr, layout); + 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() { + track(layout.size() as isize); + } + 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 -> track must never panic regardless of size + track((usize::MAX / 2) as isize); + track((usize::MAX / 2) as isize); + } + + #[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); + track(SETTLE_THRESHOLD * 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 + track(SETTLE_THRESHOLD * 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..f695cc75e58 --- /dev/null +++ b/native/core/src/execution/memory_pools/real_usage_pool.rs @@ -0,0 +1,269 @@ +// 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. +#[derive(Debug)] +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. +#[derive(Debug)] +pub(crate) struct RealUsageMemoryPool { + 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::Display for RealUsageMemoryPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "RealUsageMemoryPool(ceiling={}, inner={})", + self.ceiling, self.inner + ) + } +} + +impl RealUsageMemoryPool { + /// 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, + } + } +} + +/// 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) +} + +impl MemoryPool for RealUsageMemoryPool { + fn name(&self) -> &str { + "RealUsageMemoryPool" + } + + 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. + if self.ceiling != 0 && additional != 0 { + let real = self.balance_source.current(); + if real.saturating_add(additional) > self.ceiling { + // `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; \ + 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}; + + /// 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(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() { + // 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, 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 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 + // 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(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 + // 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" + ); + assert!( + pool.try_grow(&reservation, 1).is_err(), + "real usage over the ceiling should reject the grow" + ); + // Keep `held` alive until after the assertions above. + drop(held); + } + + #[test] + 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 b5b4da61f9f..47291658cf0 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(feature = "oom-guard")] +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..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,20 +92,14 @@ pub mod parquet; #[cfg(debug_assertions)] pub mod debug; -#[cfg(all( - not(target_env = "msvc"), - feature = "jemalloc", - not(feature = "mimalloc") -))] +#[cfg(feature = "oom-guard")] #[global_allocator] -static GLOBAL: Jemalloc = Jemalloc; +static GLOBAL: oom_guard::AccountingAllocator = + oom_guard::AccountingAllocator::new(InnerAllocator); -#[cfg(all( - feature = "mimalloc", - not(all(not(target_env = "msvc"), feature = "jemalloc")) -))] +#[cfg(not(feature = "oom-guard"))] #[global_allocator] -static GLOBAL: MiMalloc = MiMalloc; +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 43f030a7d9d..9f611e4f55b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -824,7 +824,10 @@ 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 " + + "`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") @@ -1001,6 +1004,29 @@ 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 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) + + 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) 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 }