You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/06/03 16:32:12 UTC
[PR] feat: prototype allocator-level OOM circuit breaker (OomGuard) [datafusion-comet]
andygrove opened a new pull request, #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582
## Which issue does this PR close?
Relates to #4576. This is an exploratory prototype of the RSS circuit breaker ("OomGuard") half of that issue, not a complete implementation, so it does not close it.
## Rationale for this change
Comet's memory accounting relies on voluntary `MemoryPool` reservations, which miss allocations made by Arrow buffers, join scratch space, and expression kernels. When real native memory exceeds the container limit, the OS/YARN/Kubernetes kills the entire executor JVM, losing every task and all cached data on it.
This prototype adds an executor-global, allocator-level circuit breaker. It tracks the real bytes the global allocator hands out and, when an armed, over-budget condition is detected on a query-worker thread, fails that single task with a retriable `ResourcesExhausted` error instead of letting the executor get OOM-killed. The approach adapts the `AccountingAllocator` from apache/datafusion#22626 (the byte-tracking allocator wrapper) rather than depending on it, since that code lives in DataFusion's test-only `sqllogictest` crate.
## What changes are included in this PR?
Gated behind a new `oom-guard` cargo feature; the default build is unchanged with zero added per-allocation overhead.
- `native/core/src/execution/memory_pools/oom_guard.rs` (new): `AccountingAllocator<A>` wrapping the inner global allocator; a single process-wide balance with per-thread drift settled at a 64 KiB threshold; `arm`/`disarm`/`stamp_current_thread`/`current_balance`; a typed `OomGuardPanic`, raised via `panic_any` on an armed, stamped thread that crosses the limit, with reentrancy protection so the panic's own boxing allocation does not recurse.
- `native/core/src/lib.rs`: under the `oom-guard` feature, installs the wrapper as `#[global_allocator]` over jemalloc / mimalloc / system; mutually exclusive cfgs leave the default build untouched.
- `native/core/src/execution/jni_api.rs`: stamps tokio worker threads (`on_thread_start`) and the JNI caller thread; arms the guard from config in `createPlan`; maps `OomGuardPanic` to `DataFusionError::ResourcesExhausted` at both execution boundaries (the spawned/channel path, both producer and consumer, and the busy-poll `block_on` path).
- `spark/src/main/scala/org/apache/comet/CometConf.scala`: registers `spark.comet.exec.memoryGuard.enabled` (default false) and `spark.comet.exec.memoryGuard.size` (optional; defaults to the executor off-heap size).
Known limitations / out of scope for this prototype (candidates for follow-ups):
- Executor-global granularity only; no per-task attribution or fairness.
- Only tokio workers and the JNI caller thread are stamped, so allocations on `spawn_blocking`/IO/other pools are tracked but cannot themselves trip the breaker.
- Layout-byte accounting only; no periodic resync to real jemalloc/mimalloc resident stats.
- Does not feed the live balance back into DataFusion's `MemoryPool` (the "online accounting" half of #4576).
- Process-global, last-writer-wins limit, armed per `createPlan`.
Note: commits carry `[skip ci]` intentionally while this is an early prototype.
## How are these changes tested?
- Rust unit tests in `oom_guard.rs` cover the decision/settle helpers, and that the breaker trips only on an armed, stamped thread (disarmed never trips, unstamped never trips).
- An end-to-end Rust test drives a real 64 MiB heap allocation through the installed `AccountingAllocator` and asserts an `OomGuardPanic` is raised and caught.
- Verified the build and `clippy -D warnings` across the default, `oom-guard`, and `jemalloc,oom-guard` feature combinations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM circuit breaker (OomGuard) [datafusion-comet]
Posted by "avantgardnerio (via GitHub)" <gi...@apache.org>.
avantgardnerio commented on code in PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#discussion_r3351290034
##########
native/core/src/execution/memory_pools/oom_guard.rs:
##########
@@ -0,0 +1,352 @@
+// 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);
+
+thread_local! {
+ /// Un-flushed per-thread delta.
+ static LOCAL_DRIFT: Cell<isize> = const { Cell::new(0) };
+ /// Is this a query-worker thread eligible for enforcement?
+ static STAMPED: Cell<bool> = const { Cell::new(false) };
+ /// Set while a guard panic is unwinding this thread, to avoid double-faults.
+ static UNWINDING: Cell<bool> = 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);
+ ARMED.store(true, Ordering::Relaxed);
+}
+
+/// Disarm the guard (enforcement off; tracking continues cheaply).
+#[allow(dead_code)] // used only by tests
+pub fn disarm() {
+ ARMED.store(false, 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.
+pub 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<DataFusionError> {
+ let g = panic.downcast_ref::<OomGuardPanic>()?;
+ 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
+ )))
+}
+
+/// Current process-wide balance in bytes (never reported negative).
+#[allow(dead_code)] // used only by tests
+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) {
+ 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) {
+ // panic_any boxes the payload, which re-enters this allocator and calls
+ // track() again. Set UNWINDING first so that re-entrant call short-circuits
+ // above, preventing infinite recursion / a double panic from inside alloc.
+ 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<isize> {
+ *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<A: GlobalAlloc> {
+ inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+ pub const fn new(inner: A) -> Self {
+ Self { inner }
+ }
+}
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+ 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 {
+ let new_ptr = self.inner.realloc(ptr, layout, new_size);
+ if !new_ptr.is_null() {
+ // 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);
Review Comment:
Just fixed this bug in DF. You need to panic before the realloc, otherwise the caller still has the old pointer and tries to free it on unwind and segfaults.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#discussion_r3604119540
##########
native/core/src/execution/memory_pools/oom_guard.rs:
##########
@@ -0,0 +1,352 @@
+// 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);
+
+thread_local! {
+ /// Un-flushed per-thread delta.
+ static LOCAL_DRIFT: Cell<isize> = const { Cell::new(0) };
+ /// Is this a query-worker thread eligible for enforcement?
+ static STAMPED: Cell<bool> = const { Cell::new(false) };
+ /// Set while a guard panic is unwinding this thread, to avoid double-faults.
+ static UNWINDING: Cell<bool> = 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);
+ ARMED.store(true, Ordering::Relaxed);
+}
+
+/// Disarm the guard (enforcement off; tracking continues cheaply).
+#[allow(dead_code)] // used only by tests
+pub fn disarm() {
+ ARMED.store(false, 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.
+pub 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<DataFusionError> {
+ let g = panic.downcast_ref::<OomGuardPanic>()?;
+ 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
+ )))
+}
+
+/// Current process-wide balance in bytes (never reported negative).
+#[allow(dead_code)] // used only by tests
+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) {
+ 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) {
+ // panic_any boxes the payload, which re-enters this allocator and calls
+ // track() again. Set UNWINDING first so that re-entrant call short-circuits
+ // above, preventing infinite recursion / a double panic from inside alloc.
+ 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<isize> {
+ *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<A: GlobalAlloc> {
+ inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+ pub const fn new(inner: A) -> Self {
+ Self { inner }
+ }
+}
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+ 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 {
+ let new_ptr = self.inner.realloc(ptr, layout, new_size);
+ if !new_ptr.is_null() {
+ // 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);
Review Comment:
Thanks for the pointer. The current OomGuard sequence panics through `std::alloc::handle_alloc_error` before returning the resized pointer, so the caller does not observe the new allocation and its unwind uses the original pointer. Let me pin that with a comment near the panic call so a future refactor cannot silently swap the order and re-introduce the double-free.
##########
native/core/src/execution/memory_pools/real_usage_pool.rs:
##########
@@ -0,0 +1,348 @@
+// 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
+/// wiring uses `oom_guard::current_balance`; tests inject a controllable value.
+type BalanceSource = Arc<dyn Fn() -> usize + Send + Sync>;
+
+/// 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<dyn MemoryPool>,
+ /// 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<usize>,
+ 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 RealUsagePool {
+ /// Wrap `inner` with the real-usage gate using the live OomGuard balance.
+ pub(crate) fn new(
+ inner: Arc<dyn MemoryPool>,
+ ceiling: usize,
+ fair_share: Option<usize>,
+ ) -> Self {
+ Self {
+ inner,
+ ceiling,
+ fair_share,
+ balance_source: Arc::new(oom_guard::current_balance),
+ }
+ }
+
+ /// Wrap `inner` with an explicit balance source (test seam).
+ #[cfg(test)]
+ fn with_balance_source(
+ inner: Arc<dyn MemoryPool>,
+ ceiling: usize,
+ fair_share: Option<usize>,
+ 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<usize>) -> bool {
+ match share {
+ None => true,
+ Some(s) => reserved.saturating_add(additional) > s,
+ }
+}
+
+impl MemoryPool for 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,
Review Comment:
`additional` is the extra headroom the pool grants above real jemalloc usage: on a `try_grow` call, the guard admits the request if `jemalloc_stats.allocated + requested <= configured_pool_limit + additional`. It exists so short-lived spikes (e.g. an operator's staging buffer inside a single `try_grow`) do not trip the breaker when the sustained usage would still fit. Set to 0 when you want a strict cap. I will add a doc comment on the field explaining it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "avantgardnerio (via GitHub)" <gi...@apache.org>.
avantgardnerio commented on PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#issuecomment-5330718698
I think there was talk of tracking which operator was using the memory in the pool? I don't know if anything ever came of it, so I think the most granularity you can have is a MemoryPool per query / context.
AFAIK, all you need to know about when to spill is when allocations would otherwise fail? Which the OomGuard could tell you.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove closed pull request #4582: feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting)
URL: https://github.com/apache/datafusion-comet/pull/4582
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#discussion_r3572076944
##########
native/core/src/execution/memory_pools/real_usage_pool.rs:
##########
@@ -0,0 +1,348 @@
+// 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
+/// wiring uses `oom_guard::current_balance`; tests inject a controllable value.
+type BalanceSource = Arc<dyn Fn() -> usize + Send + Sync>;
+
+/// 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<dyn MemoryPool>,
+ /// 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<usize>,
+ 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 RealUsagePool {
+ /// Wrap `inner` with the real-usage gate using the live OomGuard balance.
+ pub(crate) fn new(
+ inner: Arc<dyn MemoryPool>,
+ ceiling: usize,
+ fair_share: Option<usize>,
+ ) -> Self {
+ Self {
+ inner,
+ ceiling,
+ fair_share,
+ balance_source: Arc::new(oom_guard::current_balance),
+ }
+ }
+
+ /// Wrap `inner` with an explicit balance source (test seam).
+ #[cfg(test)]
+ fn with_balance_source(
+ inner: Arc<dyn MemoryPool>,
+ ceiling: usize,
+ fair_share: Option<usize>,
+ 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<usize>) -> bool {
+ match share {
+ None => true,
+ Some(s) => reserved.saturating_add(additional) > s,
+ }
+}
+
+impl MemoryPool for 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,
Review Comment:
how additional would be used/calculated? is it sort of extra buffer?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#issuecomment-5669369951
I'm closing this rather than continuing to iterate on it.
The honest reason is that it grew into three things at once — allocator-level
byte tracking, a cooperative gate on the memory pool, and a panic-based circuit
breaker — and only the first of those is something I'm confident enough in to
defend. Reviewing it properly against the merged state of main turned up a few
things that undercut the design as described:
The two enforcement layers don't actually layer. The argument in the description
is that the cooperative gate fires before the breaker because it tests projected
usage while the breaker tests actual usage. But the gate only runs inside
`try_grow`, and the whole premise of the feature is that memory grows through
allocations that never call `try_grow`. In exactly the case this exists for, the
breaker fires with no spill attempted.
The stamped-thread set is wider than I thought. Tokio runs `on_thread_start` on
every thread the runtime spawns, including the blocking pool, not just the
workers — so a panic raised inside a `spawn_blocking` task becomes a `JoinError`
that reaches neither of the catch sites in `executePlan`.
`LOCAL_DRIFT` leaks on thread exit. It's a plain `Cell<isize>` with no TLS
destructor, so up to 64 KiB of un-flushed drift is discarded every time a thread
dies, and blocking threads churn on a 10s idle timeout. On a long-lived executor
that's a slowly accumulating bias.
And more fundamentally: the tracked balance is layout bytes, not RSS. For a
guard whose job is preventing a container OOM kill, the gap between those two —
fragmentation, jemalloc's retained pages, mmap, anything a C dependency
allocates — is unbounded and always in the dangerous direction.
None of that is wasted. @avantgardnerio's catch that realloc must panic before
delegating, and the finding that concurrent guard panics abort with exit 133
unless serialized, are the kind of thing you only learn by running it, and
they'll carry over.
What I'd like to do instead, in order:
I've opened #5933, a draft documenting the memory model as it actually is today
— the allocator inventory, the FFI ownership asymmetry, why declared
reservations diverge from RSS, and what the container counts. Most of what I
wrote while reviewing this PR turned out to be existing behavior, and it's
useful whether or not any guard ever lands.
Then I'd like to extract just the `AccountingAllocator` as pure observability:
no gate, no panics, feature-gated, with the balance exposed next to pool
reservations in tracing. That makes the accounting gap measurable, which is what
every memory bug report needs and nobody currently has, and it's small enough to
review properly.
Whether to enforce on that number at all is a decision I'd rather make after
seeing how well it tracks RSS on real workloads, rather than before.
@comphead — you mentioned you had a repro pending on #5314. That would still be
very useful against the observability-only version; it would tell us directly
whether the balance is a trustworthy signal.
The branch stays available for anyone who wants to pull pieces from it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM circuit breaker (OomGuard) [datafusion-comet]
Posted by "cetra3 (via GitHub)" <gi...@apache.org>.
cetra3 commented on PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#issuecomment-4667144199
We've been using https://github.com/cetra3/thresher ourselves to handle something similar with datafusion. Namely, we log a trace when the threshold is reached with the heap dump of jemalloc, based upon [the example](https://github.com/cetra3/thresher/blob/main/examples/jemalloc.rs)
It's helped us so far nail down some gnarly memory allocation related bugs with our compaction process. I'm hoping to expand this to queries soon
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#issuecomment-5427543046
I have a good repro but it would depend on #5314
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4582:
URL: https://github.com/apache/datafusion-comet/pull/4582#issuecomment-5330599818
I'm going to test this approach soon
one thing, if I understand correctly this approach would help with overall accounting difference between DF memory pools and actual resident memory usage. However it wouldn't get us closer to identify threshold for spilling.
Currently it's hard(if even possible) to do fine grained analysis and identify which object/layer allocated resources through this approach. So spilling IMO still relies on DF memory pools unless @avantgardnerio as usual have an elegant memory solution :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org