heliosdb_proxy/cache/
mod.rs

1//! Query Caching Module
2//!
3//! Provides multi-tier query caching for HeliosProxy:
4//!
5//! - **L1 Hot Cache**: Per-connection, exact match, LRU eviction
6//! - **L2 Warm Cache**: Shared, normalized queries, configurable storage
7//! - **L3 Semantic Cache**: Vector similarity for AI workloads
8//!
9//! # Architecture
10//!
11//! ```text
12//!                     ┌─────────────────────────────────────────────────┐
13//!                     │                QUERY CACHE LAYER                 │
14//!                     │                                                  │
15//!   Query ───────────►│ ┌──────────────────────────────────────────────┐│
16//!                     ││ L1: Hot Cache (in-memory, <1ms)               ││
17//!                     │└──────────────────────────────────────────────┘│
18//!                     │         │ miss                                  │
19//!                     │         ▼                                       │
20//!                     │ ┌──────────────────────────────────────────────┐│
21//!                     ││ L2: Warm Cache (shared memory, <5ms)          ││
22//!                     │└──────────────────────────────────────────────┘│
23//!                     │         │ miss                                  │
24//!                     │         ▼                                       │
25//!                     │ ┌──────────────────────────────────────────────┐│
26//!                     ││ L3: Semantic Cache (vector similarity, <20ms) ││
27//!                     │└──────────────────────────────────────────────┘│
28//!                     │         │ miss                                  │
29//!                     │         ▼                                       │
30//!                     │       BACKEND                                   │
31//!                     └─────────────────────────────────────────────────┘
32//! ```
33//!
34//! # Usage
35//!
36//! ```rust,ignore
37//! use heliosdb_lite::proxy::cache::{QueryCache, CacheConfig};
38//!
39//! let config = CacheConfig::default();
40//! let cache = QueryCache::new(config);
41//!
42//! // Check cache before executing query
43//! if let Some(result) = cache.get(&query, &context).await {
44//!     return result;
45//! }
46//!
47//! // Execute query and cache result
48//! let result = execute_query(&query).await?;
49//! cache.put(&query, &context, result.clone()).await;
50//! ```
51
52pub mod config;
53pub mod hints;
54pub mod invalidation;
55pub mod l1_hot;
56pub mod l2_warm;
57pub mod l3_semantic;
58pub mod metrics;
59pub mod normalizer;
60pub mod result;
61
62// Re-exports
63pub use config::{CacheConfig, L1Config, L2Config, L3Config, StorageBackend};
64pub use hints::{parse_cache_hints, CacheHint};
65pub use invalidation::{InvalidationManager, InvalidationMode};
66pub use l1_hot::L1HotCache;
67pub use l2_warm::L2WarmCache;
68pub use l3_semantic::L3SemanticCache;
69pub use metrics::{CacheMetrics, CacheStatsLevelSnapshot, CacheStatsSnapshot};
70pub use normalizer::{NormalizedQuery, QueryNormalizer};
71pub use result::{CacheKey, CachedResult};
72
73use bytes::Bytes;
74use dashmap::DashMap;
75use std::sync::Arc;
76use std::time::{Duration, Instant};
77
78/// Query cache context (for cache key generation)
79#[derive(Debug, Clone, Hash, Eq, PartialEq)]
80pub struct CacheContext {
81    /// Database name
82    pub database: String,
83    /// Username (for RLS)
84    pub user: Option<String>,
85    /// Branch name (for HeliosDB branching)
86    pub branch: Option<String>,
87    /// Connection ID (for L1 cache)
88    pub connection_id: Option<u64>,
89}
90
91impl Default for CacheContext {
92    fn default() -> Self {
93        Self {
94            database: "default".to_string(),
95            user: None,
96            branch: None,
97            connection_id: None,
98        }
99    }
100}
101
102/// Cache lookup result
103#[derive(Debug)]
104pub enum CacheLookup {
105    /// Cache hit with result
106    Hit {
107        result: CachedResult,
108        level: CacheLevel,
109    },
110    /// Cache miss
111    Miss,
112}
113
114/// Cache level indicator
115#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub enum CacheLevel {
117    L1Hot,
118    L2Warm,
119    L3Semantic,
120}
121
122impl std::fmt::Display for CacheLevel {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        match self {
125            CacheLevel::L1Hot => write!(f, "L1"),
126            CacheLevel::L2Warm => write!(f, "L2"),
127            CacheLevel::L3Semantic => write!(f, "L3"),
128        }
129    }
130}
131
132/// Main query cache implementation
133pub struct QueryCache {
134    /// Configuration
135    config: CacheConfig,
136
137    /// L1: Per-connection hot cache (exact match)
138    l1_caches: DashMap<u64, Arc<L1HotCache>>,
139
140    /// L2: Shared normalized cache
141    l2_cache: Option<Arc<L2WarmCache>>,
142
143    /// L3: Semantic similarity cache
144    l3_cache: Option<Arc<L3SemanticCache>>,
145
146    /// Query normalizer
147    normalizer: Arc<QueryNormalizer>,
148
149    /// Cache invalidation manager
150    invalidator: Arc<InvalidationManager>,
151
152    /// Metrics collector
153    metrics: Arc<CacheMetrics>,
154
155    /// Request coalescing for cache stampede prevention
156    #[allow(dead_code)]
157    pending_requests: DashMap<CacheKey, Arc<tokio::sync::Notify>>,
158}
159
160impl QueryCache {
161    /// Create a new query cache with the given configuration
162    pub fn new(config: CacheConfig) -> Self {
163        let l2_cache = if config.l2.enabled {
164            Some(Arc::new(L2WarmCache::new(config.l2.clone())))
165        } else {
166            None
167        };
168
169        let l3_cache = if config.l3.enabled {
170            Some(Arc::new(L3SemanticCache::new(config.l3.clone())))
171        } else {
172            None
173        };
174
175        let invalidator = Arc::new(InvalidationManager::new(config.invalidation.clone()));
176
177        Self {
178            config: config.clone(),
179            l1_caches: DashMap::new(),
180            l2_cache,
181            l3_cache,
182            normalizer: Arc::new(QueryNormalizer::new()),
183            invalidator,
184            metrics: Arc::new(CacheMetrics::new()),
185            pending_requests: DashMap::new(),
186        }
187    }
188
189    /// Get or create L1 cache for a connection
190    pub fn get_l1_cache(&self, connection_id: u64) -> Arc<L1HotCache> {
191        self.l1_caches
192            .entry(connection_id)
193            .or_insert_with(|| Arc::new(L1HotCache::new(self.config.l1.clone())))
194            .clone()
195    }
196
197    /// Remove L1 cache for a connection (on disconnect)
198    pub fn remove_l1_cache(&self, connection_id: u64) {
199        self.l1_caches.remove(&connection_id);
200    }
201
202    /// Number of per-connection L1 caches currently retained. Grows by one per
203    /// connection that runs a cacheable read; must return to ~0 as connections
204    /// close (see `remove_l1_cache`). Exposed for leak observability/testing.
205    pub fn l1_cache_count(&self) -> usize {
206        self.l1_caches.len()
207    }
208
209    /// Look up a query in the cache hierarchy
210    pub async fn get(&self, query: &str, context: &CacheContext) -> CacheLookup {
211        // Parse cache hints
212        let hints = parse_cache_hints(query);
213
214        // Skip cache if hint says so
215        if hints.skip {
216            self.metrics.record_skip();
217            return CacheLookup::Miss;
218        }
219
220        let start = Instant::now();
221
222        // L1: Check hot cache (exact match)
223        if self.config.l1.enabled {
224            if let Some(conn_id) = context.connection_id {
225                let l1 = self.get_l1_cache(conn_id);
226                if let Some(result) = l1.get(query) {
227                    self.metrics.record_hit(CacheLevel::L1Hot, start.elapsed());
228                    return CacheLookup::Hit {
229                        result,
230                        level: CacheLevel::L1Hot,
231                    };
232                }
233            }
234        }
235
236        // Normalize query for L2/L3 lookup
237        let normalized = self.normalizer.normalize(query);
238        let cache_key = CacheKey::new(&normalized, context);
239
240        // L2: Check warm cache (normalized match)
241        if let Some(ref l2) = self.l2_cache {
242            if let Some(result) = l2.get(&cache_key).await {
243                self.metrics.record_hit(CacheLevel::L2Warm, start.elapsed());
244
245                // Promote to L1
246                if self.config.l1.enabled {
247                    if let Some(conn_id) = context.connection_id {
248                        let l1 = self.get_l1_cache(conn_id);
249                        l1.put(query.to_string(), result.clone());
250                    }
251                }
252
253                return CacheLookup::Hit {
254                    result,
255                    level: CacheLevel::L2Warm,
256                };
257            }
258        }
259
260        // L3: Check semantic cache (similarity match)
261        if hints.semantic_cache {
262            if let Some(ref l3) = self.l3_cache {
263                if let Some(result) = l3.get(query, context).await {
264                    self.metrics
265                        .record_hit(CacheLevel::L3Semantic, start.elapsed());
266                    return CacheLookup::Hit {
267                        result,
268                        level: CacheLevel::L3Semantic,
269                    };
270                }
271            }
272        }
273
274        self.metrics.record_miss(start.elapsed());
275        CacheLookup::Miss
276    }
277
278    /// Store a query result in the cache
279    pub async fn put(
280        &self,
281        query: &str,
282        context: &CacheContext,
283        data: Bytes,
284        row_count: usize,
285        execution_time: Duration,
286    ) {
287        // Parse cache hints
288        let hints = parse_cache_hints(query);
289
290        // Skip if hint says so
291        if hints.skip {
292            return;
293        }
294
295        // Normalize query
296        let normalized = self.normalizer.normalize(query);
297
298        // Determine TTL
299        let ttl = hints
300            .ttl
301            .unwrap_or_else(|| self.get_table_ttl(&normalized.tables));
302
303        // Check size limit
304        if data.len() > self.config.max_result_size {
305            self.metrics.record_size_exceeded();
306            return;
307        }
308
309        // Create cached result
310        let result = CachedResult {
311            data,
312            row_count,
313            cached_at: Instant::now(),
314            ttl,
315            tables: normalized.tables.clone(),
316            execution_time,
317        };
318
319        // Store in L1 (exact match)
320        if self.config.l1.enabled {
321            if let Some(conn_id) = context.connection_id {
322                let l1 = self.get_l1_cache(conn_id);
323                l1.put(query.to_string(), result.clone());
324            }
325        }
326
327        // Store in L2 (normalized)
328        if let Some(ref l2) = self.l2_cache {
329            let cache_key = CacheKey::new(&normalized, context);
330            l2.put(cache_key.clone(), result.clone()).await;
331
332            // Register for invalidation
333            for table in &normalized.tables {
334                self.invalidator.register(&cache_key, table);
335            }
336        }
337
338        // Store in L3 (semantic) if hint enabled
339        if hints.semantic_cache {
340            if let Some(ref l3) = self.l3_cache {
341                l3.put(query, context, result).await;
342            }
343        }
344
345        self.metrics.record_put();
346    }
347
348    /// Invalidate any cached results that reference a table written by `sql`.
349    /// Normalizes the (write) query to extract its tables, then drops their
350    /// cached entries.
351    pub async fn invalidate_query(&self, sql: &str) {
352        let normalized = self.normalizer.normalize(sql);
353        if !normalized.tables.is_empty() {
354            self.invalidate_tables(&normalized.tables).await;
355        }
356    }
357
358    /// Invalidate cache entries for specific tables
359    pub async fn invalidate_tables(&self, tables: &[String]) {
360        for table in tables {
361            let keys = self.invalidator.get_keys_for_table(table);
362
363            // Invalidate L2
364            if let Some(ref l2) = self.l2_cache {
365                for key in &keys {
366                    l2.remove(key).await;
367                }
368            }
369
370            self.invalidator.invalidate_table(table);
371        }
372
373        // L1 caches are invalidated on next access (TTL-based)
374        // L3 semantic cache has its own TTL handling
375
376        self.metrics.record_invalidation(tables.len());
377    }
378
379    /// Clear all caches
380    pub async fn clear(&self, levels: &[CacheLevel]) {
381        for level in levels {
382            match level {
383                CacheLevel::L1Hot => {
384                    self.l1_caches.clear();
385                }
386                CacheLevel::L2Warm => {
387                    if let Some(ref l2) = self.l2_cache {
388                        l2.clear().await;
389                    }
390                }
391                CacheLevel::L3Semantic => {
392                    if let Some(ref l3) = self.l3_cache {
393                        l3.clear().await;
394                    }
395                }
396            }
397        }
398
399        self.metrics.record_clear();
400    }
401
402    /// Get cache statistics
403    pub fn stats(&self) -> CacheStatsSnapshot {
404        self.metrics.snapshot()
405    }
406
407    /// Get configuration
408    pub fn config(&self) -> &CacheConfig {
409        &self.config
410    }
411
412    /// Get the invalidation manager (for WAL subscription)
413    pub fn invalidator(&self) -> Arc<InvalidationManager> {
414        self.invalidator.clone()
415    }
416
417    /// Get table-specific TTL or default
418    fn get_table_ttl(&self, tables: &[String]) -> Duration {
419        // Find shortest TTL among tables
420        let mut min_ttl = self.config.default_ttl;
421
422        for table in tables {
423            if let Some(table_config) = self.config.table_configs.get(table) {
424                if table_config.ttl < min_ttl {
425                    min_ttl = table_config.ttl;
426                }
427            }
428        }
429
430        min_ttl
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn test_cache_context_default() {
440        let ctx = CacheContext::default();
441        assert_eq!(ctx.database, "default");
442        assert!(ctx.user.is_none());
443        assert!(ctx.branch.is_none());
444        assert!(ctx.connection_id.is_none());
445    }
446
447    #[test]
448    fn test_cache_level_display() {
449        assert_eq!(format!("{}", CacheLevel::L1Hot), "L1");
450        assert_eq!(format!("{}", CacheLevel::L2Warm), "L2");
451        assert_eq!(format!("{}", CacheLevel::L3Semantic), "L3");
452    }
453
454    /// Per-connection L1 caches must be reclaimable so they don't leak per
455    /// session. `get_l1_cache` creates one; `remove_l1_cache` reclaims it.
456    #[test]
457    fn l1_cache_is_reclaimed_on_remove() {
458        let cache = QueryCache::new(CacheConfig::default());
459        assert_eq!(cache.l1_cache_count(), 0);
460        let _ = cache.get_l1_cache(42);
461        let _ = cache.get_l1_cache(43);
462        assert_eq!(cache.l1_cache_count(), 2);
463        cache.remove_l1_cache(42);
464        cache.remove_l1_cache(43);
465        assert_eq!(cache.l1_cache_count(), 0, "L1 caches must be reclaimed");
466    }
467
468    #[tokio::test]
469    async fn test_query_cache_creation() {
470        let config = CacheConfig::default();
471        let cache = QueryCache::new(config);
472
473        assert!(cache.config.l1.enabled);
474        assert!(cache.config.l2.enabled);
475    }
476
477    #[tokio::test]
478    async fn test_l1_cache_per_connection() {
479        let config = CacheConfig::default();
480        let cache = QueryCache::new(config);
481
482        let l1_a = cache.get_l1_cache(1);
483        let l1_b = cache.get_l1_cache(2);
484        let l1_a2 = cache.get_l1_cache(1);
485
486        // Same connection should get same cache
487        assert!(Arc::ptr_eq(&l1_a, &l1_a2));
488        // Different connections should get different caches
489        assert!(!Arc::ptr_eq(&l1_a, &l1_b));
490    }
491
492    #[tokio::test]
493    async fn test_cache_miss() {
494        let config = CacheConfig::default();
495        let cache = QueryCache::new(config);
496        let context = CacheContext::default();
497
498        let result = cache.get("SELECT * FROM users", &context).await;
499        assert!(matches!(result, CacheLookup::Miss));
500    }
501}
heliosdb_proxy/cache/mod.rs

heliosdb_proxy/cache/
mod.rs
