Skip to main content

sz_orm_core/
shadow.rs

1//! 双轨影子流量校验模块
2//!
3//! 在 ORM 查询执行时,同步通过原生路径执行相同 SQL 并比较结果集,
4//! 用于在上线重大优化时验证 ORM 行为与原生驱动的一致性。
5//!
6//! # 设计
7//!
8//! [`ShadowConnection`] 是一个装饰器(decorator),包装一个 ORM [`Connection`]
9//! 和一个用于比对的"原生" [`Connection`](通常底层是相同的 sqlx pool)。
10//! 每次 ORM 执行 `query` 时,原生连接同步执行相同 SQL,然后比较结果。
11//!
12//! 与 SKILL.md 中提到的 `tower` 中间件模式不同,sz-orm 是库而非 HTTP 服务,
13//! 装饰器模式更贴合 ORM 的 `Connection` trait 抽象。
14//!
15//! # 使用示例
16//!
17//! ```ignore
18//! use sz_orm_core::shadow::{ShadowConnection, ShadowConfig};
19//!
20//! let orm_conn = pool.acquire().await?;
21//! let raw_conn = pool.acquire().await?;
22//! let mut shadow = ShadowConnection::new(orm_conn, raw_conn, ShadowConfig::default());
23//!
24//! // 执行查询时自动双轨比对
25//! let rows = shadow.query("SELECT * FROM users").await?;
26//! ```
27
28use std::collections::HashMap;
29use std::sync::atomic::{AtomicU64, Ordering};
30use std::time::{Duration, Instant};
31
32use crate::pool::{Connection, QueryRows};
33use crate::Value;
34
35/// 影子流量比较结果
36#[derive(Debug, Clone)]
37pub struct ShadowComparison {
38    /// 执行的 SQL
39    pub sql: String,
40    /// ORM 路径执行耗时
41    pub orm_duration: Duration,
42    /// 原生路径执行耗时
43    pub raw_duration: Duration,
44    /// ORM 返回的行数
45    pub orm_rows: usize,
46    /// 原生返回的行数
47    pub raw_rows: usize,
48    /// 是否一致
49    pub consistent: bool,
50    /// 不一致时的差异描述(首条差异)
51    pub mismatch: Option<String>,
52}
53
54impl ShadowComparison {
55    /// ORM 延迟与原生延迟的比值(>1 表示 ORM 更慢)
56    pub fn latency_ratio(&self) -> f64 {
57        if self.raw_duration.is_zero() {
58            1.0
59        } else {
60            self.orm_duration.as_secs_f64() / self.raw_duration.as_secs_f64()
61        }
62    }
63}
64
65/// 影子流量累计统计
66#[derive(Debug, Default)]
67pub struct ShadowStats {
68    /// 累计比较次数
69    pub comparisons: AtomicU64,
70    /// 不一致次数
71    pub mismatches: AtomicU64,
72    /// ORM 累计耗时(微秒)
73    pub orm_total_us: AtomicU64,
74    /// 原生累计耗时(微秒)
75    pub raw_total_us: AtomicU64,
76}
77
78impl ShadowStats {
79    /// 平均 ORM 耗时(微秒)
80    pub fn avg_orm_us(&self) -> u64 {
81        let n = self.comparisons.load(Ordering::Relaxed);
82        // 使用 checked_div 替代手动零检查,避免 clippy::manual_checked_ops 警告
83        self.orm_total_us
84            .load(Ordering::Relaxed)
85            .checked_div(n)
86            .unwrap_or(0)
87    }
88
89    /// 平均原生耗时(微秒)
90    pub fn avg_raw_us(&self) -> u64 {
91        let n = self.comparisons.load(Ordering::Relaxed);
92        // 使用 checked_div 替代手动零检查,避免 clippy::manual_checked_ops 警告
93        self.raw_total_us
94            .load(Ordering::Relaxed)
95            .checked_div(n)
96            .unwrap_or(0)
97    }
98
99    /// 不一致率
100    pub fn mismatch_rate(&self) -> f64 {
101        let n = self.comparisons.load(Ordering::Relaxed);
102        if n == 0 {
103            0.0
104        } else {
105            self.mismatches.load(Ordering::Relaxed) as f64 / n as f64
106        }
107    }
108}
109
110/// 影子流量配置
111#[derive(Debug, Clone)]
112pub struct ShadowConfig {
113    /// 单侧执行超时(秒)。超时记为不一致。
114    pub timeout: Duration,
115    /// 是否仅比较行数(跳过逐字段比对,用于快速验证)
116    pub row_count_only: bool,
117    /// 最大比对行数(避免超大结果集内存爆炸)
118    pub max_compare_rows: usize,
119}
120
121impl Default for ShadowConfig {
122    fn default() -> Self {
123        Self {
124            timeout: Duration::from_secs(3),
125            row_count_only: false,
126            max_compare_rows: 10_000,
127        }
128    }
129}
130
131/// 不一致处理策略
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133pub enum MismatchAction {
134    /// 仅记录,继续返回 ORM 结果
135    Record,
136    /// 记录并返回错误
137    Error,
138    /// 记录并触发 panic(用于测试环境)
139    Panic,
140}
141
142/// 双轨影子流量校验连接装饰器
143///
144/// 包装 ORM 连接和原生连接,每次查询时同步执行两路并比较结果。
145///
146/// # 线程安全
147///
148/// `ShadowConnection` 内部的 ORM 连接和原生连接都要求 `Connection: Send + Sync`,
149/// 因此 `ShadowConnection` 自动满足 `Send + Sync`。
150pub struct ShadowConnection<C: Connection> {
151    /// ORM 路径连接
152    orm_conn: C,
153    /// 原生路径连接
154    raw_conn: C,
155    /// 配置
156    config: ShadowConfig,
157    /// 不匹配时的行为
158    on_mismatch: MismatchAction,
159    /// 累计统计
160    stats: ShadowStats,
161}
162
163impl<C: Connection> ShadowConnection<C> {
164    /// 创建影子流量连接
165    pub fn new(orm_conn: C, raw_conn: C, config: ShadowConfig) -> Self {
166        Self {
167            orm_conn,
168            raw_conn,
169            config,
170            on_mismatch: MismatchAction::Record,
171            stats: ShadowStats::default(),
172        }
173    }
174
175    /// 设置不匹配行为
176    pub fn with_mismatch_action(mut self, action: MismatchAction) -> Self {
177        self.on_mismatch = action;
178        self
179    }
180
181    /// 获取统计快照
182    pub fn stats(&self) -> &ShadowStats {
183        &self.stats
184    }
185
186    /// 执行双轨查询并比较结果
187    ///
188    /// 返回 ORM 路径的结果。如果不匹配,根据 `on_mismatch` 策略处理。
189    pub async fn query_shadow(&mut self, sql: &str) -> Result<QueryRows, crate::DbError> {
190        // 并发执行两路查询
191        let orm_fut = self.orm_conn.query(sql);
192        let raw_fut = self.raw_conn.query(sql);
193
194        let orm_start = Instant::now();
195        let orm_result = tokio::time::timeout(self.config.timeout, orm_fut).await;
196        let orm_duration = orm_start.elapsed();
197
198        let raw_start = Instant::now();
199        let raw_result = tokio::time::timeout(self.config.timeout, raw_fut).await;
200        let raw_duration = raw_start.elapsed();
201
202        // 统计累计耗时
203        self.stats
204            .orm_total_us
205            .fetch_add(orm_duration.as_micros() as u64, Ordering::Relaxed);
206        self.stats
207            .raw_total_us
208            .fetch_add(raw_duration.as_micros() as u64, Ordering::Relaxed);
209        self.stats.comparisons.fetch_add(1, Ordering::Relaxed);
210
211        // 处理超时与错误
212        let orm_rows = match orm_result {
213            Ok(Ok(rows)) => rows,
214            Ok(Err(e)) => {
215                self.stats.mismatches.fetch_add(1, Ordering::Relaxed);
216                // H-5 修复:传播 handle_mismatch 的 Result(Panic 模式下返回 Internal 错误)
217                self.handle_mismatch(ShadowComparison {
218                    sql: sql.to_string(),
219                    orm_duration,
220                    raw_duration,
221                    orm_rows: 0,
222                    raw_rows: 0,
223                    consistent: false,
224                    mismatch: Some(format!("ORM error: {}", e)),
225                })?;
226                return Err(e);
227            }
228            Err(_) => {
229                self.stats.mismatches.fetch_add(1, Ordering::Relaxed);
230                return Err(crate::DbError::ConnectionTimeout(format!(
231                    "ORM shadow path timeout after {:?}",
232                    self.config.timeout
233                )));
234            }
235        };
236
237        let raw_rows = match raw_result {
238            Ok(Ok(rows)) => rows,
239            Ok(Err(e)) => {
240                self.stats.mismatches.fetch_add(1, Ordering::Relaxed);
241                // H-5 修复:传播 handle_mismatch 的 Result(Panic 模式下返回 Internal 错误)
242                self.handle_mismatch(ShadowComparison {
243                    sql: sql.to_string(),
244                    orm_duration,
245                    raw_duration,
246                    orm_rows: orm_rows.len(),
247                    raw_rows: 0,
248                    consistent: false,
249                    mismatch: Some(format!("Raw path error: {}", e)),
250                })?;
251                // 原生路径失败不影响 ORM 路径返回
252                return Ok(orm_rows);
253            }
254            Err(_) => {
255                self.stats.mismatches.fetch_add(1, Ordering::Relaxed);
256                // H-5 修复:传播 handle_mismatch 的 Result(Panic 模式下返回 Internal 错误)
257                self.handle_mismatch(ShadowComparison {
258                    sql: sql.to_string(),
259                    orm_duration,
260                    raw_duration,
261                    orm_rows: orm_rows.len(),
262                    raw_rows: 0,
263                    consistent: false,
264                    mismatch: Some(format!("Raw path timeout after {:?}", self.config.timeout)),
265                })?;
266                return Ok(orm_rows);
267            }
268        };
269
270        // 比较结果集
271        let comparison = self.compare_rows(sql, orm_duration, raw_duration, &orm_rows, &raw_rows);
272        if !comparison.consistent {
273            self.stats.mismatches.fetch_add(1, Ordering::Relaxed);
274            // H-5 修复:传播 handle_mismatch 的 Result(Panic 模式下返回 Internal 错误)
275            self.handle_mismatch(comparison)?;
276        }
277
278        Ok(orm_rows)
279    }
280
281    /// 比较两路结果集
282    fn compare_rows(
283        &self,
284        sql: &str,
285        orm_duration: Duration,
286        raw_duration: Duration,
287        orm_rows: &[HashMap<String, Value>],
288        raw_rows: &[HashMap<String, Value>],
289    ) -> ShadowComparison {
290        // 行数比较
291        if orm_rows.len() != raw_rows.len() {
292            return ShadowComparison {
293                sql: sql.to_string(),
294                orm_duration,
295                raw_duration,
296                orm_rows: orm_rows.len(),
297                raw_rows: raw_rows.len(),
298                consistent: false,
299                mismatch: Some(format!(
300                    "Row count mismatch: ORM={} vs Raw={}",
301                    orm_rows.len(),
302                    raw_rows.len()
303                )),
304            };
305        }
306
307        // 仅比较行数时直接返回一致
308        if self.config.row_count_only {
309            return ShadowComparison {
310                sql: sql.to_string(),
311                orm_duration,
312                raw_duration,
313                orm_rows: orm_rows.len(),
314                raw_rows: raw_rows.len(),
315                consistent: true,
316                mismatch: None,
317            };
318        }
319
320        // 逐行逐字段比较(限制最大行数)
321        let limit = self.config.max_compare_rows.min(orm_rows.len());
322        for (i, (orm_row, raw_row)) in orm_rows[..limit]
323            .iter()
324            .zip(raw_rows[..limit].iter())
325            .enumerate()
326        {
327            if orm_row.len() != raw_row.len() {
328                return ShadowComparison {
329                    sql: sql.to_string(),
330                    orm_duration,
331                    raw_duration,
332                    orm_rows: orm_rows.len(),
333                    raw_rows: raw_rows.len(),
334                    consistent: false,
335                    mismatch: Some(format!(
336                        "Row {} column count mismatch: ORM={} vs Raw={}",
337                        i,
338                        orm_row.len(),
339                        raw_row.len()
340                    )),
341                };
342            }
343            for (key, orm_val) in orm_row {
344                match raw_row.get(key) {
345                    Some(raw_val) if !values_equal(orm_val, raw_val) => {
346                        return ShadowComparison {
347                            sql: sql.to_string(),
348                            orm_duration,
349                            raw_duration,
350                            orm_rows: orm_rows.len(),
351                            raw_rows: raw_rows.len(),
352                            consistent: false,
353                            mismatch: Some(format!(
354                                "Row {} column '{}' value mismatch: ORM={:?} vs Raw={:?}",
355                                i, key, orm_val, raw_val
356                            )),
357                        };
358                    }
359                    None => {
360                        return ShadowComparison {
361                            sql: sql.to_string(),
362                            orm_duration,
363                            raw_duration,
364                            orm_rows: orm_rows.len(),
365                            raw_rows: raw_rows.len(),
366                            consistent: false,
367                            mismatch: Some(format!(
368                                "Row {} column '{}' missing in raw result",
369                                i, key
370                            )),
371                        };
372                    }
373                    _ => {}
374                }
375            }
376        }
377
378        ShadowComparison {
379            sql: sql.to_string(),
380            orm_duration,
381            raw_duration,
382            orm_rows: orm_rows.len(),
383            raw_rows: raw_rows.len(),
384            consistent: true,
385            mismatch: None,
386        }
387    }
388
389    /// 根据策略处理不匹配
390    ///
391    /// 返回 `Err` 仅在 `MismatchAction::Panic` 模式下(H-5 修复:将 panic! 改为
392    /// 返回 `Result`,由调用方决定如何处理)。其他模式返回 `Ok(())`。
393    fn handle_mismatch(&self, comparison: ShadowComparison) -> Result<(), crate::DbError> {
394        tracing::warn!(
395            target: "sz_orm::shadow",
396            sql = %comparison.sql,
397            orm_rows = comparison.orm_rows,
398            raw_rows = comparison.raw_rows,
399            orm_us = ?comparison.orm_duration,
400            raw_us = ?comparison.raw_duration,
401            mismatch = ?comparison.mismatch,
402            "Shadow traffic mismatch detected"
403        );
404
405        match self.on_mismatch {
406            MismatchAction::Record => Ok(()),
407            MismatchAction::Error => Ok(()),
408            MismatchAction::Panic => Err(crate::DbError::Internal(format!(
409                "Shadow traffic mismatch: SQL={:?} mismatch={:?}",
410                comparison.sql, comparison.mismatch
411            ))),
412        }
413    }
414}
415
416/// 值相等比较(容忍类型相近的等价,如 I64(42) vs I32(42))
417fn values_equal(a: &Value, b: &Value) -> bool {
418    use Value::*;
419    match (a, b) {
420        (Null, Null) => true,
421        (Bool(x), Bool(y)) => x == y,
422        (I8(x), I8(y)) => x == y,
423        (I8(x), I16(y)) => (*x as i16) == *y,
424        (I16(x), I8(y)) => *x == (*y as i16),
425        (I16(x), I16(y)) => x == y,
426        (I16(x), I32(y)) => (*x as i32) == *y,
427        (I32(x), I16(y)) => *x == (*y as i32),
428        (I32(x), I32(y)) => x == y,
429        (I32(x), I64(y)) => (*x as i64) == *y,
430        (I64(x), I32(y)) => *x == (*y as i64),
431        (I64(x), I64(y)) => x == y,
432        (U8(x), U8(y)) => x == y,
433        (U16(x), U16(y)) => x == y,
434        (U32(x), U32(y)) => x == y,
435        (U64(x), U64(y)) => x == y,
436        (F32(x), F32(y)) => (x - y).abs() < 1e-6,
437        (F32(x), F64(y)) => ((*x as f64) - y).abs() < 1e-6,
438        (F64(x), F32(y)) => (x - (*y as f64)).abs() < 1e-6,
439        (F64(x), F64(y)) => (x - y).abs() < 1e-9,
440        (String(x), String(y)) => x == y,
441        (Bytes(x), Bytes(y)) => x == y,
442        _ => a == b,
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use super::*;
449
450    #[test]
451    fn test_values_equal_numeric_cross_type() {
452        assert!(values_equal(&Value::I32(42), &Value::I64(42)));
453        assert!(values_equal(&Value::I8(1), &Value::I16(1)));
454        assert!(!values_equal(&Value::I32(42), &Value::I64(43)));
455    }
456
457    #[test]
458    fn test_values_equal_float_tolerance() {
459        assert!(values_equal(&Value::F32(1.0), &Value::F64(1.0)));
460        assert!(!values_equal(&Value::F32(1.0), &Value::F64(2.0)));
461    }
462
463    #[test]
464    fn test_shadow_stats_mismatch_rate() {
465        let stats = ShadowStats::default();
466        assert_eq!(stats.mismatch_rate(), 0.0);
467        stats.comparisons.store(10, Ordering::Relaxed);
468        stats.mismatches.store(1, Ordering::Relaxed);
469        assert!((stats.mismatch_rate() - 0.1).abs() < 1e-9);
470    }
471}