adriangb commented on code in PR #25651: URL: https://github.com/apache/datafusion/pull/25651#discussion_r4087298298
########## 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: This function is called on every `grow`, `try_grow` and `shrink` of every pool that reports to the tracker. For the SLT atomic, that is cheap. But the module docs suggest "allocator statistics", and to get a current value from jemalloc's `stats.allocated` you must advance `epoch` through `mallctl` on each call. That can slow down queries. Can we document this requirement? ```suggestion /// 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>; ``` ########## 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: The unflushed count of a thread is lost when the thread exits. Tokio blocking-pool threads exit after they are idle, and these threads are used for local file I/O. So the error is not limited to `threads * FLUSH_BYTES`, and it can grow over a long run. Also, `allocated_bytes()` clamps at zero, so a negative bias is not visible. I understand why `UNFLUSHED` has no destructor (the destructor registration can allocate from inside the allocator). If you keep this design, please correct the comment: ```suggestion /// 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 off by up to `FLUSH_BYTES` per live thread. The unflushed /// count of a thread that exits (e.g. an idle Tokio blocking thread) is lost, /// so this error can grow during a long run. const FLUSH_BYTES: isize = 256 * 1024; ``` ########## 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: No test covers the logging logic: a log line at the threshold, and a re-arm after drift falls. The `last_logged` logic uses a CAS and has an edge case (falling drift lowers the baseline). A test would help here. One option: `observe` returns `bool` for "logged", or the tracker keeps a `logged_count`, and the tests assert on it with `with_log_threshold(1000)`. ########## 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: Drift is only sampled when a reservation changes. A query that allocates a lot but reserves little is the case this tool must find. That query is only sampled when a different reservation changes, possibly in a different file or after the memory is freed. A possible solution: make it possible to sample without a reservation change, so that the caller can also sample from a timer or from the allocator flush path. For example: ```rust /// Compare the current allocated bytes with the reserved total now, /// without a reservation change. pub fn sample(&self, source: &str) { self.observe(source, "", self.reserved()); } ``` ########## 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: These 6 files are mostly the spill tests, where the accounting is most important. After the `SET`, their reservations are not counted, but their allocations are. This makes the drift larger, and a peak can then be attributed to a different file. Is it possible to wrap the new pool too? If not, the README can say that the total is not correct while these files run. ########## 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: These fields show which reservation change took the sample. They do not show where the untracked memory came from. When files run concurrently, even the pool can be wrong: a full SLT run attributes the peak to `order.slt / ExternalSorterMerge[0]` (84 MB), but `order.slt` alone peaks at 2.4 MB. With `--test-threads 1`, the pool is correct, but the consumer is still only where the sample was taken. For example, `ParquetSink(SerializedFileWriter)` is named at the peak, but the total reserved at that moment is 144.5 KB of 43.5 MB allocated. Can the docs say this, so that nobody reads it as blame? ```suggestion /// Label of the [`DriftLoggingPool`] whose reservation change took this /// 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. This shows when /// drift was sampled, not what caused it. pub consumer: String, ``` ########## 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: With the default concurrency, this line names a file and consumer that can be wrong (see the table in the review summary). A short hint here tells the reader how to get a correct result: ```suggestion if let Some(peak) = memory_drift_tracker().and_then(|t| t.peak_drift()) { eprintln!("Peak memory drift: {peak}"); if options.test_threads > 1 { eprintln!( "Test files ran concurrently, so the file and consumer above can be wrong. \ Run with --test-threads 1 to attribute drift to one file." ); } } ``` ########## 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: With `--test-threads 1`, the results are stable and attributed to the correct file. Two full runs gave `push_down_filter_regression.slt` at 43.3 MB and 43.6 MB, and that file alone gives 42.9 MB. I suggest that we document this as the way to investigate: ````suggestion Test files run concurrently by default, so the comparison is process-wide: allocated bytes across the whole process against reservations summed across all files. In this mode, the file in the output can be wrong. To attribute drift to one file, run one file at a time: ```shell cargo test --test sqllogictests -- --test-threads 1 ``` In both modes, the consumer in the output is the reservation change that took the sample, not necessarily the code that allocated the untracked memory. Most drift at this scale is memory that `MemoryPool` does not track by design (for example, in-flight batches), so a large drift is a lead to investigate, not necessarily a bug. Files that `SET datafusion.runtime.memory_limit` replace their pool and drop out of the reserved total. ```` ########## 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: With the default 64 MB threshold, this command printed no lines in a full run, with or without `--test-threads 1`, because the largest drift was 43 MB. With a 4 MB threshold (a local change), the output was useful: it showed each row-group cycle of the `COPY` in `push_down_filter_regression.slt`. Can the threshold be an SLT option that is passed to `MemoryDriftTracker::with_log_threshold`? For example, `--memory-drift-log-threshold` / `SLT_MEMORY_DRIFT_LOG_THRESHOLD`. -- 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]
