linkrs_query/cache/
plan_cache.rs

1//! Query Plan Cache Module
2//!
3//! Provides Prepared Statement style query plan caching with support for parameterized queries.
4//!
5//! # Design objectives
6//!
7//! 1. Cache query plan parsing, validation and planning results
8//! 2. Support for parameterized queries (Prepared Statement)
9//! 3. Limit memory usage to prevent unlimited growth
10//! 4. Thread-safe, supporting highly concurrent access
11//!
12//! # Scenarios of use
13//!
14//! - Repeated execution of the same query template
15//! - Batch insert/update operations
16//! - Applications use Prepared Statements
17
18use moka::sync::Cache;
19use std::hash::Hasher;
20use std::sync::Arc;
21use std::time::{Duration, Instant};
22
23use crate::executor::streaming::plan::PhysicalPlan;
24use crate::planning::plan::execution_plan::PartitionSpec;
25use linkrs_metrics::StatsManager;
26
27use super::config::{CachePriority, PlanCacheConfig};
28use super::stats::PlanCacheStats;
29
30pub mod entry;
31pub mod key;
32pub mod params;
33pub use entry::CachedPlan;
34pub use key::{ParamPosition, PlanCacheContext, PlanCacheKey, PlanCachePutContext};
35pub use params::ParameterizedQueryHandler;
36
37pub struct QueryPlanCache {
38    /// Cache storage - using moka for high-performance concurrent access
39    cache: Cache<PlanCacheKey, Arc<CachedPlan>>,
40    /// Configuration
41    config: PlanCacheConfig,
42    /// Statistics
43    stats: Arc<PlanCacheStats>,
44    /// Stats manager for reporting cache metrics
45    stats_manager: std::sync::RwLock<Option<Arc<StatsManager>>>,
46}
47
48impl std::fmt::Debug for QueryPlanCache {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("QueryPlanCache")
51            .field("config", &self.config)
52            .field("stats", &self.stats.snapshot())
53            .finish()
54    }
55}
56
57impl QueryPlanCache {
58    /// Create a new query plan cache.
59    pub fn new(config: PlanCacheConfig) -> Self {
60        let max_weight = config.effective_max_weight();
61
62        let cache = Cache::builder()
63            .weigher(|_key, value: &Arc<CachedPlan>| -> u32 {
64                let arc_overhead = std::mem::size_of::<Arc<CachedPlan>>();
65                (value.estimate_memory() + arc_overhead) as u32
66            })
67            .max_capacity(max_weight)
68            .time_to_live(Duration::from_secs(config.ttl_config.base_ttl_seconds))
69            .build();
70
71        let stats = Arc::new(PlanCacheStats::new(config.memory_budget));
72
73        Self {
74            cache,
75            config,
76            stats,
77            stats_manager: std::sync::RwLock::new(None),
78        }
79    }
80
81    pub fn with_stats_manager(self, stats_manager: Arc<StatsManager>) -> Self {
82        if let Ok(mut guard) = self.stats_manager.write() {
83            *guard = Some(stats_manager);
84        }
85        self
86    }
87
88    /// Set the stats manager after creation
89    pub fn set_stats_manager(&self, stats_manager: Arc<StatsManager>) {
90        if let Ok(mut guard) = self.stats_manager.write() {
91            *guard = Some(stats_manager);
92        }
93    }
94
95    /// Obtaining the cached plan
96    ///
97    /// # Parameters
98    /// - `query`: The text of the query
99    ///
100    /// # Returns
101    /// - `Some(Arc<CachedPlan>)`: Cached plan
102    /// - `None`: No results were found, or there was a hash collision.
103    pub fn get(&self, query: &str) -> Option<Arc<CachedPlan>> {
104        self.get_with_space(query, None, None)
105    }
106
107    /// Look up a cached plan with space, schema, and parameter type context.
108    ///
109    /// space_name and schema_version are incorporated into the cache key
110    /// to prevent cross-space plan reuse and stale plans after DDL.
111    ///
112    /// param_type_signature is a hash of parameter *types* so that
113    /// the same query with different param types gets a different cache key.
114    ///
115    /// index_version forces replan after index DDL (use the same value
116    /// that was passed to `put_with_context` to ensure key matching).
117    pub fn get_with_space(
118        &self,
119        query: &str,
120        space_name: Option<String>,
121        schema_version: Option<u64>,
122    ) -> Option<Arc<CachedPlan>> {
123        self.get_with_full_context(query, space_name, schema_version, None, None)
124    }
125
126    /// Look up with space, schema version, AND index version.
127    ///
128    /// the index_version dimension ensures that a plan compiled with a
129    /// certain index state is not reused after index DDL.
130    pub fn get_with_full_space(
131        &self,
132        query: &str,
133        space_name: Option<String>,
134        schema_version: Option<u64>,
135        index_version: Option<u64>,
136    ) -> Option<Arc<CachedPlan>> {
137        self.get_with_full_context(query, space_name, schema_version, None, index_version)
138    }
139
140    /// Full-context cache lookup with parameter type signature and index version.
141    pub fn get_with_full_context(
142        &self,
143        query: &str,
144        space_name: Option<String>,
145        schema_version: Option<u64>,
146        param_type_signature: Option<u64>,
147        index_version: Option<u64>,
148    ) -> Option<Arc<CachedPlan>> {
149        self.get_with_context(
150            query,
151            PlanCacheContext {
152                space_name,
153                schema_version,
154                index_version,
155                param_type_signature,
156                ..Default::default()
157            },
158        )
159    }
160
161    pub fn get_with_context(
162        &self,
163        query: &str,
164        context: PlanCacheContext,
165    ) -> Option<Arc<CachedPlan>> {
166        let key = PlanCacheKey::from_query_with_context(query, context);
167
168        if let Some(plan) = self.cache.get(&key) {
169            if plan.query_template != query {
170                log::warn!(
171                    "Query plan cache hash collision detected: hash={}, expected_query={}, cached_query={}",
172                    key.hash,
173                    query,
174                    plan.query_template
175                );
176                self.stats.counters.record_miss();
177                if let Ok(ref sm_guard) = self.stats_manager.read() {
178                    if let Some(ref sm) = **sm_guard {
179                        sm.record_cache_hit_global(false);
180                    }
181                }
182                return None;
183            }
184
185            self.stats.counters.record_hit();
186            if let Ok(ref sm_guard) = self.stats_manager.read() {
187                if let Some(ref sm) = **sm_guard {
188                    sm.record_cache_hit_global(true);
189                }
190            }
191            return Some(plan);
192        }
193
194        self.stats.counters.record_miss();
195        if let Ok(ref sm_guard) = self.stats_manager.read() {
196            if let Some(ref sm) = **sm_guard {
197                sm.record_cache_hit_global(false);
198            }
199        }
200        None
201    }
202
203    /// Put the plan in the cache.
204    ///
205    /// stores an [`Arc<PhysicalPlan>`] instead of [`ExecutionPlan`].
206    ///
207    /// # Parameters
208    /// - `query`: Query text
209    /// - `plan`: Arena-based physical plan
210    /// - `param_positions`: Information about the positions of the parameters
211    pub fn put(&self, query: &str, plan: Arc<PhysicalPlan>, param_positions: Vec<ParamPosition>) {
212        self.put_with_context(query, plan, param_positions, PlanCachePutContext::default());
213    }
214
215    /// Put the plan in the cache with dependent tables.
216    pub fn put_with_tables(
217        &self,
218        query: &str,
219        plan: Arc<PhysicalPlan>,
220        param_positions: Vec<ParamPosition>,
221        dependent_tables: Vec<String>,
222    ) {
223        self.put_with_context(
224            query,
225            plan,
226            param_positions,
227            PlanCachePutContext {
228                dependent_tables,
229                ..PlanCachePutContext::default()
230            },
231        );
232    }
233
234    /// Put the plan with full context (space, schema version, index version, tables).
235    ///
236    /// space_name and schema_version are incorporated into the cache key
237    /// to prevent cross-space reuse and stale plans after DDL.
238    ///
239    /// `param_type_signature` is derived from the parameter type
240    /// declarations (not values) and is included in the cache key.
241    ///
242    /// index_version is incorporated into the cache key to force replan
243    /// after index DDL (CREATE/DROP index) even when schema_version is unchanged.
244    ///
245    /// stores [`Arc<PhysicalPlan>`] — the immutable arena plan.
246    pub fn put_with_context(
247        &self,
248        query: &str,
249        plan: Arc<PhysicalPlan>,
250        param_positions: Vec<ParamPosition>,
251        context: PlanCachePutContext,
252    ) {
253        let PlanCachePutContext {
254            dependent_tables,
255            space_name,
256            schema_version,
257            index_version,
258            is_dml,
259            is_transaction,
260            optimizer_version,
261            planning_config_hash,
262        } = context;
263        let query_bytes = query.len();
264        let param_type_sig = Self::compute_param_type_signature(&param_positions);
265        let key = PlanCacheKey::from_query_with_context(
266            query,
267            PlanCacheContext {
268                space_name,
269                schema_version,
270                index_version,
271                param_type_signature: param_type_sig,
272                optimizer_version,
273                planning_config_hash,
274            },
275        );
276
277        let priority = if self.config.priority_config.enable_priority {
278            self.calculate_priority(&plan)
279        } else {
280            CachePriority::Normal
281        };
282
283        let complexity_score = self.calculate_complexity_score(&plan);
284        let estimated_compute_cost = self.estimate_compute_cost(&plan);
285        let current_ttl = Duration::from_secs(self.config.ttl_config.base_ttl_seconds);
286
287        let cached_plan = Arc::new(CachedPlan {
288            query_template: query.to_string(),
289            plan,
290            param_positions,
291            created_at: Instant::now(),
292            last_accessed: Instant::now(),
293            access_count: 0,
294            avg_execution_time_ms: 0.0,
295            execution_count: 0,
296            priority,
297            complexity_score,
298            estimated_compute_cost,
299            current_ttl,
300            dependent_tables,
301            is_dml,
302            is_transaction,
303        });
304
305        let is_update = self.cache.contains_key(&key);
306        self.cache.insert(key, cached_plan);
307
308        if !is_update {
309            self.stats.record_query_size(query_bytes);
310        }
311
312        let current_entries = self.cache.entry_count() as usize;
313        let current_memory = self.estimate_current_memory();
314        self.stats.memory.update(current_memory, current_entries);
315    }
316
317    /// Calculate priority based on query characteristics.
318    ///
319    /// uses operator count from the arena [`PhysicalPlan`].
320    fn calculate_priority(&self, plan: &PhysicalPlan) -> CachePriority {
321        let complexity = self.calculate_complexity_score(plan);
322
323        if complexity > 1000 {
324            CachePriority::High
325        } else if complexity > 100 {
326            CachePriority::Normal
327        } else {
328            CachePriority::Low
329        }
330    }
331
332    /// Calculate complexity score from the arena [`PhysicalPlan`].
333    ///
334    /// uses operator count and fragment count as a proxy for complexity.
335    fn calculate_complexity_score(&self, plan: &PhysicalPlan) -> u32 {
336        let op_count = plan.operator_count() as u32;
337        let frag_count = plan.fragment_count() as u32;
338        (op_count * 20).max(frag_count * 10)
339    }
340
341    /// Estimate compute cost in milliseconds.
342    fn estimate_compute_cost(&self, plan: &PhysicalPlan) -> u64 {
343        let complexity = self.calculate_complexity_score(plan);
344        (complexity as u64 * 10).max(100)
345    }
346
347    /// Estimate current memory usage
348    fn estimate_current_memory(&self) -> usize {
349        self.cache
350            .iter()
351            .map(|entry| entry.1.estimate_memory())
352            .sum()
353    }
354
355    /// Look up a cached plan keyed by its physical partition layout.
356    ///
357    /// The key incorporates the layout fingerprint (ranges, source, layout
358    /// version), so a plan cached under one layout is never reused for a
359    /// different one.  This is the lookup used by callers that already know
360    /// the current layout (e.g. re-planning against a previous execution's
361    /// metadata); plain-text lookups deliberately miss partitioned entries.
362    pub fn get_with_partition(&self, query: &str, spec: &PartitionSpec) -> Option<Arc<CachedPlan>> {
363        self.get_with_partition_context(query, spec, PlanCacheContext::default())
364    }
365
366    /// Partition-keyed lookup that additionally scopes by the full
367    /// compatibility context, matching [`put_with_partition`](Self::put_with_partition).
368    pub fn get_with_partition_context(
369        &self,
370        query: &str,
371        spec: &PartitionSpec,
372        context: PlanCacheContext,
373    ) -> Option<Arc<CachedPlan>> {
374        let key = PlanCacheKey::from_query_with_partition_and_context(query, spec, context);
375
376        if let Some(plan) = self.cache.get(&key) {
377            if plan.query_template != query {
378                log::warn!(
379                    "Query plan cache hash collision detected: hash={}, expected_query={}, cached_query={}",
380                    key.hash,
381                    query,
382                    plan.query_template
383                );
384                self.stats.counters.record_miss();
385                return None;
386            }
387            self.stats.counters.record_hit();
388            if let Ok(ref sm_guard) = self.stats_manager.read() {
389                if let Some(ref sm) = **sm_guard {
390                    sm.record_cache_hit_global(true);
391                }
392            }
393            return Some(plan);
394        }
395
396        self.stats.counters.record_miss();
397        if let Ok(ref sm_guard) = self.stats_manager.read() {
398            if let Some(ref sm) = **sm_guard {
399                sm.record_cache_hit_global(false);
400            }
401        }
402        None
403    }
404
405    /// Store a plan under its physical partition layout key.
406    ///
407    /// Only partitioned plans should use this; the plain-text key is left
408    /// free so a subsequent lookup that cannot provide the layout does not
409    /// silently reuse a layout-dependent plan.
410    pub fn put_with_partition(
411        &self,
412        query: &str,
413        spec: &PartitionSpec,
414        plan: Arc<PhysicalPlan>,
415        param_positions: Vec<ParamPosition>,
416        context: PlanCachePutContext,
417    ) {
418        let PlanCachePutContext {
419            dependent_tables,
420            space_name,
421            schema_version,
422            index_version,
423            is_dml,
424            is_transaction,
425            optimizer_version,
426            planning_config_hash,
427        } = context;
428        let param_type_sig = Self::compute_param_type_signature(&param_positions);
429        let key = PlanCacheKey::from_query_with_partition_and_context(
430            query,
431            spec,
432            PlanCacheContext {
433                space_name,
434                schema_version,
435                index_version,
436                param_type_signature: param_type_sig,
437                optimizer_version,
438                planning_config_hash,
439            },
440        );
441        let priority = if self.config.priority_config.enable_priority {
442            self.calculate_priority(&plan)
443        } else {
444            CachePriority::Normal
445        };
446        let complexity_score = self.calculate_complexity_score(&plan);
447        let cached_plan = Arc::new(CachedPlan {
448            query_template: query.to_string(),
449            plan,
450            param_positions,
451            created_at: Instant::now(),
452            last_accessed: Instant::now(),
453            access_count: 0,
454            avg_execution_time_ms: 0.0,
455            execution_count: 0,
456            priority,
457            complexity_score,
458            estimated_compute_cost: self.estimate_compute_cost_from_score(complexity_score),
459            current_ttl: Duration::from_secs(self.config.ttl_config.base_ttl_seconds),
460            dependent_tables,
461            is_dml,
462            is_transaction,
463        });
464        self.cache.insert(key, cached_plan);
465        self.stats.memory.update(
466            self.estimate_current_memory(),
467            self.cache.entry_count() as usize,
468        );
469    }
470
471    /// Estimate compute cost in milliseconds from an already-computed score.
472    fn estimate_compute_cost_from_score(&self, complexity: u32) -> u64 {
473        (complexity as u64 * 10).max(100)
474    }
475
476    /// Compute a parameter type signature from param positions.
477    ///
478    /// produces a hash of the parameter *types* (not values) so that
479    /// plans with different param type declarations get different cache keys.
480    pub(crate) fn compute_param_type_signature(params: &[ParamPosition]) -> Option<u64> {
481        if params.is_empty() {
482            return None;
483        }
484        use std::hash::Hash;
485        let mut hasher = std::collections::hash_map::DefaultHasher::new();
486        for p in params {
487            p.index.hash(&mut hasher);
488            p.name.hash(&mut hasher);
489            p.expected_type.hash(&mut hasher);
490        }
491        Some(hasher.finish())
492    }
493
494    /// Record the statistics on the execution of the plan.
495    ///
496    /// # Parameter
497    /// - `query`: Query content
498    /// - `execution_time_ms`: Execution time (in milliseconds)
499    pub fn record_execution(&self, query: &str, execution_time_ms: f64) {
500        let key = PlanCacheKey::from_query(query);
501
502        if let Some(plan) = self.cache.get(&key) {
503            let alpha = 0.1;
504            let new_avg = plan.avg_execution_time_ms * (1.0 - alpha) + execution_time_ms * alpha;
505
506            let updated_plan = Arc::new(CachedPlan {
507                execution_count: plan.execution_count + 1,
508                avg_execution_time_ms: new_avg,
509                ..(*plan).clone()
510            });
511
512            self.cache.insert(key, updated_plan);
513        }
514    }
515
516    /// Record execution with space context.
517    ///
518    /// index_version must match the value used during `put_with_context` to
519    /// locate the cached plan; otherwise the keys will differ and the record
520    /// will miss.
521    ///
522    /// `param_type_signature` must match the signature used during
523    /// `put_with_context` (i.e. the hash of the parameter *types*) or the
524    /// record will miss for parameterized queries.
525    pub fn record_execution_with_space(
526        &self,
527        query: &str,
528        execution_time_ms: f64,
529        space_name: Option<String>,
530        schema_version: Option<u64>,
531        index_version: Option<u64>,
532        param_type_signature: Option<u64>,
533    ) {
534        let key = PlanCacheKey::from_query_with_full_context(
535            query,
536            space_name,
537            schema_version,
538            param_type_signature,
539            index_version,
540        );
541
542        if let Some(plan) = self.cache.get(&key) {
543            let alpha = 0.1;
544            let new_avg = plan.avg_execution_time_ms * (1.0 - alpha) + execution_time_ms * alpha;
545
546            let updated_plan = Arc::new(CachedPlan {
547                execution_count: plan.execution_count + 1,
548                avg_execution_time_ms: new_avg,
549                ..(*plan).clone()
550            });
551
552            self.cache.insert(key, updated_plan);
553        }
554    }
555
556    /// Check whether the query has been cached.
557    pub fn contains(&self, query: &str) -> bool {
558        let key = PlanCacheKey::from_query(query);
559        self.cache.contains_key(&key)
560    }
561
562    /// Invalidate the cache entry
563    pub fn invalidate(&self, query: &str) -> bool {
564        let key = PlanCacheKey::from_query(query);
565        let removed = self.cache.remove(&key).is_some();
566
567        if removed {
568            self.stats.counters.record_eviction();
569            self.update_stats();
570        }
571
572        removed
573    }
574
575    /// Invalidate all cache entries for a given space.
576    ///
577    /// called after DDL/index/schema changes to force replanning.
578    /// Iterates all entries and removes those matching the space name.
579    pub fn invalidate_space(&self, space_name: &str) -> usize {
580        // Collect keys to remove while iterating (can't mutate during iter).
581        let keys_to_remove: Vec<PlanCacheKey> = self
582            .cache
583            .iter()
584            .filter(|(k, _)| k.space_name.as_deref() == Some(space_name))
585            .map(|(k, _)| PlanCacheKey {
586                hash: k.hash,
587                query_text: k.query_text.clone(),
588                partition_fingerprint: k.partition_fingerprint,
589                space_name: k.space_name.clone(),
590                schema_version: k.schema_version,
591                param_type_signature: k.param_type_signature,
592                index_version: k.index_version,
593                optimizer_version: k.optimizer_version,
594                planning_config_hash: k.planning_config_hash,
595            })
596            .collect();
597        let removed = keys_to_remove.len();
598        for key in &keys_to_remove {
599            self.cache.remove(key);
600        }
601        if removed > 0 {
602            self.stats.counters.record_eviction();
603            self.update_stats();
604        }
605        removed
606    }
607
608    /// Get cache entries for eviction (internal use)
609    pub fn get_cache_entries(&self) -> Vec<(Arc<PlanCacheKey>, f64, usize)> {
610        self.cache
611            .iter()
612            .map(|(k, v)| {
613                let value_score = v.value_score();
614                (k.clone(), value_score, v.query_template.len())
615            })
616            .collect()
617    }
618
619    /// Increment eviction count (internal use)
620    pub fn increment_eviction_count(&self, count: u64) {
621        for _ in 0..count {
622            self.stats.counters.record_eviction();
623        }
624    }
625
626    /// Clear all caches.
627    pub fn clear(&self) {
628        self.cache.invalidate_all();
629        self.stats.reset();
630    }
631
632    /// Obtain statistical information
633    pub fn stats(&self) -> Arc<PlanCacheStats> {
634        self.stats.clone()
635    }
636
637    /// Record a hit without a cache lookup, for the Level 2 same-shape DML
638    /// memo that short-circuits before `get_with_context`. Keeps the hit rate
639    /// accounting consistent with the regular lookup path.
640    pub fn record_memo_hit(&self) {
641        self.stats.counters.record_hit();
642        if let Ok(ref sm_guard) = self.stats_manager.read() {
643            if let Some(ref sm) = **sm_guard {
644                sm.record_cache_hit_global(true);
645            }
646        }
647    }
648
649    /// Get statistics snapshot
650    pub fn stats_snapshot(&self) -> super::stats::PlanCacheStatsSnapshot {
651        self.stats.snapshot()
652    }
653
654    /// Clean up expired entries.
655    /// Note: moka handles TTL automatically, so this is a no-op
656    pub fn cleanup_expired(&self) {
657        // moka handles TTL automatically
658    }
659
660    /// Get the number of cached entries
661    pub fn len(&self) -> usize {
662        self.cache.entry_count() as usize
663    }
664
665    /// Check whether the cache is empty.
666    pub fn is_empty(&self) -> bool {
667        self.cache.entry_count() == 0
668    }
669
670    /// Get the configuration
671    pub fn config(&self) -> &PlanCacheConfig {
672        &self.config
673    }
674
675    /// Update internal statistics
676    fn update_stats(&self) {
677        let current_entries = self.cache.entry_count() as usize;
678        let current_memory = self.estimate_current_memory();
679        self.stats.memory.update(current_memory, current_entries);
680    }
681}
682
683impl Default for QueryPlanCache {
684    fn default() -> Self {
685        Self::new(PlanCacheConfig::default())
686    }
687}
688
689#[cfg(test)]
690mod tests;
linkrs_query/cache/plan_cache.rs

linkrs_query/cache/
plan_cache.rs
