## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/execution/src/memory_pool/pool.rs

    Line
    Count
    Source (jump to first uncovered line)
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    use crate::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
    19
    use datafusion_common::{resources_datafusion_err, DataFusionError, Result};
    20
    use hashbrown::HashMap;
    21
    use log::debug;
    22
    use parking_lot::Mutex;
    23
    use std::{
    24
        num::NonZeroUsize,
    25
        sync::atomic::{AtomicU64, AtomicUsize, Ordering},
    26
    };
    27
    28
    /// A [`MemoryPool`] that enforces no limit
    29
    #[derive(Debug, Default)]
    30
    pub struct UnboundedMemoryPool {
    31
        used: AtomicUsize,
    32
    }
    33
    34
    impl MemoryPool for UnboundedMemoryPool {
    35
    43.5k
        fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
    36
    43.5k
            self.used.fetch_add(additional, Ordering::Relaxed);
    37
    43.5k
        }
    38
    39
    22.2k
        fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
    40
    22.2k
            self.used.fetch_sub(shrink, Ordering::Relaxed);
    41
    22.2k
        }
    42
    43
    43.5k
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    44
    43.5k
            self.grow(reservation, additional);
    45
    43.5k
            Ok(())
    46
    43.5k
        }
    47
    48
    15
        fn reserved(&self) -> usize {
    49
    15
            self.used.load(Ordering::Relaxed)
    50
    15
        }
    51
    }
    52
    53
    /// A [`MemoryPool`] that implements a greedy first-come first-serve limit.
    54
    ///
    55
    /// This pool works well for queries that do not need to spill or have
    56
    /// a single spillable operator. See [`FairSpillPool`] if there are
    57
    /// multiple spillable operators that all will spill.
    58
    #[derive(Debug)]
    59
    pub struct GreedyMemoryPool {
    60
        pool_size: usize,
    61
        used: AtomicUsize,
    62
    }
    63
    64
    impl GreedyMemoryPool {
    65
        /// Allocate up to `limit` bytes
    66
    34
        pub fn new(pool_size: usize) -> Self {
    67
    34
            debug!("Created new GreedyMemoryPool(pool_size={pool_size})"0);
    68
    34
            Self {
    69
    34
                pool_size,
    70
    34
                used: AtomicUsize::new(0),
    71
    34
            }
    72
    34
        }
    73
    }
    74
    75
    impl MemoryPool for GreedyMemoryPool {
    76
    0
        fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
    77
    0
            self.used.fetch_add(additional, Ordering::Relaxed);
    78
    0
        }
    79
    80
    41
        fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
    81
    41
            self.used.fetch_sub(shrink, Ordering::Relaxed);
    82
    41
        }
    83
    84
    217
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    85
    217
            self.used
    86
    217
                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
    87
    217
                    let new_used = used + additional;
    88
    217
                    (new_used <= self.pool_size).then_some(new_used)
    89
    217
                })
    90
    217
                .map_err(|used| {
    91
    81
                    insufficient_capacity_err(
    92
    81
                        reservation,
    93
    81
                        additional,
    94
    81
                        self.pool_size.saturating_sub(used),
    95
    81
                    )
    96
    217
                })?81;
    97
    136
            Ok(())
    98
    217
        }
    99
    100
    1
        fn reserved(&self) -> usize {
    101
    1
            self.used.load(Ordering::Relaxed)
    102
    1
        }
    103
    }
    104
    105
    /// A [`MemoryPool`] that prevents spillable reservations from using more than
    106
    /// an even fraction of the available memory sans any unspillable reservations
    107
    /// (i.e. `(pool_size - unspillable_memory) / num_spillable_reservations`)
    108
    ///
    109
    /// This pool works best when you know beforehand the query has
    110
    /// multiple spillable operators that will likely all need to
    111
    /// spill. Sometimes it will cause spills even when there was
    112
    /// sufficient memory (reserved for other operators) to avoid doing
    113
    /// so.
    114
    ///
    115
    /// ```text
    116
    ///    ┌───────────────────────z──────────────────────z───────────────┐
    117
    ///    │                       z                      z               │
    118
    ///    │                       z                      z               │
    119
    ///    │       Spillable       z       Unspillable    z     Free      │
    120
    ///    │        Memory         z        Memory        z    Memory     │
    121
    ///    │                       z                      z               │
    122
    ///    │                       z                      z               │
    123
    ///    └───────────────────────z──────────────────────z───────────────┘
    124
    /// ```
    125
    ///
    126
    /// Unspillable memory is allocated in a first-come, first-serve fashion
    127
    #[derive(Debug)]
    128
    pub struct FairSpillPool {
    129
        /// The total memory limit
    130
        pool_size: usize,
    131
    132
        state: Mutex<FairSpillPoolState>,
    133
    }
    134
    135
    #[derive(Debug)]
    136
    struct FairSpillPoolState {
    137
        /// The number of consumers that can spill
    138
        num_spill: usize,
    139
    140
        /// The total amount of memory reserved that can be spilled
    141
        spillable: usize,
    142
    143
        /// The total amount of memory reserved by consumers that cannot spill
    144
        unspillable: usize,
    145
    }
    146
    147
    impl FairSpillPool {
    148
        /// Allocate up to `limit` bytes
    149
    12
        pub fn new(pool_size: usize) -> Self {
    150
    12
            debug!("Created new FairSpillPool(pool_size={pool_size})"0);
    151
    12
            Self {
    152
    12
                pool_size,
    153
    12
                state: Mutex::new(FairSpillPoolState {
    154
    12
                    num_spill: 0,
    155
    12
                    spillable: 0,
    156
    12
                    unspillable: 0,
    157
    12
                }),
    158
    12
            }
    159
    12
        }
    160
    }
    161
    162
    impl MemoryPool for FairSpillPool {
    163
    32
        fn register(&self, consumer: &MemoryConsumer) {
    164
    32
            if consumer.can_spill {
    165
    32
                self.state.lock().num_spill += 1;
    166
    32
            }0
    167
    32
        }
    168
    169
    32
        fn unregister(&self, consumer: &MemoryConsumer) {
    170
    32
            if consumer.can_spill {
    171
    32
                let mut state = self.state.lock();
    172
    32
                state.num_spill = state.num_spill.checked_sub(1).unwrap();
    173
    32
            }0
    174
    32
        }
    175
    176
    0
        fn grow(&self, reservation: &MemoryReservation, additional: usize) {
    177
    0
            let mut state = self.state.lock();
    178
    0
            match reservation.registration.consumer.can_spill {
    179
    0
                true => state.spillable += additional,
    180
    0
                false => state.unspillable += additional,
    181
            }
    182
    0
        }
    183
    184
    84
        fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
    185
    84
            let mut state = self.state.lock();
    186
    84
            match reservation.registration.consumer.can_spill {
    187
    84
                true => state.spillable -= shrink,
    188
    0
                false => state.unspillable -= shrink,
    189
            }
    190
    84
        }
    191
    192
    162
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    193
    162
            let mut state = self.state.lock();
    194
    162
    195
    162
            match reservation.registration.consumer.can_spill {
    196
                true => {
    197
                    // The total amount of memory available to spilling consumers
    198
    162
                    let spill_available = self.pool_size.saturating_sub(state.unspillable);
    199
    162
    200
    162
                    // No spiller may use more than their fraction of the memory available
    201
    162
                    let available = spill_available
    202
    162
                        .checked_div(state.num_spill)
    203
    162
                        .unwrap_or(spill_available);
    204
    162
    205
    162
                    if reservation.size + additional > available {
    206
    104
                        return Err(insufficient_capacity_err(
    207
    104
                            reservation,
    208
    104
                            additional,
    209
    104
                            available,
    210
    104
                        ));
    211
    58
                    }
    212
    58
                    state.spillable += additional;
    213
                }
    214
                false => {
    215
    0
                    let available = self
    216
    0
                        .pool_size
    217
    0
                        .saturating_sub(state.unspillable + state.spillable);
    218
    0
    219
    0
                    if available < additional {
    220
    0
                        return Err(insufficient_capacity_err(
    221
    0
                            reservation,
    222
    0
                            additional,
    223
    0
                            available,
    224
    0
                        ));
    225
    0
                    }
    226
    0
                    state.unspillable += additional;
    227
                }
    228
            }
    229
    58
            Ok(())
    230
    162
        }
    231
    232
    0
        fn reserved(&self) -> usize {
    233
    0
            let state = self.state.lock();
    234
    0
            state.spillable + state.unspillable
    235
    0
        }
    236
    }
    237
    238
    /// Constructs a resources error based upon the individual [`MemoryReservation`].
    239
    ///
    240
    /// The error references the `bytes already allocated` for the reservation,
    241
    /// and not the total within the collective [`MemoryPool`],
    242
    /// nor the total across multiple reservations with the same [`MemoryConsumer`].
    243
    #[inline(always)]
    244
    185
    fn insufficient_capacity_err(
    245
    185
        reservation: &MemoryReservation,
    246
    185
        additional: usize,
    247
    185
        available: usize,
    248
    185
    ) -> DataFusionError {
    249
    185
        resources_datafusion_err!("Failed to allocate additional {} bytes for {} with {} bytes already allocated for this reservation - {} bytes remain available for the total pool", additional, reservation.registration.consumer.name, reservation.size, available)
    250
    185
    }
    251
    252
    /// A [`MemoryPool`] that tracks the consumers that have
    253
    /// reserved memory within the inner memory pool.
    254
    ///
    255
    /// By tracking memory reservations more carefully this pool
    256
    /// can provide better error messages on the largest memory users
    257
    ///
    258
    /// Tracking is per hashed [`MemoryConsumer`], not per [`MemoryReservation`].
    259
    /// The same consumer can have multiple reservations.
    260
    #[derive(Debug)]
    261
    pub struct TrackConsumersPool<I> {
    262
        inner: I,
    263
        top: NonZeroUsize,
    264
        tracked_consumers: Mutex<HashMap<MemoryConsumer, AtomicU64>>,
    265
    }
    266
    267
    impl<I: MemoryPool> TrackConsumersPool<I> {
    268
        /// Creates a new [`TrackConsumersPool`].
    269
        ///
    270
        /// The `top` determines how many Top K [`MemoryConsumer`]s to include
    271
        /// in the reported [`DataFusionError::ResourcesExhausted`].
    272
    34
        pub fn new(inner: I, top: NonZeroUsize) -> Self {
    273
    34
            Self {
    274
    34
                inner,
    275
    34
                top,
    276
    34
                tracked_consumers: Default::default(),
    277
    34
            }
    278
    34
        }
    279
    280
        /// Determine if there are multiple [`MemoryConsumer`]s registered
    281
        /// which have the same name.
    282
        ///
    283
        /// This is very tied to the implementation of the memory consumer.
    284
    89
        fn has_multiple_consumers(&self, name: &String) -> bool {
    285
    89
            let consumer = MemoryConsumer::new(name);
    286
    89
            let consumer_with_spill = consumer.clone().with_can_spill(true);
    287
    89
            let guard = self.tracked_consumers.lock();
    288
    89
            guard.contains_key(&consumer) && guard.contains_key(&consumer_with_spill)83
    289
    89
        }
    290
    291
        /// The top consumers in a report string.
    292
    81
        pub fn report_top(&self, top: usize) -> String {
    293
    81
            let mut consumers = self
    294
    81
                .tracked_consumers
    295
    81
                .lock()
    296
    81
                .iter()
    297
    89
                .map(|(consumer, reserved)| {
    298
    89
                    (
    299
    89
                        (consumer.name().to_owned(), consumer.can_spill()),
    300
    89
                        reserved.load(Ordering::Acquire),
    301
    89
                    )
    302
    89
                })
    303
    81
                .collect::<Vec<_>>();
    304
    81
            consumers.sort_by(|a, b| b.1.cmp(&a.1)8); // inverse ordering
    305
    81
    306
    81
            consumers[0..std::cmp::min(top, consumers.len())]
    307
    81
                .iter()
    308
    89
                .map(|((name, can_spill), size)| {
    309
    89
                    if self.has_multiple_consumers(name) {
    310
    0
                        format!("{name}(can_spill={}) consumed {:?} bytes", can_spill, size)
    311
                    } else {
    312
    89
                        format!("{name} consumed {:?} bytes", size)
    313
                    }
    314
    89
                })
    315
    81
                .collect::<Vec<_>>()
    316
    81
                .join(", ")
    317
    81
        }
    318
    }
    319
    320
    impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
    321
    72
        fn register(&self, consumer: &MemoryConsumer) {
    322
    72
            self.inner.register(consumer);
    323
    72
    324
    72
            let mut guard = self.tracked_consumers.lock();
    325
    72
            if let Some(already_reserved0) = guard.insert(consumer.clone(), Default::default())
    326
    0
            {
    327
    0
                guard.entry_ref(consumer).and_modify(|bytes| {
    328
    0
                    bytes.fetch_add(
    329
    0
                        already_reserved.load(Ordering::Acquire),
    330
    0
                        Ordering::AcqRel,
    331
    0
                    );
    332
    0
                });
    333
    72
            }
    334
    72
        }
    335
    336
    72
        fn unregister(&self, consumer: &MemoryConsumer) {
    337
    72
            self.inner.unregister(consumer);
    338
    72
            self.tracked_consumers.lock().remove(consumer);
    339
    72
        }
    340
    341
    0
        fn grow(&self, reservation: &MemoryReservation, additional: usize) {
    342
    0
            self.inner.grow(reservation, additional);
    343
    0
            self.tracked_consumers
    344
    0
                .lock()
    345
    0
                .entry_ref(reservation.consumer())
    346
    0
                .and_modify(|bytes| {
    347
    0
                    bytes.fetch_add(additional as u64, Ordering::AcqRel);
    348
    0
                });
    349
    0
        }
    350
    351
    41
        fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
    352
    41
            self.inner.shrink(reservation, shrink);
    353
    41
            self.tracked_consumers
    354
    41
                .lock()
    355
    41
                .entry_ref(reservation.consumer())
    356
    41
                .and_modify(|bytes| {
    357
    41
                    bytes.fetch_sub(shrink as u64, Ordering::AcqRel);
    358
    41
                });
    359
    41
        }
    360
    361
    217
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    362
    217
            self.inner
    363
    217
                .try_grow(reservation, additional)
    364
    217
                .map_err(|e| match e81 {
    365
    81
                    DataFusionError::ResourcesExhausted(e) => {
    366
    81
                        // wrap OOM message in top consumers
    367
    81
                        DataFusionError::ResourcesExhausted(
    368
    81
                            provide_top_memory_consumers_to_error_msg(
    369
    81
                                e,
    370
    81
                                self.report_top(self.top.into()),
    371
    81
                            ),
    372
    81
                        )
    373
                    }
    374
    0
                    _ => e,
    375
    217
                }81)?81;
    376
    377
    136
            self.tracked_consumers
    378
    136
                .lock()
    379
    136
                .entry_ref(reservation.consumer())
    380
    136
                .and_modify(|bytes| {
    381
    136
                    bytes.fetch_add(additional as u64, Ordering::AcqRel);
    382
    136
                });
    383
    136
            Ok(())
    384
    217
        }
    385
    386
    1
        fn reserved(&self) -> usize {
    387
    1
            self.inner.reserved()
    388
    1
        }
    389
    }
    390
    391
    81
    fn provide_top_memory_consumers_to_error_msg(
    392
    81
        error_msg: String,
    393
    81
        top_consumers: String,
    394
    81
    ) -> String {
    395
    81
        format!("Additional allocation failed with top memory consumers (across reservations) as: {}. Error: {}", top_consumers, error_msg)
    396
    81
    }
    397
    398
    #[cfg(test)]
    399
    mod tests {
    400
        use super::*;
    401
        use std::sync::Arc;
    402
    403
        #[test]
    404
        fn test_fair() {
    405
            let pool = Arc::new(FairSpillPool::new(100)) as _;
    406
    407
            let mut r1 = MemoryConsumer::new("unspillable").register(&pool);
    408
            // Can grow beyond capacity of pool
    409
            r1.grow(2000);
    410
            assert_eq!(pool.reserved(), 2000);
    411
    412
            let mut r2 = MemoryConsumer::new("r2")
    413
                .with_can_spill(true)
    414
                .register(&pool);
    415
            // Can grow beyond capacity of pool
    416
            r2.grow(2000);
    417
    418
            assert_eq!(pool.reserved(), 4000);
    419
    420
            let err = r2.try_grow(1).unwrap_err().strip_backtrace();
    421
            assert_eq!(err, "Resources exhausted: Failed to allocate additional 1 bytes for r2 with 2000 bytes already allocated for this reservation - 0 bytes remain available for the total pool");
    422
    423
            let err = r2.try_grow(1).unwrap_err().strip_backtrace();
    424
            assert_eq!(err, "Resources exhausted: Failed to allocate additional 1 bytes for r2 with 2000 bytes already allocated for this reservation - 0 bytes remain available for the total pool");
    425
    426
            r1.shrink(1990);
    427
            r2.shrink(2000);
    428
    429
            assert_eq!(pool.reserved(), 10);
    430
    431
            r1.try_grow(10).unwrap();
    432
            assert_eq!(pool.reserved(), 20);
    433
    434
            // Can grow r2 to 80 as only spilling consumer
    435
            r2.try_grow(80).unwrap();
    436
            assert_eq!(pool.reserved(), 100);
    437
    438
            r2.shrink(70);
    439
    440
            assert_eq!(r1.size(), 20);
    441
            assert_eq!(r2.size(), 10);
    442
            assert_eq!(pool.reserved(), 30);
    443
    444
            let mut r3 = MemoryConsumer::new("r3")
    445
                .with_can_spill(true)
    446
                .register(&pool);
    447
    448
            let err = r3.try_grow(70).unwrap_err().strip_backtrace();
    449
            assert_eq!(err, "Resources exhausted: Failed to allocate additional 70 bytes for r3 with 0 bytes already allocated for this reservation - 40 bytes remain available for the total pool");
    450
    451
            //Shrinking r2 to zero doesn't allow a3 to allocate more than 45
    452
            r2.free();
    453
            let err = r3.try_grow(70).unwrap_err().strip_backtrace();
    454
            assert_eq!(err, "Resources exhausted: Failed to allocate additional 70 bytes for r3 with 0 bytes already allocated for this reservation - 40 bytes remain available for the total pool");
    455
    456
            // But dropping r2 does
    457
            drop(r2);
    458
            assert_eq!(pool.reserved(), 20);
    459
            r3.try_grow(80).unwrap();
    460
    461
            assert_eq!(pool.reserved(), 100);
    462
            r1.free();
    463
            assert_eq!(pool.reserved(), 80);
    464
    465
            let mut r4 = MemoryConsumer::new("s4").register(&pool);
    466
            let err = r4.try_grow(30).unwrap_err().strip_backtrace();
    467
            assert_eq!(err, "Resources exhausted: Failed to allocate additional 30 bytes for s4 with 0 bytes already allocated for this reservation - 20 bytes remain available for the total pool");
    468
        }
    469
    470
        #[test]
    471
        fn test_tracked_consumers_pool() {
    472
            let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    473
                GreedyMemoryPool::new(100),
    474
                NonZeroUsize::new(3).unwrap(),
    475
            ));
    476
    477
            // Test: use all the different interfaces to change reservation size
    478
    479
            // set r1=50, using grow and shrink
    480
            let mut r1 = MemoryConsumer::new("r1").register(&pool);
    481
            r1.grow(70);
    482
            r1.shrink(20);
    483
    484
            // set r2=15 using try_grow
    485
            let mut r2 = MemoryConsumer::new("r2").register(&pool);
    486
            r2.try_grow(15)
    487
                .expect("should succeed in memory allotment for r2");
    488
    489
            // set r3=20 using try_resize
    490
            let mut r3 = MemoryConsumer::new("r3").register(&pool);
    491
            r3.try_resize(25)
    492
                .expect("should succeed in memory allotment for r3");
    493
            r3.try_resize(20)
    494
                .expect("should succeed in memory allotment for r3");
    495
    496
            // set r4=10
    497
            // this should not be reported in top 3
    498
            let mut r4 = MemoryConsumer::new("r4").register(&pool);
    499
            r4.grow(10);
    500
    501
            // Test: reports if new reservation causes error
    502
            // using the previously set sizes for other consumers
    503
            let mut r5 = MemoryConsumer::new("r5").register(&pool);
    504
            let expected = "Additional allocation failed with top memory consumers (across reservations) as: r1 consumed 50 bytes, r3 consumed 20 bytes, r2 consumed 15 bytes. Error: Failed to allocate additional 150 bytes for r5 with 0 bytes already allocated for this reservation - 5 bytes remain available for the total pool";
    505
            let res = r5.try_grow(150);
    506
            assert!(
    507
                matches!(
    508
                    &res,
    509
                    Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    510
                ),
    511
                "should provide list of top memory consumers, instead found {:?}",
    512
                res
    513
            );
    514
        }
    515
    516
        #[test]
    517
        fn test_tracked_consumers_pool_register() {
    518
            let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    519
                GreedyMemoryPool::new(100),
    520
                NonZeroUsize::new(3).unwrap(),
    521
            ));
    522
    523
            let same_name = "foo";
    524
    525
            // Test: see error message when no consumers recorded yet
    526
            let mut r0 = MemoryConsumer::new(same_name).register(&pool);
    527
            let expected = "Additional allocation failed with top memory consumers (across reservations) as: foo consumed 0 bytes. Error: Failed to allocate additional 150 bytes for foo with 0 bytes already allocated for this reservation - 100 bytes remain available for the total pool";
    528
            let res = r0.try_grow(150);
    529
            assert!(
    530
                matches!(
    531
                    &res,
    532
                    Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    533
                ),
    534
                "should provide proper error when no reservations have been made yet, instead found {:?}", res
    535
            );
    536
    537
            // API: multiple registrations using the same hashed consumer,
    538
            // will be recognized as the same in the TrackConsumersPool.
    539
    540
            // Test: will be the same per Top Consumers reported.
    541
            r0.grow(10); // make r0=10, pool available=90
    542
            let new_consumer_same_name = MemoryConsumer::new(same_name);
    543
            let mut r1 = new_consumer_same_name.register(&pool);
    544
            // TODO: the insufficient_capacity_err() message is per reservation, not per consumer.
    545
            // a followup PR will clarify this message "0 bytes already allocated for this reservation"
    546
            let expected = "Additional allocation failed with top memory consumers (across reservations) as: foo consumed 10 bytes. Error: Failed to allocate additional 150 bytes for foo with 0 bytes already allocated for this reservation - 90 bytes remain available for the total pool";
    547
            let res = r1.try_grow(150);
    548
            assert!(
    549
                matches!(
    550
                    &res,
    551
                    Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    552
                ),
    553
                "should provide proper error with same hashed consumer (a single foo=10 bytes, available=90), instead found {:?}", res
    554
            );
    555
    556
            // Test: will accumulate size changes per consumer, not per reservation
    557
            r1.grow(20);
    558
            let expected = "Additional allocation failed with top memory consumers (across reservations) as: foo consumed 30 bytes. Error: Failed to allocate additional 150 bytes for foo with 20 bytes already allocated for this reservation - 70 bytes remain available for the total pool";
    559
            let res = r1.try_grow(150);
    560
            assert!(
    561
                matches!(
    562
                    &res,
    563
                    Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    564
                ),
    565
                "should provide proper error with same hashed consumer (a single foo=30 bytes, available=70), instead found {:?}", res
    566
            );
    567
    568
            // Test: different hashed consumer, (even with the same name),
    569
            // will be recognized as different in the TrackConsumersPool
    570
            let consumer_with_same_name_but_different_hash =
    571
                MemoryConsumer::new(same_name).with_can_spill(true);
    572
            let mut r2 = consumer_with_same_name_but_different_hash.register(&pool);
    573
            let expected = "Additional allocation failed with top memory consumers (across reservations) as: foo(can_spill=false) consumed 30 bytes, foo(can_spill=true) consumed 0 bytes. Error: Failed to allocate additional 150 bytes for foo with 0 bytes already allocated for this reservation - 70 bytes remain available for the total pool";
    574
            let res = r2.try_grow(150);
    575
            assert!(
    576
                matches!(
    577
                    &res,
    578
                    Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    579
                ),
    580
                "should provide proper error with different hashed consumer (foo(can_spill=false)=30 bytes and foo(can_spill=true)=0 bytes, available=70), instead found {:?}", res
    581
            );
    582
        }
    583
    584
        #[test]
    585
        fn test_tracked_consumers_pool_deregister() {
    586
            fn test_per_pool_type(pool: Arc<dyn MemoryPool>) {
    587
                // Baseline: see the 2 memory consumers
    588
                let mut r0 = MemoryConsumer::new("r0").register(&pool);
    589
                r0.grow(10);
    590
                let r1_consumer = MemoryConsumer::new("r1");
    591
                let mut r1 = r1_consumer.clone().register(&pool);
    592
                r1.grow(20);
    593
                let expected = "Additional allocation failed with top memory consumers (across reservations) as: r1 consumed 20 bytes, r0 consumed 10 bytes. Error: Failed to allocate additional 150 bytes for r0 with 10 bytes already allocated for this reservation - 70 bytes remain available for the total pool";
    594
                let res = r0.try_grow(150);
    595
                assert!(
    596
                    matches!(
    597
                        &res,
    598
                        Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected)
    599
                    ),
    600
                    "should provide proper error with both consumers, instead found {:?}",
    601
                    res
    602
                );
    603
    604
                // Test: unregister one
    605
                // only the remaining one should be listed
    606
                pool.unregister(&r1_consumer);
    607
                let expected_consumers = "Additional allocation failed with top memory consumers (across reservations) as: r0 consumed 10 bytes";
    608
                let res = r0.try_grow(150);
    609
                assert!(
    610
                    matches!(
    611
                        &res,
    612
                        Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected_consumers)
    613
                    ),
    614
                    "should provide proper error with only 1 consumer left registered, instead found {:?}", res
    615
                );
    616
    617
                // Test: actual message we see is the `available is 70`. When it should be `available is 90`.
    618
                // This is because the pool.shrink() does not automatically occur within the inner_pool.deregister().
    619
                let expected_70_available = "Failed to allocate additional 150 bytes for r0 with 10 bytes already allocated for this reservation - 70 bytes remain available for the total pool";
    620
                let res = r0.try_grow(150);
    621
                assert!(
    622
                    matches!(
    623
                        &res,
    624
                        Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected_70_available)
    625
                    ),
    626
                    "should find that the inner pool will still count all bytes for the deregistered consumer until the reservation is dropped, instead found {:?}", res
    627
                );
    628
    629
                // Test: the registration needs to free itself (or be dropped),
    630
                // for the proper error message
    631
                r1.free();
    632
                let expected_90_available = "Failed to allocate additional 150 bytes for r0 with 10 bytes already allocated for this reservation - 90 bytes remain available for the total pool";
    633
                let res = r0.try_grow(150);
    634
                assert!(
    635
                    matches!(
    636
                        &res,
    637
                        Err(DataFusionError::ResourcesExhausted(ref e)) if e.to_string().contains(expected_90_available)
    638
                    ),
    639
                    "should correctly account the total bytes after reservation is free, instead found {:?}", res
    640
                );
    641
            }
    642
    643
            let tracked_spill_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    644
                FairSpillPool::new(100),
    645
                NonZeroUsize::new(3).unwrap(),
    646
            ));
    647
            test_per_pool_type(tracked_spill_pool);
    648
    649
            let tracked_greedy_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    650
                GreedyMemoryPool::new(100),
    651
                NonZeroUsize::new(3).unwrap(),
    652
            ));
    653
            test_per_pool_type(tracked_greedy_pool);
    654
        }
    655
    656
        #[test]
    657
        fn test_tracked_consumers_pool_use_beyond_errors() {
    658
            let upcasted: Arc<dyn std::any::Any + Send + Sync> =
    659
                Arc::new(TrackConsumersPool::new(
    660
                    GreedyMemoryPool::new(100),
    661
                    NonZeroUsize::new(3).unwrap(),
    662
                ));
    663
            let pool: Arc<dyn MemoryPool> = Arc::clone(&upcasted)
    664
                .downcast::<TrackConsumersPool<GreedyMemoryPool>>()
    665
                .unwrap();
    666
            // set r1=20
    667
            let mut r1 = MemoryConsumer::new("r1").register(&pool);
    668
            r1.grow(20);
    669
            // set r2=15
    670
            let mut r2 = MemoryConsumer::new("r2").register(&pool);
    671
            r2.grow(15);
    672
            // set r3=45
    673
            let mut r3 = MemoryConsumer::new("r3").register(&pool);
    674
            r3.grow(45);
    675
    676
            let downcasted = upcasted
    677
                .downcast::<TrackConsumersPool<GreedyMemoryPool>>()
    678
                .unwrap();
    679
    680
            // Test: can get runtime metrics, even without an error thrown
    681
            let expected = "r3 consumed 45 bytes, r1 consumed 20 bytes";
    682
            let res = downcasted.report_top(2);
    683
            assert_eq!(
    684
                res, expected,
    685
                "should provide list of top memory consumers, instead found {:?}",
    686
                res
    687
            );
    688
        }
    689
    }