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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
