linkrs_query/cache/cte_cache/
manager.rs

1use moka::sync::Cache;
2use std::sync::Arc;
3use std::time::Duration;
4
5use crate::cache::config::CteCacheConfig;
6use crate::cache::stats::{CteCacheStats, CteCacheStatsSnapshot};
7use linkrs_metrics::StatsManager;
8
9use super::entry::CteCacheEntry;
10
11/// CTE Cache Manager
12///
13/// Managing caching of CTE (Common Table Expression) query results and ensuring thread-safe access.
14#[derive(Debug)]
15pub struct CteCacheManager {
16    /// Cache storage - using moka for high-performance concurrent access with weigher
17    cache: Cache<String, Arc<CteCacheEntry>>,
18    /// Configuration
19    config: Arc<std::sync::RwLock<CteCacheConfig>>,
20    /// Statistics
21    stats: Arc<CteCacheStats>,
22    /// Stats manager for reporting cache metrics
23    stats_manager: std::sync::RwLock<Option<Arc<StatsManager>>>,
24}
25
26impl CteCacheManager {
27    /// Create a new cache manager.
28    pub fn new() -> Self {
29        Self::with_config(CteCacheConfig::default())
30    }
31
32    /// Create using the configuration.
33    pub fn with_config(config: CteCacheConfig) -> Self {
34        let max_weight = config.max_size as u64;
35
36        let cache = Cache::builder()
37            .weigher(|_key, value: &Arc<CteCacheEntry>| -> u32 {
38                let arc_overhead = std::mem::size_of::<Arc<CteCacheEntry>>();
39                (value.estimate_memory() + arc_overhead) as u32
40            })
41            .max_capacity(max_weight)
42            .time_to_live(Duration::from_secs(config.entry_ttl_seconds))
43            .build();
44
45        let stats = Arc::new(CteCacheStats::new(config.max_size));
46
47        Self {
48            cache,
49            config: Arc::new(std::sync::RwLock::new(config.clone())),
50            stats,
51            stats_manager: std::sync::RwLock::new(None),
52        }
53    }
54
55    pub fn with_stats_manager(mut self, stats_manager: Arc<StatsManager>) -> Self {
56        self.stats_manager = std::sync::RwLock::new(Some(stats_manager));
57        self
58    }
59
60    /// Set the stats manager after creation
61    pub fn set_stats_manager(&self, stats_manager: Arc<StatsManager>) {
62        if let Ok(mut guard) = self.stats_manager.write() {
63            *guard = Some(stats_manager);
64        }
65    }
66
67    /// Obtain the configuration.
68    pub fn config(&self) -> CteCacheConfig {
69        self.config.read().expect("Config lock poisoned").clone()
70    }
71
72    /// Update the configuration.
73    pub fn set_config(&self, config: CteCacheConfig) {
74        let mut cfg = self.config.write().expect("Config lock poisoned");
75        self.stats.memory.set_max_bytes(config.max_size);
76        *cfg = config;
77    }
78
79    /// Determine whether to enable caching.
80    pub fn is_enabled(&self) -> bool {
81        self.config.read().expect("Config lock poisoned").enabled
82    }
83
84    /// Determine whether the results of the CTE (Common Table Expression) are cached.
85    ///
86    /// # Parameters
87    /// `cte_definition`: Text defining the Common Table Expression (CTE).
88    /// estimated_rows: The estimated number of rows
89    /// `is_deterministic`: Whether the CTE (Common Table Expression) is deterministic.
90    pub fn should_cache(
91        &self,
92        cte_definition: &str,
93        estimated_rows: u64,
94        is_deterministic: bool,
95    ) -> bool {
96        let config = self.config.read().expect("Config lock poisoned");
97
98        if !config.enabled {
99            return false;
100        }
101
102        if !is_deterministic {
103            return false;
104        }
105
106        if estimated_rows < config.min_row_count || estimated_rows > config.max_row_count {
107            return false;
108        }
109
110        let reuse_prob = self.predict_reuse_probability(cte_definition);
111        if reuse_prob < 0.3 {
112            return false;
113        }
114
115        true
116    }
117
118    /// Predict the probability of reuse
119    pub(super) fn predict_reuse_probability(&self, cte_definition: &str) -> f64 {
120        let cte_hash = Self::compute_hash(cte_definition);
121
122        if let Some(entry) = self.cache.get(&cte_hash) {
123            return entry.reuse_probability;
124        }
125
126        let complexity = cte_definition.len() as f64 / 100.0;
127        let base_prob = 0.5;
128        let complexity_bonus = (complexity / 10.0).min(0.3);
129
130        base_prob + complexity_bonus
131    }
132
133    /// Calculate the hash value defined by the CTE (Common Table Expression).
134    pub(super) fn compute_hash(cte_definition: &str) -> String {
135        use std::collections::hash_map::DefaultHasher;
136        use std::hash::{Hash, Hasher};
137
138        let mut hasher = DefaultHasher::new();
139        cte_definition.hash(&mut hasher);
140        format!("{:016x}", hasher.finish())
141    }
142
143    /// Store the data in the cache.
144    pub fn put(&self, cte_definition: &str, data: Vec<u8>, row_count: u64) -> Option<String> {
145        self.put_with_cost(cte_definition, data, row_count, 100)
146    }
147
148    /// Store the data in the cache with compute cost.
149    pub fn put_with_cost(
150        &self,
151        cte_definition: &str,
152        data: Vec<u8>,
153        row_count: u64,
154        compute_cost_ms: u64,
155    ) -> Option<String> {
156        let config = self.config.read().expect("Config lock poisoned");
157
158        if !config.enabled {
159            return None;
160        }
161
162        if data.len() > config.max_entry_size {
163            self.stats.counters.record_rejection();
164            return None;
165        }
166
167        drop(config);
168
169        let cte_hash = Self::compute_hash(cte_definition);
170        let entry = Arc::new(CteCacheEntry::new(
171            cte_hash.clone(),
172            cte_definition.to_string(),
173            data,
174            row_count,
175            compute_cost_ms,
176        ));
177
178        self.cache.insert(cte_hash.clone(), entry);
179
180        self.stats.counters.record_insertion();
181        self.update_stats();
182
183        Some(cte_hash)
184    }
185
186    /// Store the data in the cache with dependent tables.
187    pub fn put_with_tables(
188        &self,
189        cte_definition: &str,
190        data: Vec<u8>,
191        row_count: u64,
192        compute_cost_ms: u64,
193        dependent_tables: Vec<String>,
194    ) -> Option<String> {
195        let config = self.config.read().expect("Config lock poisoned");
196
197        if !config.enabled {
198            return None;
199        }
200
201        if data.len() > config.max_entry_size {
202            self.stats.counters.record_rejection();
203            return None;
204        }
205
206        drop(config);
207
208        let cte_hash = Self::compute_hash(cte_definition);
209        let mut entry = CteCacheEntry::new(
210            cte_hash.clone(),
211            cte_definition.to_string(),
212            data,
213            row_count,
214            compute_cost_ms,
215        );
216        entry.dependent_tables = dependent_tables;
217
218        self.cache.insert(cte_hash.clone(), Arc::new(entry));
219
220        self.stats.counters.record_insertion();
221        self.update_stats();
222
223        Some(cte_hash)
224    }
225
226    /// Evict low priority entries
227    pub fn evict_low_priority(&self, target_bytes: usize) -> usize {
228        let mut freed = 0;
229        let mut to_remove = Vec::new();
230
231        let entries: Vec<_> = self
232            .cache
233            .iter()
234            .map(|entry| {
235                let value_score = entry.1.value_score();
236                (
237                    entry.0.as_ref().clone(),
238                    value_score,
239                    entry.1.data_size,
240                    entry.1.priority,
241                )
242            })
243            .collect();
244
245        let mut entries_sorted = entries;
246        entries_sorted.sort_by(|a, b| {
247            a.1.partial_cmp(&b.1)
248                .unwrap_or(std::cmp::Ordering::Equal)
249                .then_with(|| a.3.cmp(&b.3))
250        });
251
252        for (key, _, size, _) in entries_sorted {
253            if freed >= target_bytes {
254                break;
255            }
256            to_remove.push(key);
257            freed += size;
258        }
259
260        for key in &to_remove {
261            self.cache.invalidate(key);
262        }
263
264        if freed > 0 {
265            self.stats.counters.record_eviction();
266            self.update_stats();
267        }
268
269        freed
270    }
271
272    /// Retrieve data from the cache.
273    pub fn get(&self, cte_definition: &str) -> Option<Arc<Vec<u8>>> {
274        let config = self.config.read().expect("Config lock poisoned");
275
276        if !config.enabled {
277            return None;
278        }
279
280        drop(config);
281
282        let cte_hash = Self::compute_hash(cte_definition);
283
284        if let Some(entry) = self.cache.get(&cte_hash) {
285            self.stats.counters.record_hit();
286            if let Ok(ref sm_guard) = self.stats_manager.read() {
287                if let Some(ref sm) = **sm_guard {
288                    sm.record_cache_hit_global(true);
289                }
290            }
291            Some(entry.data.clone())
292        } else {
293            self.stats.counters.record_miss();
294            if let Ok(ref sm_guard) = self.stats_manager.read() {
295                if let Some(ref sm) = **sm_guard {
296                    sm.record_cache_hit_global(false);
297                }
298            }
299            None
300        }
301    }
302
303    /// Check whether it exists in the cache.
304    pub fn contains(&self, cte_definition: &str) -> bool {
305        let cte_hash = Self::compute_hash(cte_definition);
306        self.cache.contains_key(&cte_hash)
307    }
308
309    /// Invalidate the cache entry
310    pub fn invalidate(&self, cte_definition: &str) -> bool {
311        let cte_hash = Self::compute_hash(cte_definition);
312        self.invalidate_by_hash(&cte_hash)
313    }
314
315    /// Invalidate cache entry by hash
316    pub fn invalidate_by_hash(&self, cte_hash: &str) -> bool {
317        let removed = self.cache.remove(cte_hash).is_some();
318
319        if removed {
320            self.stats.counters.record_eviction();
321            self.update_stats();
322        }
323
324        removed
325    }
326
327    /// Invalidate cache entries by table name
328    pub fn invalidate_by_table(&self, table_name: &str) -> usize {
329        let mut count = 0;
330        let keys_to_remove: Vec<String> = self
331            .cache
332            .iter()
333            .filter(|entry| entry.1.dependent_tables.iter().any(|t| t == table_name))
334            .map(|entry| entry.0.as_ref().clone())
335            .collect();
336
337        for key in keys_to_remove {
338            if self.cache.remove(&key).is_some() {
339                count += 1;
340            }
341        }
342
343        if count > 0 {
344            self.stats.counters.record_eviction();
345            self.update_stats();
346        }
347
348        count
349    }
350
351    /// Get cache entries for eviction (internal use)
352    pub fn get_cache_entries(&self) -> Vec<(String, f64, usize)> {
353        self.cache
354            .iter()
355            .map(|entry| {
356                let value_score = entry.1.value_score();
357                (entry.0.as_ref().clone(), value_score, entry.1.data_size)
358            })
359            .collect()
360    }
361
362    /// Increment eviction count (internal use)
363    pub fn increment_evicted_count(&self, count: u64) {
364        for _ in 0..count {
365            self.stats.counters.record_eviction();
366        }
367    }
368
369    /// Clear all caches.
370    pub fn clear(&self) {
371        self.cache.invalidate_all();
372        self.stats.reset();
373    }
374
375    /// Obtain statistical information
376    pub fn get_stats(&self) -> CteCacheStatsSnapshot {
377        self.update_stats();
378        self.stats.snapshot()
379    }
380
381    /// Update internal statistics
382    fn update_stats(&self) {
383        let current_entries = self.cache.entry_count() as usize;
384        let current_memory = self.estimate_current_memory();
385        self.stats.memory.update(current_memory, current_entries);
386    }
387
388    /// Estimate current memory usage
389    fn estimate_current_memory(&self) -> usize {
390        self.cache
391            .iter()
392            .map(|entry| entry.1.estimate_memory())
393            .sum()
394    }
395
396    /// Reset statistics
397    pub fn reset_stats(&self) {
398        self.stats.reset();
399    }
400
401    /// Get the current memory usage
402    pub fn current_memory(&self) -> usize {
403        self.estimate_current_memory()
404    }
405
406    /// Obtain the number of cached entries
407    pub fn entry_count(&self) -> usize {
408        self.cache.entry_count() as usize
409    }
410
411    /// Clearance of obsolete entries
412    /// Note: moka handles TTL automatically, so this is a no-op
413    pub fn cleanup_expired(&self) -> usize {
414        0
415    }
416
417    /// Get the statistics object
418    pub fn stats(&self) -> Arc<CteCacheStats> {
419        self.stats.clone()
420    }
421}
422
423impl Default for CteCacheManager {
424    fn default() -> Self {
425        Self::new()
426    }
427}
linkrs_query/cache/cte_cache/manager.rs

linkrs_query/cache/cte_cache/
manager.rs
