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    pub fn new(batch_size: u32, interval: Duration, total_timeout: Duration) -> Self {
55        Self {
56            batch_size: batch_size.max(1),
57            interval,
58            total_timeout,
59        }
60    }
61
62    pub fn with_batch_size(mut self, size: u32) -> Self {
63        self.batch_size = size.max(1);
64        self
65    }
66
67    pub fn with_interval(mut self, interval: Duration) -> Self {
68        self.interval = interval;
69        self
70    }
71
72    pub fn with_total_timeout(mut self, timeout: Duration) -> Self {
73        self.total_timeout = timeout;
74        self
75    }
76}
77
78/// 预热配置
79#[derive(Debug, Clone, Default)]
80pub struct PrewarmConfig {
81    /// 是否自动预热
82    pub auto_prewarm: bool,
83    /// 渐进式配置（None 表示一次性预热）
84    pub progressive: Option<ProgressiveConfig>,
85}
86
87impl PrewarmConfig {
88    pub fn new() -> Self {
89        Self::default()
90    }
91
92    pub fn with_auto_prewarm(mut self, enabled: bool) -> Self {
93        self.auto_prewarm = enabled;
94        self
95    }
96
97    pub fn with_progressive(mut self, config: ProgressiveConfig) -> Self {
98        self.progressive = Some(config);
99        self
100    }
101}
102
103// ============================================================================
104// 进度指标
105// ============================================================================
106
107/// 预热进度（无锁原子计数器）
108#[derive(Debug)]
109pub struct PrewarmProgress {
110    warmed: AtomicU32,
111    target: u32,
112    failed: AtomicU32,
113    elapsed_ns: AtomicU64,
114    is_completed: AtomicBool,
115}
116
117impl PrewarmProgress {
118    pub fn new(target: u32) -> Self {
119        Self {
120            warmed: AtomicU32::new(0),
121            target,
122            failed: AtomicU32::new(0),
123            elapsed_ns: AtomicU64::new(0),
124            is_completed: AtomicBool::new(false),
125        }
126    }
127
128    pub fn record_success(&self) {
129        self.warmed.fetch_add(1, Ordering::Relaxed);
130    }
131
132    pub fn record_failure(&self) {
133        self.failed.fetch_add(1, Ordering::Relaxed);
134    }
135
136    pub fn set_elapsed(&self, duration: Duration) {
137        self.elapsed_ns
138            .store(duration.as_nanos() as u64, Ordering::Relaxed);
139    }
140
141    pub fn mark_completed(&self) {
142        self.is_completed.store(true, Ordering::Release);
143    }
144
145    pub fn snapshot(&self) -> PrewarmProgressSnapshot {
146        PrewarmProgressSnapshot {
147            warmed: self.warmed.load(Ordering::Relaxed),
148            target: self.target,
149            failed: self.failed.load(Ordering::Relaxed),
150            elapsed: Duration::from_nanos(self.elapsed_ns.load(Ordering::Relaxed)),
151            is_completed: self.is_completed.load(Ordering::Acquire),
152        }
153    }
154}
155
156/// 预热进度快照
157#[derive(Debug, Clone)]
158pub struct PrewarmProgressSnapshot {
159    pub warmed: u32,
160    pub target: u32,
161    pub failed: u32,
162    pub elapsed: Duration,
163    pub is_completed: bool,
164}
165
166impl PrewarmProgressSnapshot {
167    /// 进度百分比（0.0 ~ 1.0）
168    pub fn percent(&self) -> f64 {
169        if self.target == 0 {
170            1.0
171        } else {
172            (self.warmed + self.failed) as f64 / self.target as f64
173        }
174    }
175
176    /// 是否全部成功
177    pub fn all_succeeded(&self) -> bool {
178        self.is_completed && self.failed == 0 && self.warmed == self.target
179    }
180}
181
182// ============================================================================
183// 多池统一预热汇总
184// ============================================================================
185
186/// 单个后端预热结果
187#[derive(Debug, Clone)]
188pub struct BackendPrewarmResult {
189    pub backend: String,
190    pub warmed: u32,
191    pub failed: u32,
192    pub elapsed: Duration,
193    pub errors: Vec<String>,
194}
195
196/// 多池统一预热汇总
197#[derive(Debug, Clone)]
198pub struct PrewarmSummary {
199    pub results: Vec<BackendPrewarmResult>,
200}
201
202impl PrewarmSummary {
203    pub fn new() -> Self {
204        Self {
205            results: Vec::new(),
206        }
207    }
208
209    pub fn add(&mut self, result: BackendPrewarmResult) {
210        self.results.push(result);
211    }
212
213    pub fn total_warmed(&self) -> u32 {
214        self.results.iter().map(|r| r.warmed).sum()
215    }
216
217    pub fn total_failed(&self) -> u32 {
218        self.results.iter().map(|r| r.failed).sum()
219    }
220
221    pub fn all_succeeded(&self) -> bool {
222        !self.results.is_empty() && self.results.iter().all(|r| r.failed == 0)
223    }
224}
225
226impl Default for PrewarmSummary {
227    fn default() -> Self {
228        Self::new()
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn test_prewarm_config_defaults() {
238        let config = PrewarmConfig::default();
239        assert!(!config.auto_prewarm);
240        assert!(config.progressive.is_none());
241    }
242
243    #[test]
244    fn test_prewarm_config_builders() {
245        let config = PrewarmConfig::new()
246            .with_auto_prewarm(true)
247            .with_progressive(ProgressiveConfig::default());
248        assert!(config.auto_prewarm);
249        assert!(config.progressive.is_some());
250    }
251
252    #[test]
253    fn test_progressive_config_defaults() {
254        let config = ProgressiveConfig::default();
255        assert_eq!(config.batch_size, 2);
256        assert_eq!(config.interval, Duration::from_millis(10));
257        assert_eq!(config.total_timeout, Duration::from_secs(30));
258    }
259
260    #[test]
261    fn test_progressive_config_batch_size_min_1() {
262        let config = ProgressiveConfig::new(0, Duration::from_millis(5), Duration::from_secs(10));
263        assert_eq!(config.batch_size, 1);
264    }
265
266    #[test]
267    fn test_prewarm_progress_snapshot() {
268        let progress = PrewarmProgress::new(10);
269        progress.record_success();
270        progress.record_success();
271        progress.record_failure();
272        progress.set_elapsed(Duration::from_millis(100));
273        progress.mark_completed();
274
275        let snap = progress.snapshot();
276        assert_eq!(snap.warmed, 2);
277        assert_eq!(snap.target, 10);
278        assert_eq!(snap.failed, 1);
279        assert_eq!(snap.elapsed, Duration::from_millis(100));
280        assert!(snap.is_completed);
281        assert!((snap.percent() - 0.3).abs() < 0.001);
282    }
283
284    #[test]
285    fn test_prewarm_progress_all_succeeded() {
286        let progress = PrewarmProgress::new(3);
287        progress.record_success();
288        progress.record_success();
289        progress.record_success();
290        progress.mark_completed();
291
292        let snap = progress.snapshot();
293        assert!(snap.all_succeeded());
294    }
295
296    #[test]
297    fn test_prewarm_progress_not_all_succeeded_with_failure() {
298        let progress = PrewarmProgress::new(3);
299        progress.record_success();
300        progress.record_success();
301        progress.record_failure();
302        progress.mark_completed();
303
304        let snap = progress.snapshot();
305        assert!(!snap.all_succeeded());
306    }
307
308    #[test]
309    fn test_prewarm_summary_aggregation() {
310        let mut summary = PrewarmSummary::new();
311        summary.add(BackendPrewarmResult {
312            backend: "mysql".into(),
313            warmed: 5,
314            failed: 0,
315            elapsed: Duration::from_millis(50),
316            errors: vec![],
317        });
318        summary.add(BackendPrewarmResult {
319            backend: "pg".into(),
320            warmed: 3,
321            failed: 1,
322            elapsed: Duration::from_millis(40),
323            errors: vec!["connection refused".into()],
324        });
325
326        assert_eq!(summary.total_warmed(), 8);
327        assert_eq!(summary.total_failed(), 1);
328        assert!(!summary.all_succeeded());
329    }
330
331    #[test]
332    fn test_prewarm_summary_all_succeeded() {
333        let mut summary = PrewarmSummary::new();
334        summary.add(BackendPrewarmResult {
335            backend: "mysql".into(),
336            warmed: 5,
337            failed: 0,
338            elapsed: Duration::from_millis(50),
339            errors: vec![],
340        });
341        summary.add(BackendPrewarmResult {
342            backend: "pg".into(),
343            warmed: 3,
344            failed: 0,
345            elapsed: Duration::from_millis(40),
346            errors: vec![],
347        });
348
349        assert_eq!(summary.total_warmed(), 8);
350        assert_eq!(summary.total_failed(), 0);
351        assert!(summary.all_succeeded());
352    }
353
354    #[test]
355    fn test_prewarm_summary_empty() {
356        let summary = PrewarmSummary::new();
357        assert_eq!(summary.total_warmed(), 0);
358        assert_eq!(summary.total_failed(), 0);
359        assert!(!summary.all_succeeded());
360    }
361
362    #[test]
363    fn test_progressive_config_builders() {
364        let config = ProgressiveConfig::default()
365            .with_batch_size(5)
366            .with_interval(Duration::from_millis(20))
367            .with_total_timeout(Duration::from_secs(60));
368        assert_eq!(config.batch_size, 5);
369        assert_eq!(config.interval, Duration::from_millis(20));
370        assert_eq!(config.total_timeout, Duration::from_secs(60));
371    }
372
373    #[test]
374    fn test_progressive_config_with_batch_size_min_1() {
375        let config = ProgressiveConfig::default().with_batch_size(0);
376        assert_eq!(config.batch_size, 1);
377    }
378
379    #[test]
380    fn test_progressive_config_interval_zero() {
381        let config = ProgressiveConfig::new(2, Duration::ZERO, Duration::from_secs(10));
382        assert_eq!(config.interval, Duration::ZERO);
383    }
384
385    #[test]
386    fn test_progressive_config_total_timeout_zero() {
387        let config = ProgressiveConfig::new(2, Duration::from_millis(5), Duration::ZERO);
388        assert_eq!(config.total_timeout, Duration::ZERO);
389    }
390
391    #[test]
392    fn test_prewarm_progress_percent_zero() {
393        let progress = PrewarmProgress::new(5);
394        let snap = progress.snapshot();
395        assert!((snap.percent() - 0.0).abs() < 0.001);
396    }
397
398    #[test]
399    fn test_prewarm_progress_percent_full() {
400        let progress = PrewarmProgress::new(3);
401        progress.record_success();
402        progress.record_success();
403        progress.record_success();
404        progress.mark_completed();
405        let snap = progress.snapshot();
406        assert!((snap.percent() - 1.0).abs() < 0.001);
407    }
408
409    #[test]
410    fn test_prewarm_progress_warmed_plus_failed_le_target() {
411        let progress = PrewarmProgress::new(10);
412        for _ in 0..7 {
413            progress.record_success();
414        }
415        for _ in 0..3 {
416            progress.record_failure();
417        }
418        progress.mark_completed();
419        let snap = progress.snapshot();
420        assert!(snap.warmed + snap.failed <= snap.target);
421        assert_eq!(snap.warmed + snap.failed, 10);
422    }
423
424    #[test]
425    fn test_prewarm_progress_target_zero() {
426        let progress = PrewarmProgress::new(0);
427        let snap = progress.snapshot();
428        assert_eq!(snap.target, 0);
429        assert!(
430            (snap.percent() - 1.0).abs() < 0.001,
431            "target=0 时 percent 应为 1.0"
432        );
433    }
434
435    #[test]
436    fn test_backend_prewarm_result_fields() {
437        let result = BackendPrewarmResult {
438            backend: "mysql".into(),
439            warmed: 10,
440            failed: 2,
441            elapsed: Duration::from_millis(200),
442            errors: vec!["timeout".into(), "refused".into()],
443        };
444        assert_eq!(result.backend, "mysql");
445        assert_eq!(result.warmed, 10);
446        assert_eq!(result.failed, 2);
447        assert_eq!(result.errors.len(), 2);
448    }
449
450    #[test]
451    fn test_prewarm_summary_partial_failure() {
452        let mut summary = PrewarmSummary::new();
453        summary.add(BackendPrewarmResult {
454            backend: "mysql".into(),
455            warmed: 5,
456            failed: 0,
457            elapsed: Duration::from_millis(50),
458            errors: vec![],
459        });
460        summary.add(BackendPrewarmResult {
461            backend: "oracle".into(),
462            warmed: 0,
463            failed: 3,
464            elapsed: Duration::from_millis(30),
465            errors: vec!["unreachable".into()],
466        });
467        assert_eq!(summary.total_warmed(), 5);
468        assert_eq!(summary.total_failed(), 3);
469        assert!(!summary.all_succeeded());
470        assert_eq!(summary.results.len(), 2);
471    }
472}
sz_orm_core/prewarm.rs

sz_orm_core/
prewarm.rs
