quill_sql/utils/cache/
lru_k.rs

1use super::Replacer;
2use crate::buffer::FrameId;
3use crate::error::{QuillSQLError, QuillSQLResult};
4use std::collections::{HashMap, LinkedList};
5
6#[derive(Debug)]
7struct LRUKNode {
8    k: usize,
9    // 该frame最近k次被访问的时间
10    history: LinkedList<u64>,
11    // 是否可被置换
12    is_evictable: bool,
13}
14impl LRUKNode {
15    pub fn new(k: usize) -> Self {
16        Self {
17            k,
18            history: LinkedList::new(),
19            is_evictable: false,
20        }
21    }
22    pub fn record_access(&mut self, timestamp: u64) {
23        self.history.push_back(timestamp);
24        if self.history.len() > self.k {
25            self.history.pop_front();
26        }
27    }
28}
29
30#[derive(Debug)]
31pub struct LRUKReplacer {
32    // 当前可置换的frame数
33    current_size: usize,
34    // 可置换的frame数上限
35    replacer_size: usize,
36    k: usize,
37    node_store: HashMap<FrameId, LRUKNode>,
38    // 当前时间戳（从0递增）
39    current_timestamp: u64,
40}
41
42impl Replacer for LRUKReplacer {
43    fn new(capacity: usize) -> Self {
44        // Use a default K value or make K configurable via another mechanism if needed
45        // For now, let's assume a reasonable default K, e.g., 2
46        const DEFAULT_K: usize = 2;
47        Self {
48            current_size: 0,
49            replacer_size: capacity, // Use capacity from trait
50            k: DEFAULT_K,
51            node_store: HashMap::with_capacity(capacity),
52            current_timestamp: 0,
53        }
54    }
55
56    /// Drop victim from replacer; returns frame id if available.
57    fn evict(&mut self) -> Option<FrameId> {
58        let mut victim: Option<(FrameId, u64, u64)> = None; // (frame_id, k_distance, oldest_ts)
59        for (&frame_id, node) in &self.node_store {
60            if !node.is_evictable {
61                continue;
62            }
63            let oldest = node.history.front().copied().unwrap_or(0);
64            let k_distance = if node.history.len() < self.k {
65                u64::MAX
66            } else {
67                self.current_timestamp.saturating_sub(oldest)
68            };
69            match victim {
70                None => victim = Some((frame_id, k_distance, oldest)),
71                Some((_, best_dist, best_oldest)) => {
72                    if k_distance > best_dist || (k_distance == best_dist && oldest < best_oldest) {
73                        victim = Some((frame_id, k_distance, oldest));
74                    }
75                }
76            }
77        }
78        if let Some((frame_id, _, _)) = victim {
79            self.remove(frame_id);
80            Some(frame_id)
81        } else {
82            None
83        }
84    }
85
86    // 记录frame的访问
87    fn record_access(&mut self, frame_id: FrameId) -> QuillSQLResult<()> {
88        if let Some(node) = self.node_store.get_mut(&frame_id) {
89            node.record_access(self.current_timestamp);
90            self.current_timestamp += 1;
91        } else {
92            // 创建新node
93            if self.node_store.len() >= self.replacer_size {
94                return Err(QuillSQLError::Internal(
95                    "frame size exceeds the limit".to_string(),
96                ));
97            }
98            let mut node = LRUKNode::new(self.k);
99            node.record_access(self.current_timestamp);
100            self.current_timestamp += 1;
101            self.node_store.insert(frame_id, node);
102        }
103        Ok(())
104    }
105
106    // 设置frame是否可被置换
107    fn set_evictable(&mut self, frame_id: FrameId, set_evictable: bool) -> QuillSQLResult<()> {
108        if let Some(node) = self.node_store.get_mut(&frame_id) {
109            let evictable = node.is_evictable;
110            node.is_evictable = set_evictable;
111            if set_evictable && !evictable {
112                self.current_size += 1;
113            } else if !set_evictable && evictable {
114                self.current_size -= 1;
115            }
116            Ok(())
117        } else {
118            Err(QuillSQLError::Internal("frame not found".to_string()))
119        }
120    }
121
122    // 移除frame
123    fn remove(&mut self, frame_id: FrameId) {
124        if let Some(node) = self.node_store.get(&frame_id) {
125            assert!(node.is_evictable, "frame is not evictable");
126            self.node_store.remove(&frame_id);
127            self.current_size -= 1;
128        }
129    }
130
131    // 获取当前可置换的frame数
132    fn size(&self) -> usize {
133        self.current_size
134    }
135}
136
137impl LRUKReplacer {
138    pub fn with_k(num_frames: usize, k: usize) -> Self {
139        Self {
140            current_size: 0,
141            replacer_size: num_frames,
142            k,
143            node_store: HashMap::new(),
144            current_timestamp: 0,
145        }
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::utils::cache::Replacer; // Import the trait
153
154    #[test]
155    pub fn test_lru_k_set_evictable() {
156        let mut replacer = LRUKReplacer::with_k(3, 2); // Use specific constructor
157        replacer.record_access(1).unwrap();
158        replacer.set_evictable(1, true).unwrap();
159        assert_eq!(replacer.size(), 1);
160        replacer.set_evictable(1, false).unwrap();
161        assert_eq!(replacer.size(), 0);
162    }
163
164    #[test]
165    pub fn test_lru_k_evict_all_pages_at_least_k() {
166        let mut replacer = LRUKReplacer::with_k(2, 3);
167        replacer.record_access(1).unwrap();
168        replacer.record_access(2).unwrap();
169        replacer.record_access(2).unwrap();
170        replacer.record_access(1).unwrap(); // ts=3
171        replacer.record_access(2).unwrap(); // ts=4
172        replacer.set_evictable(1, true).unwrap();
173        replacer.set_evictable(2, true).unwrap();
174        // Frame 1 history: [0, 3], k-dist = 5 - 0 = 5
175        // Frame 2 history: [1, 2, 4], k-dist = 5 - 1 = 4
176        let frame_id = replacer.evict();
177        assert_eq!(frame_id, Some(1)); // Frame 1 has larger k-distance
178    }
179
180    #[test]
181    pub fn test_lru_k_evict_some_page_less_than_k() {
182        let mut replacer = LRUKReplacer::with_k(3, 3);
183        replacer.record_access(1).unwrap(); // ts=0
184        replacer.record_access(2).unwrap(); // ts=1, history < k
185        replacer.record_access(3).unwrap(); // ts=2, history < k
186        replacer.record_access(1).unwrap(); // ts=3
187        replacer.record_access(1).unwrap(); // ts=4, history = [0, 3, 4]
188        replacer.record_access(3).unwrap(); // ts=5, history = [2, 5]
189                                            // Frame 1: history=[0,3,4], k=3, k-dist = 6-0 = 6
190                                            // Frame 2: history=[1], k=3, infinite k-dist, oldest_ts=1
191                                            // Frame 3: history=[2,5], k=3, infinite k-dist, oldest_ts=2
192        replacer.set_evictable(1, true).unwrap();
193        replacer.set_evictable(2, true).unwrap();
194        replacer.set_evictable(3, true).unwrap();
195        // Evict should prioritize finite k-dist (frame 1), BUT the logic was updated
196        // to prioritize infinite k-dist with oldest timestamp. Frame 2 (ts=1) is older than Frame 3 (ts=2).
197        let frame_id = replacer.evict();
198        assert_eq!(frame_id, Some(2));
199    }
200
201    #[test]
202    pub fn test_lru_k_test_case() {
203        let mut lru_replacer = LRUKReplacer::with_k(7, 2); // k=2
204
205        // Scenario: add six elements
206        lru_replacer.record_access(1).unwrap(); // ts=0
207        lru_replacer.record_access(2).unwrap(); // ts=1
208        lru_replacer.record_access(3).unwrap(); // ts=2
209        lru_replacer.record_access(4).unwrap(); // ts=3
210        lru_replacer.record_access(5).unwrap(); // ts=4
211        lru_replacer.record_access(6).unwrap(); // ts=5
212        lru_replacer.set_evictable(1, true).unwrap();
213        lru_replacer.set_evictable(2, true).unwrap();
214        lru_replacer.set_evictable(3, true).unwrap();
215        lru_replacer.set_evictable(4, true).unwrap();
216        lru_replacer.set_evictable(5, true).unwrap();
217        lru_replacer.set_evictable(6, false).unwrap(); // 6 not evictable initially
218        assert_eq!(5, lru_replacer.size());
219
220        // Scenario: access frame 1 again
221        lru_replacer.record_access(1).unwrap(); // ts=6, history=[0, 6]
222                                                // Frame 1: k=2, k-dist = 7-0 = 7
223                                                // Frame 2: history=[1], k=2, inf k-dist, oldest=1
224                                                // Frame 3: history=[2], k=2, inf k-dist, oldest=2
225                                                // Frame 4: history=[3], k=2, inf k-dist, oldest=3
226                                                // Frame 5: history=[4], k=2, inf k-dist, oldest=4
227
228        // Scenario: Evict three pages.
229        // Infinite distance frames first, ordered by oldest timestamp.
230        let value = lru_replacer.evict(); // Evict frame 2 (oldest inf dist)
231        assert_eq!(Some(2), value);
232        let value = lru_replacer.evict(); // Evict frame 3
233        assert_eq!(Some(3), value);
234        let value = lru_replacer.evict(); // Evict frame 4
235        assert_eq!(Some(4), value);
236        assert_eq!(lru_replacer.size(), 2); // Remaining: [1, 5] (6 is non-evictable)
237
238        // Scenario: Insert new frames 3, 4, update 5
239        lru_replacer.record_access(3).unwrap(); // ts=7, history=[7], inf k-dist, oldest=7
240        lru_replacer.record_access(4).unwrap(); // ts=8, history=[8], inf k-dist, oldest=8
241        lru_replacer.record_access(5).unwrap(); // ts=9, history=[4, 9], k-dist=10-4=6
242        lru_replacer.record_access(4).unwrap(); // ts=10, history=[8, 10], k-dist=11-8=3
243        lru_replacer.set_evictable(3, true).unwrap();
244        lru_replacer.set_evictable(4, true).unwrap();
245        assert_eq!(4, lru_replacer.size()); // Now [1, 5, 3, 4] are evictable
246                                            // Frame 1: hist=[0, 6], k=2, k-dist = 11-0 = 11
247                                            // Frame 5: hist=[4, 9], k=2, k-dist = 11-4 = 7
248                                            // Frame 3: hist=[7], k=2, inf k-dist, oldest=7
249                                            // Frame 4: hist=[8, 10], k=2, k-dist = 11-8 = 3
250
251        // Scenario: Evict next.
252        // Should be infinite dist frame 3 (oldest_ts=7)
253        let value = lru_replacer.evict();
254        assert_eq!(Some(3), value);
255        assert_eq!(3, lru_replacer.size()); // Remaining: [1, 5, 4]
256
257        // Set 6 to be evictable.
258        lru_replacer.set_evictable(6, true).unwrap();
259        assert_eq!(4, lru_replacer.size()); // Now [1, 5, 4, 6]
260                                            // Frame 6: hist=[5], k=2, inf k-dist, oldest=5
261                                            // Evict should be frame 6 (oldest inf dist)
262        let value = lru_replacer.evict();
263        assert_eq!(Some(6), value);
264        assert_eq!(3, lru_replacer.size()); // Remaining: [1, 5, 4]
265
266        // Set 1 to non-evictable
267        lru_replacer.set_evictable(1, false).unwrap();
268        assert_eq!(2, lru_replacer.size()); // Remaining evictable: [5, 4]
269                                            // Frame 5: k-dist=7
270                                            // Frame 4: k-dist=3
271                                            // Evict should be frame 5 (max k-dist)
272        let value = lru_replacer.evict();
273        assert_eq!(Some(5), value);
274        assert_eq!(1, lru_replacer.size()); // Remaining evictable: [4]
275
276        // Update access history for 1
277        lru_replacer.record_access(1).unwrap(); // ts=11, hist=[6, 11]
278        lru_replacer.record_access(1).unwrap(); // ts=12, hist=[11, 12]
279        lru_replacer.set_evictable(1, true).unwrap();
280        assert_eq!(2, lru_replacer.size()); // Evictable: [4, 1]
281                                            // Frame 4: k-dist=3
282                                            // Frame 1: k-dist=13-11=2
283                                            // Evict should be frame 4 (max k-dist)
284        let value = lru_replacer.evict();
285        assert_eq!(Some(4), value);
286
287        assert_eq!(1, lru_replacer.size()); // Remaining: [1]
288        let value = lru_replacer.evict();
289        assert_eq!(Some(1), value);
290        assert_eq!(0, lru_replacer.size());
291
292        // Empty replacer
293        assert_eq!(None, lru_replacer.evict());
294        assert_eq!(0, lru_replacer.size());
295        // Remove on non-existent frame should do nothing silently in remove(), size unchanged
296        // lru_replacer.remove(1); // Let's skip this, remove expects node to exist
297        assert_eq!(0, lru_replacer.size());
298    }
299}
quill_sql/utils/cache/lru_k.rs

quill_sql/utils/cache/
lru_k.rs
