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/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`. 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/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()?; 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 11b96ef227756..2b186c79c3516 100644 --- a/benchmarks/src/util/memory.rs +++ b/benchmarks/src/util/memory.rs @@ -15,8 +15,34 @@ // specific language governing permissions and limitations // under the License. -/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api -pub fn print_memory_stats() { +use datafusion::execution::memory_pool::MemoryPool; + +use super::PeakRecordingPool; + +/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by +/// 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(memory_pool); +} + +/// Print the peak reservation `memory_pool` has seen. +/// +/// 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(recorder.max_reserved()) + ); + } +} + +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..a3606ca0a7b7a --- /dev/null +++ b/benchmarks/src/util/memory_pool.rs @@ -0,0 +1,381 @@ +// 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::{AtomicUsize, Ordering}, + }, +}; + +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use datafusion_common::Result; + +/// 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. 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; +/// 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)?; +/// reservation.shrink(512); +/// +/// // The pool is back to empty, but the high-water mark is retained. +/// assert_eq!(pool.reserved(), 0); +/// 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, + /// 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. + max: AtomicUsize, +} + +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), + } + } + + /// 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`] 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. 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(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed); + } + + /// Add `additional` granted bytes to the running total and publish it to + /// both high-water marks. + /// + /// 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); + } +} + +impl Debug for PeakRecordingPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PeakRecordingPool") + .field("inner", &self.inner) + .field("peak", &self.peak_reserved()) + .field("max", &self.max_reserved()) + .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(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(additional); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use datafusion::execution::memory_pool::GreedyMemoryPool; + + use super::*; + + /// 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), + ))); + let pool = Arc::clone(&recording) as Arc; + (recording, pool) + } + + #[test] + fn records_high_water_mark_across_reservations() { + let (recording, pool) = 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!(recording.peak_reserved(), 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!(recording.peak_reserved(), 700); + } + + #[test] + fn failed_growth_does_not_move_the_peak() { + let (recording, pool) = 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!(recording.peak_reserved(), 600); + } + + #[test] + fn reset_clears_the_window_but_not_the_run_maximum() { + let (recording, pool) = pool(1024); + + let reservation = MemoryConsumer::new("a").register(&pool); + reservation.try_grow(800).unwrap(); + reservation.shrink(800); + + recording.reset_peak(); + assert_eq!(recording.peak_reserved(), 0); + assert_eq!(recording.max_reserved(), 800); + + reservation.try_grow(100).unwrap(); + assert_eq!(recording.peak_reserved(), 100); + 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); + 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 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))); + } + + /// 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 (recording, pool) = 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!(recording.peak_reserved(), 1024); + + // ...and dropping it releases the bytes while the peak is retained. + drop(reservation); + assert_eq!(pool.reserved(), 0); + assert_eq!(recording.peak_reserved(), 1024); + } +} diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs index 6dc11c0f425bd..43855ea468ef5 100644 --- a/benchmarks/src/util/mod.rs +++ b/benchmarks/src/util/mod.rs @@ -18,9 +18,11 @@ //! 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; 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..6c63ceec6423c 100644 --- a/benchmarks/src/util/run.rs +++ b/benchmarks/src/util/run.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +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}; @@ -22,6 +24,7 @@ use serde_json::Value; use std::{ collections::HashMap, path::Path, + sync::Arc, time::{Duration, SystemTime}, }; @@ -91,6 +94,16 @@ pub struct BenchQuery { #[serde(serialize_with = "serialize_start_time")] start_time: SystemTime, success: bool, + /// Peak [`MemoryPool`] reservation observed while running this query, in + /// 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")] + pool_peak_bytes: Option, } /// Internal representation of a single benchmark query iteration result. pub struct QueryResult { @@ -102,6 +115,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 { @@ -117,15 +134,44 @@ 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. + if let Some(recorder) = self.peak_recorder() { + recorder.reset_peak(); + } 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; @@ -135,10 +181,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 }) + .push(QueryIter { elapsed, row_count }); + self.queries[idx].pool_peak_bytes = pool_peak_bytes; } else { panic!("no cases existed yet"); } @@ -159,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"); } @@ -182,3 +236,87 @@ 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 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(); + 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")); + } +}