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/09/23 15:53:30 UTC
[PR] feat: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
andygrove opened a new pull request, #25651:
URL: https://github.com/apache/datafusion/pull/25651
## Which issue does this PR close?
- Part of #25650.
## Rationale for this change
`MemoryPool` accounting only covers memory that operators explicitly reserve, so a process can be OOM killed while the pool reports plenty of headroom. As a first step, this makes the gap visible: it logs the difference between what pools have reserved and what's actually allocated, and it's on by default in SLT. It only logs and never fails a test, unlike #22626 which was reverted in #22860.
## What changes are included in this PR?
- `datafusion-execution`: `MemoryDriftTracker` and `DriftLoggingPool`. The pool wraps any `MemoryPool` and reports reservation changes to a tracker, which can be shared by many pools. The tracker compares the reserved total against allocated bytes from a caller-supplied `Fn() -> usize`, since DataFusion doesn't choose the global allocator. It logs at `info` each time positive drift rises by 64 MB, naming the pool and consumer, and records the peak.
- `sqllogictest`: a `CountingAllocator` global allocator, with every test file's pool wrapped against one process-wide tracker. The peak drift is printed at the end of the run. Disable it with `--memory-drift false`.
Test files run concurrently, so the comparison is process-wide rather than per file. Files that `SET datafusion.runtime.memory_limit` replace their pool and drop out of the reserved total (6 of 504 files).
Example output from a full local run:
```
Peak memory drift: drift=84.3 MB allocated=85.4 MB reserved=1176.3 KB pool=order.slt consumer=ExternalSorterMerge[0]
```
The allocator batches its counts per thread, so SLT runtime is unchanged locally (7s both with and without it).
## Are these changes tested?
Yes. There are unit tests for the tracker and pool, plus a doc test. I also ran the full SLT suite locally with drift logging on and off.
## Are there any user-facing changes?
New public types in `datafusion_execution::memory_pool`. No changes to existing APIs.
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
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: 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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #25651:
URL: https://github.com/apache/datafusion/pull/25651#issuecomment-5818187407
@andygrove please fix the doc format
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #25651:
URL: https://github.com/apache/datafusion/pull/25651
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
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: 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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #25651:
URL: https://github.com/apache/datafusion/pull/25651#discussion_r4087990709
##########
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:
This is close to a copy of `PeakRecordingPool` in `peak_recording.rs`: the same pass-through of every `MemoryPool` method, the same `Display` passthrough, the same "`inner` is expected to be empty" caveat, and near-identical delegation and failed-growth tests. Each pass-through wrapper also has to remember to forward any method added to `MemoryPool` later, because the defaulted ones (`register`, `unregister`, `memory_limit`) silently fall back to the default.
Could `PeakRecordingPool` report to the tracker instead, e.g. `PeakRecordingPool::new(pool).with_drift_tracker(tracker, label)`? That removes this type, its `Debug`/`Display` impls and the two duplicated tests. `DriftLoggingPool::tracker()` also has no caller.
##########
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:
`self.log_threshold as isize` wraps for values above `isize::MAX`, so `with_log_threshold(usize::MAX)` or `--memory-drift-log-threshold 18446744073709551615` logs on every reservation change instead of never (checked with a copy of `observe`). `saturating_add_unsigned` avoids the cast:
```suggestion
let rose = untracked >= last.saturating_add_unsigned(self.log_threshold);
```
##########
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:
A cheap way to avoid the loss without a TLS destructor: tokio runs the `on_thread_stop` hook on every worker and blocking thread just before it exits, outside the allocator, so the runner can flush there.
```rust
/// Flush this thread's unflushed count into `ALLOCATED`.
pub fn flush_thread_allocations() {
let _ = UNFLUSHED.try_with(|unflushed| {
ALLOCATED.fetch_add(unflushed.replace(0), Ordering::Relaxed);
});
}
```
```rust
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.on_thread_stop(flush_thread_allocations)
.build()?
```
In a standalone check with the same counting scheme, 40 `spawn_blocking` tasks each returned a 100 KB `Vec`, and their threads then went idle and exited. Without the hook, 21 KB of the 4000 KB was counted. With it, 3919 KB was counted.
Dropping `COUNTING` and counting from process start would also remove the other source of negative bias, which is memory allocated before `enable_memory_drift_logging` and freed after it. The PR description reports no runtime change from counting, so `--memory-drift false` would only need to skip the pool wrapping.
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #25651:
URL: https://github.com/apache/datafusion/pull/25651#issuecomment-5804216635
@comphead I addressed the smaller points in f79c1f984:
- `allocated_bytes` is no longer public, and `CountingAllocator` no longer takes a type parameter; it always uses `System`.
- I merged the negative-drift test into the threshold test, and the `tracker()` test helper now takes a threshold.
- The Substrait round-trip engine now calls `rewrap_replaced_pool` too. With the call, `information_schema.slt` and `set_variable.slt` log 40 and 5 drift lines in round-trip mode. Without it, they log 37 and 2.
- The `--memory-drift` help text now names `RUST_LOG=datafusion_execution::memory_pool=info` and says that only rises of `--memory-drift-log-threshold` bytes are logged.
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #25651:
URL: https://github.com/apache/datafusion/pull/25651#issuecomment-5800115157
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/25651?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `87.67606%` with `35 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 82.45%. Comparing base ([`7570366`](https://app.codecov.io/gh/apache/datafusion/commit/7570366fd929daf9ced744bb8397686b50565b18?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`44c87fb`](https://app.codecov.io/gh/apache/datafusion/commit/44c87fb8511960e32cdbf167869b5c0de56c9352?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 3 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/25651?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [datafusion/execution/src/memory\_pool/drift.rs](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&filepath=datafusion%2Fexecution%2Fsrc%2Fmemory_pool%2Fdrift.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9leGVjdXRpb24vc3JjL21lbW9yeV9wb29sL2RyaWZ0LnJz) | 87.79% | [23 Missing and 3 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [datafusion/sqllogictest/src/memory\_drift.rs](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&filepath=datafusion%2Fsqllogictest%2Fsrc%2Fmemory_drift.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9zcWxsb2dpY3Rlc3Qvc3JjL21lbW9yeV9kcmlmdC5ycw==) | 87.93% | [4 Missing and 3 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [datafusion/sqllogictest/bin/sqllogictests.rs](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&filepath=datafusion%2Fsqllogictest%2Fbin%2Fsqllogictests.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9zcWxsb2dpY3Rlc3QvYmluL3NxbGxvZ2ljdGVzdHMucnM=) | 66.66% | [0 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25651?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #25651 +/- ##
==========================================
- Coverage 82.45% 82.45% -0.01%
==========================================
Files 1140 1142 +2
Lines 436589 436872 +283
Branches 436589 436872 +283
==========================================
+ Hits 359996 360229 +233
- Misses 54838 54877 +39
- Partials 21755 21766 +11
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/25651?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
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: 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: log drift between MemoryPool reservations and allocated bytes in SLT [datafusion]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #25651:
URL: https://github.com/apache/datafusion/pull/25651#issuecomment-5818236879
@comphead Fixed in 4cbebae19. The `MemoryPool` doc links in `drift.rs` broke when I removed the import. `cargo doc --document-private-items --no-deps --workspace` with `-D warnings` now passes locally.
--
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