dynamo_mocker/cache/
hash_cache.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::common::evictor::LRUEvictor;
5use dynamo_tokens::blocks::UniqueBlock;
6use std::collections::HashMap;
7
8/// Hash-based KV cache with O(1) block lookups, maintaining active (ref-counted) and
9/// inactive (LRU-evictable) pools.
10pub struct HashCache {
11    active_blocks: HashMap<UniqueBlock, usize>,
12    inactive_blocks: LRUEvictor<UniqueBlock>,
13    max_capacity: usize,
14}
15
16impl HashCache {
17    /// Create a new HashCache with the given maximum block capacity.
18    pub fn new(max_capacity: usize) -> Self {
19        Self {
20            active_blocks: HashMap::new(),
21            inactive_blocks: LRUEvictor::default(),
22            max_capacity,
23        }
24    }
25
26    /// Get the reference count of an active block, if it exists.
27    pub fn get_active_ref_count(&self, block: &UniqueBlock) -> Option<usize> {
28        self.active_blocks.get(block).copied()
29    }
30
31    /// Increment the reference count of an active block. Returns the new count.
32    pub fn increment_ref(&mut self, block: &UniqueBlock) -> usize {
33        let ref_count = self
34            .active_blocks
35            .get_mut(block)
36            .expect("block must be active to increment ref");
37        *ref_count += 1;
38        *ref_count
39    }
40
41    /// Decrement the reference count of an active block. Returns the new count.
42    pub fn decrement_ref(&mut self, block: &UniqueBlock) -> usize {
43        let ref_count = self
44            .active_blocks
45            .get_mut(block)
46            .expect("block must be active to decrement ref");
47        *ref_count -= 1;
48        *ref_count
49    }
50
51    /// Insert a block into the active pool with the given reference count.
52    pub fn insert_active(&mut self, block: UniqueBlock, ref_count: usize) {
53        self.active_blocks.insert(block, ref_count);
54    }
55
56    /// Remove a block from the active pool. Returns the reference count, or None if not found.
57    pub fn remove_active(&mut self, block: &UniqueBlock) -> Option<usize> {
58        self.active_blocks.remove(block)
59    }
60
61    /// Check if a block is in the active pool.
62    pub fn contains_active(&self, block: &UniqueBlock) -> bool {
63        self.active_blocks.contains_key(block)
64    }
65
66    /// Insert a block into the inactive pool (LRU order).
67    pub fn insert_inactive(&mut self, block: UniqueBlock) {
68        self.inactive_blocks.insert(block);
69    }
70
71    /// Remove a block from the inactive pool. Returns true if it was found.
72    pub fn remove_inactive(&mut self, block: &UniqueBlock) -> bool {
73        self.inactive_blocks.remove(block)
74    }
75
76    /// Evict the least-recently-used block from the inactive pool.
77    pub fn evict_inactive(&mut self) -> Option<UniqueBlock> {
78        self.inactive_blocks.evict()
79    }
80
81    /// Check if a block is in the inactive pool.
82    pub fn contains_inactive(&self, block: &UniqueBlock) -> bool {
83        self.inactive_blocks.contains(block)
84    }
85
86    /// Check if a block exists in either active or inactive pool.
87    pub fn contains(&self, block: &UniqueBlock) -> bool {
88        self.active_blocks.contains_key(block) || self.inactive_blocks.contains(block)
89    }
90
91    /// Move block from active to inactive (ref_count reached 0).
92    pub fn deactivate(&mut self, block: &UniqueBlock) {
93        debug_assert!(
94            self.active_blocks.contains_key(block),
95            "deactivate called on non-active block"
96        );
97        debug_assert!(
98            !self.inactive_blocks.contains(block),
99            "deactivate called on already-inactive block"
100        );
101        self.active_blocks.remove(block);
102        self.inactive_blocks.insert(block.clone());
103    }
104
105    /// Move block from inactive to active with ref_count=1. Returns true if found.
106    pub fn reactivate(&mut self, block: &UniqueBlock) -> bool {
107        if self.inactive_blocks.remove(block) {
108            self.active_blocks.insert(block.clone(), 1);
109            true
110        } else {
111            false
112        }
113    }
114
115    /// Check if total blocks (active + inactive) has reached max_capacity.
116    pub fn is_at_capacity(&self) -> bool {
117        self.active_blocks.len() + self.inactive_blocks.len() >= self.max_capacity
118    }
119
120    /// Get the number of active blocks.
121    pub fn num_active(&self) -> usize {
122        self.active_blocks.len()
123    }
124
125    /// Get the number of inactive blocks.
126    pub fn num_inactive(&self) -> usize {
127        self.inactive_blocks.len()
128    }
129
130    /// Get the maximum block capacity.
131    pub fn max_capacity(&self) -> usize {
132        self.max_capacity
133    }
134
135    /// Get the current capacity (active + inactive blocks).
136    pub fn current_capacity(&self) -> usize {
137        self.active_blocks.len() + self.inactive_blocks.len()
138    }
139
140    /// Iterate over active block keys.
141    pub fn active_keys(&self) -> impl Iterator<Item = &UniqueBlock> {
142        self.active_blocks.keys()
143    }
144
145    /// Iterate over inactive block keys.
146    pub fn inactive_keys(&self) -> impl Iterator<Item = &UniqueBlock> {
147        self.inactive_blocks.keys()
148    }
149
150    /// Direct access to active blocks map (for tests that check ref counts).
151    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
152        &self.active_blocks
153    }
154}
dynamo_mocker/cache/hash_cache.rs

dynamo_mocker/cache/
hash_cache.rs
