1use std::collections::{BTreeMap, BTreeSet};
4
5use serde::{Deserialize, Serialize};
6
7use eredu_core::{cache::CacheBlockId, residency::CacheEvictionPolicy};
8
9#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
11pub struct MutableCacheTail {
12 pub bytes: u64,
14 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#[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 pub const fn new() -> Self {
41 Self {
42 access_clock: 0,
43 blocks: BTreeMap::new(),
44 tails: BTreeMap::new(),
45 }
46 }
47
48 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 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 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 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 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 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 pub fn is_leased(&self, id: &CacheBlockId) -> Result<bool, CacheLifecycleError> {
169 self.lease_count(id).map(|leases| leases != 0)
170 }
171
172 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 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 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 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 pub fn set_tail(&mut self, layer: usize, tail: MutableCacheTail) -> Option<MutableCacheTail> {
247 self.tails.insert(layer, tail)
248 }
249
250 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 pub fn tail(&self, layer: usize) -> Option<MutableCacheTail> {
264 self.tails.get(&layer).copied()
265 }
266
267 pub fn tails(&self) -> impl Iterator<Item = (usize, MutableCacheTail)> + '_ {
269 self.tails.iter().map(|(layer, tail)| (*layer, *tail))
270 }
271
272 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#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
327pub enum CacheLifecycleError {
328 #[error("duplicate cache block {0:?}")]
330 DuplicateBlock(CacheBlockId),
331 #[error("missing cache block {0:?}")]
333 MissingBlock(CacheBlockId),
334 #[error("cache block replacement contains a duplicate removal")]
336 DuplicateRemoval,
337 #[error("cache block is leased by active attention: {0:?}")]
339 BlockLeased(CacheBlockId),
340 #[error("cache block {id:?} has lease count {actual}, expected {expected}")]
342 UnexpectedLeaseCount {
343 id: CacheBlockId,
345 expected: usize,
347 actual: usize,
349 },
350 #[error("cache block lease count overflowed: {0:?}")]
352 LeaseOverflow(CacheBlockId),
353 #[error("cache block lease count underflowed: {0:?}")]
355 LeaseUnderflow(CacheBlockId),
356 #[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}