sz_orm_core/
pool_elastic.rs

1//! v6.7.0 连接池弹性：动态扩缩容 + 多级熔断 + 健康检查 + 连接预热。
2//!
3//! `PoolElasticController` 周期采集池 stats，waiting > threshold 扩容，idle > threshold 缩容。
4//! `TieredCircuitBreaker` 三级熔断（连接→节点→全局），`ConnectionHealthChecker` 连续 3 次失败才剔除。
5
6use std::sync::Mutex;
7use std::time::{Duration, Instant};
8
9use serde::{Deserialize, Serialize};
10
11use crate::circuit_breaker::{CircuitBreaker, CircuitState, DefaultCircuitBreaker};
12
13// ============================================================================
14// 配置
15// ============================================================================
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ElasticConfig {
19    pub min_connections: u32,
20    pub max_connections: u32,
21    pub scale_up_threshold: u32,
22    pub scale_down_threshold: u32,
23    pub scale_interval: Duration,
24    pub health_check_interval: Duration,
25    pub health_check_sql: String,
26}
27
28impl Default for ElasticConfig {
29    fn default() -> Self {
30        Self {
31            min_connections: 5,
32            max_connections: 50,
33            scale_up_threshold: 10,
34            scale_down_threshold: 20,
35            scale_interval: Duration::from_secs(1),
36            health_check_interval: Duration::from_secs(30),
37            health_check_sql: "SELECT 1".to_string(),
38        }
39    }
40}
41
42impl ElasticConfig {
43    pub fn validate(&self) -> Result<(), String> {
44        if self.min_connections < 1 || self.min_connections > 100 {
45            return Err("min_connections 必须在 [1,100]".to_string());
46        }
47        if self.max_connections < self.min_connections || self.max_connections > 1000 {
48            return Err("max_connections 必须在 [min,1000]".to_string());
49        }
50        if self.scale_interval < Duration::from_secs(1)
51            || self.scale_interval > Duration::from_secs(60)
52        {
53            return Err("scale_interval 必须在 [1s,60s]".to_string());
54        }
55        if self.health_check_interval < Duration::from_secs(10)
56            || self.health_check_interval > Duration::from_secs(300)
57        {
58            return Err("health_check_interval 必须在 [10s,300s]".to_string());
59        }
60        Ok(())
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub enum ScaleReason {
66    HighLoad,
67    LowLoad,
68    HealthCheckFailure,
69}
70
71#[derive(Debug, Clone)]
72pub struct ScaleEvent {
73    pub from: u32,
74    pub to: u32,
75    pub reason: ScaleReason,
76    pub timestamp: Instant,
77    pub elapsed: Duration,
78}
79
80// ============================================================================
81// 多级熔断器
82// ============================================================================
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
85pub enum CircuitTier {
86    Connection,
87    Node,
88    Global,
89}
90
91pub struct TieredCircuitBreaker {
92    connection_level: Mutex<DefaultCircuitBreaker>,
93    node_level: Mutex<DefaultCircuitBreaker>,
94    global_level: Mutex<DefaultCircuitBreaker>,
95}
96
97impl TieredCircuitBreaker {
98    pub fn new(failure_threshold: usize, reset_timeout: Duration) -> Self {
99        Self {
100            connection_level: Mutex::new(DefaultCircuitBreaker::new(
101                failure_threshold,
102                reset_timeout,
103            )),
104            node_level: Mutex::new(DefaultCircuitBreaker::new(failure_threshold, reset_timeout)),
105            global_level: Mutex::new(DefaultCircuitBreaker::new(failure_threshold, reset_timeout)),
106        }
107    }
108
109    pub fn can_execute(&self, tier: CircuitTier) -> bool {
110        let cb = match tier {
111            CircuitTier::Connection => &self.connection_level,
112            CircuitTier::Node => &self.node_level,
113            CircuitTier::Global => &self.global_level,
114        };
115        cb.lock().unwrap().can_execute()
116    }
117
118    pub fn record_success(&self, tier: CircuitTier) {
119        let cb = match tier {
120            CircuitTier::Connection => &self.connection_level,
121            CircuitTier::Node => &self.node_level,
122            CircuitTier::Global => &self.global_level,
123        };
124        cb.lock().unwrap().record_success();
125    }
126
127    pub fn record_failure(&self, tier: CircuitTier) {
128        let cb = match tier {
129            CircuitTier::Connection => &self.connection_level,
130            CircuitTier::Node => &self.node_level,
131            CircuitTier::Global => &self.global_level,
132        };
133        cb.lock().unwrap().record_failure();
134    }
135
136    pub fn state(&self, tier: CircuitTier) -> CircuitState {
137        let cb = match tier {
138            CircuitTier::Connection => &self.connection_level,
139            CircuitTier::Node => &self.node_level,
140            CircuitTier::Global => &self.global_level,
141        };
142        cb.lock().unwrap().state()
143    }
144
145    pub fn can_execute_any(&self) -> bool {
146        self.can_execute(CircuitTier::Global)
147            && self.can_execute(CircuitTier::Node)
148            && self.can_execute(CircuitTier::Connection)
149    }
150}
151
152// ============================================================================
153// 健康检查器
154// ============================================================================
155
156pub struct ConnectionHealthChecker {
157    failure_counts: Mutex<std::collections::HashMap<String, u32>>,
158    failure_threshold: u32,
159}
160
161impl ConnectionHealthChecker {
162    pub fn new(failure_threshold: u32) -> Self {
163        Self {
164            failure_counts: Mutex::new(std::collections::HashMap::new()),
165            failure_threshold,
166        }
167    }
168
169    pub fn record_failure(&self, conn_id: &str) -> bool {
170        let mut counts = self.failure_counts.lock().unwrap();
171        let count = counts.entry(conn_id.to_string()).or_insert(0);
172        *count += 1;
173        *count >= self.failure_threshold
174    }
175
176    pub fn record_success(&self, conn_id: &str) {
177        let mut counts = self.failure_counts.lock().unwrap();
178        counts.insert(conn_id.to_string(), 0);
179    }
180
181    pub fn should_evict(&self, conn_id: &str) -> bool {
182        let counts = self.failure_counts.lock().unwrap();
183        *counts.get(conn_id).unwrap_or(&0) >= self.failure_threshold
184    }
185
186    pub fn failure_count(&self, conn_id: &str) -> u32 {
187        let counts = self.failure_counts.lock().unwrap();
188        *counts.get(conn_id).unwrap_or(&0)
189    }
190}
191
192// ============================================================================
193// 池统计
194// ============================================================================
195
196#[derive(Debug, Clone, Default)]
197pub struct PoolStats {
198    pub total_connections: u32,
199    pub idle_connections: u32,
200    pub waiting_requests: u32,
201}
202
203impl PoolStats {
204    pub fn active_connections(&self) -> u32 {
205        self.total_connections.saturating_sub(self.idle_connections)
206    }
207}
208
209// ============================================================================
210// 弹性控制器
211// ============================================================================
212
213pub struct PoolElasticController {
214    config: ElasticConfig,
215    circuit_breaker: TieredCircuitBreaker,
216    health_checker: ConnectionHealthChecker,
217    scale_events: Mutex<Vec<ScaleEvent>>,
218    current_size: Mutex<u32>,
219}
220
221impl PoolElasticController {
222    pub fn new(config: ElasticConfig) -> Self {
223        Self {
224            config,
225            circuit_breaker: TieredCircuitBreaker::new(5, Duration::from_secs(30)),
226            health_checker: ConnectionHealthChecker::new(3),
227            scale_events: Mutex::new(Vec::new()),
228            current_size: Mutex::new(5),
229        }
230    }
231
232    pub fn should_scale_up(stats: &PoolStats, config: &ElasticConfig) -> bool {
233        stats.waiting_requests > config.scale_up_threshold
234    }
235
236    pub fn should_scale_down(stats: &PoolStats, config: &ElasticConfig) -> bool {
237        stats.idle_connections > config.scale_down_threshold
238    }
239
240    pub fn evaluate_and_scale(&self, stats: &PoolStats) -> Option<ScaleEvent> {
241        let mut current = self.current_size.lock().unwrap();
242        let start = Instant::now();
243
244        if Self::should_scale_up(stats, &self.config) {
245            let target = (*current + stats.waiting_requests).min(self.config.max_connections);
246            if target > *current {
247                let event = ScaleEvent {
248                    from: *current,
249                    to: target,
250                    reason: ScaleReason::HighLoad,
251                    timestamp: Instant::now(),
252                    elapsed: start.elapsed(),
253                };
254                *current = target;
255                self.scale_events.lock().unwrap().push(event.clone());
256                return Some(event);
257            }
258        }
259
260        if Self::should_scale_down(stats, &self.config) {
261            let target = (*current / 2).max(self.config.min_connections);
262            if target < *current {
263                let event = ScaleEvent {
264                    from: *current,
265                    to: target,
266                    reason: ScaleReason::LowLoad,
267                    timestamp: Instant::now(),
268                    elapsed: start.elapsed(),
269                };
270                *current = target;
271                self.scale_events.lock().unwrap().push(event.clone());
272                return Some(event);
273            }
274        }
275
276        None
277    }
278
279    pub fn scale_events(&self) -> Vec<ScaleEvent> {
280        self.scale_events.lock().unwrap().clone()
281    }
282
283    pub fn current_size(&self) -> u32 {
284        *self.current_size.lock().unwrap()
285    }
286
287    pub fn circuit_breaker(&self) -> &TieredCircuitBreaker {
288        &self.circuit_breaker
289    }
290
291    pub fn health_checker(&self) -> &ConnectionHealthChecker {
292        &self.health_checker
293    }
294
295    pub fn config(&self) -> &ElasticConfig {
296        &self.config
297    }
298
299    pub fn prewarm(&self, target: u32) -> u32 {
300        let mut current = self.current_size.lock().unwrap();
301        let target = target
302            .min(self.config.max_connections)
303            .max(self.config.min_connections);
304        *current = target;
305        target
306    }
307}
308
309// ============================================================================
310// v6.8.0 PERF-POOL-01：连接池 IO 复用
311// ============================================================================
312
313/// 连接池 IO 复用通道
314///
315/// 在连接上复用 prepared statement 通道，减少重复 prepare 开销。
316/// 首次执行 SQL 时 prepare 并缓存句柄，后续执行直接复用缓存句柄。
317///
318/// 需启用 `pool-io-reuse` feature。
319#[cfg(feature = "pool-io-reuse")]
320pub struct IoReuseChannel {
321    cache: crate::prepared_cache::PreparedStatementCache,
322}
323
324#[cfg(feature = "pool-io-reuse")]
325impl IoReuseChannel {
326    /// 创建 IO 复用通道，内部持有 `PreparedStatementCache`
327    ///
328    /// `max_size_per_conn` 为每连接最大缓存句柄数（默认 256）
329    #[must_use]
330    pub fn new(max_size_per_conn: usize) -> Self {
331        Self {
332            cache: crate::prepared_cache::PreparedStatementCache::new(max_size_per_conn),
333        }
334    }
335
336    /// 复用 prepared statement 通道执行查询
337    ///
338    /// 流程：
339    /// 1. 查找缓存句柄 → 命中则直接执行返回
340    /// 2. 未命中 → 调用 `prepare_fn` 获取执行闭包 → 缓存 → 执行返回
341    ///
342    /// # 参数
343    /// - `conn_id`: 连接唯一标识
344    /// - `sql`: SQL 文本
345    /// - `params`: 参数列表
346    /// - `tables`: 涉及的表名列表（用于表级失效索引）
347    /// - `prepare_fn`: 首次执行时的 prepare 闭包，返回执行函数
348    pub async fn execute_reuse(
349        &self,
350        conn_id: crate::prepared_cache::ConnId,
351        sql: &str,
352        params: &[crate::value::Value],
353        tables: Vec<String>,
354        prepare_fn: impl FnOnce() -> crate::prepared_cache::ExecuteFn,
355    ) -> Result<crate::pool::QueryRows, crate::error::DbError> {
356        use crate::prepared_cache::PreparedLookup;
357
358        match self.cache.get_or_prepare(conn_id, sql, params).await? {
359            PreparedLookup::Hit(rows) => Ok(rows),
360            PreparedLookup::Miss => {
361                let execute_fn = prepare_fn();
362                self.cache
363                    .store_handle(conn_id, sql, tables, std::sync::Arc::clone(&execute_fn));
364                execute_fn(params).await
365            }
366        }
367    }
368
369    /// 返回缓存统计快照
370    pub fn stats(&self) -> crate::prepared_cache::PreparedStatementCacheStatsSnapshot {
371        self.cache.stats()
372    }
373
374    /// 失效连接级缓存（连接关闭时调用）
375    pub fn invalidate_conn(&self, conn_id: crate::prepared_cache::ConnId) {
376        self.cache.invalidate_conn(conn_id);
377    }
378
379    /// 失效表级缓存（表结构变更时调用）
380    pub fn invalidate_table(&self, table: &str) {
381        self.cache.invalidate_table(table);
382    }
383}
384
385// ============================================================================
386// 测试
387// ============================================================================
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    #[test]
394    fn elastic_config_default_valid() {
395        let config = ElasticConfig::default();
396        assert!(config.validate().is_ok());
397    }
398
399    #[test]
400    fn elastic_config_invalid_min() {
401        let mut config = ElasticConfig::default();
402        config.min_connections = 0;
403        assert!(config.validate().is_err());
404    }
405
406    #[test]
407    fn elastic_config_invalid_max() {
408        let mut config = ElasticConfig::default();
409        config.max_connections = 3;
410        assert!(config.validate().is_err());
411    }
412
413    #[test]
414    fn tiered_circuit_independent() {
415        let cb = TieredCircuitBreaker::new(2, Duration::from_secs(60));
416        assert!(cb.can_execute(CircuitTier::Connection));
417        assert!(cb.can_execute(CircuitTier::Node));
418        assert!(cb.can_execute(CircuitTier::Global));
419
420        cb.record_failure(CircuitTier::Connection);
421        cb.record_failure(CircuitTier::Connection);
422        assert_eq!(cb.state(CircuitTier::Connection), CircuitState::Open);
423        assert_eq!(cb.state(CircuitTier::Node), CircuitState::Closed);
424        assert_eq!(cb.state(CircuitTier::Global), CircuitState::Closed);
425    }
426
427    #[test]
428    fn tiered_circuit_global_blocks_all() {
429        let cb = TieredCircuitBreaker::new(1, Duration::from_secs(60));
430        cb.record_failure(CircuitTier::Global);
431        assert_eq!(cb.state(CircuitTier::Global), CircuitState::Open);
432        assert!(!cb.can_execute_any());
433    }
434
435    #[test]
436    fn tiered_circuit_half_open_recovery() {
437        let cb = TieredCircuitBreaker::new(1, Duration::from_millis(10));
438        cb.record_failure(CircuitTier::Connection);
439        assert_eq!(cb.state(CircuitTier::Connection), CircuitState::Open);
440        std::thread::sleep(Duration::from_millis(20));
441        assert!(cb.can_execute(CircuitTier::Connection));
442        assert_eq!(cb.state(CircuitTier::Connection), CircuitState::HalfOpen);
443        cb.record_success(CircuitTier::Connection);
444        assert_eq!(cb.state(CircuitTier::Connection), CircuitState::Closed);
445    }
446
447    #[test]
448    fn health_check_single_failure_no_evict() {
449        let checker = ConnectionHealthChecker::new(3);
450        assert!(!checker.record_failure("conn-1"));
451        assert!(!checker.should_evict("conn-1"));
452    }
453
454    #[test]
455    fn health_check_three_failures_evict() {
456        let checker = ConnectionHealthChecker::new(3);
457        checker.record_failure("conn-1");
458        checker.record_failure("conn-1");
459        assert!(checker.record_failure("conn-1"));
460        assert!(checker.should_evict("conn-1"));
461    }
462
463    #[test]
464    fn health_check_success_resets() {
465        let checker = ConnectionHealthChecker::new(3);
466        checker.record_failure("conn-1");
467        checker.record_failure("conn-1");
468        checker.record_success("conn-1");
469        assert_eq!(checker.failure_count("conn-1"), 0);
470        assert!(!checker.should_evict("conn-1"));
471    }
472
473    #[test]
474    fn scale_up_on_high_load() {
475        let controller = PoolElasticController::new(ElasticConfig::default());
476        let stats = PoolStats {
477            total_connections: 10,
478            idle_connections: 0,
479            waiting_requests: 20,
480        };
481        let event = controller.evaluate_and_scale(&stats);
482        assert!(event.is_some());
483        let event = event.unwrap();
484        assert_eq!(event.reason, ScaleReason::HighLoad);
485        assert!(event.to > event.from);
486    }
487
488    #[test]
489    fn scale_down_on_low_load() {
490        let config = ElasticConfig {
491            min_connections: 2,
492            max_connections: 50,
493            scale_up_threshold: 10,
494            scale_down_threshold: 5,
495            scale_interval: Duration::from_secs(1),
496            health_check_interval: Duration::from_secs(30),
497            health_check_sql: "SELECT 1".to_string(),
498        };
499        let controller = PoolElasticController::new(config);
500        *controller.current_size.lock().unwrap() = 20;
501        let stats = PoolStats {
502            total_connections: 20,
503            idle_connections: 18,
504            waiting_requests: 0,
505        };
506        let event = controller.evaluate_and_scale(&stats);
507        assert!(event.is_some());
508        let event = event.unwrap();
509        assert_eq!(event.reason, ScaleReason::LowLoad);
510        assert!(event.to < event.from);
511    }
512
513    #[test]
514    fn scale_respects_max() {
515        let config = ElasticConfig {
516            min_connections: 1,
517            max_connections: 15,
518            scale_up_threshold: 5,
519            scale_down_threshold: 20,
520            scale_interval: Duration::from_secs(1),
521            health_check_interval: Duration::from_secs(30),
522            health_check_sql: "SELECT 1".to_string(),
523        };
524        let controller = PoolElasticController::new(config);
525        *controller.current_size.lock().unwrap() = 10;
526        let stats = PoolStats {
527            total_connections: 10,
528            idle_connections: 0,
529            waiting_requests: 100,
530        };
531        let event = controller.evaluate_and_scale(&stats).unwrap();
532        assert_eq!(event.to, 15, "不应超过 max_connections");
533    }
534
535    #[test]
536    fn scale_respects_min() {
537        let config = ElasticConfig {
538            min_connections: 5,
539            max_connections: 50,
540            scale_up_threshold: 10,
541            scale_down_threshold: 3,
542            scale_interval: Duration::from_secs(1),
543            health_check_interval: Duration::from_secs(30),
544            health_check_sql: "SELECT 1".to_string(),
545        };
546        let controller = PoolElasticController::new(config);
547        *controller.current_size.lock().unwrap() = 8;
548        let stats = PoolStats {
549            total_connections: 8,
550            idle_connections: 7,
551            waiting_requests: 0,
552        };
553        let event = controller.evaluate_and_scale(&stats).unwrap();
554        assert_eq!(event.to, 5, "不应低于 min_connections");
555    }
556
557    #[test]
558    fn no_scale_when_stable() {
559        let controller = PoolElasticController::new(ElasticConfig::default());
560        let stats = PoolStats {
561            total_connections: 10,
562            idle_connections: 5,
563            waiting_requests: 3,
564        };
565        let event = controller.evaluate_and_scale(&stats);
566        assert!(event.is_none());
567    }
568
569    #[test]
570    fn prewarm_to_min() {
571        let config = ElasticConfig {
572            min_connections: 5,
573            max_connections: 50,
574            scale_up_threshold: 10,
575            scale_down_threshold: 20,
576            scale_interval: Duration::from_secs(1),
577            health_check_interval: Duration::from_secs(30),
578            health_check_sql: "SELECT 1".to_string(),
579        };
580        let controller = PoolElasticController::new(config);
581        let warmed = controller.prewarm(5);
582        assert_eq!(warmed, 5);
583        assert_eq!(controller.current_size(), 5);
584    }
585
586    #[test]
587    fn prewarm_clamps_to_max() {
588        let config = ElasticConfig {
589            min_connections: 1,
590            max_connections: 10,
591            scale_up_threshold: 10,
592            scale_down_threshold: 20,
593            scale_interval: Duration::from_secs(1),
594            health_check_interval: Duration::from_secs(30),
595            health_check_sql: "SELECT 1".to_string(),
596        };
597        let controller = PoolElasticController::new(config);
598        let warmed = controller.prewarm(100);
599        assert_eq!(warmed, 10);
600    }
601
602    #[test]
603    fn scale_events_recorded() {
604        let controller = PoolElasticController::new(ElasticConfig::default());
605        let stats = PoolStats {
606            total_connections: 5,
607            idle_connections: 0,
608            waiting_requests: 20,
609        };
610        controller.evaluate_and_scale(&stats);
611        assert_eq!(controller.scale_events().len(), 1);
612    }
613
614    #[test]
615    fn wiring_public_api() {
616        let controller = PoolElasticController::new(ElasticConfig::default());
617        assert!(controller.config().validate().is_ok());
618        assert!(controller
619            .circuit_breaker()
620            .can_execute(CircuitTier::Global));
621        assert_eq!(controller.current_size(), 5);
622    }
623}
sz_orm_core/pool_elastic.rs

sz_orm_core/
pool_elastic.rs
