Fix test_tracked_consumers_pool_deregister

Author: erenavsarogullari

datafusion/execution/src/memory_pool/pool.rs

@@ -423,7 +423,6 @@ pub struct TrackConsumersPool<I> {
 
 impl<I: MemoryPool> Display for TrackConsumersPool<I> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        // let tracked_consumers = self.tracked_consumers.lock();
         write!(
             f,
             "{}",
@@ -482,6 +481,11 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
         }
     }
 
+    /// Returns a reference to the wrapped inner [`MemoryPool`].
+    pub fn inner(&self) -> &I {
+        &self.inner
+    }
+
     /// Returns a snapshot of all currently tracked consumers.
     pub fn metrics(&self) -> Vec<MemoryConsumerMetrics> {
         self.tracked_consumers
@@ -624,7 +628,7 @@ fn provide_top_memory_consumers_to_error_msg(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use insta::{Settings, allow_duplicates, assert_snapshot};
+    use insta::{Settings, allow_duplicates, assert_snapshot, with_settings};
     use std::sync::Arc;
 
     fn make_settings() -> Settings {
@@ -824,72 +828,80 @@ mod tests {
 
     #[test]
     fn test_tracked_consumers_pool_deregister() {
-        fn test_per_pool_type(pool: Arc<dyn MemoryPool>) {
-            // Baseline: see the 2 memory consumers
-            let setting = make_settings();
-            let _bound = setting.bind_to_scope();
-            let r0 = MemoryConsumer::new("r0").register(&pool);
-            r0.grow(10);
-            let r1_consumer = MemoryConsumer::new("r1");
-            let r1 = r1_consumer.register(&pool);
-            r1.grow(20);
-
-            let res = r0.try_grow(150);
-            assert!(res.is_err());
-            let error = res.unwrap_err().strip_backtrace();
-            allow_duplicates!(assert_snapshot!(error, @r"
-            Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
-              r1#[ID](can spill: false) consumed 20.0 B, peak 20.0 B,
-              r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
-            Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 70.0 B remain available for the total memory pool: greedy(used: 30.0 B, pool_size: 100.0 B)
-            "));
-
-            // Test: unregister one
-            // only the remaining one should be listed
-            drop(r1);
-            let res = r0.try_grow(150);
-            assert!(res.is_err());
-            let error = res.unwrap_err().strip_backtrace();
-            allow_duplicates!(assert_snapshot!(error, @r"
-            Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
-              r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
-            Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: fair(pool_size: 100.0 B)
-            "));
-
-            // Test: actual message we see is the `available is 70`. When it should be `available is 90`.
-            // This is because the pool.shrink() does not automatically occur within the inner_pool.deregister().
-            let res = r0.try_grow(150);
-            assert!(res.is_err());
-            let error = res.unwrap_err().strip_backtrace();
-            allow_duplicates!(assert_snapshot!(error, @r"
-            Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
-              r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
-            Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: fair(pool_size: 100.0 B)
-            "));
-
-            // Test: the registration needs to free itself (or be dropped),
-            // for the proper error message
-            let res = r0.try_grow(150);
-            assert!(res.is_err());
-            let error = res.unwrap_err().strip_backtrace();
-            allow_duplicates!(assert_snapshot!(error, @r"
-            Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
-              r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
-            Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: fair(pool_size: 100.0 B)
-            "));
+        fn test_per_pool_type<P: MemoryPool + 'static>(pool: Arc<TrackConsumersPool<P>>) {
+            // `snapshot_suffix` ties each insta snapshot to this pool's inner backend; filters
+            // normalize inner pool `Display` so fair vs greedy share the same `@` reference text.
+            with_settings!({
+                snapshot_suffix => pool.inner().name().to_string(),
+                filters => vec![
+                    (
+                        r"([^\s]+)\#\d+\(can spill: (true|false)\)",
+                        "$1#[ID](can spill: $2)",
+                    ),
+                    (
+                        r"for the total memory pool: [^\n]+",
+                        "for the total memory pool: [INNER_POOL]",
+                    ),
+                ],
+            }, {
+                let memory_pool: Arc<dyn MemoryPool> = Arc::<TrackConsumersPool<P>>::clone(&pool);
+                let r0 = MemoryConsumer::new("r0").register(&memory_pool);
+                r0.grow(10);
+                let r1 = MemoryConsumer::new("r1").register(&memory_pool);
+                r1.grow(20);
+
+                // Baseline: see the 2 memory consumers
+                let error = r0.try_grow(150).unwrap_err().strip_backtrace();
+                assert_snapshot!(error, @r"
+                Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
+                  r1#[ID](can spill: false) consumed 20.0 B, peak 20.0 B,
+                  r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
+                Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 70.0 B remain available for the total memory pool: [INNER_POOL]
+                ");
+
+                // Test: unregister one — only the remaining consumer should be listed
+                drop(r1);
+                let error = r0.try_grow(150).unwrap_err().strip_backtrace();
+                assert_snapshot!(error, @r"
+                Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
+                  r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
+                Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: [INNER_POOL]
+                ");
+
+                // Test: actual message we see is the `available is 70`. When it should be `available is 90`.
+                // This is because the pool.shrink() does not automatically occur within the inner_pool.deregister().
+                let error = r0.try_grow(150).unwrap_err().strip_backtrace();
+                assert_snapshot!(error, @r"
+                Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
+                  r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
+                Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: [INNER_POOL]
+                ");
+
+                // Test: the registration needs to free itself (or be dropped),
+                // for the proper error message
+                let error = r0.try_grow(150).unwrap_err().strip_backtrace();
+                assert_snapshot!(error, @r"
+                Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
+                  r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
+                Error: Failed to allocate additional 150.0 B for r0 with 10.0 B already allocated for this reservation - 90.0 B remain available for the total memory pool: [INNER_POOL]
+                ");
+                }
+            );
         }
 
-        let tracked_spill_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
-            FairSpillPool::new(100),
-            NonZeroUsize::new(3).unwrap(),
-        ));
-        test_per_pool_type(tracked_spill_pool);
+        allow_duplicates! {
+            let tracked_spill_pool = Arc::new(TrackConsumersPool::new(
+                FairSpillPool::new(100),
+                NonZeroUsize::new(3).unwrap(),
+            ));
+            test_per_pool_type(tracked_spill_pool);
 
-        let tracked_greedy_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
-            GreedyMemoryPool::new(100),
-            NonZeroUsize::new(3).unwrap(),
-        ));
-        test_per_pool_type(tracked_greedy_pool);
+            let tracked_greedy_pool = Arc::new(TrackConsumersPool::new(
+                GreedyMemoryPool::new(100),
+                NonZeroUsize::new(3).unwrap(),
+            ));
+            test_per_pool_type(tracked_greedy_pool);
+        }
     }
 
     #[test]