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
}