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}