1use crate::collections::map::HashMap;
2#[cfg(any(test, debug_assertions))]
3use crate::slot::{AnchorState, PayloadAnchorLifecycle, SlotInvariantError};
4#[cfg(any(test, debug_assertions))]
5use crate::SlotTable;
6use crate::{
7 slot::{DetachedSubtree, GroupKey, NodeLifecycle},
8 ScopeId,
9};
10use std::cmp::Ordering;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
16pub enum RetentionMode {
17 #[default]
18 DisposeWhenInactive,
19 RetainWhenInactive,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct RetentionBudget {
24 pub max_retained_subtrees: Option<usize>,
25 pub max_retained_bytes: Option<usize>,
26 pub max_age_passes: Option<u64>,
27}
28
29impl RetentionBudget {
30 pub const UNBOUNDED: Self = Self {
31 max_retained_subtrees: None,
32 max_retained_bytes: None,
33 max_age_passes: None,
34 };
35}
36
37impl Default for RetentionBudget {
38 fn default() -> Self {
39 Self::UNBOUNDED
40 }
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
44pub enum RetentionEvictionPolicy {
45 #[default]
46 LeastRecentlyDetached,
47 LeastRecentlyRestored,
48 LargestFirst,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct RetentionPolicy {
53 pub budget: RetentionBudget,
54 pub eviction: RetentionEvictionPolicy,
55}
56
57impl RetentionPolicy {
58 pub const UNBOUNDED: Self = Self {
59 budget: RetentionBudget::UNBOUNDED,
60 eviction: RetentionEvictionPolicy::LeastRecentlyDetached,
61 };
62}
63
64impl Default for RetentionPolicy {
65 fn default() -> Self {
66 Self::UNBOUNDED
67 }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
71pub(crate) struct RetainKey {
72 pub(crate) parent_scope: Option<ScopeId>,
73 pub(crate) key: GroupKey,
74}
75
76pub(crate) struct RetainedGroup {
77 pub(crate) subtree: DetachedSubtree,
78 detached_pass: u64,
79 detached_order: u64,
80 last_restored_order: u64,
81}
82
83impl RetainedGroup {
84 fn node_count(&self) -> usize {
85 self.subtree.node_count()
86 }
87
88 fn payload_count(&self) -> usize {
89 self.subtree.payload_count()
90 }
91
92 fn scope_count(&self) -> usize {
93 self.subtree.scope_count()
94 }
95
96 fn anchor_count(&self) -> usize {
97 self.subtree.anchor_count()
98 }
99
100 fn heap_bytes(&self) -> usize {
101 self.subtree.heap_bytes()
102 }
103}
104
105#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
106pub(crate) struct RetentionDebugStats {
107 pub(crate) subtree_count: usize,
108 pub(crate) group_count: usize,
109 pub(crate) payload_count: usize,
110 pub(crate) node_count: usize,
111 pub(crate) scope_count: usize,
112 pub(crate) anchor_count: usize,
113 pub(crate) heap_bytes: usize,
114 pub(crate) evictions_total: usize,
115}
116
117pub(crate) struct RetentionManager {
118 groups: HashMap<RetainKey, RetainedGroup>,
119 restored_at_by_key: HashMap<RetainKey, u64>,
120 policy: RetentionPolicy,
121 pass_clock: u64,
122 operation_clock: u64,
123 evictions_total: usize,
124}
125
126impl Default for RetentionManager {
127 fn default() -> Self {
128 Self::new(RetentionPolicy::default())
129 }
130}
131
132impl RetentionManager {
133 pub(crate) fn new(policy: RetentionPolicy) -> Self {
134 Self {
135 groups: HashMap::default(),
136 restored_at_by_key: HashMap::default(),
137 policy,
138 pass_clock: 0,
139 operation_clock: 0,
140 evictions_total: 0,
141 }
142 }
143
144 pub(crate) fn set_policy(&mut self, policy: RetentionPolicy) {
145 self.policy = policy;
146 }
147
148 pub(crate) fn take(&mut self, key: RetainKey) -> Option<DetachedSubtree> {
149 let retained = self.groups.get(&key)?;
150 retained
151 .subtree
152 .assert_root_key(key.key, "retention restore");
153 retained
154 .subtree
155 .assert_root_parent_detached("retention restore");
156 retained
157 .subtree
158 .assert_node_lifecycle(NodeLifecycle::RetainedDetached, "retention restore");
159 let restored_order = self.tick_operation();
160 let retained = self
161 .groups
162 .remove(&key)
163 .expect("validated retained subtree key must still exist");
164 self.restored_at_by_key.insert(key, restored_order);
165 Some(retained.subtree)
166 }
167
168 pub(crate) fn take_after_restore_preflight(
169 &mut self,
170 key: RetainKey,
171 preflight: impl FnOnce(&DetachedSubtree),
172 ) -> Option<DetachedSubtree> {
173 preflight(&self.groups.get(&key)?.subtree);
174 self.take(key)
175 }
176
177 pub(crate) fn insert(
178 &mut self,
179 key: RetainKey,
180 mut subtree: DetachedSubtree,
181 ) -> Vec<DetachedSubtree> {
182 assert!(
183 !self.groups.contains_key(&key),
184 "retained subtree key collision: parent_scope={:?} key={:?}",
185 key.parent_scope,
186 key.key,
187 );
188 subtree.assert_root_key(key.key, "retention insert");
189 subtree.assert_root_parent_detached("retention insert");
190 let detached_order = self.tick_operation();
191 let last_restored_order = self.restored_at_by_key.remove(&key).unwrap_or_default();
192 subtree.mark_nodes_retained_detached();
193 subtree.assert_node_lifecycle(NodeLifecycle::RetainedDetached, "retention insert");
194 self.groups.insert(
195 key,
196 RetainedGroup {
197 subtree,
198 detached_pass: self.pass_clock,
199 detached_order,
200 last_restored_order,
201 },
202 );
203 self.evict_to_budget()
204 }
205
206 pub(crate) fn is_empty(&self) -> bool {
207 self.groups.is_empty()
208 }
209
210 pub(crate) fn debug_stats(&self) -> RetentionDebugStats {
211 RetentionDebugStats {
212 subtree_count: self.groups.len(),
213 group_count: self
214 .groups
215 .values()
216 .map(|retained| retained.subtree.group_count())
217 .sum(),
218 payload_count: self.groups.values().map(RetainedGroup::payload_count).sum(),
219 node_count: self.groups.values().map(RetainedGroup::node_count).sum(),
220 scope_count: self.groups.values().map(RetainedGroup::scope_count).sum(),
221 anchor_count: self.groups.values().map(RetainedGroup::anchor_count).sum(),
222 heap_bytes: self.groups.values().map(RetainedGroup::heap_bytes).sum(),
223 evictions_total: self.evictions_total,
224 }
225 }
226
227 pub(crate) fn into_subtrees(self) -> Vec<DetachedSubtree> {
228 self.groups
229 .into_values()
230 .map(|retained| retained.subtree)
231 .collect()
232 }
233
234 pub(crate) fn subtrees(&self) -> impl Iterator<Item = &DetachedSubtree> + '_ {
235 self.groups.values().map(|retained| &retained.subtree)
236 }
237
238 #[cfg(test)]
239 pub(crate) fn subtrees_mut(&mut self) -> impl Iterator<Item = &mut DetachedSubtree> + '_ {
240 self.groups
241 .values_mut()
242 .map(|retained| &mut retained.subtree)
243 }
244
245 #[cfg(any(test, debug_assertions))]
246 pub(crate) fn validate(&self, table: &SlotTable) -> Result<(), SlotInvariantError> {
247 for (key, retained) in &self.groups {
248 let subtree = &retained.subtree;
249 if subtree.root_key() != key.key {
250 return Err(SlotInvariantError::RetainedRootKeyMismatch {
251 parent_scope: key.parent_scope,
252 expected: key.key,
253 actual: subtree.root_key(),
254 });
255 }
256
257 if subtree.root_parent_anchor().is_valid() {
258 return Err(SlotInvariantError::RetainedRootHasActiveParent {
259 root_key: subtree.root_key(),
260 parent_anchor: subtree.root_parent_anchor(),
261 });
262 }
263
264 subtree.validate_detached()?;
265
266 for anchor in subtree.group_anchors() {
267 match table.anchor_state(anchor) {
268 Some(AnchorState::Detached) => {}
269 Some(AnchorState::Active(active_index)) => {
270 return Err(SlotInvariantError::RetainedSubtreeAnchorStillActive {
271 root_key: subtree.root_key(),
272 anchor,
273 active_index,
274 });
275 }
276 actual => {
277 return Err(SlotInvariantError::RetainedAnchorStateMismatch {
278 root_key: subtree.root_key(),
279 anchor,
280 actual,
281 });
282 }
283 }
284 }
285
286 for scope_id in subtree.scope_ids_iter() {
287 if let Some(active_anchor) = table.scope_index_anchor(scope_id) {
288 return Err(SlotInvariantError::RetainedScopeStillActive {
289 root_key: subtree.root_key(),
290 scope_id,
291 active_anchor,
292 });
293 }
294 }
295
296 for payload_anchor in subtree.payload_anchors() {
297 match table.payload_anchor_lifecycle(payload_anchor) {
298 Some(PayloadAnchorLifecycle::Detached) => {}
299 Some(PayloadAnchorLifecycle::Active) => {
300 let (active_owner, active_index) = table
301 .payload_anchor_active_location(payload_anchor)
302 .expect("active retained payload anchor must expose its location");
303 return Err(SlotInvariantError::RetainedPayloadAnchorStillActive {
304 root_key: subtree.root_key(),
305 payload_anchor,
306 active_owner,
307 active_index,
308 });
309 }
310 actual => {
311 return Err(SlotInvariantError::RetainedPayloadAnchorStateMismatch {
312 root_key: subtree.root_key(),
313 payload_anchor,
314 actual,
315 });
316 }
317 }
318 }
319
320 for (node_id, lifecycle) in subtree.node_states() {
321 if lifecycle != NodeLifecycle::RetainedDetached {
322 return Err(SlotInvariantError::RetainedNodeLifecycleMismatch {
323 root_key: subtree.root_key(),
324 node_id,
325 actual: lifecycle,
326 });
327 }
328 }
329 }
330 Ok(())
331 }
332
333 #[cfg(any(test, debug_assertions))]
334 pub(crate) fn debug_verify(&self, table: &SlotTable) {
335 if crate::slot_validation_diagnostics_enabled() {
336 if let Err(err) = self.validate(table) {
337 panic!("retention invariant violation: {err:?}");
338 }
339 }
340 }
341
342 pub(crate) fn evictions_total(&self) -> usize {
343 self.evictions_total
344 }
345
346 pub(crate) fn advance_pass(&mut self) -> Vec<DetachedSubtree> {
347 self.pass_clock = self.pass_clock.saturating_add(1);
348 self.evict_to_budget()
349 }
350
351 fn tick_operation(&mut self) -> u64 {
352 self.operation_clock = self.operation_clock.saturating_add(1);
353 self.operation_clock
354 }
355
356 fn evict_to_budget(&mut self) -> Vec<DetachedSubtree> {
357 let mut evicted = Vec::new();
358 while let Some(key) = self.budget_eviction_key() {
359 let Some(retained) = self.groups.remove(&key) else {
360 break;
361 };
362 self.restored_at_by_key.remove(&key);
363 self.evictions_total = self.evictions_total.saturating_add(1);
364 retained
365 .subtree
366 .assert_root_key(key.key, "retention eviction");
367 retained
368 .subtree
369 .assert_root_parent_detached("retention eviction");
370 retained
371 .subtree
372 .assert_node_lifecycle(NodeLifecycle::RetainedDetached, "retention eviction");
373 evicted.push(retained.subtree);
374 }
375 evicted
376 }
377
378 fn budget_eviction_key(&self) -> Option<RetainKey> {
379 if self.groups.is_empty() {
380 return None;
381 }
382
383 if let Some(max_age_passes) = self.policy.budget.max_age_passes {
384 if let Some(key) = self.age_eviction_key(max_age_passes) {
385 return Some(key);
386 }
387 }
388
389 let over_count = self
390 .policy
391 .budget
392 .max_retained_subtrees
393 .is_some_and(|max| self.groups.len() > max);
394 let over_bytes = self
395 .policy
396 .budget
397 .max_retained_bytes
398 .is_some_and(|max| self.retained_heap_bytes() > max);
399
400 (over_count || over_bytes)
401 .then(|| self.policy_eviction_key())
402 .flatten()
403 }
404
405 fn age_eviction_key(&self, max_age_passes: u64) -> Option<RetainKey> {
406 self.groups
407 .iter()
408 .filter(|(_, retained)| {
409 self.pass_clock.saturating_sub(retained.detached_pass) > max_age_passes
410 })
411 .min_by(|(left_key, left), (right_key, right)| {
412 left.detached_pass
413 .cmp(&right.detached_pass)
414 .then_with(|| left.detached_order.cmp(&right.detached_order))
415 .then_with(|| retain_key_cmp(left_key, right_key))
416 })
417 .map(|(key, _)| *key)
418 }
419
420 fn policy_eviction_key(&self) -> Option<RetainKey> {
421 match self.policy.eviction {
422 RetentionEvictionPolicy::LeastRecentlyDetached => self.least_recently_detached_key(),
423 RetentionEvictionPolicy::LeastRecentlyRestored => self.least_recently_restored_key(),
424 RetentionEvictionPolicy::LargestFirst => self.largest_first_key(),
425 }
426 }
427
428 fn least_recently_detached_key(&self) -> Option<RetainKey> {
429 self.groups
430 .iter()
431 .min_by(|(left_key, left), (right_key, right)| {
432 left.detached_order
433 .cmp(&right.detached_order)
434 .then_with(|| retain_key_cmp(left_key, right_key))
435 })
436 .map(|(key, _)| *key)
437 }
438
439 fn least_recently_restored_key(&self) -> Option<RetainKey> {
440 self.groups
441 .iter()
442 .min_by(|(left_key, left), (right_key, right)| {
443 left.last_restored_order
444 .cmp(&right.last_restored_order)
445 .then_with(|| left.detached_order.cmp(&right.detached_order))
446 .then_with(|| retain_key_cmp(left_key, right_key))
447 })
448 .map(|(key, _)| *key)
449 }
450
451 fn largest_first_key(&self) -> Option<RetainKey> {
452 self.groups
453 .iter()
454 .max_by(|(left_key, left), (right_key, right)| {
455 left.heap_bytes()
456 .cmp(&right.heap_bytes())
457 .then_with(|| retain_key_cmp(left_key, right_key))
458 })
459 .map(|(key, _)| *key)
460 }
461
462 fn retained_heap_bytes(&self) -> usize {
463 self.groups.values().map(RetainedGroup::heap_bytes).sum()
464 }
465}
466
467fn retain_key_cmp(left: &RetainKey, right: &RetainKey) -> Ordering {
468 (
469 left.parent_scope,
470 left.key.static_key,
471 left.key.explicit_key,
472 left.key.ordinal,
473 )
474 .cmp(&(
475 right.parent_scope,
476 right.key.static_key,
477 right.key.explicit_key,
478 right.key.ordinal,
479 ))
480}
481
482#[cfg(test)]
483mod tests {
484 use super::*;
485
486 #[test]
487 fn retention_budget_default_is_unbounded() {
488 assert_eq!(RetentionBudget::default(), RetentionBudget::UNBOUNDED);
489 assert_eq!(RetentionBudget::default().max_retained_subtrees, None);
490 assert_eq!(RetentionBudget::default().max_retained_bytes, None);
491 assert_eq!(RetentionBudget::default().max_age_passes, None);
492 }
493
494 #[test]
495 fn retention_policy_default_uses_unbounded_budget_and_detach_lru() {
496 assert_eq!(RetentionPolicy::default(), RetentionPolicy::UNBOUNDED);
497 assert_eq!(
498 RetentionPolicy::default().budget,
499 RetentionBudget::UNBOUNDED
500 );
501 assert_eq!(
502 RetentionPolicy::default().eviction,
503 RetentionEvictionPolicy::LeastRecentlyDetached
504 );
505 }
506
507 #[test]
508 fn retention_budget_can_express_all_limits() {
509 let budget = RetentionBudget {
510 max_retained_subtrees: Some(3),
511 max_retained_bytes: Some(4096),
512 max_age_passes: Some(5),
513 };
514
515 assert_eq!(budget.max_retained_subtrees, Some(3));
516 assert_eq!(budget.max_retained_bytes, Some(4096));
517 assert_eq!(budget.max_age_passes, Some(5));
518 }
519}