Skip to main content

datafusion_execution/memory_pool/
mod.rs

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//! [`MemoryPool`] for memory management during query execution, [`proxy`] for
19//! help with allocation accounting.
20
21use datafusion_common::{Result, internal_datafusion_err};
22use std::any::Any;
23use std::fmt::Display;
24use std::hash::{Hash, Hasher};
25use std::{cmp::Ordering, sync::Arc, sync::atomic};
26
27mod peak_recording;
28mod pool;
29
30#[cfg(feature = "arrow_buffer_pool")]
31pub mod arrow;
32
33pub mod proxy {
34    pub use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt};
35}
36
37pub use datafusion_common::{
38    human_readable_count, human_readable_duration, human_readable_size, units,
39};
40pub use peak_recording::*;
41pub use pool::*;
42
43/// Tracks and potentially limits memory use across operators during execution.
44///
45/// # Memory Management Overview
46///
47/// DataFusion is a streaming query engine, processing most queries without
48/// buffering the entire input. Most operators require a fixed amount of memory
49/// based on the schema and target batch size. However, certain operations such
50/// as sorting and grouping/joining, require buffering intermediate results,
51/// which can require memory proportional to the number of input rows.
52///
53/// Rather than tracking all allocations, DataFusion takes a pragmatic approach:
54/// Intermediate memory used as data streams through the system is not accounted
55/// (it assumed to be "small") but the large consumers of memory must register
56/// and constrain their use. This design trades off the additional code
57/// complexity of memory tracking with limiting resource usage.
58///
59/// When limiting memory with a `MemoryPool` you should typically reserve some
60/// overhead (e.g. 10%) for the "small" memory allocations that are not tracked.
61///
62/// # Memory Management Design
63///
64/// As explained above, DataFusion's design ONLY limits operators that require
65/// "large" amounts of memory (proportional to number of input rows), such as
66/// `GroupByHashExec`. It does NOT track and limit memory used internally by
67/// other operators such as `DataSourceExec` or the `RecordBatch`es that flow
68/// between operators. Furthermore, operators should not reserve memory for the
69/// batches they produce. Instead, if a consumer operator needs to hold batches
70/// from its producers in memory for an extended period, it is the consumer
71/// operator's responsibility to reserve the necessary memory for those batches.
72///
73/// In order to avoid allocating memory until the OS or the container system
74/// kills the process, DataFusion `ExecutionPlan`s (operators) that consume
75/// large amounts of memory must first request their desired allocation from a
76/// [`MemoryPool`] before allocating more.  The request is typically managed via
77/// a  [`MemoryReservation`] and [`MemoryConsumer`].
78///
79/// If the allocation is successful, the operator should proceed and allocate
80/// the desired memory. If the allocation fails, the operator must either first
81/// free memory (e.g. by spilling to local disk) and try again, or error.
82///
83/// Note that a `MemoryPool` can be shared by concurrently executing plans,
84/// which can be used to control memory usage in a multi-tenant system.
85///
86/// # How MemoryPool works by example
87///
88/// Scenario 1:
89/// For `Filter` operator, `RecordBatch`es will stream through it, so it
90/// don't have to keep track of memory usage through [`MemoryPool`].
91///
92/// Scenario 2:
93/// For `CrossJoin` operator, if the input size gets larger, the intermediate
94/// state will also grow. So `CrossJoin` operator will use [`MemoryPool`] to
95/// limit the memory usage.
96/// 2.1 `CrossJoin` operator has read a new batch, asked memory pool for
97/// additional memory. Memory pool updates the usage and returns success.
98/// 2.2 `CrossJoin` has read another batch, and tries to reserve more memory
99/// again, memory pool does not have enough memory. Since `CrossJoin` operator
100/// has not implemented spilling, it will stop execution and return an error.
101///
102/// Scenario 3:
103/// For `Aggregate` operator, its intermediate states will also accumulate as
104/// the input size gets larger, but with spilling capability. When it tries to
105/// reserve more memory from the memory pool, and the memory pool has already
106/// reached the memory limit, it will return an error. Then, `Aggregate`
107/// operator will spill the intermediate buffers to disk, and release memory
108/// from the memory pool, and continue to retry memory reservation.
109///
110/// # Related Structs
111///
112/// To better understand memory management in DataFusion, here are the key structs
113/// and their relationships:
114///
115/// - [`MemoryConsumer`]: A named allocation traced by a particular operator. If an
116///   execution is parallelized, and there are multiple partitions of the same
117///   operator, each partition will have a separate `MemoryConsumer`.
118/// - `SharedRegistration`: A registration of a `MemoryConsumer` with a `MemoryPool`.
119///   `SharedRegistration` and `MemoryPool` have a many-to-one relationship. `MemoryPool`
120///   implementation can decide how to allocate memory based on the registered consumers.
121///   (e.g. `FairSpillPool` will try to share available memory evenly among all registered
122///   consumers)
123/// - [`MemoryReservation`]: Each `MemoryConsumer`/operator can have multiple
124///   `MemoryReservation`s for different internal data structures. The relationship
125///   between `MemoryConsumer` and `MemoryReservation` is one-to-many. This design
126///   enables cleaner operator implementations:
127///   - Different `MemoryReservation`s can be used for different purposes
128///   - `MemoryReservation` follows RAII principles - to release a reservation,
129///     simply drop the `MemoryReservation` object. When all `MemoryReservation`s
130///     for a `SharedRegistration` are dropped, the `SharedRegistration` is dropped
131///     when its reference count reaches zero, automatically unregistering the
132///     `MemoryConsumer` from the `MemoryPool`.
133///
134/// ## Relationship Diagram
135///
136/// ```text
137/// ┌──────────────────┐     ┌──────────────────┐
138/// │MemoryReservation │     │MemoryReservation │
139/// └───┬──────────────┘     └──────────────────┘ ......
140///     │belongs to                    │
141///     │      ┌───────────────────────┘           │  │
142///     │      │                                   │  │
143///     ▼      ▼                                   ▼  ▼
144/// ┌────────────────────────┐       ┌────────────────────────┐
145/// │   SharedRegistration   │       │   SharedRegistration   │
146/// │   ┌────────────────┐   │       │   ┌────────────────┐   │
147/// │   │                │   │       │   │                │   │
148/// │   │ MemoryConsumer │   │       │   │ MemoryConsumer │   │
149/// │   │                │   │       │   │                │   │
150/// │   └────────────────┘   │       │   └────────────────┘   │
151/// └────────────┬───────────┘       └────────────┬───────────┘
152///              │                                │
153///              │                        register│into
154///              │                                │
155///              └─────────────┐   ┌──────────────┘
156///                            │   │
157///                            ▼   ▼
158///    ╔═══════════════════════════════════════════════════╗
159///    ║                                                   ║
160///    ║                    MemoryPool                     ║
161///    ║                                                   ║
162///    ╚═══════════════════════════════════════════════════╝
163/// ```
164///
165/// For example, there are two parallel partitions of an operator X: each partition
166/// corresponds to a `MemoryConsumer` in the above diagram. Inside each partition of
167/// operator X, there are typically several `MemoryReservation`s - one for each
168/// internal data structure that needs memory tracking (e.g., 1 reservation for the hash
169/// table, and 1 reservation for buffered input, etc.).
170///
171/// # Implementing `MemoryPool`
172///
173/// You can implement a custom allocation policy by implementing the
174/// [`MemoryPool`] trait and configuring a `SessionContext` appropriately.
175/// However, DataFusion comes with the following simple memory pool implementations that
176/// handle many common cases:
177///
178/// * [`UnboundedMemoryPool`]: no memory limits (the default)
179///
180/// * [`GreedyMemoryPool`]: Limits memory usage to a fixed size using a "first
181///   come first served" policy
182///
183/// * [`FairSpillPool`]: Limits memory usage to a fixed size, allocating memory
184///   to all spilling operators fairly
185///
186/// * [`TrackConsumersPool`]: Wraps another [`MemoryPool`] and tracks consumers,
187///   providing better error messages on the largest memory users.
188pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display {
189    /// Return pool name
190    fn name(&self) -> &str;
191
192    /// Registers a new [`MemoryConsumer`]
193    ///
194    /// Note: Subsequent calls to [`Self::grow`] must be made to reserve memory
195    fn register(&self, _consumer: &MemoryConsumer) {}
196
197    /// Records the destruction of a [`MemoryReservation`] with [`MemoryConsumer`]
198    ///
199    /// Note: Prior calls to [`Self::shrink`] must be made to free any reserved memory
200    fn unregister(&self, _consumer: &MemoryConsumer) {}
201
202    /// Infallibly grow the provided `reservation` by `additional` bytes
203    ///
204    /// This must always succeed
205    fn grow(&self, reservation: &MemoryReservation, additional: usize);
206
207    /// Infallibly shrink the provided `reservation` by `shrink` bytes
208    fn shrink(&self, reservation: &MemoryReservation, shrink: usize);
209
210    /// Attempt to grow the provided `reservation` by `additional` bytes
211    ///
212    /// On error the `allocation` will not be increased in size
213    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()>;
214
215    /// Return the total amount of memory reserved
216    fn reserved(&self) -> usize;
217
218    /// Return the memory limit of the pool
219    ///
220    /// The default implementation of `MemoryPool::memory_limit`
221    /// will return `MemoryLimit::Unknown`.
222    /// If you are using your custom memory pool, but have the requirement to
223    /// know the memory usage limit of the pool, please implement this method
224    /// to return it(`Memory::Finite(limit)`).
225    fn memory_limit(&self) -> MemoryLimit {
226        MemoryLimit::Unknown
227    }
228}
229
230impl dyn MemoryPool {
231    /// Returns `true` if this pool is of type `T`.
232    pub fn is<T: MemoryPool>(&self) -> bool {
233        (self as &dyn Any).is::<T>()
234    }
235
236    /// Attempts to downcast this pool to a concrete type `T`.
237    pub fn downcast_ref<T: MemoryPool>(&self) -> Option<&T> {
238        (self as &dyn Any).downcast_ref()
239    }
240}
241
242/// Memory limit of `MemoryPool`
243pub enum MemoryLimit {
244    Infinite,
245    /// Bounded memory limit in bytes.
246    Finite(usize),
247    Unknown,
248}
249
250/// A memory consumer is a named allocation traced by a particular
251/// [`MemoryReservation`] in a [`MemoryPool`]. All allocations are registered to
252/// a particular `MemoryConsumer`;
253///
254/// Each `MemoryConsumer` is identifiable by a process-unique id, and is therefore not cloneable,
255/// If you want a clone of a `MemoryConsumer`, you should look into [`MemoryConsumer::clone_with_new_id`],
256/// but note that this `MemoryConsumer` may be treated as a separate entity based on the used pool,
257/// and is only guaranteed to share the name and inner properties.
258///
259/// For help with allocation accounting, see the [`proxy`] module.
260///
261/// [proxy]: datafusion_common::utils::proxy
262#[derive(Debug)]
263pub struct MemoryConsumer {
264    name: String,
265    can_spill: bool,
266    id: usize,
267}
268
269impl PartialEq for MemoryConsumer {
270    fn eq(&self, other: &Self) -> bool {
271        let is_same_id = self.id == other.id;
272
273        #[cfg(debug_assertions)]
274        if is_same_id {
275            assert_eq!(self.name, other.name);
276            assert_eq!(self.can_spill, other.can_spill);
277        }
278
279        is_same_id
280    }
281}
282
283impl Eq for MemoryConsumer {}
284
285impl Hash for MemoryConsumer {
286    fn hash<H: Hasher>(&self, state: &mut H) {
287        self.id.hash(state);
288        self.name.hash(state);
289        self.can_spill.hash(state);
290    }
291}
292
293impl MemoryConsumer {
294    fn new_unique_id() -> usize {
295        static ID: atomic::AtomicUsize = atomic::AtomicUsize::new(0);
296        ID.fetch_add(1, atomic::Ordering::Relaxed)
297    }
298
299    /// Create a new empty [`MemoryConsumer`] that can be grown using [`MemoryReservation`]
300    pub fn new(name: impl Into<String>) -> Self {
301        Self {
302            name: name.into(),
303            can_spill: false,
304            id: Self::new_unique_id(),
305        }
306    }
307
308    /// Returns a clone of this [`MemoryConsumer`] with a new unique id,
309    /// which can be registered with a [`MemoryPool`],
310    /// This new consumer is separate from the original.
311    pub fn clone_with_new_id(&self) -> Self {
312        Self {
313            name: self.name.clone(),
314            can_spill: self.can_spill,
315            id: Self::new_unique_id(),
316        }
317    }
318
319    /// Return the unique id of this [`MemoryConsumer`]
320    pub fn id(&self) -> usize {
321        self.id
322    }
323
324    /// Set whether this allocation can be spilled to disk
325    pub fn with_can_spill(self, can_spill: bool) -> Self {
326        Self { can_spill, ..self }
327    }
328
329    /// Returns true if this allocation can spill to disk
330    pub fn can_spill(&self) -> bool {
331        self.can_spill
332    }
333
334    /// Returns the name associated with this allocation
335    pub fn name(&self) -> &str {
336        &self.name
337    }
338
339    /// Registers this [`MemoryConsumer`] with the provided [`MemoryPool`] returning
340    /// a [`MemoryReservation`] that can be used to grow or shrink the memory reservation
341    pub fn register(self, pool: &Arc<dyn MemoryPool>) -> MemoryReservation {
342        pool.register(&self);
343        MemoryReservation {
344            registration: Arc::new(SharedRegistration {
345                pool: Arc::clone(pool),
346                consumer: self,
347            }),
348            size: atomic::AtomicUsize::new(0),
349        }
350    }
351}
352
353/// A registration of a [`MemoryConsumer`] with a [`MemoryPool`].
354///
355/// Calls [`MemoryPool::unregister`] on drop to return any memory to
356/// the underlying pool.
357#[derive(Debug)]
358struct SharedRegistration {
359    pool: Arc<dyn MemoryPool>,
360    consumer: MemoryConsumer,
361}
362
363impl Drop for SharedRegistration {
364    fn drop(&mut self) {
365        self.pool.unregister(&self.consumer);
366    }
367}
368
369/// A [`MemoryReservation`] tracks an individual reservation of a
370/// number of bytes of memory in a [`MemoryPool`] that is freed back
371/// to the pool on drop.
372///
373/// The reservation can be grown or shrunk over time.
374#[derive(Debug)]
375pub struct MemoryReservation {
376    registration: Arc<SharedRegistration>,
377    size: atomic::AtomicUsize,
378}
379
380impl MemoryReservation {
381    /// Returns the size of this reservation in bytes
382    pub fn size(&self) -> usize {
383        self.size.load(atomic::Ordering::Relaxed)
384    }
385
386    /// Returns [MemoryConsumer] for this [MemoryReservation]
387    pub fn consumer(&self) -> &MemoryConsumer {
388        &self.registration.consumer
389    }
390
391    /// Frees all bytes from this reservation back to the underlying
392    /// pool, returning the number of bytes freed.
393    pub fn free(&self) -> usize {
394        let size = self.size.swap(0, atomic::Ordering::Relaxed);
395        if size != 0 {
396            self.registration.pool.shrink(self, size);
397        }
398        size
399    }
400
401    /// Frees `capacity` bytes from this reservation
402    ///
403    /// # Panics
404    ///
405    /// Panics if `capacity` exceeds [`Self::size`]
406    pub fn shrink(&self, capacity: usize) {
407        self.size
408            .fetch_update(
409                atomic::Ordering::Relaxed,
410                atomic::Ordering::Relaxed,
411                |prev| prev.checked_sub(capacity),
412            )
413            .unwrap_or_else(|prev| {
414                panic!("Cannot free the capacity {capacity} out of allocated size {prev}")
415            });
416        self.registration.pool.shrink(self, capacity);
417    }
418
419    /// Tries to free `capacity` bytes from this reservation
420    /// if `capacity` does not exceed [`Self::size`].
421    /// Returns new reservation size,
422    /// or error if shrinking capacity is more than allocated size.
423    pub fn try_shrink(&self, capacity: usize) -> Result<usize> {
424        let prev = self
425            .size
426            .fetch_update(
427                atomic::Ordering::Relaxed,
428                atomic::Ordering::Relaxed,
429                |prev| prev.checked_sub(capacity),
430            )
431            .map_err(|prev| {
432                internal_datafusion_err!(
433                    "Cannot free the capacity {capacity} out of allocated size {prev}"
434                )
435            })?;
436
437        self.registration.pool.shrink(self, capacity);
438        Ok(prev - capacity)
439    }
440
441    /// Sets the size of this reservation to `capacity`
442    pub fn resize(&self, capacity: usize) {
443        let size = self.size.load(atomic::Ordering::Relaxed);
444        match capacity.cmp(&size) {
445            Ordering::Greater => self.grow(capacity - size),
446            Ordering::Less => self.shrink(size - capacity),
447            _ => {}
448        }
449    }
450
451    /// Try to set the size of this reservation to `capacity`
452    pub fn try_resize(&self, capacity: usize) -> Result<()> {
453        let size = self.size.load(atomic::Ordering::Relaxed);
454        match capacity.cmp(&size) {
455            Ordering::Greater => self.try_grow(capacity - size)?,
456            Ordering::Less => {
457                self.try_shrink(size - capacity)?;
458            }
459            _ => {}
460        };
461        Ok(())
462    }
463
464    /// Increase the size of this reservation by `capacity` bytes
465    pub fn grow(&self, capacity: usize) {
466        self.registration.pool.grow(self, capacity);
467        self.size.fetch_add(capacity, atomic::Ordering::Relaxed);
468    }
469
470    /// Try to increase the size of this reservation by `capacity`
471    /// bytes, returning error if there is insufficient capacity left
472    /// in the pool.
473    pub fn try_grow(&self, capacity: usize) -> Result<()> {
474        self.registration.pool.try_grow(self, capacity)?;
475        self.size.fetch_add(capacity, atomic::Ordering::Relaxed);
476        Ok(())
477    }
478
479    /// Splits off `capacity` bytes from this [`MemoryReservation`]
480    /// into a new [`MemoryReservation`] with the same
481    /// [`MemoryConsumer`].
482    ///
483    /// This can be useful to free part of this reservation with RAAI
484    /// style dropping
485    ///
486    /// # Panics
487    ///
488    /// Panics if `capacity` exceeds [`Self::size`]
489    pub fn split(&self, capacity: usize) -> MemoryReservation {
490        self.size
491            .fetch_update(
492                atomic::Ordering::Relaxed,
493                atomic::Ordering::Relaxed,
494                |prev| prev.checked_sub(capacity),
495            )
496            .unwrap();
497        Self {
498            size: atomic::AtomicUsize::new(capacity),
499            registration: Arc::clone(&self.registration),
500        }
501    }
502
503    /// Returns a new empty [`MemoryReservation`] with the same [`MemoryConsumer`]
504    pub fn new_empty(&self) -> Self {
505        Self {
506            size: atomic::AtomicUsize::new(0),
507            registration: Arc::clone(&self.registration),
508        }
509    }
510
511    /// Splits off all the bytes from this [`MemoryReservation`] into
512    /// a new [`MemoryReservation`] with the same [`MemoryConsumer`]
513    pub fn take(&mut self) -> MemoryReservation {
514        self.split(self.size.load(atomic::Ordering::Relaxed))
515    }
516}
517
518impl Drop for MemoryReservation {
519    fn drop(&mut self) {
520        self.free();
521    }
522}
523
524#[cfg(test)]
525mod tests {
526    use super::*;
527
528    #[test]
529    fn test_id_uniqueness() {
530        let mut ids = std::collections::HashSet::new();
531        for _ in 0..100 {
532            let consumer = MemoryConsumer::new("test");
533            assert!(ids.insert(consumer.id())); // Ensures unique insertion
534        }
535    }
536
537    #[test]
538    fn test_memory_pool_underflow() {
539        let pool = Arc::new(GreedyMemoryPool::new(50)) as _;
540        let a1 = MemoryConsumer::new("a1").register(&pool);
541        assert_eq!(pool.reserved(), 0);
542
543        a1.grow(100);
544        assert_eq!(pool.reserved(), 100);
545
546        assert_eq!(a1.free(), 100);
547        assert_eq!(pool.reserved(), 0);
548
549        a1.try_grow(100).unwrap_err();
550        assert_eq!(pool.reserved(), 0);
551
552        a1.try_grow(30).unwrap();
553        assert_eq!(pool.reserved(), 30);
554
555        let a2 = MemoryConsumer::new("a2").register(&pool);
556        a2.try_grow(25).unwrap_err();
557        assert_eq!(pool.reserved(), 30);
558
559        drop(a1);
560        assert_eq!(pool.reserved(), 0);
561
562        a2.try_grow(25).unwrap();
563        assert_eq!(pool.reserved(), 25);
564    }
565
566    #[test]
567    fn test_split() {
568        let pool = Arc::new(GreedyMemoryPool::new(50)) as _;
569        let r1 = MemoryConsumer::new("r1").register(&pool);
570
571        r1.try_grow(20).unwrap();
572        assert_eq!(r1.size(), 20);
573        assert_eq!(pool.reserved(), 20);
574
575        // take 5 from r1, should still have same reservation split
576        let r2 = r1.split(5);
577        assert_eq!(r1.size(), 15);
578        assert_eq!(r2.size(), 5);
579        assert_eq!(pool.reserved(), 20);
580
581        // dropping r1 frees 15 but retains 5 as they have the same consumer
582        drop(r1);
583        assert_eq!(r2.size(), 5);
584        assert_eq!(pool.reserved(), 5);
585    }
586
587    #[test]
588    fn test_new_empty() {
589        let pool = Arc::new(GreedyMemoryPool::new(50)) as _;
590        let r1 = MemoryConsumer::new("r1").register(&pool);
591
592        r1.try_grow(20).unwrap();
593        let r2 = r1.new_empty();
594        r2.try_grow(5).unwrap();
595
596        assert_eq!(r1.size(), 20);
597        assert_eq!(r2.size(), 5);
598        assert_eq!(pool.reserved(), 25);
599    }
600
601    #[test]
602    fn test_take() {
603        let pool = Arc::new(GreedyMemoryPool::new(50)) as _;
604        let mut r1 = MemoryConsumer::new("r1").register(&pool);
605
606        r1.try_grow(20).unwrap();
607        let r2 = r1.take();
608        r2.try_grow(5).unwrap();
609
610        assert_eq!(r1.size(), 0);
611        assert_eq!(r2.size(), 25);
612        assert_eq!(pool.reserved(), 25);
613
614        // r1 can still grow again
615        r1.try_grow(3).unwrap();
616        assert_eq!(r1.size(), 3);
617        assert_eq!(r2.size(), 25);
618        assert_eq!(pool.reserved(), 28);
619    }
620
621    #[test]
622    fn test_downcast() {
623        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(50));
624
625        assert!(pool.is::<GreedyMemoryPool>());
626        assert!(!pool.is::<UnboundedMemoryPool>());
627
628        let greedy: &GreedyMemoryPool = pool.downcast_ref().unwrap();
629        assert_eq!(greedy.reserved(), 0);
630        assert!(pool.downcast_ref::<UnboundedMemoryPool>().is_none());
631    }
632
633    #[test]
634    fn test_try_shrink() {
635        let pool = Arc::new(GreedyMemoryPool::new(100)) as _;
636        let r1 = MemoryConsumer::new("r1").register(&pool);
637
638        r1.try_grow(50).unwrap();
639        assert_eq!(r1.size(), 50);
640        assert_eq!(pool.reserved(), 50);
641
642        // Successful shrink returns new size and frees pool memory
643        let new_size = r1.try_shrink(30).unwrap();
644        assert_eq!(new_size, 20);
645        assert_eq!(r1.size(), 20);
646        assert_eq!(pool.reserved(), 20);
647
648        // Freed pool memory is now available to other consumers
649        let r2 = MemoryConsumer::new("r2").register(&pool);
650        r2.try_grow(80).unwrap();
651        assert_eq!(pool.reserved(), 100);
652
653        // Shrinking more than allocated fails without changing state
654        let err = r1.try_shrink(25);
655        assert!(err.is_err());
656        assert_eq!(r1.size(), 20);
657        assert_eq!(pool.reserved(), 100);
658
659        // Shrink to exactly zero
660        let new_size = r1.try_shrink(20).unwrap();
661        assert_eq!(new_size, 0);
662        assert_eq!(r1.size(), 0);
663        assert_eq!(pool.reserved(), 80);
664    }
665}