andygrove commented on code in PR #25651: URL: https://github.com/apache/datafusion/pull/25651#discussion_r4087847969
########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,434 @@ +// 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. +pub type AllocatedBytesFn = Arc<dyn Fn() -> usize + Send + Sync>; Review Comment: Done in 956f9a0cb, using your wording. ########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,434 @@ +// 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. +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`] that made the observation. + pub pool: String, + /// Consumer whose reservation change triggered the observation. + 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() + } + + 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); + } + + fn observe(&self, pool: &str, consumer: &str, reserved: usize) { + let allocated = (self.allocated)(); Review Comment: Added `MemoryDriftTracker::sample(source)` with a test in 956f9a0cb. I did not add a timer in the SLT runner yet: timer samples would all be recorded under one source name, so the peak's file and consumer would say even less. Happy to add it in a follow-up if it's useful. ########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,434 @@ +// 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. +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`] that made the observation. + pub pool: String, + /// Consumer whose reservation change triggered the observation. + 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() + } + + 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); + } + + fn observe(&self, pool: &str, consumer: &str, reserved: usize) { + 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; + } + let updated = self + .last_logged + .compare_exchange(last, untracked, Ordering::Relaxed, Ordering::Relaxed) + .is_ok(); + if updated && rose { Review Comment: `observe` now returns `true` when it logs. New tests in 956f9a0cb cover logging at the threshold, re-arming after drift falls (the baseline-lowering case), and negative drift counting as zero. ########## 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: Kept the design and used your comment in 956f9a0cb. ########## 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. Review Comment: Fixed in 956f9a0cb: after each statement, the SLT runner wraps the pool again if a statement replaced it (checked with `pool.is::<DriftLoggingPool>()`). When `aggregate_memory_spill.slt` runs alone, queries after its `SET datafusion.runtime.memory_limit = '1M'` still log with non-zero reserved bytes. ########## datafusion/execution/src/memory_pool/drift.rs: ########## @@ -0,0 +1,434 @@ +// 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. +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`] that made the observation. + pub pool: String, + /// Consumer whose reservation change triggered the observation. + pub consumer: String, Review Comment: Done in 956f9a0cb, using your wording, plus a mention of the new `sample` source. ########## datafusion/sqllogictest/bin/sqllogictests.rs: ########## @@ -386,6 +396,10 @@ async fn run_tests() -> Result<()> { HumanDuration(start.elapsed()) ))?; + if let Some(peak) = memory_drift_tracker().and_then(|t| t.peak_drift()) { + eprintln!("Peak memory drift: {peak}"); + } Review Comment: Added in 956f9a0cb. ########## datafusion/sqllogictest/README.md: ########## @@ -101,6 +101,29 @@ SLT_TIMING_SUMMARY=1 cargo test --test sqllogictests SLT_TIMING_DEBUG_SLOW_FILES=1 cargo test --test sqllogictests ``` +### Memory drift + +The runner compares the bytes reserved in `MemoryPool`s with the bytes actually +allocated, to find operators whose memory is not tracked by the pool (see +[#25650](https://github.com/apache/datafusion/issues/25650)). It is enabled by +default and prints the largest drift seen at the end of the run. It only logs +and never fails a test. + +Test files run concurrently, so the comparison is process-wide: allocated bytes +across the whole process against reservations summed across all files. Files +that `SET datafusion.runtime.memory_limit` replace their pool and drop out of +the reserved total. Review Comment: Added in 956f9a0cb. I also removed the sentence about `SET` files dropping out of the reserved total, because the runner now wraps the replacement pool too. ########## datafusion/sqllogictest/README.md: ########## @@ -101,6 +101,29 @@ SLT_TIMING_SUMMARY=1 cargo test --test sqllogictests SLT_TIMING_DEBUG_SLOW_FILES=1 cargo test --test sqllogictests ``` +### Memory drift + +The runner compares the bytes reserved in `MemoryPool`s with the bytes actually +allocated, to find operators whose memory is not tracked by the pool (see +[#25650](https://github.com/apache/datafusion/issues/25650)). It is enabled by +default and prints the largest drift seen at the end of the run. It only logs +and never fails a test. + +Test files run concurrently, so the comparison is process-wide: allocated bytes +across the whole process against reservations summed across all files. Files +that `SET datafusion.runtime.memory_limit` replace their pool and drop out of +the reserved total. + +```shell +# Log each 64 MB rise in drift, with the file and consumer that triggered it +RUST_LOG=datafusion_execution::memory_pool=info cargo test --test sqllogictests +``` Review Comment: Added `--memory-drift-log-threshold` / `SLT_MEMORY_DRIFT_LOG_THRESHOLD` (in bytes) in 956f9a0cb. It is passed to `with_log_threshold`, and the README has a 4 MB example with `--test-threads 1`. -- 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]
