## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

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

    Line
    Count
    Source
    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::{
    19
        human_readable_size, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
    20
    };
    21
    use datafusion_common::HashMap;
    22
    use datafusion_common::{resources_datafusion_err, DataFusionError, Result};
    23
    use log::debug;
    24
    use parking_lot::Mutex;
    25
    use std::{
    26
        num::NonZeroUsize,
    27
        sync::atomic::{AtomicUsize, Ordering},
    28
    };
    29
    30
    /// A [`MemoryPool`] that enforces no limit
    31
    #[derive(Debug, Default)]
    32
    pub struct UnboundedMemoryPool {
    33
        used: AtomicUsize,
    34
    }
    35
    36
    impl MemoryPool for UnboundedMemoryPool {
    37
    0
        fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
    38
    0
            self.used.fetch_add(additional, Ordering::Relaxed);
    39
    0
        }
    40
    41
    0
        fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
    42
    0
            self.used.fetch_sub(shrink, Ordering::Relaxed);
    43
    0
        }
    44
    45
    0
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    46
    0
            self.grow(reservation, additional);
    47
    0
            Ok(())
    48
    0
        }
    49
    50
    0
        fn reserved(&self) -> usize {
    51
    0
            self.used.load(Ordering::Relaxed)
    52
    0
        }
    53
    54
    0
        fn memory_limit(&self) -> MemoryLimit {
    55
    0
            MemoryLimit::Infinite
    56
    0
        }
    57
    }
    58
    59
    /// A [`MemoryPool`] that implements a greedy first-come first-serve limit.
    60
    ///
    61
    /// This pool works well for queries that do not need to spill or have
    62
    /// a single spillable operator. See [`FairSpillPool`] if there are
    63
    /// multiple spillable operators that all will spill.
    64
    #[derive(Debug)]
    65
    pub struct GreedyMemoryPool {
    66
        pool_size: usize,
    67
        used: AtomicUsize,
    68
    }
    69
    70
    impl GreedyMemoryPool {
    71
        /// Create a new pool that can allocate up to `pool_size` bytes
    72
    0
        pub fn new(pool_size: usize) -> Self {
    73
    0
            debug!("Created new GreedyMemoryPool(pool_size={pool_size})");
    74
    0
            Self {
    75
    0
                pool_size,
    76
    0
                used: AtomicUsize::new(0),
    77
    0
            }
    78
    0
        }
    79
    }
    80
    81
    impl MemoryPool for GreedyMemoryPool {
    82
    0
        fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
    83
    0
            self.used.fetch_add(additional, Ordering::Relaxed);
    84
    0
        }
    85
    86
    0
        fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
    87
    0
            self.used.fetch_sub(shrink, Ordering::Relaxed);
    88
    0
        }
    89
    90
    0
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    91
    0
            self.used
    92
    0
                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
    93
    0
                    let new_used = used + additional;
    94
    0
                    (new_used <= self.pool_size).then_some(new_used)
    95
    0
                })
    96
    0
                .map_err(|used| {
    97
    0
                    insufficient_capacity_err(
    98
    0
                        reservation,
    99
    0
                        additional,
    100
    0
                        self.pool_size.saturating_sub(used),
    101
                    )
    102
    0
                })?;
    103
    0
            Ok(())
    104
    0
        }
    105
    106
    0
        fn reserved(&self) -> usize {
    107
    0
            self.used.load(Ordering::Relaxed)
    108
    0
        }
    109
    110
    0
        fn memory_limit(&self) -> MemoryLimit {
    111
    0
            MemoryLimit::Finite(self.pool_size)
    112
    0
        }
    113
    }
    114
    115
    /// A [`MemoryPool`] that prevents spillable reservations from using more than
    116
    /// an even fraction of the available memory sans any unspillable reservations
    117
    /// (i.e. `(pool_size - unspillable_memory) / num_spillable_reservations`)
    118
    ///
    119
    /// This pool works best when you know beforehand the query has
    120
    /// multiple spillable operators that will likely all need to
    121
    /// spill. Sometimes it will cause spills even when there was
    122
    /// sufficient memory (reserved for other operators) to avoid doing
    123
    /// so.
    124
    ///
    125
    /// ```text
    126
    ///    ┌───────────────────────z──────────────────────z───────────────┐
    127
    ///    │                       z                      z               │
    128
    ///    │                       z                      z               │
    129
    ///    │       Spillable       z       Unspillable    z     Free      │
    130
    ///    │        Memory         z        Memory        z    Memory     │
    131
    ///    │                       z                      z               │
    132
    ///    │                       z                      z               │
    133
    ///    └───────────────────────z──────────────────────z───────────────┘
    134
    /// ```
    135
    ///
    136
    /// Unspillable memory is allocated in a first-come, first-serve fashion
    137
    #[derive(Debug)]
    138
    pub struct FairSpillPool {
    139
        /// The total memory limit
    140
        pool_size: usize,
    141
    142
        state: Mutex<FairSpillPoolState>,
    143
    }
    144
    145
    #[derive(Debug)]
    146
    struct FairSpillPoolState {
    147
        /// The number of consumers that can spill
    148
        num_spill: usize,
    149
    150
        /// The total amount of memory reserved that can be spilled
    151
        spillable: usize,
    152
    153
        /// The total amount of memory reserved by consumers that cannot spill
    154
        unspillable: usize,
    155
    }
    156
    157
    impl FairSpillPool {
    158
        /// Allocate up to `limit` bytes
    159
    0
        pub fn new(pool_size: usize) -> Self {
    160
    0
            debug!("Created new FairSpillPool(pool_size={pool_size})");
    161
    0
            Self {
    162
    0
                pool_size,
    163
    0
                state: Mutex::new(FairSpillPoolState {
    164
    0
                    num_spill: 0,
    165
    0
                    spillable: 0,
    166
    0
                    unspillable: 0,
    167
    0
                }),
    168
    0
            }
    169
    0
        }
    170
    }
    171
    172
    impl MemoryPool for FairSpillPool {
    173
    0
        fn register(&self, consumer: &MemoryConsumer) {
    174
    0
            if consumer.can_spill {
    175
    0
                self.state.lock().num_spill += 1;
    176
    0
            }
    177
    0
        }
    178
    179
    0
        fn unregister(&self, consumer: &MemoryConsumer) {
    180
    0
            if consumer.can_spill {
    181
    0
                let mut state = self.state.lock();
    182
    0
                state.num_spill = state.num_spill.checked_sub(1).unwrap();
    183
    0
            }
    184
    0
        }
    185
    186
    0
        fn grow(&self, reservation: &MemoryReservation, additional: usize) {
    187
    0
            let mut state = self.state.lock();
    188
    0
            match reservation.registration.consumer.can_spill {
    189
    0
                true => state.spillable += additional,
    190
    0
                false => state.unspillable += additional,
    191
            }
    192
    0
        }
    193
    194
    0
        fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
    195
    0
            let mut state = self.state.lock();
    196
    0
            match reservation.registration.consumer.can_spill {
    197
    0
                true => state.spillable -= shrink,
    198
    0
                false => state.unspillable -= shrink,
    199
            }
    200
    0
        }
    201
    202
    0
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    203
    0
            let mut state = self.state.lock();
    204
    205
    0
            match reservation.registration.consumer.can_spill {
    206
                true => {
    207
                    // The total amount of memory available to spilling consumers
    208
    0
                    let spill_available = self.pool_size.saturating_sub(state.unspillable);
    209
    210
                    // No spiller may use more than their fraction of the memory available
    211
    0
                    let available = spill_available
    212
    0
                        .checked_div(state.num_spill)
    213
    0
                        .unwrap_or(spill_available);
    214
    215
    0
                    if reservation.size + additional > available {
    216
    0
                        return Err(insufficient_capacity_err(
    217
    0
                            reservation,
    218
    0
                            additional,
    219
    0
                            available,
    220
    0
                        ));
    221
    0
                    }
    222
    0
                    state.spillable += additional;
    223
                }
    224
                false => {
    225
    0
                    let available = self
    226
    0
                        .pool_size
    227
    0
                        .saturating_sub(state.unspillable + state.spillable);
    228
    229
    0
                    if available < additional {
    230
    0
                        return Err(insufficient_capacity_err(
    231
    0
                            reservation,
    232
    0
                            additional,
    233
    0
                            available,
    234
    0
                        ));
    235
    0
                    }
    236
    0
                    state.unspillable += additional;
    237
                }
    238
            }
    239
    0
            Ok(())
    240
    0
        }
    241
    242
    0
        fn reserved(&self) -> usize {
    243
    0
            let state = self.state.lock();
    244
    0
            state.spillable + state.unspillable
    245
    0
        }
    246
    247
    0
        fn memory_limit(&self) -> MemoryLimit {
    248
    0
            MemoryLimit::Finite(self.pool_size)
    249
    0
        }
    250
    }
    251
    252
    /// Constructs a resources error based upon the individual [`MemoryReservation`].
    253
    ///
    254
    /// The error references the `bytes already allocated` for the reservation,
    255
    /// and not the total within the collective [`MemoryPool`],
    256
    /// nor the total across multiple reservations with the same [`MemoryConsumer`].
    257
    #[inline(always)]
    258
    0
    fn insufficient_capacity_err(
    259
    0
        reservation: &MemoryReservation,
    260
    0
        additional: usize,
    261
    0
        available: usize,
    262
    0
    ) -> DataFusionError {
    263
    0
        resources_datafusion_err!("Failed to allocate additional {} for {} with {} already allocated for this reservation - {} remain available for the total pool",
    264
    0
        human_readable_size(additional), reservation.registration.consumer.name, human_readable_size(reservation.size), human_readable_size(available))
    265
    0
    }
    266
    267
    #[derive(Debug)]
    268
    struct TrackedConsumer {
    269
        name: String,
    270
        can_spill: bool,
    271
        reserved: AtomicUsize,
    272
        peak: AtomicUsize,
    273
    }
    274
    275
    impl TrackedConsumer {
    276
        /// Shorthand to return the currently reserved value
    277
    0
        fn reserved(&self) -> usize {
    278
    0
            self.reserved.load(Ordering::Relaxed)
    279
    0
        }
    280
    281
        /// Return the peak value
    282
    0
        fn peak(&self) -> usize {
    283
    0
            self.peak.load(Ordering::Relaxed)
    284
    0
        }
    285
    286
        /// Grows the tracked consumer's reserved size,
    287
        /// should be called after the pool has successfully performed the grow().
    288
    0
        fn grow(&self, additional: usize) {
    289
    0
            self.reserved.fetch_add(additional, Ordering::Relaxed);
    290
    0
            self.peak.fetch_max(self.reserved(), Ordering::Relaxed);
    291
    0
        }
    292
    293
        /// Reduce the tracked consumer's reserved size,
    294
        /// should be called after the pool has successfully performed the shrink().
    295
    0
        fn shrink(&self, shrink: usize) {
    296
    0
            self.reserved.fetch_sub(shrink, Ordering::Relaxed);
    297
    0
        }
    298
    }
    299
    300
    /// A [`MemoryPool`] that tracks the consumers that have
    301
    /// reserved memory within the inner memory pool.
    302
    ///
    303
    /// By tracking memory reservations more carefully this pool
    304
    /// can provide better error messages on the largest memory users
    305
    /// when memory allocation fails.
    306
    ///
    307
    /// Tracking is per hashed [`MemoryConsumer`], not per [`MemoryReservation`].
    308
    /// The same consumer can have multiple reservations.
    309
    ///
    310
    /// # Automatic Usage via [`RuntimeEnvBuilder`]
    311
    ///
    312
    /// The easiest way to use `TrackConsumersPool` is via
    313
    /// [`RuntimeEnvBuilder::with_memory_limit()`].
    314
    ///
    315
    /// [`RuntimeEnvBuilder`]: crate::runtime_env::RuntimeEnvBuilder
    316
    /// [`RuntimeEnvBuilder::with_memory_limit()`]: crate::runtime_env::RuntimeEnvBuilder::with_memory_limit
    317
    ///
    318
    /// # Usage Examples
    319
    ///
    320
    /// For more examples of using `TrackConsumersPool`, see the [memory_pool_tracking.rs] example
    321
    ///
    322
    /// [memory_pool_tracking.rs]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/memory_pool_tracking.rs
    323
    /// [memory_pool_execution_plan.rs]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/memory_pool_execution_plan.rs
    324
    #[derive(Debug)]
    325
    pub struct TrackConsumersPool<I> {
    326
        /// The wrapped memory pool that actually handles reservation logic
    327
        inner: I,
    328
        /// The amount of consumers to report(ordered top to bottom by reservation size)
    329
        top: NonZeroUsize,
    330
        /// Maps consumer_id --> TrackedConsumer
    331
        tracked_consumers: Mutex<HashMap<usize, TrackedConsumer>>,
    332
    }
    333
    334
    impl<I: MemoryPool> TrackConsumersPool<I> {
    335
        /// Creates a new [`TrackConsumersPool`].
    336
        ///
    337
        /// # Arguments
    338
        /// * `inner` - The underlying memory pool that handles actual memory allocation
    339
        /// * `top` - The number of top memory consumers to include in error messages
    340
        ///
    341
        /// # Note
    342
        /// In most cases, you should use [`RuntimeEnvBuilder::with_memory_limit()`](crate::runtime_env::RuntimeEnvBuilder::with_memory_limit)
    343
        /// instead of creating this pool manually, as it automatically sets up tracking with
    344
        /// sensible defaults (top 5 consumers).
    345
        ///
    346
        /// # Example
    347
        ///
    348
        /// ```rust
    349
        /// use datafusion_execution::memory_pool::{
    350
        ///     FairSpillPool, GreedyMemoryPool, TrackConsumersPool,
    351
        /// };
    352
        /// use std::num::NonZeroUsize;
    353
        ///
    354
        /// // Create with a greedy pool backend, reporting top 3 consumers in error messages
    355
        /// let tracked_greedy = TrackConsumersPool::new(
    356
        ///     GreedyMemoryPool::new(1024 * 1024), // 1MB limit
    357
        ///     NonZeroUsize::new(3).unwrap(),
    358
        /// );
    359
        ///
    360
        /// // Create with a fair spill pool backend, reporting top 5 consumers in error messages
    361
        /// let tracked_fair = TrackConsumersPool::new(
    362
        ///     FairSpillPool::new(2 * 1024 * 1024), // 2MB limit
    363
        ///     NonZeroUsize::new(5).unwrap(),
    364
        /// );
    365
        /// ```
    366
        ///
    367
        /// # Impact on Error Messages
    368
        ///
    369
        /// The `top` determines how many Top K [`MemoryConsumer`]s to include
    370
        /// in the reported [`DataFusionError::ResourcesExhausted`].
    371
    0
        pub fn new(inner: I, top: NonZeroUsize) -> Self {
    372
    0
            Self {
    373
    0
                inner,
    374
    0
                top,
    375
    0
                tracked_consumers: Default::default(),
    376
    0
            }
    377
    0
        }
    378
    379
        /// Returns a formatted string with the top memory consumers.
    380
    0
        pub fn report_top(&self, top: usize) -> String {
    381
    0
            let mut consumers = self
    382
    0
                .tracked_consumers
    383
    0
                .lock()
    384
    0
                .iter()
    385
    0
                .map(|(consumer_id, tracked_consumer)| {
    386
    0
                    (
    387
    0
                        (
    388
    0
                            *consumer_id,
    389
    0
                            tracked_consumer.name.to_owned(),
    390
    0
                            tracked_consumer.can_spill,
    391
    0
                            tracked_consumer.peak(),
    392
    0
                        ),
    393
    0
                        tracked_consumer.reserved(),
    394
    0
                    )
    395
    0
                })
    396
    0
                .collect::<Vec<_>>();
    397
    0
            consumers.sort_by(|a, b| b.1.cmp(&a.1)); // inverse ordering
    398
    399
    0
            consumers[0..std::cmp::min(top, consumers.len())]
    400
    0
                .iter()
    401
    0
                .map(|((id, name, can_spill, peak), size)| {
    402
    0
                    format!(
    403
    0
                        "  {name}#{id}(can spill: {can_spill}) consumed {}, peak {}",
    404
    0
                        human_readable_size(*size),
    405
    0
                        human_readable_size(*peak),
    406
                    )
    407
    0
                })
    408
    0
                .collect::<Vec<_>>()
    409
    0
                .join(",\n")
    410
    0
                + "."
    411
    0
        }
    412
    }
    413
    414
    impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
    415
    0
        fn register(&self, consumer: &MemoryConsumer) {
    416
    0
            self.inner.register(consumer);
    417
    418
    0
            let mut guard = self.tracked_consumers.lock();
    419
    0
            let existing = guard.insert(
    420
    0
                consumer.id(),
    421
    0
                TrackedConsumer {
    422
    0
                    name: consumer.name().to_string(),
    423
    0
                    can_spill: consumer.can_spill(),
    424
    0
                    reserved: Default::default(),
    425
    0
                    peak: Default::default(),
    426
    0
                },
    427
    0
            );
    428
    429
    0
            debug_assert!(
    430
    0
                existing.is_none(),
    431
    0
                "Registered was called twice on the same consumer"
    432
            );
    433
    0
        }
    434
    435
    0
        fn unregister(&self, consumer: &MemoryConsumer) {
    436
    0
            self.inner.unregister(consumer);
    437
    0
            self.tracked_consumers.lock().remove(&consumer.id());
    438
    0
        }
    439
    440
    0
        fn grow(&self, reservation: &MemoryReservation, additional: usize) {
    441
    0
            self.inner.grow(reservation, additional);
    442
    0
            self.tracked_consumers
    443
    0
                .lock()
    444
    0
                .entry(reservation.consumer().id())
    445
    0
                .and_modify(|tracked_consumer| {
    446
    0
                    tracked_consumer.grow(additional);
    447
    0
                });
    448
    0
        }
    449
    450
    0
        fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
    451
    0
            self.inner.shrink(reservation, shrink);
    452
    0
            self.tracked_consumers
    453
    0
                .lock()
    454
    0
                .entry(reservation.consumer().id())
    455
    0
                .and_modify(|tracked_consumer| {
    456
    0
                    tracked_consumer.shrink(shrink);
    457
    0
                });
    458
    0
        }
    459
    460
    0
        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
    461
    0
            self.inner
    462
    0
                .try_grow(reservation, additional)
    463
    0
                .map_err(|e| match e {
    464
    0
                    DataFusionError::ResourcesExhausted(e) => {
    465
                        // wrap OOM message in top consumers
    466
    0
                        DataFusionError::ResourcesExhausted(
    467
    0
                            provide_top_memory_consumers_to_error_msg(
    468
    0
                                &reservation.consumer().name,
    469
    0
                                &e,
    470
    0
                                &self.report_top(self.top.into()),
    471
    0
                            ),
    472
    0
                        )
    473
                    }
    474
    0
                    _ => e,
    475
    0
                })?;
    476
    477
    0
            self.tracked_consumers
    478
    0
                .lock()
    479
    0
                .entry(reservation.consumer().id())
    480
    0
                .and_modify(|tracked_consumer| {
    481
    0
                    tracked_consumer.grow(additional);
    482
    0
                });
    483
    0
            Ok(())
    484
    0
        }
    485
    486
    0
        fn reserved(&self) -> usize {
    487
    0
            self.inner.reserved()
    488
    0
        }
    489
    490
    0
        fn memory_limit(&self) -> MemoryLimit {
    491
    0
            self.inner.memory_limit()
    492
    0
        }
    493
    }
    494
    495
    0
    fn provide_top_memory_consumers_to_error_msg(
    496
    0
        consumer_name: &str,
    497
    0
        error_msg: &str,
    498
    0
        top_consumers: &str,
    499
    0
    ) -> String {
    500
    0
        format!("Additional allocation failed for {consumer_name} with top memory consumers (across reservations) as:\n{top_consumers}\nError: {error_msg}")
    501
    0
    }
    502
    503
    #[cfg(test)]
    504
    mod tests {
    505
        use super::*;
    506
        use insta::{allow_duplicates, assert_snapshot, Settings};
    507
        use std::sync::Arc;
    508
    509
        fn make_settings() -> Settings {
    510
            let mut settings = Settings::clone_current();
    511
            settings.add_filter(
    512
                r"([^\s]+)\#\d+\(can spill: (true|false)\)",
    513
                "$1#[ID](can spill: $2)",
    514
            );
    515
            settings
    516
        }
    517
    518
        #[test]
    519
        fn test_fair() {
    520
            let pool = Arc::new(FairSpillPool::new(100)) as _;
    521
    522
            let mut r1 = MemoryConsumer::new("unspillable").register(&pool);
    523
            // Can grow beyond capacity of pool
    524
            r1.grow(2000);
    525
            assert_eq!(pool.reserved(), 2000);
    526
    527
            let mut r2 = MemoryConsumer::new("r2")
    528
                .with_can_spill(true)
    529
                .register(&pool);
    530
            // Can grow beyond capacity of pool
    531
            r2.grow(2000);
    532
    533
            assert_eq!(pool.reserved(), 4000);
    534
    535
            let err = r2.try_grow(1).unwrap_err().strip_backtrace();
    536
            assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 1.0 B for r2 with 2000.0 B already allocated for this reservation - 0.0 B remain available for the total pool");
    537
    538
            let err = r2.try_grow(1).unwrap_err().strip_backtrace();
    539
            assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 1.0 B for r2 with 2000.0 B already allocated for this reservation - 0.0 B remain available for the total pool");
    540
    541
            r1.shrink(1990);
    542
            r2.shrink(2000);
    543
    544
            assert_eq!(pool.reserved(), 10);
    545
    546
            r1.try_grow(10).unwrap();
    547
            assert_eq!(pool.reserved(), 20);
    548
    549
            // Can grow r2 to 80 as only spilling consumer
    550
            r2.try_grow(80).unwrap();
    551
            assert_eq!(pool.reserved(), 100);
    552
    553
            r2.shrink(70);
    554
    555
            assert_eq!(r1.size(), 20);
    556
            assert_eq!(r2.size(), 10);
    557
            assert_eq!(pool.reserved(), 30);
    558
    559
            let mut r3 = MemoryConsumer::new("r3")
    560
                .with_can_spill(true)
    561
                .register(&pool);
    562
    563
            let err = r3.try_grow(70).unwrap_err().strip_backtrace();
    564
            assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 70.0 B for r3 with 0.0 B already allocated for this reservation - 40.0 B remain available for the total pool");
    565
    566
            //Shrinking r2 to zero doesn't allow a3 to allocate more than 45
    567
            r2.free();
    568
            let err = r3.try_grow(70).unwrap_err().strip_backtrace();
    569
            assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 70.0 B for r3 with 0.0 B already allocated for this reservation - 40.0 B remain available for the total pool");
    570
    571
            // But dropping r2 does
    572
            drop(r2);
    573
            assert_eq!(pool.reserved(), 20);
    574
            r3.try_grow(80).unwrap();
    575
    576
            assert_eq!(pool.reserved(), 100);
    577
            r1.free();
    578
            assert_eq!(pool.reserved(), 80);
    579
    580
            let mut r4 = MemoryConsumer::new("s4").register(&pool);
    581
            let err = r4.try_grow(30).unwrap_err().strip_backtrace();
    582
            assert_snapshot!(err, @"Resources exhausted: Failed to allocate additional 30.0 B for s4 with 0.0 B already allocated for this reservation - 20.0 B remain available for the total pool");
    583
        }
    584
    585
        #[test]
    586
        fn test_tracked_consumers_pool() {
    587
            let setting = make_settings();
    588
            let _bound = setting.bind_to_scope();
    589
            let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    590
                GreedyMemoryPool::new(100),
    591
                NonZeroUsize::new(3).unwrap(),
    592
            ));
    593
    594
            // Test: use all the different interfaces to change reservation size
    595
    596
            // set r1=50, using grow and shrink
    597
            let mut r1 = MemoryConsumer::new("r1").register(&pool);
    598
            r1.grow(50);
    599
            r1.grow(20);
    600
            r1.shrink(20);
    601
    602
            // set r2=15 using try_grow
    603
            let mut r2 = MemoryConsumer::new("r2").register(&pool);
    604
            r2.try_grow(15)
    605
                .expect("should succeed in memory allotment for r2");
    606
    607
            // set r3=20 using try_resize
    608
            let mut r3 = MemoryConsumer::new("r3").register(&pool);
    609
            r3.try_resize(25)
    610
                .expect("should succeed in memory allotment for r3");
    611
            r3.try_resize(20)
    612
                .expect("should succeed in memory allotment for r3");
    613
    614
            // set r4=10
    615
            // this should not be reported in top 3
    616
            let mut r4 = MemoryConsumer::new("r4").register(&pool);
    617
            r4.grow(10);
    618
    619
            // Test: reports if new reservation causes error
    620
            // using the previously set sizes for other consumers
    621
            let mut r5 = MemoryConsumer::new("r5").register(&pool);
    622
            let res = r5.try_grow(150);
    623
            assert!(res.is_err());
    624
            let error = res.unwrap_err().strip_backtrace();
    625
            assert_snapshot!(error, @r"
    626
            Resources exhausted: Additional allocation failed for r5 with top memory consumers (across reservations) as:
    627
              r1#[ID](can spill: false) consumed 50.0 B, peak 70.0 B,
    628
              r3#[ID](can spill: false) consumed 20.0 B, peak 25.0 B,
    629
              r2#[ID](can spill: false) consumed 15.0 B, peak 15.0 B.
    630
            Error: Failed to allocate additional 150.0 B for r5 with 0.0 B already allocated for this reservation - 5.0 B remain available for the total pool
    631
            ");
    632
        }
    633
    634
        #[test]
    635
        fn test_tracked_consumers_pool_register() {
    636
            let setting = make_settings();
    637
            let _bound = setting.bind_to_scope();
    638
            let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    639
                GreedyMemoryPool::new(100),
    640
                NonZeroUsize::new(3).unwrap(),
    641
            ));
    642
    643
            let same_name = "foo";
    644
    645
            // Test: see error message when no consumers recorded yet
    646
            let mut r0 = MemoryConsumer::new(same_name).register(&pool);
    647
            let res = r0.try_grow(150);
    648
            assert!(res.is_err());
    649
            let error = res.unwrap_err().strip_backtrace();
    650
            assert_snapshot!(error, @r"
    651
            Resources exhausted: Additional allocation failed for foo with top memory consumers (across reservations) as:
    652
              foo#[ID](can spill: false) consumed 0.0 B, peak 0.0 B.
    653
            Error: Failed to allocate additional 150.0 B for foo with 0.0 B already allocated for this reservation - 100.0 B remain available for the total pool
    654
            ");
    655
    656
            // API: multiple registrations using the same hashed consumer,
    657
            // will be recognized *differently* in the TrackConsumersPool.
    658
    659
            r0.grow(10); // make r0=10, pool available=90
    660
            let new_consumer_same_name = MemoryConsumer::new(same_name);
    661
            let mut r1 = new_consumer_same_name.register(&pool);
    662
            // TODO: the insufficient_capacity_err() message is per reservation, not per consumer.
    663
            // a followup PR will clarify this message "0 bytes already allocated for this reservation"
    664
            let res = r1.try_grow(150);
    665
            assert!(res.is_err());
    666
            let error = res.unwrap_err().strip_backtrace();
    667
            assert_snapshot!(error, @r"
    668
            Resources exhausted: Additional allocation failed for foo with top memory consumers (across reservations) as:
    669
              foo#[ID](can spill: false) consumed 10.0 B, peak 10.0 B,
    670
              foo#[ID](can spill: false) consumed 0.0 B, peak 0.0 B.
    671
            Error: Failed to allocate additional 150.0 B for foo with 0.0 B already allocated for this reservation - 90.0 B remain available for the total pool
    672
            ");
    673
    674
            // Test: will accumulate size changes per consumer, not per reservation
    675
            r1.grow(20);
    676
    677
            let res = r1.try_grow(150);
    678
            assert!(res.is_err());
    679
            let error = res.unwrap_err().strip_backtrace();
    680
            assert_snapshot!(error, @r"
    681
            Resources exhausted: Additional allocation failed for foo with top memory consumers (across reservations) as:
    682
              foo#[ID](can spill: false) consumed 20.0 B, peak 20.0 B,
    683
              foo#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
    684
            Error: Failed to allocate additional 150.0 B for foo with 20.0 B already allocated for this reservation - 70.0 B remain available for the total pool
    685
            ");
    686
    687
            // Test: different hashed consumer, (even with the same name),
    688
            // will be recognized as different in the TrackConsumersPool
    689
            let consumer_with_same_name_but_different_hash =
    690
                MemoryConsumer::new(same_name).with_can_spill(true);
    691
            let mut r2 = consumer_with_same_name_but_different_hash.register(&pool);
    692
            let res = r2.try_grow(150);
    693
            assert!(res.is_err());
    694
            let error = res.unwrap_err().strip_backtrace();
    695
            assert_snapshot!(error, @r"
    696
            Resources exhausted: Additional allocation failed for foo with top memory consumers (across reservations) as:
    697
              foo#[ID](can spill: false) consumed 20.0 B, peak 20.0 B,
    698
              foo#[ID](can spill: false) consumed 10.0 B, peak 10.0 B,
    699
              foo#[ID](can spill: true) consumed 0.0 B, peak 0.0 B.
    700
            Error: Failed to allocate additional 150.0 B for foo with 0.0 B already allocated for this reservation - 70.0 B remain available for the total pool
    701
            ");
    702
        }
    703
    704
        #[test]
    705
        fn test_tracked_consumers_pool_deregister() {
    706
            fn test_per_pool_type(pool: Arc<dyn MemoryPool>) {
    707
                // Baseline: see the 2 memory consumers
    708
                let setting = make_settings();
    709
                let _bound = setting.bind_to_scope();
    710
                let mut r0 = MemoryConsumer::new("r0").register(&pool);
    711
                r0.grow(10);
    712
                let r1_consumer = MemoryConsumer::new("r1");
    713
                let mut r1 = r1_consumer.register(&pool);
    714
                r1.grow(20);
    715
    716
                let res = r0.try_grow(150);
    717
                assert!(res.is_err());
    718
                let error = res.unwrap_err().strip_backtrace();
    719
                allow_duplicates!(assert_snapshot!(error, @r"
    720
                    Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
    721
                      r1#[ID](can spill: false) consumed 20.0 B, peak 20.0 B,
    722
                      r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
    723
                    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 pool
    724
                    "));
    725
    726
                // Test: unregister one
    727
                // only the remaining one should be listed
    728
                drop(r1);
    729
                let res = r0.try_grow(150);
    730
                assert!(res.is_err());
    731
                let error = res.unwrap_err().strip_backtrace();
    732
                allow_duplicates!(assert_snapshot!(error, @r"
    733
                    Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
    734
                      r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
    735
                    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 pool
    736
                    "));
    737
    738
                // Test: actual message we see is the `available is 70`. When it should be `available is 90`.
    739
                // This is because the pool.shrink() does not automatically occur within the inner_pool.deregister().
    740
                let res = r0.try_grow(150);
    741
                assert!(res.is_err());
    742
                let error = res.unwrap_err().strip_backtrace();
    743
                allow_duplicates!(assert_snapshot!(error, @r"
    744
                    Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
    745
                      r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
    746
                    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 pool
    747
                    "));
    748
    749
                // Test: the registration needs to free itself (or be dropped),
    750
                // for the proper error message
    751
                let res = r0.try_grow(150);
    752
                assert!(res.is_err());
    753
                let error = res.unwrap_err().strip_backtrace();
    754
                allow_duplicates!(assert_snapshot!(error, @r"
    755
                    Resources exhausted: Additional allocation failed for r0 with top memory consumers (across reservations) as:
    756
                      r0#[ID](can spill: false) consumed 10.0 B, peak 10.0 B.
    757
                    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 pool
    758
                    "));
    759
            }
    760
    761
            let tracked_spill_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    762
                FairSpillPool::new(100),
    763
                NonZeroUsize::new(3).unwrap(),
    764
            ));
    765
            test_per_pool_type(tracked_spill_pool);
    766
    767
            let tracked_greedy_pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
    768
                GreedyMemoryPool::new(100),
    769
                NonZeroUsize::new(3).unwrap(),
    770
            ));
    771
            test_per_pool_type(tracked_greedy_pool);
    772
        }
    773
    774
        #[test]
    775
        fn test_tracked_consumers_pool_use_beyond_errors() {
    776
            let setting = make_settings();
    777
            let _bound = setting.bind_to_scope();
    778
            let upcasted: Arc<dyn std::any::Any + Send + Sync> =
    779
                Arc::new(TrackConsumersPool::new(
    780
                    GreedyMemoryPool::new(100),
    781
                    NonZeroUsize::new(3).unwrap(),
    782
                ));
    783
            let pool: Arc<dyn MemoryPool> = Arc::clone(&upcasted)
    784
                .downcast::<TrackConsumersPool<GreedyMemoryPool>>()
    785
                .unwrap();
    786
            // set r1=20
    787
            let mut r1 = MemoryConsumer::new("r1").register(&pool);
    788
            r1.grow(20);
    789
            // set r2=15
    790
            let mut r2 = MemoryConsumer::new("r2").register(&pool);
    791
            r2.grow(15);
    792
            // set r3=45
    793
            let mut r3 = MemoryConsumer::new("r3").register(&pool);
    794
            r3.grow(45);
    795
    796
            let downcasted = upcasted
    797
                .downcast::<TrackConsumersPool<GreedyMemoryPool>>()
    798
                .unwrap();
    799
    800
            // Test: can get runtime metrics, even without an error thrown
    801
            let res = downcasted.report_top(2);
    802
            assert_snapshot!(res, @r"
    803
            r3#[ID](can spill: false) consumed 45.0 B, peak 45.0 B,
    804
            r1#[ID](can spill: false) consumed 20.0 B, peak 20.0 B.
    805
            ");
    806
        }
    807
    }