andygrove commented on code in PR #25651: URL: https://github.com/apache/datafusion/pull/25651#discussion_r4088063561
########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,517 @@ +// 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. + +//! Logs the drift between what [`MemoryPool`]s have reserved and what the +//! process has actually allocated. +//! +//! [`MemoryPool`] accounting is voluntary: only allocations that an operator +//! explicitly reserves are counted. Anything else (in-flight batches, kernel +//! scratch space, untracked buffers) is invisible to the pool, so a process +//! can run out of memory while its pool reports plenty of headroom. See +//! <https://github.com/apache/datafusion/issues/25650>. +//! +//! This module only observes. Nothing here changes how memory is granted or +//! limited; it logs the gap so that operators which under-report can be found. +//! +//! DataFusion is a library and does not choose the global allocator, so the +//! allocated byte count is supplied by the caller, e.g. from a counting +//! [`GlobalAlloc`](std::alloc::GlobalAlloc) wrapper or allocator statistics. + +use std::{ + fmt::{Debug, Display, Formatter}, + sync::{ + Arc, + atomic::{AtomicIsize, AtomicUsize, Ordering}, + }, +}; + +use datafusion_common::{Result, human_readable_size}; +use parking_lot::Mutex; + +use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}; + +/// Returns the number of bytes currently allocated by the process. +/// +/// This is called on every reservation change (`grow`, `try_grow` and +/// `shrink`) of every pool that reports to the tracker, so it must be cheap, +/// e.g. a single atomic load. Do not read allocator statistics that need a +/// refresh on each call (such as jemalloc's `epoch`). +pub type AllocatedBytesFn = Arc<dyn Fn() -> usize + Send + Sync>; + +/// Default rise in drift, in bytes, needed before another line is logged. +pub const DEFAULT_DRIFT_LOG_THRESHOLD: usize = 64 * 1024 * 1024; + +/// Compares allocated bytes against the total reserved by every +/// [`DriftLoggingPool`] that reports to it. +/// +/// A single tracker can be shared by many pools, e.g. one pool per +/// `SessionContext` in a process running several at once. The reserved total +/// is then summed across all of them, which is what has to be compared with a +/// process-wide allocated byte count. +/// +/// Drift is `allocated - reserved`. A line is logged at `info` level each time +/// drift rises by at least the log threshold, naming the pool and consumer +/// whose reservation change triggered the check. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion_execution::memory_pool::{ +/// # MemoryConsumer, MemoryDriftTracker, MemoryPool, DriftLoggingPool, UnboundedMemoryPool, +/// # }; +/// // A real caller would read a counting allocator or allocator stats here. +/// let tracker = Arc::new(MemoryDriftTracker::new(Arc::new(|| 10_000))); +/// let pool: Arc<dyn MemoryPool> = Arc::new(DriftLoggingPool::new( +/// Arc::new(UnboundedMemoryPool::default()), +/// Arc::clone(&tracker), +/// "example", +/// )); +/// +/// let reservation = MemoryConsumer::new("op").register(&pool); +/// reservation.grow(4_000); +/// +/// assert_eq!(tracker.reserved(), 4_000); +/// assert_eq!(tracker.peak_drift().unwrap().drift, 6_000); +/// ``` +pub struct MemoryDriftTracker { + allocated: AllocatedBytesFn, + log_threshold: usize, + /// Total reserved across every pool reporting to this tracker. + reserved: AtomicUsize, + /// Positive drift at the time of the last logged line. + last_logged: AtomicIsize, + /// Largest drift seen. + peak_drift: AtomicIsize, + /// Where the largest drift was seen. + peak: Mutex<Option<DriftSample>>, +} + +/// One observation of drift, recorded by [`MemoryDriftTracker`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DriftSample { + /// Label of the [`DriftLoggingPool`] whose reservation change took this + /// sample, or the source passed to [`MemoryDriftTracker::sample`]. When + /// several pools share a tracker, the untracked memory can come from any + /// of them. + pub pool: String, + /// Consumer whose reservation change took this sample (empty for + /// [`MemoryDriftTracker::sample`]). This shows when drift was sampled, + /// not what caused it. + pub consumer: String, + /// Bytes reserved across all pools reporting to the tracker. + pub reserved: usize, + /// Bytes allocated, as reported by the tracker's [`AllocatedBytesFn`]. + pub allocated: usize, + /// `allocated - reserved`. + pub drift: isize, +} + +impl Display for DriftSample { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let drift = if self.drift < 0 { + format!("-{}", human_readable_size(self.drift.unsigned_abs())) + } else { + human_readable_size(self.drift as usize) + }; + write!( + f, + "drift={drift} allocated={} reserved={} pool={} consumer={}", + human_readable_size(self.allocated), + human_readable_size(self.reserved), + self.pool, + self.consumer, + ) + } +} + +impl MemoryDriftTracker { + /// Create a tracker that reads allocated bytes from `allocated`, logging + /// every [`DEFAULT_DRIFT_LOG_THRESHOLD`] of drift. + pub fn new(allocated: AllocatedBytesFn) -> Self { + Self { + allocated, + log_threshold: DEFAULT_DRIFT_LOG_THRESHOLD, + reserved: AtomicUsize::new(0), + last_logged: AtomicIsize::new(0), + peak_drift: AtomicIsize::new(isize::MIN), + peak: Mutex::new(None), + } + } + + /// Log a line each time drift rises by `log_threshold` bytes. + pub fn with_log_threshold(mut self, log_threshold: usize) -> Self { + self.log_threshold = log_threshold; + self + } + + /// Bytes currently reserved across all pools reporting to this tracker. + pub fn reserved(&self) -> usize { + self.reserved.load(Ordering::Relaxed) + } + + /// The largest drift seen so far, if any reservation has been made. + pub fn peak_drift(&self) -> Option<DriftSample> { + self.peak.lock().clone() + } + + /// Compare the current allocated bytes with the reserved total now, + /// without a reservation change. + /// + /// Drift is otherwise only sampled when a reservation changes, so memory + /// allocated by code that reserves little can go unseen until some other + /// reservation changes. Calling this periodically, e.g. from a timer, + /// closes that gap. `source` is recorded as the pool label. + pub fn sample(&self, source: &str) { + self.observe(source, "", self.reserved()); + } + + fn grew(&self, pool: &str, consumer: &str, additional: usize) { + let reserved = + self.reserved.fetch_add(additional, Ordering::Relaxed) + additional; + self.observe(pool, consumer, reserved); + } + + fn shrank(&self, pool: &str, consumer: &str, shrink: usize) { + let reserved = self.reserved.fetch_sub(shrink, Ordering::Relaxed) - shrink; + self.observe(pool, consumer, reserved); + } + + /// Returns `true` if a line was logged. + fn observe(&self, pool: &str, consumer: &str, reserved: usize) -> bool { + let allocated = (self.allocated)(); + let drift = allocated as isize - reserved as isize; + + let sample = || DriftSample { + pool: pool.to_string(), + consumer: consumer.to_string(), + reserved, + allocated, + drift, + }; + + // Lock-free check first so the lock is only taken for a new peak. + if self.peak_drift.fetch_max(drift, Ordering::Relaxed) < drift { + let mut peak = self.peak.lock(); + if peak.as_ref().is_none_or(|p| drift > p.drift) { + *peak = Some(sample()); + } + } + + // Only rising positive drift is logged: that is untracked memory, which + // is what leads to OOM kills. Negative drift (e.g. an operator + // reserving ahead of allocating) counts as zero. Falling drift quietly + // lowers the baseline so the next rise is seen. + let untracked = drift.max(0); + let last = self.last_logged.load(Ordering::Relaxed); + let rose = untracked >= last.saturating_add(self.log_threshold as isize); + if !rose && untracked >= last { + return false; + } + let updated = self + .last_logged + .compare_exchange(last, untracked, Ordering::Relaxed, Ordering::Relaxed) + .is_ok(); + let logged = updated && rose; + if logged { + log::info!("memory drift: {}", sample()); + } + logged + } +} + +impl Debug for MemoryDriftTracker { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MemoryDriftTracker") + .field("log_threshold", &self.log_threshold) + .field("reserved", &self.reserved()) + .field("peak", &self.peak_drift()) + .finish() + } +} + +/// Wraps a [`MemoryPool`], reporting every reservation change to a +/// [`MemoryDriftTracker`]. +/// +/// Every method delegates to the wrapped pool, so wrapping does not change how +/// memory is granted, limited, or reported. As with other wrappers, +/// downcasting the pool finds this wrapper rather than the pool it wraps. +pub struct DriftLoggingPool { Review Comment: Done in f79c1f984: I removed `DriftLoggingPool` and added `PeakRecordingPool::with_drift_tracker(tracker, label)`. The duplicate impls and tests and the unused `tracker()` accessor are gone. ########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,517 @@ +// 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. + +//! Logs the drift between what [`MemoryPool`]s have reserved and what the +//! process has actually allocated. +//! +//! [`MemoryPool`] accounting is voluntary: only allocations that an operator +//! explicitly reserves are counted. Anything else (in-flight batches, kernel +//! scratch space, untracked buffers) is invisible to the pool, so a process +//! can run out of memory while its pool reports plenty of headroom. See +//! <https://github.com/apache/datafusion/issues/25650>. +//! +//! This module only observes. Nothing here changes how memory is granted or +//! limited; it logs the gap so that operators which under-report can be found. +//! +//! DataFusion is a library and does not choose the global allocator, so the +//! allocated byte count is supplied by the caller, e.g. from a counting +//! [`GlobalAlloc`](std::alloc::GlobalAlloc) wrapper or allocator statistics. + +use std::{ + fmt::{Debug, Display, Formatter}, + sync::{ + Arc, + atomic::{AtomicIsize, AtomicUsize, Ordering}, + }, +}; + +use datafusion_common::{Result, human_readable_size}; +use parking_lot::Mutex; + +use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}; + +/// Returns the number of bytes currently allocated by the process. +/// +/// This is called on every reservation change (`grow`, `try_grow` and +/// `shrink`) of every pool that reports to the tracker, so it must be cheap, +/// e.g. a single atomic load. Do not read allocator statistics that need a +/// refresh on each call (such as jemalloc's `epoch`). +pub type AllocatedBytesFn = Arc<dyn Fn() -> usize + Send + Sync>; + +/// Default rise in drift, in bytes, needed before another line is logged. +pub const DEFAULT_DRIFT_LOG_THRESHOLD: usize = 64 * 1024 * 1024; + +/// Compares allocated bytes against the total reserved by every +/// [`DriftLoggingPool`] that reports to it. +/// +/// A single tracker can be shared by many pools, e.g. one pool per +/// `SessionContext` in a process running several at once. The reserved total +/// is then summed across all of them, which is what has to be compared with a +/// process-wide allocated byte count. +/// +/// Drift is `allocated - reserved`. A line is logged at `info` level each time +/// drift rises by at least the log threshold, naming the pool and consumer +/// whose reservation change triggered the check. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion_execution::memory_pool::{ +/// # MemoryConsumer, MemoryDriftTracker, MemoryPool, DriftLoggingPool, UnboundedMemoryPool, +/// # }; +/// // A real caller would read a counting allocator or allocator stats here. +/// let tracker = Arc::new(MemoryDriftTracker::new(Arc::new(|| 10_000))); +/// let pool: Arc<dyn MemoryPool> = Arc::new(DriftLoggingPool::new( +/// Arc::new(UnboundedMemoryPool::default()), +/// Arc::clone(&tracker), +/// "example", +/// )); +/// +/// let reservation = MemoryConsumer::new("op").register(&pool); +/// reservation.grow(4_000); +/// +/// assert_eq!(tracker.reserved(), 4_000); +/// assert_eq!(tracker.peak_drift().unwrap().drift, 6_000); +/// ``` +pub struct MemoryDriftTracker { + allocated: AllocatedBytesFn, + log_threshold: usize, + /// Total reserved across every pool reporting to this tracker. + reserved: AtomicUsize, + /// Positive drift at the time of the last logged line. + last_logged: AtomicIsize, + /// Largest drift seen. + peak_drift: AtomicIsize, + /// Where the largest drift was seen. + peak: Mutex<Option<DriftSample>>, +} + +/// One observation of drift, recorded by [`MemoryDriftTracker`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DriftSample { + /// Label of the [`DriftLoggingPool`] whose reservation change took this + /// sample, or the source passed to [`MemoryDriftTracker::sample`]. When + /// several pools share a tracker, the untracked memory can come from any + /// of them. + pub pool: String, + /// Consumer whose reservation change took this sample (empty for + /// [`MemoryDriftTracker::sample`]). This shows when drift was sampled, + /// not what caused it. + pub consumer: String, + /// Bytes reserved across all pools reporting to the tracker. + pub reserved: usize, + /// Bytes allocated, as reported by the tracker's [`AllocatedBytesFn`]. + pub allocated: usize, + /// `allocated - reserved`. + pub drift: isize, +} + +impl Display for DriftSample { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let drift = if self.drift < 0 { + format!("-{}", human_readable_size(self.drift.unsigned_abs())) + } else { + human_readable_size(self.drift as usize) + }; + write!( + f, + "drift={drift} allocated={} reserved={} pool={} consumer={}", + human_readable_size(self.allocated), + human_readable_size(self.reserved), + self.pool, + self.consumer, + ) + } +} + +impl MemoryDriftTracker { + /// Create a tracker that reads allocated bytes from `allocated`, logging + /// every [`DEFAULT_DRIFT_LOG_THRESHOLD`] of drift. + pub fn new(allocated: AllocatedBytesFn) -> Self { + Self { + allocated, + log_threshold: DEFAULT_DRIFT_LOG_THRESHOLD, + reserved: AtomicUsize::new(0), + last_logged: AtomicIsize::new(0), + peak_drift: AtomicIsize::new(isize::MIN), + peak: Mutex::new(None), + } + } + + /// Log a line each time drift rises by `log_threshold` bytes. + pub fn with_log_threshold(mut self, log_threshold: usize) -> Self { + self.log_threshold = log_threshold; + self + } + + /// Bytes currently reserved across all pools reporting to this tracker. + pub fn reserved(&self) -> usize { + self.reserved.load(Ordering::Relaxed) + } + + /// The largest drift seen so far, if any reservation has been made. + pub fn peak_drift(&self) -> Option<DriftSample> { + self.peak.lock().clone() + } + + /// Compare the current allocated bytes with the reserved total now, + /// without a reservation change. + /// + /// Drift is otherwise only sampled when a reservation changes, so memory + /// allocated by code that reserves little can go unseen until some other + /// reservation changes. Calling this periodically, e.g. from a timer, + /// closes that gap. `source` is recorded as the pool label. + pub fn sample(&self, source: &str) { + self.observe(source, "", self.reserved()); + } + + fn grew(&self, pool: &str, consumer: &str, additional: usize) { + let reserved = + self.reserved.fetch_add(additional, Ordering::Relaxed) + additional; + self.observe(pool, consumer, reserved); + } + + fn shrank(&self, pool: &str, consumer: &str, shrink: usize) { + let reserved = self.reserved.fetch_sub(shrink, Ordering::Relaxed) - shrink; + self.observe(pool, consumer, reserved); + } + + /// Returns `true` if a line was logged. + fn observe(&self, pool: &str, consumer: &str, reserved: usize) -> bool { + let allocated = (self.allocated)(); + let drift = allocated as isize - reserved as isize; + + let sample = || DriftSample { + pool: pool.to_string(), + consumer: consumer.to_string(), + reserved, + allocated, + drift, + }; + + // Lock-free check first so the lock is only taken for a new peak. + if self.peak_drift.fetch_max(drift, Ordering::Relaxed) < drift { + let mut peak = self.peak.lock(); + if peak.as_ref().is_none_or(|p| drift > p.drift) { + *peak = Some(sample()); + } + } + + // Only rising positive drift is logged: that is untracked memory, which + // is what leads to OOM kills. Negative drift (e.g. an operator + // reserving ahead of allocating) counts as zero. Falling drift quietly + // lowers the baseline so the next rise is seen. + let untracked = drift.max(0); + let last = self.last_logged.load(Ordering::Relaxed); + let rose = untracked >= last.saturating_add(self.log_threshold as isize); Review Comment: Fixed in f79c1f984 with `saturating_add_unsigned`. I added a test that a `usize::MAX` threshold never logs. ########## datafusion/sqllogictest/src/memory_drift.rs: ########## @@ -0,0 +1,154 @@ +// 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. + +//! Logs drift between `MemoryPool` reservations and actual allocations while +//! running sqllogictests. See <https://github.com/apache/datafusion/issues/25650>. +//! +//! Test files run concurrently in one process, so allocations cannot be split +//! per file. Instead every file's pool reports to one process-wide +//! [`MemoryDriftTracker`], which compares the sum of all reservations with the +//! bytes counted by [`CountingAllocator`]. +//! +//! Files that `SET datafusion.runtime.memory_limit` replace their pool, so +//! their reservations after that point are not included in the total. +//! +//! This only logs. It never fails a test. + +use std::{ + alloc::{GlobalAlloc, Layout, System}, + cell::Cell, + sync::{ + Arc, OnceLock, + atomic::{AtomicBool, AtomicIsize, Ordering}, + }, +}; + +use datafusion::execution::memory_pool::{ + DriftLoggingPool, MemoryDriftTracker, MemoryPool, +}; + +static ALLOCATED: AtomicIsize = AtomicIsize::new(0); +static COUNTING: AtomicBool = AtomicBool::new(false); +static TRACKER: OnceLock<Arc<MemoryDriftTracker>> = OnceLock::new(); + +/// A [`GlobalAlloc`] that counts the bytes currently allocated through it, +/// delegating the allocation itself to `A`. +/// +/// Counts requested sizes, so allocator overhead and memory retained by the +/// allocator are not included. Counting is off until +/// [`enable_memory_drift_logging`] is called. +pub struct CountingAllocator<A = System> { + inner: A, +} + +impl<A> CountingAllocator<A> { + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +/// Per-thread count is flushed to [`ALLOCATED`] once it moves this far, so +/// threads do not contend on one atomic for every allocation. The global count +/// is therefore accurate to within `threads * FLUSH_BYTES`. +const FLUSH_BYTES: isize = 256 * 1024; Review Comment: Done in f79c1f984: `flush_thread_allocations` is registered as the runtime's `on_thread_stop` hook, and I removed `COUNTING`, so counting starts with the process. `--memory-drift false` now only skips wrapping the pools. -- 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]
