sz_orm_core/
prewarm.rs

1//! # 连接池预热增强（v3.2.0）
2//!
3//! 提供自动预热、渐进式分批策略、预热进度可观测、多池统一预热。
4//!
5//! ## Feature Gate
6//!
7//! 本模块仅在 `auto-prewarm` feature 启用时编译。
8//! 手动预热 API（`Pool::prewarm()`）保持不变，向后兼容。
9//!
10//! ## 使用示例
11//!
12//! ```ignore
13//! use sz_orm_core::prewarm::{PrewarmConfig, ProgressiveConfig};
14//! use sz_orm_core::pool::{PoolConfigBuilder, Pool};
15//!
16//! let config = PoolConfigBuilder::new()
17//!     .max_size(20)
18//!     .min_idle(5)
19//!     .auto_prewarm(true)
20//!     .progressive_prewarm(ProgressiveConfig::default())
21//!     .build();
22//! // Pool::new 自动后台预热，Pool::new_async 等待预热完成
23//! ```
24
25use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
26use std::time::Duration;
27
28// ============================================================================
29// 配置结构体
30// ============================================================================
31
32/// 渐进式预热配置
33#[derive(Debug, Clone)]
34pub struct ProgressiveConfig {
35    /// 每批创建连接数
36    pub batch_size: u32,
37    /// 批间隔（避免瞬时冲击 DB）
38    pub interval: Duration,
39    /// 总超时
40    pub total_timeout: Duration,
41}
42
43impl Default for ProgressiveConfig {
44    fn default() -> Self {
45        Self {
46            batch_size: 2,
47            interval: Duration::from_millis(10),
48            total_timeout: Duration::from_secs(30),
49        }
50    }
51}
52
53impl ProgressiveConfig {
54    /// 创建渐进式预热配置
55    pub fn new(batch_size: u32, interval: Duration, total_timeout: Duration) -> Self {
56        Self {
57            batch_size: batch_size.max(1),
58            interval,
59            total_timeout,
60        }
61    }
62
63    /// 设置每批创建连接数
64    pub fn with_batch_size(mut self, size: u32) -> Self {
65        self.batch_size = size.max(1);
66        self
67    }
68
69    /// 设置批间隔
70    pub fn with_interval(mut self, interval: Duration) -> Self {
71        self.interval = interval;
72        self
73    }
74
75    /// 设置总超时
76    pub fn with_total_timeout(mut self, timeout: Duration) -> Self {
77        self.total_timeout = timeout;
78        self
79    }
80}
81
82/// 预热配置
83#[derive(Debug, Clone, Default)]
84pub struct PrewarmConfig {
85    /// 是否自动预热
86    pub auto_prewarm: bool,
87    /// 渐进式配置（None 表示一次性预热）
88    pub progressive: Option<ProgressiveConfig>,
89}
90
91impl PrewarmConfig {
92    /// 创建默认预热配置
93    pub fn new() -> Self {
94        Self::default()
95    }
96
97    /// 设置是否自动预热
98    pub fn with_auto_prewarm(mut self, enabled: bool) -> Self {
99        self.auto_prewarm = enabled;
100        self
101    }
102
103    /// 设置渐进式配置
104    pub fn with_progressive(mut self, config: ProgressiveConfig) -> Self {
105        self.progressive = Some(config);
106        self
107    }
108}
109
110// ============================================================================
111// 进度指标
112// ============================================================================
113
114/// 预热进度（无锁原子计数器）
115#[derive(Debug)]
116pub struct PrewarmProgress {
117    warmed: AtomicU32,
118    target: u32,
119    failed: AtomicU32,
120    elapsed_ns: AtomicU64,
121    is_completed: AtomicBool,
122}
123
124impl PrewarmProgress {
125    /// 创建预热进度实例
126    pub fn new(target: u32) -> Self {
127        Self {
128            warmed: AtomicU32::new(0),
129            target,
130            failed: AtomicU32::new(0),
131            elapsed_ns: AtomicU64::new(0),
132            is_completed: AtomicBool::new(false),
133        }
134    }
135
136    /// 记录一次成功预热
137    pub fn record_success(&self) {
138        self.warmed.fetch_add(1, Ordering::Relaxed);
139    }
140
141    /// 记录一次失败预热
142    pub fn record_failure(&self) {
143        self.failed.fetch_add(1, Ordering::Relaxed);
144    }
145
146    /// 设置已耗时
147    pub fn set_elapsed(&self, duration: Duration) {
148        self.elapsed_ns
149            .store(duration.as_nanos() as u64, Ordering::Relaxed);
150    }
151
152    /// 标记预热完成
153    pub fn mark_completed(&self) {
154        self.is_completed.store(true, Ordering::Release);
155    }
156
157    /// 获取进度快照
158    pub fn snapshot(&self) -> PrewarmProgressSnapshot {
159        PrewarmProgressSnapshot {
160            warmed: self.warmed.load(Ordering::Relaxed),
161            target: self.target,
162            failed: self.failed.load(Ordering::Relaxed),
163            elapsed: Duration::from_nanos(self.elapsed_ns.load(Ordering::Relaxed)),
164            is_completed: self.is_completed.load(Ordering::Acquire),
165        }
166    }
167}
168
169/// 预热进度快照
170#[derive(Debug, Clone)]
171pub struct PrewarmProgressSnapshot {
172    /// 已成功预热数
173    pub warmed: u32,
174    /// 目标连接数
175    pub target: u32,
176    /// 失败次数
177    pub failed: u32,
178    /// 已耗时
179    pub elapsed: Duration,
180    /// 是否已完成
181    pub is_completed: bool,
182}
183
184impl PrewarmProgressSnapshot {
185    /// 进度百分比（0.0 ~ 1.0）
186    pub fn percent(&self) -> f64 {
187        if self.target == 0 {
188            1.0
189        } else {
190            (self.warmed + self.failed) as f64 / self.target as f64
191        }
192    }
193
194    /// 是否全部成功
195    pub fn all_succeeded(&self) -> bool {
196        self.is_completed && self.failed == 0 && self.warmed == self.target
197    }
198}
199
200// ============================================================================
201// 多池统一预热汇总
202// ============================================================================
203
204/// 单个后端预热结果
205#[derive(Debug, Clone)]
206pub struct BackendPrewarmResult {
207    /// 后端名称
208    pub backend: String,
209    /// 已成功预热数
210    pub warmed: u32,
211    /// 失败次数
212    pub failed: u32,
213    /// 已耗时
214    pub elapsed: Duration,
215    /// 错误信息列表
216    pub errors: Vec<String>,
217}
218
219/// 多池统一预热汇总
220#[derive(Debug, Clone)]
221pub struct PrewarmSummary {
222    /// 各后端预热结果
223    pub results: Vec<BackendPrewarmResult>,
224}
225
226impl PrewarmSummary {
227    /// 创建空的预热汇总
228    pub fn new() -> Self {
229        Self {
230            results: Vec::new(),
231        }
232    }
233
234    /// 添加一个后端的预热结果
235    pub fn add(&mut self, result: BackendPrewarmResult) {
236        self.results.push(result);
237    }
238
239    /// 所有后端成功预热总数
240    pub fn total_warmed(&self) -> u32 {
241        self.results.iter().map(|r| r.warmed).sum()
242    }
243
244    /// 所有后端失败总数
245    pub fn total_failed(&self) -> u32 {
246        self.results.iter().map(|r| r.failed).sum()
247    }
248
249    /// 是否全部成功
250    pub fn all_succeeded(&self) -> bool {
251        !self.results.is_empty() && self.results.iter().all(|r| r.failed == 0)
252    }
253}
254
255impl Default for PrewarmSummary {
256    fn default() -> Self {
257        Self::new()
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn test_prewarm_config_defaults() {
267        let config = PrewarmConfig::default();
268        assert!(!config.auto_prewarm);
269        assert!(config.progressive.is_none());
270    }
271
272    #[test]
273    fn test_prewarm_config_builders() {
274        let config = PrewarmConfig::new()
275            .with_auto_prewarm(true)
276            .with_progressive(ProgressiveConfig::default());
277        assert!(config.auto_prewarm);
278        assert!(config.progressive.is_some());
279    }
280
281    #[test]
282    fn test_progressive_config_defaults() {
283        let config = ProgressiveConfig::default();
284        assert_eq!(config.batch_size, 2);
285        assert_eq!(config.interval, Duration::from_millis(10));
286        assert_eq!(config.total_timeout, Duration::from_secs(30));
287    }
288
289    #[test]
290    fn test_progressive_config_batch_size_min_1() {
291        let config = ProgressiveConfig::new(0, Duration::from_millis(5), Duration::from_secs(10));
292        assert_eq!(config.batch_size, 1);
293    }
294
295    #[test]
296    fn test_prewarm_progress_snapshot() {
297        let progress = PrewarmProgress::new(10);
298        progress.record_success();
299        progress.record_success();
300        progress.record_failure();
301        progress.set_elapsed(Duration::from_millis(100));
302        progress.mark_completed();
303
304        let snap = progress.snapshot();
305        assert_eq!(snap.warmed, 2);
306        assert_eq!(snap.target, 10);
307        assert_eq!(snap.failed, 1);
308        assert_eq!(snap.elapsed, Duration::from_millis(100));
309        assert!(snap.is_completed);
310        assert!((snap.percent() - 0.3).abs() < 0.001);
311    }
312
313    #[test]
314    fn test_prewarm_progress_all_succeeded() {
315        let progress = PrewarmProgress::new(3);
316        progress.record_success();
317        progress.record_success();
318        progress.record_success();
319        progress.mark_completed();
320
321        let snap = progress.snapshot();
322        assert!(snap.all_succeeded());
323    }
324
325    #[test]
326    fn test_prewarm_progress_not_all_succeeded_with_failure() {
327        let progress = PrewarmProgress::new(3);
328        progress.record_success();
329        progress.record_success();
330        progress.record_failure();
331        progress.mark_completed();
332
333        let snap = progress.snapshot();
334        assert!(!snap.all_succeeded());
335    }
336
337    #[test]
338    fn test_prewarm_summary_aggregation() {
339        let mut summary = PrewarmSummary::new();
340        summary.add(BackendPrewarmResult {
341            backend: "mysql".into(),
342            warmed: 5,
343            failed: 0,
344            elapsed: Duration::from_millis(50),
345            errors: vec![],
346        });
347        summary.add(BackendPrewarmResult {
348            backend: "pg".into(),
349            warmed: 3,
350            failed: 1,
351            elapsed: Duration::from_millis(40),
352            errors: vec!["connection refused".into()],
353        });
354
355        assert_eq!(summary.total_warmed(), 8);
356        assert_eq!(summary.total_failed(), 1);
357        assert!(!summary.all_succeeded());
358    }
359
360    #[test]
361    fn test_prewarm_summary_all_succeeded() {
362        let mut summary = PrewarmSummary::new();
363        summary.add(BackendPrewarmResult {
364            backend: "mysql".into(),
365            warmed: 5,
366            failed: 0,
367            elapsed: Duration::from_millis(50),
368            errors: vec![],
369        });
370        summary.add(BackendPrewarmResult {
371            backend: "pg".into(),
372            warmed: 3,
373            failed: 0,
374            elapsed: Duration::from_millis(40),
375            errors: vec![],
376        });
377
378        assert_eq!(summary.total_warmed(), 8);
379        assert_eq!(summary.total_failed(), 0);
380        assert!(summary.all_succeeded());
381    }
382
383    #[test]
384    fn test_prewarm_summary_empty() {
385        let summary = PrewarmSummary::new();
386        assert_eq!(summary.total_warmed(), 0);
387        assert_eq!(summary.total_failed(), 0);
388        assert!(!summary.all_succeeded());
389    }
390
391    #[test]
392    fn test_progressive_config_builders() {
393        let config = ProgressiveConfig::default()
394            .with_batch_size(5)
395            .with_interval(Duration::from_millis(20))
396            .with_total_timeout(Duration::from_secs(60));
397        assert_eq!(config.batch_size, 5);
398        assert_eq!(config.interval, Duration::from_millis(20));
399        assert_eq!(config.total_timeout, Duration::from_secs(60));
400    }
401
402    #[test]
403    fn test_progressive_config_with_batch_size_min_1() {
404        let config = ProgressiveConfig::default().with_batch_size(0);
405        assert_eq!(config.batch_size, 1);
406    }
407
408    #[test]
409    fn test_progressive_config_interval_zero() {
410        let config = ProgressiveConfig::new(2, Duration::ZERO, Duration::from_secs(10));
411        assert_eq!(config.interval, Duration::ZERO);
412    }
413
414    #[test]
415    fn test_progressive_config_total_timeout_zero() {
416        let config = ProgressiveConfig::new(2, Duration::from_millis(5), Duration::ZERO);
417        assert_eq!(config.total_timeout, Duration::ZERO);
418    }
419
420    #[test]
421    fn test_prewarm_progress_percent_zero() {
422        let progress = PrewarmProgress::new(5);
423        let snap = progress.snapshot();
424        assert!((snap.percent() - 0.0).abs() < 0.001);
425    }
426
427    #[test]
428    fn test_prewarm_progress_percent_full() {
429        let progress = PrewarmProgress::new(3);
430        progress.record_success();
431        progress.record_success();
432        progress.record_success();
433        progress.mark_completed();
434        let snap = progress.snapshot();
435        assert!((snap.percent() - 1.0).abs() < 0.001);
436    }
437
438    #[test]
439    fn test_prewarm_progress_warmed_plus_failed_le_target() {
440        let progress = PrewarmProgress::new(10);
441        for _ in 0..7 {
442            progress.record_success();
443        }
444        for _ in 0..3 {
445            progress.record_failure();
446        }
447        progress.mark_completed();
448        let snap = progress.snapshot();
449        assert!(snap.warmed + snap.failed <= snap.target);
450        assert_eq!(snap.warmed + snap.failed, 10);
451    }
452
453    #[test]
454    fn test_prewarm_progress_target_zero() {
455        let progress = PrewarmProgress::new(0);
456        let snap = progress.snapshot();
457        assert_eq!(snap.target, 0);
458        assert!(
459            (snap.percent() - 1.0).abs() < 0.001,
460            "target=0 时 percent 应为 1.0"
461        );
462    }
463
464    #[test]
465    fn test_backend_prewarm_result_fields() {
466        let result = BackendPrewarmResult {
467            backend: "mysql".into(),
468            warmed: 10,
469            failed: 2,
470            elapsed: Duration::from_millis(200),
471            errors: vec!["timeout".into(), "refused".into()],
472        };
473        assert_eq!(result.backend, "mysql");
474        assert_eq!(result.warmed, 10);
475        assert_eq!(result.failed, 2);
476        assert_eq!(result.errors.len(), 2);
477    }
478
479    #[test]
480    fn test_prewarm_summary_partial_failure() {
481        let mut summary = PrewarmSummary::new();
482        summary.add(BackendPrewarmResult {
483            backend: "mysql".into(),
484            warmed: 5,
485            failed: 0,
486            elapsed: Duration::from_millis(50),
487            errors: vec![],
488        });
489        summary.add(BackendPrewarmResult {
490            backend: "oracle".into(),
491            warmed: 0,
492            failed: 3,
493            elapsed: Duration::from_millis(30),
494            errors: vec!["unreachable".into()],
495        });
496        assert_eq!(summary.total_warmed(), 5);
497        assert_eq!(summary.total_failed(), 3);
498        assert!(!summary.all_succeeded());
499        assert_eq!(summary.results.len(), 2);
500    }
501}
sz_orm_core/prewarm.rs

sz_orm_core/
prewarm.rs
