From bd5f16feaa3b29754ebe611f9313ce55c537435e Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 29 Jul 2026 18:11:02 +0000 Subject: [PATCH 1/7] Report peak MemoryPool reservation per query in benchmarks 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 left untracked. The `MemoryPool` docs therefore advise reserving "some overhead (e.g. 10%)" on top of the configured limit. Nothing measures what that overhead actually is. `MemoryPool::reserved` is a live value that has usually fallen back to zero by the time a query finishes, so the peak is never observed, and benchmarks report peak RSS with nothing to compare it against. Add `PeakRecordingPool`, a delegating `MemoryPool` wrapper that records the high-water mark of `reserved()`. `CommonOpt::runtime_env_builder` installs it around the pool it already builds, so every benchmark that runs with a memory limit reports the peak without further changes. - `BenchQuery` gains `pool_peak_bytes`, reset per case by `start_new_case`, so each query gets its own reading rather than inheriting the high-water mark of the queries before it. - `print_memory_stats` prints the run-wide peak after the existing mimalloc line, leaving `mem_profile`'s output parser untouched. This is measurement only. Nothing enforces a relationship between the pool's accounting and actual allocation, and no threshold or CI check is added. The wrapper is installed only when a memory limit is configured, since without one there is no pool to wrap. `pool_peak_bytes` is omitted from the results JSON in that case, leaving the output byte-identical to before for runs without a limit. Co-Authored-By: Claude Opus 5 --- benchmarks/src/util/memory.rs | 24 ++- benchmarks/src/util/memory_pool.rs | 290 +++++++++++++++++++++++++++++ benchmarks/src/util/mod.rs | 4 + benchmarks/src/util/options.rs | 5 +- benchmarks/src/util/run.rs | 17 +- 5 files changed, 337 insertions(+), 3 deletions(-) create mode 100644 benchmarks/src/util/memory_pool.rs diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs index 11b96ef227756..c70c687c7baaa 100644 --- a/benchmarks/src/util/memory.rs +++ b/benchmarks/src/util/memory.rs @@ -15,8 +15,30 @@ // specific language governing permissions and limitations // under the License. -/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api +/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by +/// the peak `MemoryPool` reservation when a memory limit was configured. pub fn print_memory_stats() { + print_allocator_stats(); + print_pool_stats(); +} + +/// Print the peak `MemoryPool` reservation for the run. +/// +/// Prints nothing when the benchmark ran without a memory limit, since no pool +/// was installed to record. Comparing this against the peak RSS above shows how +/// much of a run's memory the pool actually accounted for — DataFusion only +/// tracks the "large" allocations that scale with input size, so the two are +/// expected to differ. +fn print_pool_stats() { + if let Some(peak) = super::max_pool_reserved() { + println!( + "Peak pool reserved: {}", + datafusion_common::human_readable_size(peak) + ); + } +} + +fn print_allocator_stats() { #[cfg(all(feature = "mimalloc", feature = "mimalloc_extended"))] { use datafusion_common::human_readable_size; diff --git a/benchmarks/src/util/memory_pool.rs b/benchmarks/src/util/memory_pool.rs new file mode 100644 index 0000000000000..17c83d7cae517 --- /dev/null +++ b/benchmarks/src/util/memory_pool.rs @@ -0,0 +1,290 @@ +// 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. +//! +//! [`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 { + 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 { + 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 = +/// 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, +} + +impl PeakRecordingPool { + /// Wrap `inner`, recording its peak reservation from here on. + pub fn new(inner: Arc) -> Self { + RECORDING.store(true, Ordering::Relaxed); + Self { inner } + } + + /// The wrapped pool. + pub fn inner(&self) -> &Arc { + &self.inner + } + + /// Publish the pool's current reservation to both high-water marks. + /// + /// Called after any operation that can raise `reserved()`. Reading the + /// total rather than accumulating deltas keeps this correct when the + /// wrapped pool declines or adjusts a request. + fn record(&self) { + let reserved = self.inner.reserved(); + PEAK_RESERVED.fetch_max(reserved, Ordering::Relaxed); + MAX_RESERVED.fetch_max(reserved, Ordering::Relaxed); + } +} + +impl Debug for PeakRecordingPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PeakRecordingPool") + .field("inner", &self.inner) + .field("peak", &PEAK_RESERVED.load(Ordering::Relaxed)) + .finish() + } +} + +impl Display for PeakRecordingPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // Deferring to the wrapped pool keeps `SHOW ALL`-style output and error + // messages identical to running without the wrapper. + Display::fmt(&self.inner, f) + } +} + +impl MemoryPool for PeakRecordingPool { + fn name(&self) -> &str { + self.inner.name() + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record(); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Mutex, MutexGuard}; + + use datafusion::execution::memory_pool::GreedyMemoryPool; + + use super::*; + + /// The high-water marks are process-wide, so these tests would otherwise + /// clobber each other when the test harness runs them in parallel. + static TEST_LOCK: Mutex<()> = Mutex::new(()); + + /// Take the lock and hand back a freshly reset pool. The guard is returned + /// so it stays held for the body of the test. + fn pool(limit: usize) -> (Arc, MutexGuard<'static, ()>) { + let guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let pool: Arc = Arc::new(PeakRecordingPool::new(Arc::new( + GreedyMemoryPool::new(limit), + ))); + reset_peak_pool_reserved(); + (pool, guard) + } + + #[test] + fn records_high_water_mark_across_reservations() { + let (pool, _guard) = pool(1024); + + let a = MemoryConsumer::new("a").register(&pool); + let b = MemoryConsumer::new("b").register(&pool); + + a.try_grow(300).unwrap(); + b.try_grow(400).unwrap(); + // Peak of the sum, not the largest single reservation. + assert_eq!(peak_pool_reserved(), Some(700)); + + a.shrink(300); + b.try_grow(100).unwrap(); + + // Falling back below the peak leaves it untouched, and the later growth + // does not reach it. + assert_eq!(pool.reserved(), 500); + assert_eq!(peak_pool_reserved(), Some(700)); + } + + #[test] + fn failed_growth_does_not_move_the_peak() { + let (pool, _guard) = pool(1024); + + let reservation = MemoryConsumer::new("a").register(&pool); + reservation.try_grow(600).unwrap(); + reservation + .try_grow(600) + .expect_err("should exceed the 1024 byte pool"); + + assert_eq!(peak_pool_reserved(), Some(600)); + } + + #[test] + fn reset_clears_the_window_but_not_the_run_maximum() { + let (pool, _guard) = pool(1024); + + let reservation = MemoryConsumer::new("a").register(&pool); + reservation.try_grow(800).unwrap(); + reservation.shrink(800); + + reset_peak_pool_reserved(); + assert_eq!(peak_pool_reserved(), Some(0)); + assert!(max_pool_reserved().unwrap() >= 800); + + reservation.try_grow(100).unwrap(); + assert_eq!(peak_pool_reserved(), Some(100)); + } + + #[test] + fn delegates_limit_and_name_to_the_wrapped_pool() { + let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let inner: Arc = Arc::new(GreedyMemoryPool::new(4096)); + let wrapped = PeakRecordingPool::new(Arc::clone(&inner)); + + assert_eq!(wrapped.name(), inner.name()); + assert_eq!(wrapped.to_string(), inner.to_string()); + assert!(matches!(wrapped.memory_limit(), MemoryLimit::Finite(4096))); + } +} diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs index 6dc11c0f425bd..c9ffdbc079e92 100644 --- a/benchmarks/src/util/mod.rs +++ b/benchmarks/src/util/mod.rs @@ -18,9 +18,13 @@ //! Shared benchmark utilities pub mod latency_object_store; mod memory; +mod memory_pool; mod options; mod run; pub use memory::print_memory_stats; +pub use memory_pool::{ + PeakRecordingPool, max_pool_reserved, peak_pool_reserved, reset_peak_pool_reserved, +}; pub use options::CommonOpt; pub use run::{BenchQuery, BenchmarkRun, QueryResult}; diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs index a3e6d2a4c5538..c744d0bf31c7f 100644 --- a/benchmarks/src/util/options.rs +++ b/benchmarks/src/util/options.rs @@ -30,7 +30,7 @@ use datafusion::{ use datafusion_common::{DataFusionError, Result}; use object_store::local::LocalFileSystem; -use super::latency_object_store::LatencyObjectStore; +use super::{latency_object_store::LatencyObjectStore, memory_pool::PeakRecordingPool}; // Common benchmark options (don't use doc comments otherwise this doc // shows up in help files) @@ -125,6 +125,9 @@ impl CommonOpt { ))); } }; + // Record the peak reservation so benchmarks can report it next to + // peak RSS. Purely observational: every call is delegated. + let pool: Arc = Arc::new(PeakRecordingPool::new(pool)); rt_builder = rt_builder .with_memory_pool(pool) .with_disk_manager_builder(DiskManagerBuilder::default()); diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs index df17674e62961..712ff059cc84e 100644 --- a/benchmarks/src/util/run.rs +++ b/benchmarks/src/util/run.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use super::memory_pool::{peak_pool_reserved, reset_peak_pool_reserved}; use datafusion::{DATAFUSION_VERSION, error::Result}; use datafusion_common::utils::get_available_parallelism; use serde::{Serialize, Serializer}; @@ -91,6 +92,13 @@ pub struct BenchQuery { #[serde(serialize_with = "serialize_start_time")] start_time: SystemTime, success: bool, + /// Peak [`MemoryPool`] reservation observed while running this query, in + /// bytes. `None` (and omitted from the JSON) when the benchmark ran without + /// a memory limit, since there is no pool to record. + /// + /// [`MemoryPool`]: datafusion::execution::memory_pool::MemoryPool + #[serde(skip_serializing_if = "Option::is_none")] + pool_peak_bytes: Option, } /// Internal representation of a single benchmark query iteration result. pub struct QueryResult { @@ -121,11 +129,15 @@ impl BenchmarkRun { } /// begin a new case. iterations added after this will be included in the new case pub fn start_new_case(&mut self, id: &str) { + // Give this query its own memory pool reading rather than inheriting + // the high-water mark of the queries that ran before it. + reset_peak_pool_reserved(); self.queries.push(BenchQuery { query: id.to_owned(), iterations: vec![], start_time: SystemTime::now(), success: true, + pool_peak_bytes: None, }); if let Some(c) = self.current_case.as_mut() { *c += 1; @@ -138,7 +150,10 @@ impl BenchmarkRun { if let Some(idx) = self.current_case { self.queries[idx] .iterations - .push(QueryIter { elapsed, row_count }) + .push(QueryIter { elapsed, row_count }); + // The peak is not reset between iterations, so this ends up holding + // the largest reservation seen across all of them. + self.queries[idx].pool_peak_bytes = peak_pool_reserved(); } else { panic!("no cases existed yet"); } From 5bb2c363ea558bb7b1269e657432ca8809a457d5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:25:02 -0500 Subject: [PATCH 2/7] Cover Arrow-side reservations reaching the peak recorder `ArrowMemoryPool` implements Arrow's `MemoryPool` by growing a DataFusion `MemoryReservation` against the pool it wraps, so a buffer claimed through it reaches `PeakRecordingPool::grow` and is recorded like any other reservation. Nothing in DataFusion claims buffers today, but that makes the peak follow the accounting as it changes rather than fixing it to the current set of manually tracked consumers. That property was assumed rather than tested. Add a test that builds an `ArrowMemoryPool` over the recording pool and asserts an Arrow-side reservation both raises the peak and releases on drop. `arrow-buffer/pool` and `datafusion-execution/arrow_buffer_pool` are enabled as dev-dependencies only, so the benchmark binaries are built with exactly the features they were before. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 2 ++ benchmarks/Cargo.toml | 5 +++++ benchmarks/src/util/memory_pool.rs | 36 ++++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 6f776e5e965f9..d5c0c144e9d6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1774,6 +1774,7 @@ name = "datafusion-benchmarks" version = "54.1.0" dependencies = [ "arrow", + "arrow-buffer", "async-trait", "bytes", "clap", @@ -1781,6 +1782,7 @@ dependencies = [ "datafusion", "datafusion-common", "datafusion-common-runtime", + "datafusion-execution", "datafusion-proto", "env_logger", "futures", diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 5dae70761f9a7..282b27e48101d 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -64,6 +64,11 @@ tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] } tokio-util = { version = "0.7.17" } [dev-dependencies] +# `pool`/`arrow_buffer_pool` are enabled only for tests, so the benchmark +# binaries are built exactly as before. They let `memory_pool`'s tests cover +# Arrow-side reservations reaching the pool via `ArrowMemoryPool`. +arrow-buffer = { workspace = true, features = ["pool"] } +datafusion-execution = { workspace = true, features = ["arrow_buffer_pool"] } datafusion-proto = { workspace = true, features = ["parquet"] } tempfile = { workspace = true } diff --git a/benchmarks/src/util/memory_pool.rs b/benchmarks/src/util/memory_pool.rs index 17c83d7cae517..4438a6c31310d 100644 --- a/benchmarks/src/util/memory_pool.rs +++ b/benchmarks/src/util/memory_pool.rs @@ -33,6 +33,12 @@ //! 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::{ @@ -287,4 +293,34 @@ mod tests { assert_eq!(wrapped.to_string(), inner.to_string()); assert!(matches!(wrapped.memory_limit(), MemoryLimit::Finite(4096))); } + + /// Arrow-side reservations reach the recorder too. + /// + /// [`ArrowMemoryPool`] implements Arrow's `MemoryPool` by growing a + /// DataFusion [`MemoryReservation`] against the pool it wraps, so a buffer + /// claimed through it lands in `grow` here. Nothing in DataFusion claims + /// buffers yet (see apache/datafusion#22898), but when something does, the + /// bytes show up in this peak without further changes — as long as the + /// adapter is built from the `RuntimeEnv`'s pool, which is the wrapped one. + /// This test pins that. + #[test] + fn records_reservations_arriving_through_the_arrow_adapter() { + use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait; + use datafusion_execution::memory_pool::arrow::ArrowMemoryPool; + + let (pool, _guard) = pool(4096); + + let arrow_pool = + ArrowMemoryPool::new(Arc::clone(&pool), MemoryConsumer::new("arrow")); + let reservation = arrow_pool.reserve(1024); + + // The Arrow-side reservation is visible as DataFusion pool usage... + assert_eq!(pool.reserved(), 1024); + assert_eq!(peak_pool_reserved(), Some(1024)); + + // ...and dropping it releases the bytes while the peak is retained. + drop(reservation); + assert_eq!(pool.reserved(), 0); + assert_eq!(peak_pool_reserved(), Some(1024)); + } } From 277076c1a1735e2f92d2aa1e354932008c6c8f66 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:23:10 -0500 Subject: [PATCH 3/7] Hold the peak reservation on the pool instead of in statics The high-water marks lived in process-wide statics, which meant the recorded peak could not be attributed to a particular pool, the `Debug` impl reported a number that was not the instance's, and the unit tests needed a mutex to keep from clobbering each other. Move both marks onto `PeakRecordingPool` and reach them through `RuntimeEnv::memory_pool`, which `dyn MemoryPool::downcast_ref` already supports. `BenchmarkRun::set_memory_pool` takes the pool the queries run against, so benchmarks that build a runtime per query now report each query against the pool it actually ran on. No change to what is measured or emitted. --- benchmarks/src/bin/benchmark_runner.rs | 5 +- benchmarks/src/clickbench.rs | 3 +- benchmarks/src/dict.rs | 1 + benchmarks/src/h2o.rs | 3 +- benchmarks/src/hj.rs | 1 + benchmarks/src/imdb/run.rs | 13 +- benchmarks/src/nlj.rs | 1 + benchmarks/src/smj.rs | 1 + benchmarks/src/sort_pushdown.rs | 13 +- benchmarks/src/sort_tpch.rs | 14 +- benchmarks/src/sql_benchmark_runner.rs | 2 +- benchmarks/src/tpcds/run.rs | 3 +- benchmarks/src/tpch/run.rs | 3 +- benchmarks/src/util/memory.rs | 28 ++-- benchmarks/src/util/memory_pool.rs | 204 ++++++++++++++----------- benchmarks/src/util/mod.rs | 4 +- benchmarks/src/util/run.rs | 111 +++++++++++++- 17 files changed, 284 insertions(+), 126 deletions(-) diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs index 71b3b1b9e0a87..c7a16086c9677 100644 --- a/benchmarks/src/bin/benchmark_runner.rs +++ b/benchmarks/src/bin/benchmark_runner.rs @@ -299,6 +299,9 @@ async fn run_simple_benchmark( let case_name = benchmark_case_name(benchmark); + // Each case gets its own `SessionContext`, so hand over its pool before the + // case starts. + run.set_memory_pool(&ctx.runtime_env().memory_pool); run.start_new_case(&case_name); for iteration in 0..config.common.iterations { @@ -312,7 +315,7 @@ async fn run_simple_benchmark( run.write_iter(elapsed, row_count); } - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(()) } diff --git a/benchmarks/src/clickbench.rs b/benchmarks/src/clickbench.rs index 70aaeb7d2d192..a2e65aa5618a9 100644 --- a/benchmarks/src/clickbench.rs +++ b/benchmarks/src/clickbench.rs @@ -213,6 +213,7 @@ impl RunOpt { self.register_hits(&ctx).await?; let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_path = get_query_path(&self.queries_path, query_id); let Some(sql) = get_query_sql(&query_path)? else { @@ -278,7 +279,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/dict.rs b/benchmarks/src/dict.rs index f8451715ea81e..e04b5f816adcc 100644 --- a/benchmarks/src/dict.rs +++ b/benchmarks/src/dict.rs @@ -333,6 +333,7 @@ impl RunOpt { let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query = &DICTIONARY_QUERIES[query_id - 1]; diff --git a/benchmarks/src/h2o.rs b/benchmarks/src/h2o.rs index 8b6e04932cb39..feb4bf2fa11ce 100644 --- a/benchmarks/src/h2o.rs +++ b/benchmarks/src/h2o.rs @@ -109,6 +109,7 @@ impl RunOpt { let iterations = self.common.iterations; let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { benchmark_run.start_new_case(&format!("Query {query_id}")); let sql = queries.get_query(query_id)?; @@ -131,7 +132,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); if self.common.debug { ctx.sql(sql) diff --git a/benchmarks/src/hj.rs b/benchmarks/src/hj.rs index 7d33bc3aa9e50..c74c82f41f699 100644 --- a/benchmarks/src/hj.rs +++ b/benchmarks/src/hj.rs @@ -518,6 +518,7 @@ impl RunOpt { } let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; diff --git a/benchmarks/src/imdb/run.rs b/benchmarks/src/imdb/run.rs index e0e302e466840..5fa8f2b23c569 100644 --- a/benchmarks/src/imdb/run.rs +++ b/benchmarks/src/imdb/run.rs @@ -295,7 +295,7 @@ impl RunOpt { let mut benchmark_run = BenchmarkRun::new(); for query_id in query_range { benchmark_run.start_new_case(&format!("Query {query_id}")); - let query_run = self.benchmark_query(query_id).await?; + let query_run = self.benchmark_query(query_id, &mut benchmark_run).await?; for iter in query_run { benchmark_run.write_iter(iter.elapsed, iter.row_count); } @@ -304,7 +304,13 @@ impl RunOpt { Ok(()) } - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let mut config = self .common .config()? @@ -314,6 +320,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -348,7 +355,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/nlj.rs b/benchmarks/src/nlj.rs index 361cc35ec200c..485ee069d1bba 100644 --- a/benchmarks/src/nlj.rs +++ b/benchmarks/src/nlj.rs @@ -211,6 +211,7 @@ impl RunOpt { let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; // Convert 1-based to 0-based index diff --git a/benchmarks/src/smj.rs b/benchmarks/src/smj.rs index 3d173b7116e2b..9282f72c2fab6 100644 --- a/benchmarks/src/smj.rs +++ b/benchmarks/src/smj.rs @@ -550,6 +550,7 @@ impl RunOpt { let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; // Convert 1-based to 0-based index diff --git a/benchmarks/src/sort_pushdown.rs b/benchmarks/src/sort_pushdown.rs index 86f1c0f5c1119..77f889e702e3d 100644 --- a/benchmarks/src/sort_pushdown.rs +++ b/benchmarks/src/sort_pushdown.rs @@ -137,7 +137,7 @@ impl RunOpt { for query_id in query_ids { benchmark_run.start_new_case(&format!("{query_id}")); - let query_results = self.benchmark_query(query_id).await; + let query_results = self.benchmark_query(query_id, &mut benchmark_run).await; match query_results { Ok(query_results) => { for iter in query_results { @@ -156,7 +156,13 @@ impl RunOpt { Ok(()) } - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let sql = self.load_query(query_id)?; let config = self.common.config()?; @@ -168,6 +174,7 @@ impl RunOpt { .with_default_features() .build(); let ctx = SessionContext::from(state); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); self.register_tables(&ctx).await?; @@ -191,7 +198,7 @@ impl RunOpt { let avg = millis.iter().sum::() / millis.len() as f64; println!("Query {query_id} avg time: {avg:.2} ms"); - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/sort_tpch.rs b/benchmarks/src/sort_tpch.rs index 2182d1a383633..d5f81c04a3ba4 100644 --- a/benchmarks/src/sort_tpch.rs +++ b/benchmarks/src/sort_tpch.rs @@ -187,7 +187,7 @@ impl RunOpt { for query_id in query_range { benchmark_run.start_new_case(&format!("{query_id}")); - let query_results = self.benchmark_query(query_id).await; + let query_results = self.benchmark_query(query_id, &mut benchmark_run).await; match query_results { Ok(query_results) => { for iter in query_results { @@ -207,7 +207,14 @@ impl RunOpt { } /// Benchmark query `query_id` in `SORT_QUERIES` - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let config = self.common.config()?; let rt = self.common.build_runtime()?; let state = SessionStateBuilder::new() @@ -216,6 +223,7 @@ impl RunOpt { .with_default_features() .build(); let ctx = SessionContext::from(state); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -250,7 +258,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/sql_benchmark_runner.rs b/benchmarks/src/sql_benchmark_runner.rs index f1cb3ad2f71b4..420e780c645aa 100644 --- a/benchmarks/src/sql_benchmark_runner.rs +++ b/benchmarks/src/sql_benchmark_runner.rs @@ -110,7 +110,7 @@ fn run_criterion_benchmark( match result { Ok(()) => { - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(()) } Err(payload) => Err(panic_payload_to_error(payload.as_ref())), diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs index 2e0274c935de3..3eaaf172c0f16 100644 --- a/benchmarks/src/tpcds/run.rs +++ b/benchmarks/src/tpcds/run.rs @@ -226,6 +226,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -290,7 +291,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs index 422bcec9ea066..47edfbac4b5a7 100644 --- a/benchmarks/src/tpch/run.rs +++ b/benchmarks/src/tpch/run.rs @@ -137,6 +137,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; let scale_factor = self.scale_factor()?; @@ -208,7 +209,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs index c70c687c7baaa..2b186c79c3516 100644 --- a/benchmarks/src/util/memory.rs +++ b/benchmarks/src/util/memory.rs @@ -15,25 +15,29 @@ // specific language governing permissions and limitations // under the License. +use datafusion::execution::memory_pool::MemoryPool; + +use super::PeakRecordingPool; + /// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by -/// the peak `MemoryPool` reservation when a memory limit was configured. -pub fn print_memory_stats() { +/// the peak reservation of `memory_pool` when a memory limit was configured. +pub fn print_memory_stats(memory_pool: &dyn MemoryPool) { print_allocator_stats(); - print_pool_stats(); + print_pool_stats(memory_pool); } -/// Print the peak `MemoryPool` reservation for the run. +/// Print the peak reservation `memory_pool` has seen. /// -/// Prints nothing when the benchmark ran without a memory limit, since no pool -/// was installed to record. Comparing this against the peak RSS above shows how -/// much of a run's memory the pool actually accounted for — DataFusion only -/// tracks the "large" allocations that scale with input size, so the two are -/// expected to differ. -fn print_pool_stats() { - if let Some(peak) = super::max_pool_reserved() { +/// Prints nothing when the benchmark ran without a memory limit, since no +/// [`PeakRecordingPool`] was installed to record. Comparing this against the +/// peak RSS above shows how much of a run's memory the pool actually accounted +/// for — DataFusion only tracks the "large" allocations that scale with input +/// size, so the two are expected to differ. +fn print_pool_stats(memory_pool: &dyn MemoryPool) { + if let Some(recorder) = PeakRecordingPool::from_pool(memory_pool) { println!( "Peak pool reserved: {}", - datafusion_common::human_readable_size(peak) + datafusion_common::human_readable_size(recorder.max_reserved()) ); } } diff --git a/benchmarks/src/util/memory_pool.rs b/benchmarks/src/util/memory_pool.rs index 4438a6c31310d..3877f2eb6946b 100644 --- a/benchmarks/src/util/memory_pool.rs +++ b/benchmarks/src/util/memory_pool.rs @@ -45,7 +45,7 @@ use std::{ fmt::{Debug, Display, Formatter}, sync::{ Arc, - atomic::{AtomicBool, AtomicUsize, Ordering}, + atomic::{AtomicUsize, Ordering}, }, }; @@ -54,70 +54,27 @@ use datafusion::execution::memory_pool::{ }; 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 { - 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 { - 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. +/// memory is granted, limited, or reported. The one thing it does change is +/// downcasting: `rt.memory_pool.downcast_ref::()` now finds this +/// wrapper instead of the pool it wraps. Nothing in the benchmarks relies on +/// that, and [`Self::from_pool`] uses the same mechanism to find the recorder. +/// +/// Both high-water marks are held per instance, so a benchmark that builds a +/// fresh runtime per query gets a reading scoped to that query without any +/// coordination. /// /// # 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 = -/// Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024)))); -/// reset_peak_pool_reserved(); +/// # use datafusion_benchmarks::util::PeakRecordingPool; +/// let recording = Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024)))); +/// let pool: Arc = Arc::clone(&recording) as _; /// /// let reservation = MemoryConsumer::new("example").register(&pool); /// reservation.try_grow(512)?; @@ -125,23 +82,63 @@ pub fn reset_peak_pool_reserved() { /// /// // 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)); +/// assert_eq!(recording.peak_reserved(), 512); +/// +/// // The recorder can also be recovered from the pool it was installed as. +/// assert_eq!(PeakRecordingPool::from_pool(&*pool).unwrap().peak_reserved(), 512); /// # Ok::<(), datafusion_common::DataFusionError>(()) /// ``` pub struct PeakRecordingPool { inner: Arc, + /// High-water mark since the last [`PeakRecordingPool::reset_peak`]. + peak: AtomicUsize, + /// High-water mark since this pool was created. Never reset. + max: AtomicUsize, } impl PeakRecordingPool { /// Wrap `inner`, recording its peak reservation from here on. pub fn new(inner: Arc) -> Self { - RECORDING.store(true, Ordering::Relaxed); - Self { inner } + Self { + inner, + peak: AtomicUsize::new(0), + max: AtomicUsize::new(0), + } } - /// The wrapped pool. - pub fn inner(&self) -> &Arc { - &self.inner + /// The recorder installed as `pool`, if there is one. + /// + /// Returns `None` whenever a benchmark runs without a memory limit, since + /// [`CommonOpt::runtime_env_builder`] only installs the wrapper alongside a + /// pool it has a limit for. + /// + /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder + pub fn from_pool(pool: &dyn MemoryPool) -> Option<&Self> { + pool.downcast_ref::() + } + + /// Peak reservation, in bytes, since the last [`Self::reset_peak`]. + pub fn peak_reserved(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + /// Peak reservation, in bytes, since this pool was created. + /// + /// Unlike [`Self::peak_reserved`] this is never reset, so it reports the + /// peak across every query that shared this pool. + pub fn max_reserved(&self) -> usize { + self.max.load(Ordering::Relaxed) + } + + /// Reset the value returned by [`Self::peak_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(&self) { + self.peak.store(0, Ordering::Relaxed); } /// Publish the pool's current reservation to both high-water marks. @@ -151,8 +148,8 @@ impl PeakRecordingPool { /// wrapped pool declines or adjusts a request. fn record(&self) { let reserved = self.inner.reserved(); - PEAK_RESERVED.fetch_max(reserved, Ordering::Relaxed); - MAX_RESERVED.fetch_max(reserved, Ordering::Relaxed); + self.peak.fetch_max(reserved, Ordering::Relaxed); + self.max.fetch_max(reserved, Ordering::Relaxed); } } @@ -160,7 +157,8 @@ impl Debug for PeakRecordingPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("PeakRecordingPool") .field("inner", &self.inner) - .field("peak", &PEAK_RESERVED.load(Ordering::Relaxed)) + .field("peak", &self.peak_reserved()) + .field("max", &self.max_reserved()) .finish() } } @@ -212,30 +210,23 @@ impl MemoryPool for PeakRecordingPool { #[cfg(test)] mod tests { - use std::sync::{Mutex, MutexGuard}; - use datafusion::execution::memory_pool::GreedyMemoryPool; use super::*; - /// The high-water marks are process-wide, so these tests would otherwise - /// clobber each other when the test harness runs them in parallel. - static TEST_LOCK: Mutex<()> = Mutex::new(()); - - /// Take the lock and hand back a freshly reset pool. The guard is returned - /// so it stays held for the body of the test. - fn pool(limit: usize) -> (Arc, MutexGuard<'static, ()>) { - let guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); - let pool: Arc = Arc::new(PeakRecordingPool::new(Arc::new( + /// A recording pool over a `GreedyMemoryPool`, returned both as the + /// recorder (to read the marks) and as the pool reservations register with. + fn pool(limit: usize) -> (Arc, Arc) { + let recording = Arc::new(PeakRecordingPool::new(Arc::new( GreedyMemoryPool::new(limit), ))); - reset_peak_pool_reserved(); - (pool, guard) + let pool = Arc::clone(&recording) as Arc; + (recording, pool) } #[test] fn records_high_water_mark_across_reservations() { - let (pool, _guard) = pool(1024); + let (recording, pool) = pool(1024); let a = MemoryConsumer::new("a").register(&pool); let b = MemoryConsumer::new("b").register(&pool); @@ -243,7 +234,7 @@ mod tests { a.try_grow(300).unwrap(); b.try_grow(400).unwrap(); // Peak of the sum, not the largest single reservation. - assert_eq!(peak_pool_reserved(), Some(700)); + assert_eq!(recording.peak_reserved(), 700); a.shrink(300); b.try_grow(100).unwrap(); @@ -251,12 +242,12 @@ mod tests { // Falling back below the peak leaves it untouched, and the later growth // does not reach it. assert_eq!(pool.reserved(), 500); - assert_eq!(peak_pool_reserved(), Some(700)); + assert_eq!(recording.peak_reserved(), 700); } #[test] fn failed_growth_does_not_move_the_peak() { - let (pool, _guard) = pool(1024); + let (recording, pool) = pool(1024); let reservation = MemoryConsumer::new("a").register(&pool); reservation.try_grow(600).unwrap(); @@ -264,28 +255,59 @@ mod tests { .try_grow(600) .expect_err("should exceed the 1024 byte pool"); - assert_eq!(peak_pool_reserved(), Some(600)); + assert_eq!(recording.peak_reserved(), 600); } #[test] fn reset_clears_the_window_but_not_the_run_maximum() { - let (pool, _guard) = pool(1024); + let (recording, pool) = pool(1024); let reservation = MemoryConsumer::new("a").register(&pool); reservation.try_grow(800).unwrap(); reservation.shrink(800); - reset_peak_pool_reserved(); - assert_eq!(peak_pool_reserved(), Some(0)); - assert!(max_pool_reserved().unwrap() >= 800); + recording.reset_peak(); + assert_eq!(recording.peak_reserved(), 0); + assert_eq!(recording.max_reserved(), 800); reservation.try_grow(100).unwrap(); - assert_eq!(peak_pool_reserved(), Some(100)); + assert_eq!(recording.peak_reserved(), 100); + assert_eq!(recording.max_reserved(), 800); + } + + #[test] + fn marks_are_per_instance() { + let (one, one_pool) = pool(1024); + let (two, _two_pool) = pool(1024); + + MemoryConsumer::new("a") + .register(&one_pool) + .try_grow(512) + .unwrap(); + + assert_eq!(one.peak_reserved(), 512); + assert_eq!(two.peak_reserved(), 0); + } + + #[test] + fn is_recoverable_from_the_pool_it_is_installed_as() { + let (recording, pool) = pool(1024); + + MemoryConsumer::new("a") + .register(&pool) + .try_grow(512) + .unwrap(); + + let found = PeakRecordingPool::from_pool(&*pool).expect("recorder installed"); + assert_eq!(found.peak_reserved(), recording.peak_reserved()); + + // A pool with no recorder in front of it reports nothing. + let plain: Arc = Arc::new(GreedyMemoryPool::new(1024)); + assert!(PeakRecordingPool::from_pool(&*plain).is_none()); } #[test] fn delegates_limit_and_name_to_the_wrapped_pool() { - let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let inner: Arc = Arc::new(GreedyMemoryPool::new(4096)); let wrapped = PeakRecordingPool::new(Arc::clone(&inner)); @@ -308,7 +330,7 @@ mod tests { use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait; use datafusion_execution::memory_pool::arrow::ArrowMemoryPool; - let (pool, _guard) = pool(4096); + let (recording, pool) = pool(4096); let arrow_pool = ArrowMemoryPool::new(Arc::clone(&pool), MemoryConsumer::new("arrow")); @@ -316,11 +338,11 @@ mod tests { // The Arrow-side reservation is visible as DataFusion pool usage... assert_eq!(pool.reserved(), 1024); - assert_eq!(peak_pool_reserved(), Some(1024)); + assert_eq!(recording.peak_reserved(), 1024); // ...and dropping it releases the bytes while the peak is retained. drop(reservation); assert_eq!(pool.reserved(), 0); - assert_eq!(peak_pool_reserved(), Some(1024)); + assert_eq!(recording.peak_reserved(), 1024); } } diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs index c9ffdbc079e92..43855ea468ef5 100644 --- a/benchmarks/src/util/mod.rs +++ b/benchmarks/src/util/mod.rs @@ -23,8 +23,6 @@ mod options; mod run; pub use memory::print_memory_stats; -pub use memory_pool::{ - PeakRecordingPool, max_pool_reserved, peak_pool_reserved, reset_peak_pool_reserved, -}; +pub use memory_pool::PeakRecordingPool; pub use options::CommonOpt; pub use run::{BenchQuery, BenchmarkRun, QueryResult}; diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs index 712ff059cc84e..40e97f2247376 100644 --- a/benchmarks/src/util/run.rs +++ b/benchmarks/src/util/run.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -use super::memory_pool::{peak_pool_reserved, reset_peak_pool_reserved}; +use super::memory_pool::PeakRecordingPool; +use datafusion::execution::memory_pool::MemoryPool; use datafusion::{DATAFUSION_VERSION, error::Result}; use datafusion_common::utils::get_available_parallelism; use serde::{Serialize, Serializer}; @@ -23,6 +24,7 @@ use serde_json::Value; use std::{ collections::HashMap, path::Path, + sync::Arc, time::{Duration, SystemTime}, }; @@ -110,6 +112,10 @@ pub struct BenchmarkRun { context: RunContext, queries: Vec, current_case: Option, + /// The pool queries run against, when one was handed over with + /// [`BenchmarkRun::set_memory_pool`]. Only read through + /// [`BenchmarkRun::peak_recorder`]. + memory_pool: Option>, } impl Default for BenchmarkRun { @@ -125,13 +131,38 @@ impl BenchmarkRun { context: RunContext::new(), queries: vec![], current_case: None, + memory_pool: None, } } + + /// Report the peak reservation of `memory_pool` alongside each query. + /// + /// Call this with the pool of the [`RuntimeEnv`] the queries run against. + /// Has no effect unless a [`PeakRecordingPool`] is installed, which + /// [`CommonOpt::runtime_env_builder`] does whenever a memory limit is + /// configured; without one `pool_peak_bytes` is omitted from the results. + /// + /// Benchmarks that build a runtime per query should call this each time, so + /// each query reports against the pool it actually ran on. + /// + /// [`RuntimeEnv`]: datafusion::execution::runtime_env::RuntimeEnv + /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder + pub fn set_memory_pool(&mut self, memory_pool: &Arc) { + self.memory_pool = Some(Arc::clone(memory_pool)); + } + + /// The recorder in front of the pool set by [`Self::set_memory_pool`]. + fn peak_recorder(&self) -> Option<&PeakRecordingPool> { + PeakRecordingPool::from_pool(self.memory_pool.as_deref()?) + } + /// begin a new case. iterations added after this will be included in the new case pub fn start_new_case(&mut self, id: &str) { // Give this query its own memory pool reading rather than inheriting // the high-water mark of the queries that ran before it. - reset_peak_pool_reserved(); + if let Some(recorder) = self.peak_recorder() { + recorder.reset_peak(); + } self.queries.push(BenchQuery { query: id.to_owned(), iterations: vec![], @@ -147,13 +178,14 @@ impl BenchmarkRun { } /// Write a new iteration to the current case pub fn write_iter(&mut self, elapsed: Duration, row_count: usize) { + // The peak is not reset between iterations, so this ends up holding the + // largest reservation seen across all of them. + let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved); if let Some(idx) = self.current_case { self.queries[idx] .iterations .push(QueryIter { elapsed, row_count }); - // The peak is not reset between iterations, so this ends up holding - // the largest reservation seen across all of them. - self.queries[idx].pool_peak_bytes = peak_pool_reserved(); + self.queries[idx].pool_peak_bytes = pool_peak_bytes; } else { panic!("no cases existed yet"); } @@ -197,3 +229,72 @@ impl BenchmarkRun { Ok(()) } } + +#[cfg(test)] +mod tests { + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer}; + + use super::*; + + fn recording_pool(limit: usize) -> Arc { + Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new( + limit, + )))) + } + + #[test] + fn each_case_reports_its_own_peak() { + let pool = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&pool); + + run.start_new_case("q1"); + let reservation = MemoryConsumer::new("q1").register(&pool); + reservation.try_grow(600).unwrap(); + run.write_iter(Duration::from_millis(1), 1); + drop(reservation); + + // The second case must not inherit the first case's high-water mark. + run.start_new_case("q2"); + let reservation = MemoryConsumer::new("q2").register(&pool); + reservation.try_grow(100).unwrap(); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(600)); + assert_eq!(run.queries[1].pool_peak_bytes, Some(100)); + } + + #[test] + fn a_later_pool_replaces_an_earlier_one() { + let first = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&first); + MemoryConsumer::new("q1") + .register(&first) + .try_grow(600) + .unwrap(); + + // Benchmarks that build a runtime per query hand over the new pool + // before the next case; the reading follows it. + let second = recording_pool(1024); + run.set_memory_pool(&second); + run.start_new_case("q2"); + MemoryConsumer::new("q2") + .register(&second) + .try_grow(100) + .unwrap(); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(100)); + } + + #[test] + fn the_peak_is_omitted_without_a_recording_pool() { + let mut run = BenchmarkRun::new(); + run.start_new_case("q1"); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, None); + assert!(!run.to_json().contains("pool_peak_bytes")); + } +} From f008bce482838c47c61343293d2dd4bae1af2228 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:23:10 -0500 Subject: [PATCH 4/7] Track the peak from deltas instead of re-reading reserved() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Recording the peak by calling `inner.reserved()` after every grow put the wrapped pool's bookkeeping on the allocation path. `--mem-pool-type` defaults to `fair`, and `FairSpillPool::reserved()` takes the same state lock `try_grow` just released, so every accounted allocation acquired it twice — in a harness whose job is measuring time, on exactly the spilling benchmarks where the number matters. Keep a running total in the wrapper instead. It stays exact because the trait grants exactly what is asked for: `grow` is infallible and `try_grow` either grants `additional` or leaves the reservation alone. Resetting now sets the mark to what is currently reserved rather than to zero, so a query that starts with data already held reports that as its floor instead of under-reporting until the next grow. --- benchmarks/src/util/memory_pool.rs | 57 +++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 12 deletions(-) diff --git a/benchmarks/src/util/memory_pool.rs b/benchmarks/src/util/memory_pool.rs index 3877f2eb6946b..a3606ca0a7b7a 100644 --- a/benchmarks/src/util/memory_pool.rs +++ b/benchmarks/src/util/memory_pool.rs @@ -90,6 +90,9 @@ use datafusion_common::Result; /// ``` pub struct PeakRecordingPool { inner: Arc, + /// Running total of everything granted through this wrapper, kept so the + /// peak can be maintained without asking `inner` for its total. + reserved: AtomicUsize, /// High-water mark since the last [`PeakRecordingPool::reset_peak`]. peak: AtomicUsize, /// High-water mark since this pool was created. Never reset. @@ -98,9 +101,13 @@ pub struct PeakRecordingPool { impl PeakRecordingPool { /// Wrap `inner`, recording its peak reservation from here on. + /// + /// `inner` is expected to be empty: the running total starts at zero, so + /// anything reserved before wrapping is not counted. pub fn new(inner: Arc) -> Self { Self { inner, + reserved: AtomicUsize::new(0), peak: AtomicUsize::new(0), max: AtomicUsize::new(0), } @@ -130,24 +137,33 @@ impl PeakRecordingPool { self.max.load(Ordering::Relaxed) } - /// Reset the value returned by [`Self::peak_reserved`], so the next reading - /// covers only what follows. + /// Reset the value returned by [`Self::peak_reserved`] to what is reserved + /// right now, so the next reading covers only what follows. /// /// [`BenchmarkRun::start_new_case`] calls this, giving each benchmark query - /// its own reading. + /// its own reading. Anything still held when a query starts — data the + /// benchmark loaded up front, say — stays in the reading, since the query + /// runs with those bytes reserved. /// /// [`BenchmarkRun::start_new_case`]: super::BenchmarkRun::start_new_case pub fn reset_peak(&self) { - self.peak.store(0, Ordering::Relaxed); + self.peak + .store(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed); } - /// Publish the pool's current reservation to both high-water marks. + /// Add `additional` granted bytes to the running total and publish it to + /// both high-water marks. /// - /// Called after any operation that can raise `reserved()`. Reading the - /// total rather than accumulating deltas keeps this correct when the - /// wrapped pool declines or adjusts a request. - fn record(&self) { - let reserved = self.inner.reserved(); + /// Accumulating deltas rather than reading [`MemoryPool::reserved`] keeps + /// the wrapped pool's own bookkeeping off this path: `FairSpillPool` takes + /// its state lock to answer `reserved()`, which would double the lock + /// traffic of every accounted allocation in the benchmark being measured. + /// The total stays exact because the trait grants exactly what is asked + /// for — `grow` is infallible and `try_grow` either grants `additional` or + /// returns an error, leaving the reservation untouched. + fn record(&self, additional: usize) { + let reserved = + self.reserved.fetch_add(additional, Ordering::Relaxed) + additional; self.peak.fetch_max(reserved, Ordering::Relaxed); self.max.fetch_max(reserved, Ordering::Relaxed); } @@ -186,16 +202,17 @@ impl MemoryPool for PeakRecordingPool { fn grow(&self, reservation: &MemoryReservation, additional: usize) { self.inner.grow(reservation, additional); - self.record(); + self.record(additional); } fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { self.inner.shrink(reservation, shrink); + self.reserved.fetch_sub(shrink, Ordering::Relaxed); } fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { self.inner.try_grow(reservation, additional)?; - self.record(); + self.record(additional); Ok(()) } @@ -275,6 +292,22 @@ mod tests { assert_eq!(recording.max_reserved(), 800); } + #[test] + fn reset_keeps_what_is_still_reserved() { + let (recording, pool) = pool(1024); + + // Something a benchmark loaded up front and holds across queries. + let held = MemoryConsumer::new("held").register(&pool); + held.try_grow(300).unwrap(); + + recording.reset_peak(); + assert_eq!(recording.peak_reserved(), 300); + + let query = MemoryConsumer::new("query").register(&pool); + query.try_grow(200).unwrap(); + assert_eq!(recording.peak_reserved(), 500); + } + #[test] fn marks_are_per_instance() { let (one, one_pool) = pool(1024); From 3e93ce0f7c68b91c5fdce86f85bc7b1116ca6514 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:23:10 -0500 Subject: [PATCH 5/7] Record the pool peak for failed queries too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `pool_peak_bytes` was only assigned in `write_iter`, which a query that failed never reaches. Its absence was documented as meaning "no memory limit was configured", so an out-of-memory query — the case where the peak is most worth seeing — serialized identically to a run with no pool at all. Assign it in `mark_failed` as well. --- benchmarks/src/util/run.rs | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs index 40e97f2247376..6c63ceec6423c 100644 --- a/benchmarks/src/util/run.rs +++ b/benchmarks/src/util/run.rs @@ -95,8 +95,11 @@ pub struct BenchQuery { start_time: SystemTime, success: bool, /// Peak [`MemoryPool`] reservation observed while running this query, in - /// bytes. `None` (and omitted from the JSON) when the benchmark ran without - /// a memory limit, since there is no pool to record. + /// bytes. Recorded for failed queries too, since a query that ran out of + /// memory is one whose peak is worth seeing. + /// + /// `None` (and omitted from the JSON) only when the benchmark ran without a + /// memory limit, since there is then no pool to record. /// /// [`MemoryPool`]: datafusion::execution::memory_pool::MemoryPool #[serde(skip_serializing_if = "Option::is_none")] @@ -206,8 +209,12 @@ impl BenchmarkRun { /// Mark current query pub fn mark_failed(&mut self) { + // A query that failed under a memory limit wrote no iteration, so this + // is the only chance to record what it had reserved when it gave up. + let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved); if let Some(idx) = self.current_case { self.queries[idx].success = false; + self.queries[idx].pool_peak_bytes = pool_peak_bytes; } else { unreachable!("Cannot mark failure: no current case"); } @@ -288,6 +295,21 @@ mod tests { assert_eq!(run.queries[0].pool_peak_bytes, Some(100)); } + #[test] + fn a_failed_query_still_reports_its_peak() { + let pool = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&pool); + + run.start_new_case("q1"); + let reservation = MemoryConsumer::new("q1").register(&pool); + reservation.try_grow(600).unwrap(); + // No `write_iter`: the query failed before completing an iteration. + run.mark_failed(); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(600)); + } + #[test] fn the_peak_is_omitted_without_a_recording_pool() { let mut run = BenchmarkRun::new(); From 6dce8d61db6d5e9aae79f357b1eae632d121217d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:23:10 -0500 Subject: [PATCH 6/7] Report the pool peak from external_aggr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `external_aggr` builds its pool directly instead of going through `CommonOpt::runtime_env_builder`, so it was the one benchmark that never got a recorder installed — despite being the benchmark that exists to run queries under a memory limit. Wrap the pool it builds and hand it to the run, so each (query, limit) case reports its peak like the rest. --- benchmarks/src/bin/external_aggr.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/benchmarks/src/bin/external_aggr.rs b/benchmarks/src/bin/external_aggr.rs index 42f25c2cb010c..226a619192ac9 100644 --- a/benchmarks/src/bin/external_aggr.rs +++ b/benchmarks/src/bin/external_aggr.rs @@ -39,7 +39,9 @@ use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::{collect, displayable}; use datafusion::prelude::*; -use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, QueryResult}; +use datafusion_benchmarks::util::{ + BenchmarkRun, CommonOpt, PeakRecordingPool, QueryResult, +}; use datafusion_common::instant::Instant; use datafusion_common::utils::get_available_parallelism; use datafusion_common::{DEFAULT_PARQUET_EXTENSION, exec_err}; @@ -169,7 +171,7 @@ impl ExternalAggrConfig { )); let query_results = self - .benchmark_query(query_id, mem_limit, mem_pool_type) + .benchmark_query(query_id, mem_limit, mem_pool_type, &mut benchmark_run) .await?; for iter in query_results { benchmark_run.write_iter(iter.elapsed, iter.row_count); @@ -182,11 +184,15 @@ impl ExternalAggrConfig { } /// Benchmark query `query_id` in `AGGR_QUERIES` + /// + /// `benchmark_run` is handed this query's runtime, which is built here + /// because each query runs under its own memory limit. async fn benchmark_query( &self, query_id: usize, mem_limit: u64, mem_pool_type: &str, + benchmark_run: &mut BenchmarkRun, ) -> Result> { let query_name = format!("Q{query_id}({})", human_readable_size(mem_limit as usize)); @@ -198,6 +204,12 @@ impl ExternalAggrConfig { return exec_err!("Invalid memory pool type: {}", mem_pool_type); } }; + // This benchmark builds its pool directly rather than going through + // `CommonOpt::runtime_env_builder`, so it has to install the recorder + // itself to report a peak. + let memory_pool: Arc = + Arc::new(PeakRecordingPool::new(memory_pool)); + benchmark_run.set_memory_pool(&memory_pool); let runtime_env = RuntimeEnvBuilder::new() .with_memory_pool(memory_pool) .build_arc()?; From 885d16d8c81bbacbc875527853ac817fa77f03cb Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:23:10 -0500 Subject: [PATCH 7/7] Document pool_peak_bytes for benchmark authors The README's "Collecting data" section is the contract benchmarks follow to produce results JSON, and it did not mention the new field or the `set_memory_pool` call that populates it. --- benchmarks/README.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/benchmarks/README.md b/benchmarks/README.md index 34c67e5151ba1..b6a7705cf94e3 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -483,6 +483,14 @@ Your benchmark should create and use an instance of `BenchmarkRun` defined in `b - Call its `start_new_case` method with a string that will appear in the "Query" column of the compare output. - Use `write_iter` to record elapsed times for the behavior you're benchmarking. +- Call `set_memory_pool` with the `RuntimeEnv`'s memory pool (`ctx.runtime_env().memory_pool`), + and again for each new runtime if your benchmark builds one per query. Each case then reports a + `pool_peak_bytes` field: the peak `MemoryPool` reservation reached while running it, which is the + largest value across that case's iterations. The field is omitted when the benchmark runs without + `--memory-limit`, since no pool is installed to record. Comparing it against the peak RSS printed + by `print_memory_stats` shows how much of the run's memory the pool actually accounted for; the + pool only tracks the "large" allocations that scale with input size, so the two are expected to + differ. - When all cases are done, call the `BenchmarkRun`'s `maybe_write_json` method, giving it the value of the `--output` structopt field on `RunOpt`.