Skip to main content

eredu_runtime/cache/
lifecycle.rs

1//! Backend-neutral ownership for live cache blocks and mutable tails.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use serde::{Deserialize, Serialize};
6
7use eredu_core::{cache::CacheBlockId, residency::CacheEvictionPolicy};
8
9/// Device-resident mutable state that has not yet become an immutable block.
10#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
11pub struct MutableCacheTail {
12    /// Concrete bytes owned by the backend on the execution device.
13    pub bytes: u64,
14    /// Exclusive logical token frontier represented by the tail.
15    pub end: i64,
16}
17
18#[derive(Debug, Clone, Copy, Eq, PartialEq)]
19struct BlockLifecycle {
20    leases: usize,
21    access_count: u64,
22    last_access: u64,
23    protected_prefix: bool,
24}
25
26/// Canonical logical ownership state for one backend cache session.
27///
28/// Backends keep concrete arrays, buffers, files, and completion objects in a
29/// separate storage map keyed by [`CacheBlockId`]. This catalog is the sole
30/// owner of leases, access order, protected-prefix status, and mutable tails.
31#[derive(Debug, Default)]
32pub struct CacheBlockLifecycle {
33    access_clock: u64,
34    blocks: BTreeMap<CacheBlockId, BlockLifecycle>,
35    tails: BTreeMap<usize, MutableCacheTail>,
36}
37
38impl CacheBlockLifecycle {
39    /// Creates an empty cache-session lifecycle.
40    pub const fn new() -> Self {
41        Self {
42            access_clock: 0,
43            blocks: BTreeMap::new(),
44            tails: BTreeMap::new(),
45        }
46    }
47
48    /// Registers a newly materialized immutable block.
49    pub fn insert(
50        &mut self,
51        id: CacheBlockId,
52        protected_prefix: bool,
53    ) -> Result<(), CacheLifecycleError> {
54        if self.blocks.contains_key(&id) {
55            return Err(CacheLifecycleError::DuplicateBlock(id));
56        }
57        let last_access = self.tick()?;
58        self.blocks.insert(
59            id,
60            BlockLifecycle {
61                leases: 0,
62                access_count: 0,
63                last_access,
64                protected_prefix,
65            },
66        );
67        Ok(())
68    }
69
70    /// Removes an unleased block from logical ownership.
71    pub fn remove(&mut self, id: &CacheBlockId) -> Result<(), CacheLifecycleError> {
72        if self.is_leased(id)? {
73            return Err(CacheLifecycleError::BlockLeased(id.clone()));
74        }
75        self.blocks.remove(id);
76        Ok(())
77    }
78
79    /// Atomically replaces a set of blocks and updates one mutable tail.
80    ///
81    /// The expected lease count makes a caller-owned truncation lease explicit:
82    /// validation completes before any logical state is changed.
83    pub fn replace(
84        &mut self,
85        removals: &[(CacheBlockId, usize)],
86        replacement: Option<(CacheBlockId, bool)>,
87        tail_layer: usize,
88        tail: MutableCacheTail,
89    ) -> Result<(), CacheLifecycleError> {
90        let removal_ids = removals.iter().map(|(id, _)| id).collect::<BTreeSet<_>>();
91        if removal_ids.len() != removals.len() {
92            return Err(CacheLifecycleError::DuplicateRemoval);
93        }
94        for (id, leases) in removals {
95            self.require_lease_count(id, *leases)?;
96        }
97        if let Some((id, _)) = &replacement {
98            if self.blocks.contains_key(id) && !removal_ids.contains(id) {
99                return Err(CacheLifecycleError::DuplicateBlock(id.clone()));
100            }
101        }
102        let replacement_access = if replacement.is_some() {
103            Some(self.tick()?)
104        } else {
105            None
106        };
107
108        for (id, _) in removals {
109            self.blocks.remove(id);
110        }
111        if let Some((id, protected_prefix)) = replacement {
112            self.blocks.insert(
113                id,
114                BlockLifecycle {
115                    leases: 0,
116                    access_count: 0,
117                    last_access: replacement_access.expect("replacement access was allocated"),
118                    protected_prefix,
119                },
120            );
121        }
122        self.tails.insert(tail_layer, tail);
123        Ok(())
124    }
125
126    /// Acquires one exact logical lease and records demand access.
127    pub fn acquire(&mut self, id: &CacheBlockId) -> Result<(), CacheLifecycleError> {
128        let leases = self
129            .blocks
130            .get(id)
131            .ok_or_else(|| CacheLifecycleError::MissingBlock(id.clone()))?
132            .leases
133            .checked_add(1)
134            .ok_or_else(|| CacheLifecycleError::LeaseOverflow(id.clone()))?;
135        let clock = self.tick()?;
136        let block = self
137            .blocks
138            .get_mut(id)
139            .expect("cache block was validated before advancing the access clock");
140        block.leases = leases;
141        block.access_count = block.access_count.saturating_add(1);
142        block.last_access = clock;
143        Ok(())
144    }
145
146    /// Releases one exact logical lease.
147    pub fn release(&mut self, id: &CacheBlockId) -> Result<(), CacheLifecycleError> {
148        let block = self
149            .blocks
150            .get_mut(id)
151            .ok_or_else(|| CacheLifecycleError::MissingBlock(id.clone()))?;
152        block.leases = block
153            .leases
154            .checked_sub(1)
155            .ok_or_else(|| CacheLifecycleError::LeaseUnderflow(id.clone()))?;
156        Ok(())
157    }
158
159    /// Returns the exact lease count for a block.
160    pub fn lease_count(&self, id: &CacheBlockId) -> Result<usize, CacheLifecycleError> {
161        self.blocks
162            .get(id)
163            .map(|block| block.leases)
164            .ok_or_else(|| CacheLifecycleError::MissingBlock(id.clone()))
165    }
166
167    /// Returns whether a block currently has any logical owner.
168    pub fn is_leased(&self, id: &CacheBlockId) -> Result<bool, CacheLifecycleError> {
169        self.lease_count(id).map(|leases| leases != 0)
170    }
171
172    /// Returns the first leased block in stable identity order.
173    pub fn first_leased(&self) -> Option<&CacheBlockId> {
174        self.blocks
175            .iter()
176            .find_map(|(id, block)| (block.leases != 0).then_some(id))
177    }
178
179    /// Returns whether the block is protected as an immutable prefix.
180    pub fn is_protected_prefix(&self, id: &CacheBlockId) -> Result<bool, CacheLifecycleError> {
181        self.blocks
182            .get(id)
183            .map(|block| block.protected_prefix)
184            .ok_or_else(|| CacheLifecycleError::MissingBlock(id.clone()))
185    }
186
187    /// Selects one eviction victim from backend-supplied physically eligible IDs.
188    pub fn eviction_candidate(
189        &self,
190        candidates: impl IntoIterator<Item = CacheBlockId>,
191        required: Option<&CacheBlockId>,
192        recent_per_layer: usize,
193        policy: CacheEvictionPolicy,
194    ) -> Result<Option<CacheBlockId>, CacheLifecycleError> {
195        let candidates = candidates.into_iter().collect::<BTreeSet<_>>();
196        for id in &candidates {
197            if !self.blocks.contains_key(id) {
198                return Err(CacheLifecycleError::MissingBlock(id.clone()));
199            }
200        }
201        let recent = recent_ids(&candidates, recent_per_layer);
202        Ok(candidates
203            .into_iter()
204            .filter(|id| {
205                let block = self.blocks.get(id).expect("candidate was validated");
206                block.leases == 0
207                    && !block.protected_prefix
208                    && required != Some(id)
209                    && !recent.contains(id)
210            })
211            .min_by_key(|id| {
212                let block = self.blocks.get(id).expect("candidate was validated");
213                match policy {
214                    CacheEvictionPolicy::LeastRecentlyUsed => {
215                        (block.last_access, block.access_count, id.clone())
216                    }
217                    CacheEvictionPolicy::LeastFrequentlyUsed => {
218                        (block.access_count, block.last_access, id.clone())
219                    }
220                }
221            }))
222    }
223
224    /// Counts unprotected blocks retained by a per-layer recent window.
225    pub fn recent_protection_counts(
226        &self,
227        candidates: impl IntoIterator<Item = CacheBlockId>,
228        recent_per_layer: usize,
229    ) -> Result<BTreeMap<usize, u64>, CacheLifecycleError> {
230        let candidates = candidates
231            .into_iter()
232            .filter_map(|id| match self.blocks.get(&id) {
233                Some(block) if !block.protected_prefix => Some(Ok(id)),
234                Some(_) => None,
235                None => Some(Err(CacheLifecycleError::MissingBlock(id))),
236            })
237            .collect::<Result<BTreeSet<_>, _>>()?;
238        let mut counts = BTreeMap::new();
239        for id in recent_ids(&candidates, recent_per_layer) {
240            *counts.entry(id.global_layer).or_default() += 1;
241        }
242        Ok(counts)
243    }
244
245    /// Replaces one layer's mutable tail, returning the prior state for rollback.
246    pub fn set_tail(&mut self, layer: usize, tail: MutableCacheTail) -> Option<MutableCacheTail> {
247        self.tails.insert(layer, tail)
248    }
249
250    /// Restores a prior tail after a failed backend admission.
251    pub fn restore_tail(&mut self, layer: usize, tail: Option<MutableCacheTail>) {
252        match tail {
253            Some(tail) => {
254                self.tails.insert(layer, tail);
255            }
256            None => {
257                self.tails.remove(&layer);
258            }
259        }
260    }
261
262    /// Returns one layer's mutable tail.
263    pub fn tail(&self, layer: usize) -> Option<MutableCacheTail> {
264        self.tails.get(&layer).copied()
265    }
266
267    /// Iterates mutable tails in stable layer order.
268    pub fn tails(&self) -> impl Iterator<Item = (usize, MutableCacheTail)> + '_ {
269        self.tails.iter().map(|(layer, tail)| (*layer, *tail))
270    }
271
272    /// Clears all blocks and tails only when no lease is active.
273    pub fn clear(&mut self) -> Result<(), CacheLifecycleError> {
274        if let Some(id) = self.first_leased() {
275            return Err(CacheLifecycleError::BlockLeased(id.clone()));
276        }
277        self.blocks.clear();
278        self.tails.clear();
279        Ok(())
280    }
281
282    fn require_lease_count(
283        &self,
284        id: &CacheBlockId,
285        expected: usize,
286    ) -> Result<(), CacheLifecycleError> {
287        let actual = self.lease_count(id)?;
288        if actual == expected {
289            Ok(())
290        } else {
291            Err(CacheLifecycleError::UnexpectedLeaseCount {
292                id: id.clone(),
293                expected,
294                actual,
295            })
296        }
297    }
298
299    fn tick(&mut self) -> Result<u64, CacheLifecycleError> {
300        self.access_clock = self
301            .access_clock
302            .checked_add(1)
303            .ok_or(CacheLifecycleError::AccessClockOverflow)?;
304        Ok(self.access_clock)
305    }
306}
307
308fn recent_ids(candidates: &BTreeSet<CacheBlockId>, limit: usize) -> BTreeSet<CacheBlockId> {
309    if limit == 0 {
310        return BTreeSet::new();
311    }
312    let mut by_layer = BTreeMap::<usize, Vec<&CacheBlockId>>::new();
313    for id in candidates {
314        by_layer.entry(id.global_layer).or_default().push(id);
315    }
316    by_layer
317        .into_values()
318        .flat_map(|mut ids| {
319            ids.sort_unstable_by(|left, right| right.start.cmp(&left.start).then(right.cmp(left)));
320            ids.into_iter().take(limit).cloned().collect::<Vec<_>>()
321        })
322        .collect()
323}
324
325/// Invalid backend-neutral cache ownership transition.
326#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
327pub enum CacheLifecycleError {
328    /// A block identity was registered more than once.
329    #[error("duplicate cache block {0:?}")]
330    DuplicateBlock(CacheBlockId),
331    /// A requested block is not registered.
332    #[error("missing cache block {0:?}")]
333    MissingBlock(CacheBlockId),
334    /// A removal list repeated one identity.
335    #[error("cache block replacement contains a duplicate removal")]
336    DuplicateRemoval,
337    /// An operation required an unleased block.
338    #[error("cache block is leased by active attention: {0:?}")]
339    BlockLeased(CacheBlockId),
340    /// A transactional replacement observed a different lease count.
341    #[error("cache block {id:?} has lease count {actual}, expected {expected}")]
342    UnexpectedLeaseCount {
343        /// Stable block identity.
344        id: CacheBlockId,
345        /// Required exact ownership count.
346        expected: usize,
347        /// Observed exact ownership count.
348        actual: usize,
349    },
350    /// Lease acquisition exceeded addressable ownership.
351    #[error("cache block lease count overflowed: {0:?}")]
352    LeaseOverflow(CacheBlockId),
353    /// A backend released a lease it did not own.
354    #[error("cache block lease count underflowed: {0:?}")]
355    LeaseUnderflow(CacheBlockId),
356    /// The stable access clock exhausted its range.
357    #[error("cache block access clock overflowed")]
358    AccessClockOverflow,
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use eredu_core::cache::CacheRepresentation;
365
366    fn id(layer: usize, start: i64) -> CacheBlockId {
367        CacheBlockId {
368            session_id: 1,
369            global_layer: layer,
370            representation: CacheRepresentation::KeyValue,
371            start,
372            end: start + 1,
373            rank: None,
374        }
375    }
376
377    #[test]
378    fn leases_are_exact_and_block_destructive_transitions() {
379        let block = id(0, 0);
380        let mut lifecycle = CacheBlockLifecycle::new();
381        lifecycle.insert(block.clone(), false).unwrap();
382        lifecycle.acquire(&block).unwrap();
383        assert_eq!(lifecycle.lease_count(&block).unwrap(), 1);
384        assert!(matches!(
385            lifecycle.remove(&block),
386            Err(CacheLifecycleError::BlockLeased(_))
387        ));
388        lifecycle.release(&block).unwrap();
389        lifecycle.remove(&block).unwrap();
390        assert_eq!(
391            lifecycle.release(&block),
392            Err(CacheLifecycleError::MissingBlock(block))
393        );
394    }
395
396    #[test]
397    fn rejected_acquisition_does_not_advance_access_order() {
398        let block = id(0, 0);
399        let missing = id(0, 1);
400        let mut lifecycle = CacheBlockLifecycle::new();
401        lifecycle.insert(block.clone(), false).unwrap();
402        let access_clock = lifecycle.access_clock;
403
404        assert_eq!(
405            lifecycle.acquire(&missing),
406            Err(CacheLifecycleError::MissingBlock(missing))
407        );
408        assert_eq!(lifecycle.access_clock, access_clock);
409
410        lifecycle.blocks.get_mut(&block).unwrap().leases = usize::MAX;
411        assert_eq!(
412            lifecycle.acquire(&block),
413            Err(CacheLifecycleError::LeaseOverflow(block))
414        );
415        assert_eq!(lifecycle.access_clock, access_clock);
416    }
417
418    #[test]
419    fn replacement_and_tail_update_are_atomic() {
420        let first = id(0, 0);
421        let second = id(0, 1);
422        let replacement = CacheBlockId {
423            end: 1,
424            ..first.clone()
425        };
426        let mut lifecycle = CacheBlockLifecycle::new();
427        lifecycle.insert(first.clone(), false).unwrap();
428        lifecycle.insert(second.clone(), false).unwrap();
429        lifecycle.acquire(&first).unwrap();
430        lifecycle
431            .replace(
432                &[(first.clone(), 1), (second.clone(), 0)],
433                Some((replacement.clone(), true)),
434                0,
435                MutableCacheTail { bytes: 0, end: 1 },
436            )
437            .unwrap();
438        assert_eq!(lifecycle.lease_count(&replacement).unwrap(), 0);
439        assert!(lifecycle.is_protected_prefix(&replacement).unwrap());
440        assert_eq!(
441            lifecycle.tail(0),
442            Some(MutableCacheTail { bytes: 0, end: 1 })
443        );
444    }
445
446    #[test]
447    fn eviction_is_deterministic_and_respects_owners_and_windows() {
448        let oldest = id(0, 0);
449        let leased = id(0, 1);
450        let recent = id(0, 2);
451        let protected = id(1, 0);
452        let mut lifecycle = CacheBlockLifecycle::new();
453        for (block, prefix) in [
454            (oldest.clone(), false),
455            (leased.clone(), false),
456            (recent.clone(), false),
457            (protected.clone(), true),
458        ] {
459            lifecycle.insert(block, prefix).unwrap();
460        }
461        lifecycle.acquire(&leased).unwrap();
462        let candidates = [oldest.clone(), leased, recent, protected];
463        assert_eq!(
464            lifecycle
465                .eviction_candidate(candidates, None, 1, CacheEvictionPolicy::LeastRecentlyUsed,)
466                .unwrap(),
467            Some(oldest)
468        );
469    }
470
471    #[test]
472    fn mutable_tail_rollback_restores_exact_state() {
473        let mut lifecycle = CacheBlockLifecycle::new();
474        let first = MutableCacheTail { bytes: 8, end: 2 };
475        assert_eq!(lifecycle.set_tail(3, first), None);
476        let prior = lifecycle.set_tail(3, MutableCacheTail { bytes: 16, end: 4 });
477        lifecycle.restore_tail(3, prior);
478        assert_eq!(lifecycle.tail(3), Some(first));
479    }
480}