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 <noreply@anthropic.com>

Author: adriangb

Cargo.lock

@@ -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 }
 

benchmarks/Cargo.toml

@@ -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));
+    }
 }