Skip to main content

mnemosyne_arena/scratch/pool/
manage.rs

1//! Lifecycle management for [`super::ScratchPool`]: provision-aware
2//! `release`, provision reset, warm-up helpers, and full-reclaim.
3
4use super::super::aligned_vec::AlignedVec;
5use super::super::element::ScratchElement;
6use super::{MAX_POOL_SLOTS, ScratchPool};
7
8impl<T: ScratchElement> ScratchPool<T> {
9    /// Reclaims every slot's storage above its recorded provision.
10    ///
11    /// With [`with_scratch_bounded`](Self::with_scratch_bounded) as the only
12    /// entry point, a provision is the largest request ever seen at that
13    /// depth; the slot keeps capacity for it (warm reuse stays allocation-free)
14    /// and surrenders everything above — growth headroom included, so the
15    /// retained steady state is exactly the working set. Slots whose provision
16    /// is zero are dropped entirely. A slot is reclaimed only when its depth
17    /// is idle; busy slots are skipped, never torn from under a live borrow.
18    ///
19    /// The intended quiescent rhythm: run transforms normally, then call this
20    /// when the workload idles — not on every `with_scratch` exit, which would
21    /// reintroduce the churn the pool exists to remove. Provisions persist, so
22    /// repeated release/idle cycles converge; a smaller *steady-state* working
23    /// set needs [`reset`](Self::reset).
24    ///
25    /// Returns the per-slot capacities after reclamation. A slot that is
26    /// currently borrowed reports its provision instead: its live capacity is
27    /// not observable without deriving a reference under an existing exclusive
28    /// borrow, the same aliasing [`capacity`](Self::capacity) exists to avoid.
29    pub fn release(&self) -> [usize; MAX_POOL_SLOTS] {
30        let mut capacities = [0usize; MAX_POOL_SLOTS];
31        for (idx, slot) in self.slots.iter().enumerate() {
32            let provision = self.provisions[idx].get();
33            if self.borrow_depth.get() > idx as u8 {
34                capacities[idx] = provision;
35                continue;
36            }
37            // SAFETY: exclusive access — the depth guard above proved this
38            // slot's nesting level is not on the stack, so no borrow of the
39            // slot can be live.
40            let vec = unsafe { &mut *slot.get() };
41            if provision == 0 {
42                if vec.capacity() != 0 {
43                    // Drop returns the allocation, then the sentinel lands.
44                    *vec = AlignedVec::dangling();
45                    self.slot_capacities[idx].set(0);
46                }
47            } else if vec.capacity() > provision {
48                vec.shrink_to(provision);
49                // No `clear`: reuse is explicitly not re-zeroed (see
50                // `with_scratch`), and zeroing here would put an O(n)
51                // memset on the warm path — the exact churn the pool
52                // exists to remove. The next growth re-zeros its new
53                // range as always.
54                self.slot_capacities[idx].set(vec.capacity());
55            }
56            capacities[idx] = vec.capacity();
57        }
58        capacities
59    }
60
61    /// Clears the recorded provisions so a later [`release`](Self::release)
62    /// reclaims every slot entirely.
63    ///
64    /// For a full working-set changeover (a consumer tearing down one workload
65    /// and starting another): reset, then run the new workload through
66    /// [`with_scratch_bounded`](Self::with_scratch_bounded), then release.
67    /// Slots that are not idle keep their buffers; the next release sees their
68    /// cleared provisions and reclaims them.
69    pub fn reset(&self) {
70        for provision in &self.provisions {
71            provision.set(0);
72        }
73    }
74
75    /// Ensures the primary slot has capacity for at least `min_capacity`
76    /// elements, growing it if necessary. No-op when the pool is borrowed.
77    #[inline]
78    pub fn prewarm(&self, min_capacity: usize) {
79        if self.borrow_depth.get() != 0 {
80            return;
81        }
82        // SAFETY: borrow_depth == 0 so no live exclusive reference to slot 0.
83        let vec = unsafe { &mut *self.slots[0].get() };
84        if vec.capacity() < min_capacity {
85            vec.ensure_len(min_capacity);
86            self.slot_capacities[0].set(vec.capacity());
87        }
88    }
89
90    /// Prewarms multiple slots in one call.
91    ///
92    /// `sizes[i]` specifies the minimum capacity for slot `i` (depth `i`).
93    /// Out-of-range indices or a borrowed slot are silently skipped.
94    /// Entries of `0` skip that slot.
95    #[inline]
96    pub fn preload(&self, sizes: &[usize]) {
97        if self.borrow_depth.get() != 0 {
98            return;
99        }
100        for (idx, &min_cap) in sizes.iter().enumerate().take(MAX_POOL_SLOTS) {
101            if min_cap == 0 {
102                continue;
103            }
104            // SAFETY: borrow_depth == 0; idx < MAX_POOL_SLOTS.
105            let vec = unsafe { &mut *self.slots[idx].get() };
106            if vec.capacity() < min_cap {
107                vec.ensure_len(min_cap);
108                self.slot_capacities[idx].set(vec.capacity());
109            }
110        }
111    }
112
113    /// Releases all slot allocations when not borrowed. No-op when borrowed.
114    #[inline]
115    pub fn shrink_all_slots(&self) {
116        if self.borrow_depth.get() != 0 {
117            return;
118        }
119        for (i, slot) in self.slots.iter().enumerate() {
120            // SAFETY: borrow_depth == 0 — no live exclusive references.
121            let vec = unsafe { &mut *slot.get() };
122            *vec = AlignedVec::dangling();
123            self.slot_capacities[i].set(0);
124        }
125    }
126}