adriangb commented on code in PR #23985: URL: https://github.com/apache/datafusion/pull/23985#discussion_r3683508774
########## benchmarks/src/util/memory_pool.rs: ########## @@ -0,0 +1,326 @@ +// 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. + +//! Records the peak [`MemoryPool`] reservation reached during a benchmark. +//! +//! DataFusion's [`MemoryPool`] deliberately accounts for only the "large" +//! allocations that scale with input size; intermediate batches flowing between +//! operators are assumed to be small and are left untracked. The [`MemoryPool`] +//! documentation therefore advises reserving "some overhead (e.g. 10%)" on top +//! of the configured limit. +//! +//! Nothing reports what that overhead actually is, because the peak reservation +//! itself is never recorded — [`MemoryPool::reserved`] is a live value that has +//! usually fallen back to zero by the time a query finishes. This module records +//! the high-water mark so benchmarks can emit it alongside the peak RSS that +//! [`print_memory_stats`] already prints, making the gap between the two +//! measurable. +//! +//! This is measurement only: nothing here enforces a relationship between the +//! two numbers. +//! +//! What lands in the peak is whatever the pool accounts for, so this follows +//! the accounting rather than fixing it in place. Arrow-side reservations made +//! through `ArrowMemoryPool` are included, because that adapter grows a +//! DataFusion reservation against the pool it wraps; nothing claims buffers +//! today, but the peak picks it up when something does. +//! +//! [`print_memory_stats`]: super::print_memory_stats + +use std::{ + fmt::{Debug, Display, Formatter}, + sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, +}; + +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_common::Result; + +/// High-water mark since the last [`reset_peak_pool_reserved`]. +static PEAK_RESERVED: AtomicUsize = AtomicUsize::new(0); + +/// High-water mark since the process started. Never reset. +static MAX_RESERVED: AtomicUsize = AtomicUsize::new(0); + +/// Whether a [`PeakRecordingPool`] has ever been constructed, used to +/// distinguish "no pool was recording" from "the pool peaked at zero bytes". +static RECORDING: AtomicBool = AtomicBool::new(false); + +/// Peak [`MemoryPool`] reservation, in bytes, since the last call to +/// [`reset_peak_pool_reserved`]. +/// +/// Returns `None` if no [`PeakRecordingPool`] has been installed, which is the +/// case whenever a benchmark runs without a memory limit. +pub fn peak_pool_reserved() -> Option<usize> { + RECORDING + .load(Ordering::Relaxed) + .then(|| PEAK_RESERVED.load(Ordering::Relaxed)) +} + +/// Peak [`MemoryPool`] reservation, in bytes, since the process started. +/// +/// Unlike [`peak_pool_reserved`] this is never reset, so it reports the peak +/// across every query in a run. Returns `None` if no [`PeakRecordingPool`] has +/// been installed. +pub fn max_pool_reserved() -> Option<usize> { + RECORDING + .load(Ordering::Relaxed) + .then(|| MAX_RESERVED.load(Ordering::Relaxed)) +} + +/// Reset the value returned by [`peak_pool_reserved`], so the next reading +/// covers only what follows. +/// +/// [`BenchmarkRun::start_new_case`] calls this, giving each benchmark query its +/// own reading. +/// +/// [`BenchmarkRun::start_new_case`]: super::BenchmarkRun::start_new_case +pub fn reset_peak_pool_reserved() { + PEAK_RESERVED.store(0, Ordering::Relaxed); +} + +/// Wraps a [`MemoryPool`], recording the high-water mark of +/// [`MemoryPool::reserved`] as reservations come and go. +/// +/// Every method delegates to the wrapped pool, so wrapping does not change how +/// memory is granted, limited, or reported. Peaks are published to the +/// process-wide counters read by [`peak_pool_reserved`] and +/// [`max_pool_reserved`] rather than held per instance, so callers can read +/// them without threading a handle through the benchmark. The benchmarks run +/// one query at a time, so a process-wide counter attributes cleanly. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; +/// # use datafusion_benchmarks::util::{ +/// # PeakRecordingPool, peak_pool_reserved, reset_peak_pool_reserved, +/// # }; +/// let pool: Arc<dyn MemoryPool> = +/// Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024)))); +/// reset_peak_pool_reserved(); +/// +/// let reservation = MemoryConsumer::new("example").register(&pool); +/// reservation.try_grow(512)?; +/// reservation.shrink(512); +/// +/// // The pool is back to empty, but the high-water mark is retained. +/// assert_eq!(pool.reserved(), 0); +/// assert_eq!(peak_pool_reserved(), Some(512)); +/// # Ok::<(), datafusion_common::DataFusionError>(()) +/// ``` +pub struct PeakRecordingPool { + inner: Arc<dyn MemoryPool>, +} Review Comment: Good call, implemented in 277076c. Led to a lot of cleanup. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
