1use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11#[derive(Clone, Debug, Hash, PartialEq, Eq)]
13struct PlanKey {
14 size: usize,
15 forward: bool,
16}
17
18#[derive(Clone)]
23struct CachedPlan {
24 size: usize,
25 forward: bool,
26 last_used: Instant,
27 usage_count: usize,
28}
29
30pub struct PlanCache {
32 cache: Arc<Mutex<HashMap<PlanKey, CachedPlan>>>,
33 max_entries: usize,
34 max_age: Duration,
35 enabled: Arc<Mutex<bool>>,
36 hit_count: Arc<Mutex<u64>>,
37 miss_count: Arc<Mutex<u64>>,
38}
39
40impl PlanCache {
41 pub fn new() -> Self {
43 Self {
44 cache: Arc::new(Mutex::new(HashMap::new())),
45 max_entries: 128,
46 max_age: Duration::from_secs(3600), enabled: Arc::new(Mutex::new(true)),
48 hit_count: Arc::new(Mutex::new(0)),
49 miss_count: Arc::new(Mutex::new(0)),
50 }
51 }
52
53 pub fn with_config(max_entries: usize, max_age: Duration) -> Self {
55 Self {
56 cache: Arc::new(Mutex::new(HashMap::new())),
57 max_entries,
58 max_age,
59 enabled: Arc::new(Mutex::new(true)),
60 hit_count: Arc::new(Mutex::new(0)),
61 miss_count: Arc::new(Mutex::new(0)),
62 }
63 }
64
65 pub fn set_enabled(&self, enabled: bool) {
67 *self.enabled.lock().expect("Operation failed") = enabled;
68 }
69
70 pub fn is_enabled(&self) -> bool {
72 *self.enabled.lock().expect("Operation failed")
73 }
74
75 pub fn clear(&self) {
77 if let Ok(mut cache) = self.cache.lock() {
78 cache.clear();
79 }
80 }
81
82 pub fn get_stats(&self) -> CacheStats {
84 let hit_count = *self.hit_count.lock().expect("Operation failed");
85 let miss_count = *self.miss_count.lock().expect("Operation failed");
86 let total_requests = hit_count + miss_count;
87 let hit_rate = if total_requests > 0 {
88 hit_count as f64 / total_requests as f64
89 } else {
90 0.0
91 };
92
93 let size = self.cache.lock().map(|c| c.len()).unwrap_or(0);
94
95 CacheStats {
96 hit_count,
97 miss_count,
98 hit_rate,
99 size,
100 max_size: self.max_entries,
101 }
102 }
103
104 pub fn track_plan_usage(&self, size: usize, forward: bool) {
109 if !*self.enabled.lock().expect("Operation failed") {
110 return;
111 }
112
113 let key = PlanKey { size, forward };
114
115 if let Ok(mut cache) = self.cache.lock() {
117 if let Some(cached) = cache.get_mut(&key) {
118 if cached.last_used.elapsed() <= self.max_age {
120 cached.last_used = Instant::now();
121 cached.usage_count += 1;
122 *self.hit_count.lock().expect("Operation failed") += 1;
123 return;
124 } else {
125 cache.remove(&key);
127 }
128 }
129 }
130
131 *self.miss_count.lock().expect("Operation failed") += 1;
133
134 if let Ok(mut cache) = self.cache.lock() {
136 if cache.len() >= self.max_entries {
138 self.evict_old_entries(&mut cache);
139 }
140
141 cache.insert(
142 key,
143 CachedPlan {
144 size,
145 forward,
146 last_used: Instant::now(),
147 usage_count: 1,
148 },
149 );
150 }
151 }
152
153 fn evict_old_entries(&self, cache: &mut HashMap<PlanKey, CachedPlan>) {
155 cache.retain(|_, v| v.last_used.elapsed() <= self.max_age);
157
158 while cache.len() >= self.max_entries {
160 if let Some((key_to_remove_, _)) = cache
161 .iter()
162 .min_by_key(|(_, v)| (v.last_used, v.usage_count))
163 .map(|(k, v)| (k.clone(), v.clone()))
164 {
165 cache.remove(&key_to_remove_);
166 } else {
167 break;
168 }
169 }
170 }
171
172 pub fn precompute_common_sizes(&self, sizes: &[usize]) {
177 for &size in sizes {
178 self.track_plan_usage(size, true);
180 self.track_plan_usage(size, false);
181 }
182 }
183}
184
185impl Default for PlanCache {
186 fn default() -> Self {
187 Self::new()
188 }
189}
190
191#[derive(Debug, Clone)]
193pub struct CacheStats {
194 pub hit_count: u64,
195 pub miss_count: u64,
196 pub hit_rate: f64,
197 pub size: usize,
198 pub max_size: usize,
199}
200
201impl std::fmt::Display for CacheStats {
202 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
203 write!(
204 f,
205 "Cache Stats: {} hits, {} misses ({:.1}% hit rate), {}/{} entries",
206 self.hit_count,
207 self.miss_count,
208 self.hit_rate * 100.0,
209 self.size,
210 self.max_size
211 )
212 }
213}
214
215static GLOBAL_PLAN_CACHE: std::sync::OnceLock<PlanCache> = std::sync::OnceLock::new();
217
218#[allow(dead_code)]
220pub fn get_global_cache() -> &'static PlanCache {
221 GLOBAL_PLAN_CACHE.get_or_init(PlanCache::new)
222}
223
224#[allow(dead_code)]
226pub fn init_global_cache(max_entries: usize, max_age: Duration) -> Result<(), &'static str> {
227 GLOBAL_PLAN_CACHE
228 .set(PlanCache::with_config(max_entries, max_age))
229 .map_err(|_| "Global plan cache already initialized")
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 #[test]
237 fn test_plan_cache_basic() {
238 let cache = PlanCache::new();
239
240 cache.track_plan_usage(128, true);
242 cache.track_plan_usage(128, true);
243
244 let stats = cache.get_stats();
246 assert_eq!(stats.hit_count, 1);
247 assert_eq!(stats.miss_count, 1);
248 }
249
250 #[test]
251 fn test_cache_eviction() {
252 let cache = PlanCache::with_config(2, Duration::from_secs(3600));
253
254 cache.track_plan_usage(64, true);
256 cache.track_plan_usage(128, true);
257
258 cache.track_plan_usage(256, true);
260
261 let stats = cache.get_stats();
262 assert_eq!(stats.size, 2);
263 }
264
265 #[test]
266 fn test_cache_disabled() {
267 let cache = PlanCache::new();
268 cache.set_enabled(false);
269
270 cache.track_plan_usage(128, true);
272 cache.track_plan_usage(128, true);
273
274 let stats = cache.get_stats();
276 assert_eq!(stats.hit_count, 0);
277 assert_eq!(stats.miss_count, 0); }
279}