sz_orm_core/
prepared_cache.rs

1//! v6.5.0 PreparedStatement 句柄缓存
2//!
3//! 缓存数据库侧 PreparedStatement 句柄，按连接分桶（句柄绑定连接），
4//! 跨连接共享 LRU 淘汰顺序。与既有 `PlanCache`（应用侧 AST 缓存）职责分离。
5//!
6//! 特性：
7//! - 按连接 ID 分桶（SQLx prepared statement 绑定连接）
8//! - bucket 级别 LRU 淘汰（max_size_per_conn 上限）
9//! - 表级精确失效（table_index 索引）
10//! - 连接级失效（invalidate_conn，连接关闭时调用）
11//! - 原子统计（命中/未命中/淘汰/降级/失效）
12//! - 句柄以闭包形式存储（避免在 sz-orm-core 暴露 SQLx 类型）
13
14use std::collections::HashMap;
15use std::hash::{Hash, Hasher};
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::sync::Arc;
18use std::time::Instant;
19
20use parking_lot::RwLock;
21
22use crate::error::DbError;
23use crate::pool::QueryRows;
24use crate::value::Value;
25
26/// 连接唯一标识（连接生命周期内不变）
27pub type ConnId = u64;
28
29/// 句柄执行闭包类型（Arc 包装，支持从 RwLock 中 clone 出来后释放锁再 await）
30pub type ExecuteFn = Arc<
31    dyn Fn(
32            &[Value],
33        ) -> std::pin::Pin<
34            Box<dyn std::future::Future<Output = Result<QueryRows, DbError>> + Send + 'static>,
35        > + Send
36        + Sync,
37>;
38
39/// 缓存查找结果
40#[derive(Debug)]
41pub enum PreparedLookup {
42    /// 缓存命中，已执行并返回结果
43    Hit(QueryRows),
44    /// 缓存未命中，调用方需 prepare 后调用 `store_handle`
45    Miss,
46}
47
48/// PreparedStatement 句柄缓存
49///
50/// 按连接 ID 分桶存储句柄，bucket 级别 LRU 淘汰。
51/// 句柄以闭包形式存储，避免在 sz-orm-core 暴露 SQLx 类型。
52pub struct PreparedStatementCache {
53    /// 按连接 ID 分桶
54    buckets: RwLock<HashMap<ConnId, PreparedStatementBucket>>,
55    /// 统计
56    stats: PreparedStatementCacheStats,
57    /// 每连接最大缓存数
58    max_size_per_conn: usize,
59}
60
61/// 连接级句柄桶
62struct PreparedStatementBucket {
63    /// SQL hash → 条目
64    entries: HashMap<u64, PreparedStatementEntry>,
65    /// SQL hash → 最后使用时间（LRU 淘汰依据）
66    lru: HashMap<u64, Instant>,
67    /// 表名 → SQL hash 列表（表级失效索引）
68    table_index: HashMap<String, Vec<u64>>,
69}
70
71/// 单个句柄条目
72struct PreparedStatementEntry {
73    /// 句柄执行闭包
74    execute_fn: ExecuteFn,
75}
76
77/// PreparedStatement 缓存统计（原子计数器）
78pub struct PreparedStatementCacheStats {
79    hits: AtomicU64,
80    misses: AtomicU64,
81    evictions: AtomicU64,
82    degradations: AtomicU64,
83    invalidations: AtomicU64,
84}
85
86/// 统计快照（用于读取一致性视图）
87#[derive(Debug, Clone)]
88pub struct PreparedStatementCacheStatsSnapshot {
89    /// 命中次数
90    pub hits: u64,
91    /// 未命中次数
92    pub misses: u64,
93    /// LRU 淘汰次数
94    pub evictions: u64,
95    /// 降级次数（缓存内部错误时降级直接执行）
96    pub degradations: u64,
97    /// 失效次数（表级 + 连接级）
98    pub invalidations: u64,
99    /// 当前缓存条目总数
100    pub size: usize,
101    /// 命中率（0.0 ~ 1.0）
102    pub hit_rate: f64,
103}
104
105impl PreparedStatementCacheStats {
106    fn new() -> Self {
107        Self {
108            hits: AtomicU64::new(0),
109            misses: AtomicU64::new(0),
110            evictions: AtomicU64::new(0),
111            degradations: AtomicU64::new(0),
112            invalidations: AtomicU64::new(0),
113        }
114    }
115
116    fn record_hit(&self) {
117        self.hits.fetch_add(1, Ordering::Relaxed);
118    }
119
120    fn record_miss(&self) {
121        self.misses.fetch_add(1, Ordering::Relaxed);
122    }
123
124    fn record_eviction(&self) {
125        self.evictions.fetch_add(1, Ordering::Relaxed);
126    }
127
128    fn record_degradation(&self) {
129        self.degradations.fetch_add(1, Ordering::Relaxed);
130    }
131
132    fn record_invalidation(&self, count: usize) {
133        self.invalidations
134            .fetch_add(count as u64, Ordering::Relaxed);
135    }
136}
137
138impl Default for PreparedStatementCacheStats {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144impl PreparedStatementCache {
145    /// 创建新缓存，指定每连接最大缓存数
146    ///
147    /// `max_size_per_conn` 默认 256（与 design.md 2.1.2 配置项一致）
148    #[must_use]
149    pub fn new(max_size_per_conn: usize) -> Self {
150        Self {
151            buckets: RwLock::new(HashMap::new()),
152            stats: PreparedStatementCacheStats::new(),
153            max_size_per_conn,
154        }
155    }
156
157    /// 查找或执行缓存句柄
158    ///
159    /// - 命中：调用 `execute_fn(params)` 返回 `PreparedLookup::Hit(rows)`
160    /// - 未命中：返回 `PreparedLookup::Miss`，调用方需 prepare 后调用 `store_handle`
161    /// - 执行错误：返回 `Err(DbError)`，调用方应降级直接执行
162    pub async fn get_or_prepare(
163        &self,
164        conn_id: ConnId,
165        sql: &str,
166        params: &[Value],
167    ) -> Result<PreparedLookup, DbError> {
168        let sql_hash = compute_sql_hash(sql);
169
170        let execute_fn = {
171            let buckets = self.buckets.read();
172            if let Some(bucket) = buckets.get(&conn_id) {
173                if let Some(entry) = bucket.entries.get(&sql_hash) {
174                    Some(Arc::clone(&entry.execute_fn))
175                } else {
176                    None
177                }
178            } else {
179                None
180            }
181        };
182
183        if let Some(execute_fn) = execute_fn {
184            match execute_fn(params).await {
185                Ok(rows) => {
186                    self.stats.record_hit();
187                    self.touch_lru(conn_id, sql_hash);
188                    return Ok(PreparedLookup::Hit(rows));
189                }
190                Err(e) => {
191                    self.stats.record_degradation();
192                    return Err(e);
193                }
194            }
195        }
196
197        self.stats.record_miss();
198        Ok(PreparedLookup::Miss)
199    }
200
201    /// 存储句柄（prepare 后调用）
202    ///
203    /// 若 bucket 已满，LRU 淘汰最久未使用的条目。
204    pub fn store_handle(
205        &self,
206        conn_id: ConnId,
207        sql: &str,
208        tables: Vec<String>,
209        execute_fn: ExecuteFn,
210    ) {
211        let sql_hash = compute_sql_hash(sql);
212
213        let mut buckets = self.buckets.write();
214        let bucket = buckets
215            .entry(conn_id)
216            .or_insert_with(PreparedStatementBucket::new);
217
218        if bucket.entries.len() >= self.max_size_per_conn && !bucket.entries.contains_key(&sql_hash)
219        {
220            if let Some(evict_hash) = bucket.find_lru_key() {
221                bucket.entries.remove(&evict_hash);
222                bucket.lru.remove(&evict_hash);
223                for table_list in bucket.table_index.values_mut() {
224                    table_list.retain(|&h| h != evict_hash);
225                }
226                self.stats.record_eviction();
227            }
228        }
229
230        for table in &tables {
231            bucket
232                .table_index
233                .entry(table.clone())
234                .or_default()
235                .push(sql_hash);
236        }
237
238        bucket.lru.insert(sql_hash, Instant::now());
239        bucket
240            .entries
241            .insert(sql_hash, PreparedStatementEntry { execute_fn });
242    }
243
244    /// 表级失效：清除该表关联的所有句柄
245    ///
246    /// 返回失效条目数
247    pub fn invalidate_table(&self, table: &str) -> usize {
248        let mut buckets = self.buckets.write();
249        let mut total = 0;
250
251        for bucket in buckets.values_mut() {
252            if let Some(hashes) = bucket.table_index.remove(table) {
253                for hash in &hashes {
254                    bucket.entries.remove(hash);
255                    bucket.lru.remove(hash);
256                }
257                total += hashes.len();
258            }
259        }
260
261        self.stats.record_invalidation(total);
262        total
263    }
264
265    /// 连接级失效：清除该连接整个 bucket
266    ///
267    /// 连接关闭时调用，返回失效条目数
268    pub fn invalidate_conn(&self, conn_id: ConnId) -> usize {
269        let mut buckets = self.buckets.write();
270        if let Some(bucket) = buckets.remove(&conn_id) {
271            let count = bucket.entries.len();
272            self.stats.record_invalidation(count);
273            count
274        } else {
275            0
276        }
277    }
278
279    /// 统计快照
280    #[must_use]
281    pub fn stats(&self) -> PreparedStatementCacheStatsSnapshot {
282        let buckets = self.buckets.read();
283        let size: usize = buckets.values().map(|b| b.entries.len()).sum();
284        let hits = self.stats.hits.load(Ordering::Relaxed);
285        let misses = self.stats.misses.load(Ordering::Relaxed);
286        let evictions = self.stats.evictions.load(Ordering::Relaxed);
287        let degradations = self.stats.degradations.load(Ordering::Relaxed);
288        let invalidations = self.stats.invalidations.load(Ordering::Relaxed);
289        let total = hits + misses;
290        let hit_rate = if total == 0 {
291            0.0
292        } else {
293            hits as f64 / total as f64
294        };
295
296        PreparedStatementCacheStatsSnapshot {
297            hits,
298            misses,
299            evictions,
300            degradations,
301            invalidations,
302            size,
303            hit_rate,
304        }
305    }
306
307    /// 更新 LRU 时间戳
308    fn touch_lru(&self, conn_id: ConnId, sql_hash: u64) {
309        let mut buckets = self.buckets.write();
310        if let Some(bucket) = buckets.get_mut(&conn_id) {
311            bucket.lru.insert(sql_hash, Instant::now());
312        }
313    }
314
315    /// 当前条目数（调试/测试用）
316    #[must_use]
317    pub fn len(&self) -> usize {
318        let buckets = self.buckets.read();
319        buckets.values().map(|b| b.entries.len()).sum()
320    }
321
322    /// 是否为空
323    #[must_use]
324    pub fn is_empty(&self) -> bool {
325        self.len() == 0
326    }
327}
328
329impl PreparedStatementBucket {
330    fn new() -> Self {
331        Self {
332            entries: HashMap::new(),
333            lru: HashMap::new(),
334            table_index: HashMap::new(),
335        }
336    }
337
338    fn find_lru_key(&self) -> Option<u64> {
339        self.lru.iter().min_by_key(|(_, &t)| t).map(|(&h, _)| h)
340    }
341}
342
343impl Default for PreparedStatementCache {
344    fn default() -> Self {
345        Self::new(256)
346    }
347}
348
349/// 计算 SQL hash（归一化后）
350fn compute_sql_hash(sql: &str) -> u64 {
351    let normalized = sql.trim().to_lowercase();
352    let mut hasher = std::collections::hash_map::DefaultHasher::new();
353    normalized.hash(&mut hasher);
354    hasher.finish()
355}
356
357/// 从 SQL 中提取表名（简单字符串扫描，不依赖 sqlparser）
358///
359/// 扫描 FROM / JOIN / INTO / UPDATE 后面的表名
360#[must_use]
361pub fn extract_tables_simple(sql: &str) -> Vec<String> {
362    let lower = sql.to_lowercase();
363    let keywords = ["from", "join", "into", "update", "table"];
364    let mut tables = Vec::new();
365    let tokens: Vec<&str> = lower.split_whitespace().collect();
366
367    for (i, token) in tokens.iter().enumerate() {
368        if keywords.contains(token) {
369            if let Some(table) = tokens.get(i + 1) {
370                let clean = table
371                    .trim_matches(|c: char| c == ',' || c == ';' || c == '(' || c == ')')
372                    .trim();
373                if !clean.is_empty()
374                    && !clean.starts_with('?')
375                    && !clean.starts_with('$')
376                    && !clean.starts_with('(')
377                {
378                    tables.push(clean.to_string());
379                }
380            }
381        }
382    }
383
384    tables.sort();
385    tables.dedup();
386    tables
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    fn make_execute_fn() -> ExecuteFn {
394        Arc::new(|_params: &[Value]| {
395            Box::pin(async {
396                Ok(vec![{
397                    let mut row = std::collections::HashMap::new();
398                    row.insert("id".to_string(), Value::I64(1));
399                    row
400                }])
401            })
402        })
403    }
404
405    fn make_failing_execute_fn() -> ExecuteFn {
406        Arc::new(|_params: &[Value]| {
407            Box::pin(async { Err(DbError::QueryError("handle execution failed".into())) })
408        })
409    }
410
411    #[tokio::test]
412    async fn test_cache_hit() {
413        let cache = PreparedStatementCache::new(256);
414        let conn_id = 1;
415        let sql = "SELECT * FROM users WHERE id = ?";
416
417        let result = cache.get_or_prepare(conn_id, sql, &[]).await.unwrap();
418        assert!(matches!(result, PreparedLookup::Miss));
419
420        cache.store_handle(conn_id, sql, vec!["users".into()], make_execute_fn());
421
422        let result = cache.get_or_prepare(conn_id, sql, &[]).await.unwrap();
423        assert!(matches!(result, PreparedLookup::Hit(_)));
424
425        let stats = cache.stats();
426        assert_eq!(stats.hits, 1);
427        assert_eq!(stats.misses, 1);
428    }
429
430    #[tokio::test]
431    async fn test_cache_miss() {
432        let cache = PreparedStatementCache::new(256);
433        let result = cache
434            .get_or_prepare(1, "SELECT * FROM users", &[])
435            .await
436            .unwrap();
437        assert!(matches!(result, PreparedLookup::Miss));
438        assert_eq!(cache.stats().misses, 1);
439        assert_eq!(cache.stats().hits, 0);
440    }
441
442    #[tokio::test]
443    async fn test_lru_eviction() {
444        let cache = PreparedStatementCache::new(2);
445        let conn_id = 1;
446
447        cache.store_handle(
448            conn_id,
449            "SELECT * FROM a",
450            vec!["a".into()],
451            make_execute_fn(),
452        );
453        std::thread::sleep(std::time::Duration::from_millis(10));
454        cache.store_handle(
455            conn_id,
456            "SELECT * FROM b",
457            vec!["b".into()],
458            make_execute_fn(),
459        );
460        std::thread::sleep(std::time::Duration::from_millis(10));
461
462        let _ = cache.get_or_prepare(conn_id, "SELECT * FROM a", &[]).await;
463        std::thread::sleep(std::time::Duration::from_millis(10));
464
465        cache.store_handle(
466            conn_id,
467            "SELECT * FROM c",
468            vec!["c".into()],
469            make_execute_fn(),
470        );
471
472        let stats = cache.stats();
473        assert!(stats.evictions >= 1);
474        assert_eq!(cache.len(), 2);
475    }
476
477    #[tokio::test]
478    async fn test_invalidate_table() {
479        let cache = PreparedStatementCache::new(256);
480        let conn_id = 1;
481
482        cache.store_handle(
483            conn_id,
484            "SELECT * FROM users",
485            vec!["users".into()],
486            make_execute_fn(),
487        );
488        cache.store_handle(
489            conn_id,
490            "SELECT * FROM orders",
491            vec!["orders".into()],
492            make_execute_fn(),
493        );
494
495        let count = cache.invalidate_table("users");
496        assert_eq!(count, 1);
497        assert_eq!(cache.len(), 1);
498
499        let result = cache
500            .get_or_prepare(conn_id, "SELECT * FROM users", &[])
501            .await
502            .unwrap();
503        assert!(matches!(result, PreparedLookup::Miss));
504    }
505
506    #[tokio::test]
507    async fn test_invalidate_conn() {
508        let cache = PreparedStatementCache::new(256);
509
510        cache.store_handle(1, "SELECT * FROM a", vec!["a".into()], make_execute_fn());
511        cache.store_handle(1, "SELECT * FROM b", vec!["b".into()], make_execute_fn());
512        cache.store_handle(2, "SELECT * FROM c", vec!["c".into()], make_execute_fn());
513
514        let count = cache.invalidate_conn(1);
515        assert_eq!(count, 2);
516        assert_eq!(cache.len(), 1);
517
518        let result = cache
519            .get_or_prepare(1, "SELECT * FROM a", &[])
520            .await
521            .unwrap();
522        assert!(matches!(result, PreparedLookup::Miss));
523    }
524
525    #[tokio::test]
526    async fn test_capacity_limit() {
527        let cache = PreparedStatementCache::new(3);
528        let conn_id = 1;
529
530        for i in 0..5 {
531            cache.store_handle(
532                conn_id,
533                &format!("SELECT * FROM table_{i}"),
534                vec![format!("table_{i}")],
535                make_execute_fn(),
536            );
537        }
538
539        assert_eq!(cache.len(), 3);
540        assert!(cache.stats().evictions >= 2);
541    }
542
543    #[tokio::test]
544    async fn test_concurrent_safety() {
545        use std::sync::Arc;
546        let cache = Arc::new(PreparedStatementCache::new(256));
547        let mut handles = Vec::new();
548
549        for thread_id in 0..4 {
550            let cache = cache.clone();
551            handles.push(tokio::spawn(async move {
552                for i in 0..10 {
553                    let sql = format!("SELECT * FROM t_{thread_id}_{i}");
554                    let _ = cache.get_or_prepare(thread_id, &sql, &[]).await;
555                    cache.store_handle(thread_id, &sql, vec![], make_execute_fn());
556                }
557            }));
558        }
559
560        for h in handles {
561            h.await.unwrap();
562        }
563
564        assert_eq!(cache.len(), 40);
565    }
566
567    #[tokio::test]
568    async fn test_degradation_on_handle_error() {
569        let cache = PreparedStatementCache::new(256);
570        let conn_id = 1;
571        let sql = "SELECT * FROM users";
572
573        cache.store_handle(conn_id, sql, vec![], make_failing_execute_fn());
574
575        let result = cache.get_or_prepare(conn_id, sql, &[]).await;
576        assert!(result.is_err());
577        assert_eq!(cache.stats().degradations, 1);
578    }
579
580    #[tokio::test]
581    async fn test_handle_not_exposed() {
582        let cache = PreparedStatementCache::new(256);
583        cache.store_handle(1, "SELECT 1", vec![], make_execute_fn());
584
585        let buckets = cache.buckets.read();
586        let bucket = buckets.get(&1).unwrap();
587        assert!(bucket.entries.contains_key(&compute_sql_hash("SELECT 1")));
588    }
589
590    #[tokio::test]
591    async fn test_normalization_hit() {
592        let cache = PreparedStatementCache::new(256);
593        let conn_id = 1;
594
595        cache.store_handle(conn_id, "SELECT * FROM users", vec![], make_execute_fn());
596
597        let result = cache
598            .get_or_prepare(conn_id, "select * from users", &[])
599            .await
600            .unwrap();
601        assert!(
602            matches!(result, PreparedLookup::Hit(_)),
603            "大小写归一化应命中"
604        );
605    }
606
607    #[tokio::test]
608    async fn test_cross_conn_isolation() {
609        let cache = PreparedStatementCache::new(256);
610        let sql = "SELECT * FROM users";
611
612        cache.store_handle(1, sql, vec![], make_execute_fn());
613
614        let result = cache.get_or_prepare(2, sql, &[]).await.unwrap();
615        assert!(matches!(result, PreparedLookup::Miss), "句柄不跨连接复用");
616    }
617
618    #[tokio::test]
619    async fn test_stats_accuracy() {
620        let cache = PreparedStatementCache::new(256);
621
622        cache.store_handle(1, "SELECT * FROM a", vec!["a".into()], make_execute_fn());
623        cache.store_handle(1, "SELECT * FROM b", vec!["b".into()], make_execute_fn());
624
625        let _ = cache.get_or_prepare(1, "SELECT * FROM a", &[]).await;
626        let _ = cache.get_or_prepare(1, "SELECT * FROM a", &[]).await;
627        let _ = cache.get_or_prepare(1, "SELECT * FROM b", &[]).await;
628        let _ = cache.get_or_prepare(1, "SELECT * FROM c", &[]).await;
629
630        cache.invalidate_table("a");
631
632        let stats = cache.stats();
633        assert_eq!(stats.hits, 3);
634        assert_eq!(stats.misses, 1);
635        assert_eq!(stats.invalidations, 1);
636        assert_eq!(stats.size, 1);
637    }
638
639    #[test]
640    fn test_extract_tables_simple() {
641        let tables =
642            extract_tables_simple("SELECT * FROM users JOIN orders ON users.id = orders.user_id");
643        assert!(tables.contains(&"users".to_string()));
644        assert!(tables.contains(&"orders".to_string()));
645
646        let tables = extract_tables_simple("INSERT INTO logs VALUES (1)");
647        assert!(tables.contains(&"logs".to_string()));
648
649        let tables = extract_tables_simple("UPDATE accounts SET balance = 0");
650        assert!(tables.contains(&"accounts".to_string()));
651    }
652}
sz_orm_core/prepared_cache.rs

sz_orm_core/
prepared_cache.rs
