use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct ProgressiveConfig {
pub batch_size: u32,
pub interval: Duration,
pub total_timeout: Duration,
}
impl Default for ProgressiveConfig {
fn default() -> Self {
Self {
batch_size: 2,
interval: Duration::from_millis(10),
total_timeout: Duration::from_secs(30),
}
}
}
impl ProgressiveConfig {
pub fn new(batch_size: u32, interval: Duration, total_timeout: Duration) -> Self {
Self {
batch_size: batch_size.max(1),
interval,
total_timeout,
}
}
pub fn with_batch_size(mut self, size: u32) -> Self {
self.batch_size = size.max(1);
self
}
pub fn with_interval(mut self, interval: Duration) -> Self {
self.interval = interval;
self
}
pub fn with_total_timeout(mut self, timeout: Duration) -> Self {
self.total_timeout = timeout;
self
}
}
#[derive(Debug, Clone, Default)]
pub struct PrewarmConfig {
pub auto_prewarm: bool,
pub progressive: Option<ProgressiveConfig>,
}
impl PrewarmConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_auto_prewarm(mut self, enabled: bool) -> Self {
self.auto_prewarm = enabled;
self
}
pub fn with_progressive(mut self, config: ProgressiveConfig) -> Self {
self.progressive = Some(config);
self
}
}
#[derive(Debug)]
pub struct PrewarmProgress {
warmed: AtomicU32,
target: u32,
failed: AtomicU32,
elapsed_ns: AtomicU64,
is_completed: AtomicBool,
}
impl PrewarmProgress {
pub fn new(target: u32) -> Self {
Self {
warmed: AtomicU32::new(0),
target,
failed: AtomicU32::new(0),
elapsed_ns: AtomicU64::new(0),
is_completed: AtomicBool::new(false),
}
}
pub fn record_success(&self) {
self.warmed.fetch_add(1, Ordering::Relaxed);
}
pub fn record_failure(&self) {
self.failed.fetch_add(1, Ordering::Relaxed);
}
pub fn set_elapsed(&self, duration: Duration) {
self.elapsed_ns
.store(duration.as_nanos() as u64, Ordering::Relaxed);
}
pub fn mark_completed(&self) {
self.is_completed.store(true, Ordering::Release);
}
pub fn snapshot(&self) -> PrewarmProgressSnapshot {
PrewarmProgressSnapshot {
warmed: self.warmed.load(Ordering::Relaxed),
target: self.target,
failed: self.failed.load(Ordering::Relaxed),
elapsed: Duration::from_nanos(self.elapsed_ns.load(Ordering::Relaxed)),
is_completed: self.is_completed.load(Ordering::Acquire),
}
}
}
#[derive(Debug, Clone)]
pub struct PrewarmProgressSnapshot {
pub warmed: u32,
pub target: u32,
pub failed: u32,
pub elapsed: Duration,
pub is_completed: bool,
}
impl PrewarmProgressSnapshot {
pub fn percent(&self) -> f64 {
if self.target == 0 {
1.0
} else {
(self.warmed + self.failed) as f64 / self.target as f64
}
}
pub fn all_succeeded(&self) -> bool {
self.is_completed && self.failed == 0 && self.warmed == self.target
}
}
#[derive(Debug, Clone)]
pub struct BackendPrewarmResult {
pub backend: String,
pub warmed: u32,
pub failed: u32,
pub elapsed: Duration,
pub errors: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct PrewarmSummary {
pub results: Vec<BackendPrewarmResult>,
}
impl PrewarmSummary {
pub fn new() -> Self {
Self {
results: Vec::new(),
}
}
pub fn add(&mut self, result: BackendPrewarmResult) {
self.results.push(result);
}
pub fn total_warmed(&self) -> u32 {
self.results.iter().map(|r| r.warmed).sum()
}
pub fn total_failed(&self) -> u32 {
self.results.iter().map(|r| r.failed).sum()
}
pub fn all_succeeded(&self) -> bool {
!self.results.is_empty() && self.results.iter().all(|r| r.failed == 0)
}
}
impl Default for PrewarmSummary {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prewarm_config_defaults() {
let config = PrewarmConfig::default();
assert!(!config.auto_prewarm);
assert!(config.progressive.is_none());
}
#[test]
fn test_prewarm_config_builders() {
let config = PrewarmConfig::new()
.with_auto_prewarm(true)
.with_progressive(ProgressiveConfig::default());
assert!(config.auto_prewarm);
assert!(config.progressive.is_some());
}
#[test]
fn test_progressive_config_defaults() {
let config = ProgressiveConfig::default();
assert_eq!(config.batch_size, 2);
assert_eq!(config.interval, Duration::from_millis(10));
assert_eq!(config.total_timeout, Duration::from_secs(30));
}
#[test]
fn test_progressive_config_batch_size_min_1() {
let config = ProgressiveConfig::new(0, Duration::from_millis(5), Duration::from_secs(10));
assert_eq!(config.batch_size, 1);
}
#[test]
fn test_prewarm_progress_snapshot() {
let progress = PrewarmProgress::new(10);
progress.record_success();
progress.record_success();
progress.record_failure();
progress.set_elapsed(Duration::from_millis(100));
progress.mark_completed();
let snap = progress.snapshot();
assert_eq!(snap.warmed, 2);
assert_eq!(snap.target, 10);
assert_eq!(snap.failed, 1);
assert_eq!(snap.elapsed, Duration::from_millis(100));
assert!(snap.is_completed);
assert!((snap.percent() - 0.3).abs() < 0.001);
}
#[test]
fn test_prewarm_progress_all_succeeded() {
let progress = PrewarmProgress::new(3);
progress.record_success();
progress.record_success();
progress.record_success();
progress.mark_completed();
let snap = progress.snapshot();
assert!(snap.all_succeeded());
}
#[test]
fn test_prewarm_progress_not_all_succeeded_with_failure() {
let progress = PrewarmProgress::new(3);
progress.record_success();
progress.record_success();
progress.record_failure();
progress.mark_completed();
let snap = progress.snapshot();
assert!(!snap.all_succeeded());
}
#[test]
fn test_prewarm_summary_aggregation() {
let mut summary = PrewarmSummary::new();
summary.add(BackendPrewarmResult {
backend: "mysql".into(),
warmed: 5,
failed: 0,
elapsed: Duration::from_millis(50),
errors: vec![],
});
summary.add(BackendPrewarmResult {
backend: "pg".into(),
warmed: 3,
failed: 1,
elapsed: Duration::from_millis(40),
errors: vec!["connection refused".into()],
});
assert_eq!(summary.total_warmed(), 8);
assert_eq!(summary.total_failed(), 1);
assert!(!summary.all_succeeded());
}
#[test]
fn test_prewarm_summary_all_succeeded() {
let mut summary = PrewarmSummary::new();
summary.add(BackendPrewarmResult {
backend: "mysql".into(),
warmed: 5,
failed: 0,
elapsed: Duration::from_millis(50),
errors: vec![],
});
summary.add(BackendPrewarmResult {
backend: "pg".into(),
warmed: 3,
failed: 0,
elapsed: Duration::from_millis(40),
errors: vec![],
});
assert_eq!(summary.total_warmed(), 8);
assert_eq!(summary.total_failed(), 0);
assert!(summary.all_succeeded());
}
#[test]
fn test_prewarm_summary_empty() {
let summary = PrewarmSummary::new();
assert_eq!(summary.total_warmed(), 0);
assert_eq!(summary.total_failed(), 0);
assert!(!summary.all_succeeded());
}
#[test]
fn test_progressive_config_builders() {
let config = ProgressiveConfig::default()
.with_batch_size(5)
.with_interval(Duration::from_millis(20))
.with_total_timeout(Duration::from_secs(60));
assert_eq!(config.batch_size, 5);
assert_eq!(config.interval, Duration::from_millis(20));
assert_eq!(config.total_timeout, Duration::from_secs(60));
}
#[test]
fn test_progressive_config_with_batch_size_min_1() {
let config = ProgressiveConfig::default().with_batch_size(0);
assert_eq!(config.batch_size, 1);
}
#[test]
fn test_progressive_config_interval_zero() {
let config = ProgressiveConfig::new(2, Duration::ZERO, Duration::from_secs(10));
assert_eq!(config.interval, Duration::ZERO);
}
#[test]
fn test_progressive_config_total_timeout_zero() {
let config = ProgressiveConfig::new(2, Duration::from_millis(5), Duration::ZERO);
assert_eq!(config.total_timeout, Duration::ZERO);
}
#[test]
fn test_prewarm_progress_percent_zero() {
let progress = PrewarmProgress::new(5);
let snap = progress.snapshot();
assert!((snap.percent() - 0.0).abs() < 0.001);
}
#[test]
fn test_prewarm_progress_percent_full() {
let progress = PrewarmProgress::new(3);
progress.record_success();
progress.record_success();
progress.record_success();
progress.mark_completed();
let snap = progress.snapshot();
assert!((snap.percent() - 1.0).abs() < 0.001);
}
#[test]
fn test_prewarm_progress_warmed_plus_failed_le_target() {
let progress = PrewarmProgress::new(10);
for _ in 0..7 {
progress.record_success();
}
for _ in 0..3 {
progress.record_failure();
}
progress.mark_completed();
let snap = progress.snapshot();
assert!(snap.warmed + snap.failed <= snap.target);
assert_eq!(snap.warmed + snap.failed, 10);
}
#[test]
fn test_prewarm_progress_target_zero() {
let progress = PrewarmProgress::new(0);
let snap = progress.snapshot();
assert_eq!(snap.target, 0);
assert!(
(snap.percent() - 1.0).abs() < 0.001,
"target=0 时 percent 应为 1.0"
);
}
#[test]
fn test_backend_prewarm_result_fields() {
let result = BackendPrewarmResult {
backend: "mysql".into(),
warmed: 10,
failed: 2,
elapsed: Duration::from_millis(200),
errors: vec!["timeout".into(), "refused".into()],
};
assert_eq!(result.backend, "mysql");
assert_eq!(result.warmed, 10);
assert_eq!(result.failed, 2);
assert_eq!(result.errors.len(), 2);
}
#[test]
fn test_prewarm_summary_partial_failure() {
let mut summary = PrewarmSummary::new();
summary.add(BackendPrewarmResult {
backend: "mysql".into(),
warmed: 5,
failed: 0,
elapsed: Duration::from_millis(50),
errors: vec![],
});
summary.add(BackendPrewarmResult {
backend: "oracle".into(),
warmed: 0,
failed: 3,
elapsed: Duration::from_millis(30),
errors: vec!["unreachable".into()],
});
assert_eq!(summary.total_warmed(), 5);
assert_eq!(summary.total_failed(), 3);
assert!(!summary.all_succeeded());
assert_eq!(summary.results.len(), 2);
}
}