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