use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use tokio::sync::Mutex;
#[async_trait]
pub trait Warmer: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str {
"Cache warmer"
}
fn timeout(&self) -> Duration {
Duration::from_secs(30)
}
async fn warm(&self) -> Result<(), WarmupError>;
}
#[derive(Debug, Clone)]
pub enum WarmupError {
Io(String),
Serialize(String),
Cache(String),
Database(String),
Custom(String),
}
impl std::fmt::Display for WarmupError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(msg) => write!(f, "IO error: {}", msg),
Self::Serialize(msg) => write!(f, "Serialize error: {}", msg),
Self::Cache(msg) => write!(f, "Cache error: {}", msg),
Self::Database(msg) => write!(f, "Database error: {}", msg),
Self::Custom(msg) => write!(f, "{}", msg),
}
}
}
impl std::error::Error for WarmupError {}
impl From<std::io::Error> for WarmupError {
fn from(e: std::io::Error) -> Self {
Self::Io(e.to_string())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WarmupStatus {
Success,
Failed,
Timeout,
Skipped,
}
impl std::fmt::Display for WarmupStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Success => write!(f, "success"),
Self::Failed => write!(f, "failed"),
Self::Timeout => write!(f, "timeout"),
Self::Skipped => write!(f, "skipped"),
}
}
}
#[derive(Debug, Clone)]
pub struct WarmupItem {
pub name: String,
pub description: String,
pub status: WarmupStatus,
pub duration_ms: u64,
pub error: Option<String>,
}
impl WarmupItem {
pub fn success(
name: impl Into<String>,
description: impl Into<String>,
duration_ms: u64,
) -> Self {
Self {
name: name.into(),
description: description.into(),
status: WarmupStatus::Success,
duration_ms,
error: None,
}
}
pub fn failed(
name: impl Into<String>,
description: impl Into<String>,
duration_ms: u64,
error: impl Into<String>,
) -> Self {
Self {
name: name.into(),
description: description.into(),
status: WarmupStatus::Failed,
duration_ms,
error: Some(error.into()),
}
}
pub fn timeout(
name: impl Into<String>,
description: impl Into<String>,
duration_ms: u64,
) -> Self {
Self {
name: name.into(),
description: description.into(),
status: WarmupStatus::Timeout,
duration_ms,
error: Some(format!("Timed out after {}ms", duration_ms)),
}
}
pub fn is_success(&self) -> bool {
self.status == WarmupStatus::Success
}
}
#[derive(Debug, Clone, Default)]
pub struct WarmupReport {
pub items: Vec<WarmupItem>,
pub total_duration_ms: u64,
}
impl WarmupReport {
pub fn new() -> Self {
Self::default()
}
pub fn add_item(&mut self, item: WarmupItem) {
self.items.push(item);
}
pub fn success_count(&self) -> usize {
self.items.iter().filter(|i| i.is_success()).count()
}
pub fn failed_count(&self) -> usize {
self.items
.iter()
.filter(|i| i.status == WarmupStatus::Failed)
.count()
}
pub fn timeout_count(&self) -> usize {
self.items
.iter()
.filter(|i| i.status == WarmupStatus::Timeout)
.count()
}
pub fn total_count(&self) -> usize {
self.items.len()
}
pub fn all_success(&self) -> bool {
!self.items.is_empty() && self.failed_count() == 0 && self.timeout_count() == 0
}
pub fn summary(&self) -> String {
let mut s = String::new();
s.push_str(&format!(
"Cache warmup: {}/{} succeeded, {} failed, {} timeout (total {}ms)\n",
self.success_count(),
self.total_count(),
self.failed_count(),
self.timeout_count(),
self.total_duration_ms
));
for item in &self.items {
s.push_str(&format!(
" - {:<20} {:<10} {}ms",
item.name, item.status, item.duration_ms
));
if let Some(err) = &item.error {
s.push_str(&format!(" ({})", err));
}
s.push('\n');
}
s
}
}
pub struct WarmupPipeline {
warmers: Vec<Box<dyn Warmer>>,
global_timeout: Duration,
}
impl WarmupPipeline {
pub fn new() -> Self {
Self {
warmers: Vec::new(),
global_timeout: Duration::from_secs(300),
}
}
pub fn register(&mut self, warmer: Box<dyn Warmer>) -> &mut Self {
self.warmers.push(warmer);
self
}
pub fn with_global_timeout(mut self, timeout: Duration) -> Self {
self.global_timeout = timeout;
self
}
pub fn count(&self) -> usize {
self.warmers.len()
}
pub fn names(&self) -> Vec<&str> {
self.warmers.iter().map(|w| w.name()).collect()
}
pub async fn warm_all(&self) -> WarmupReport {
let mut report = WarmupReport::new();
let total_start = Instant::now();
for warmer in &self.warmers {
let item = self.warm_one(warmer.as_ref()).await;
report.add_item(item);
}
report.total_duration_ms = total_start.elapsed().as_millis() as u64;
report
}
pub async fn warm_all_parallel(&self) -> WarmupReport {
let total_start = Instant::now();
let futures: Vec<_> = self
.warmers
.iter()
.map(|w| self.warm_one_async(w.as_ref()))
.collect();
let items = futures::future::join_all(futures).await;
let mut report = WarmupReport::new();
for item in items {
report.add_item(item);
}
report.total_duration_ms = total_start.elapsed().as_millis() as u64;
report
}
async fn warm_one(&self, warmer: &dyn Warmer) -> WarmupItem {
self.warm_one_async(warmer).await
}
async fn warm_one_async(&self, warmer: &dyn Warmer) -> WarmupItem {
let name = warmer.name().to_string();
let description = warmer.description().to_string();
let timeout = warmer.timeout();
let start = Instant::now();
let result = tokio::time::timeout(timeout, warmer.warm()).await;
let duration_ms = start.elapsed().as_millis() as u64;
match result {
Ok(Ok(())) => WarmupItem::success(name, description, duration_ms),
Ok(Err(e)) => WarmupItem::failed(name, description, duration_ms, e.to_string()),
Err(_) => WarmupItem::timeout(name, description, duration_ms),
}
}
pub async fn warm_one_by_name(&self, name: &str) -> WarmupReport {
let mut report = WarmupReport::new();
let total_start = Instant::now();
if let Some(warmer) = self.warmers.iter().find(|w| w.name() == name) {
let item = self.warm_one(warmer.as_ref()).await;
report.add_item(item);
} else {
report.add_item(WarmupItem {
name: name.to_string(),
description: "Not found".to_string(),
status: WarmupStatus::Skipped,
duration_ms: 0,
error: Some(format!("Warmer '{}' not registered", name)),
});
}
report.total_duration_ms = total_start.elapsed().as_millis() as u64;
report
}
}
impl Default for WarmupPipeline {
fn default() -> Self {
Self::new()
}
}
pub struct DeploymentHook {
pipeline: Arc<Mutex<WarmupPipeline>>,
}
impl DeploymentHook {
pub fn new(pipeline: WarmupPipeline) -> Self {
Self {
pipeline: Arc::new(Mutex::new(pipeline)),
}
}
pub async fn pre_warmup(&self) -> WarmupReport {
let pipeline = self.pipeline.lock().await;
pipeline.warm_all().await
}
pub async fn post_deploy(&self) -> WarmupReport {
let pipeline = self.pipeline.lock().await;
pipeline.warm_all().await
}
pub async fn rollback(&self) -> WarmupReport {
let pipeline = self.pipeline.lock().await;
pipeline.warm_all().await
}
pub fn pipeline(&self) -> Arc<Mutex<WarmupPipeline>> {
self.pipeline.clone()
}
}
pub struct NoopWarmer {
name: String,
delay_ms: u64,
should_fail: bool,
}
impl NoopWarmer {
pub fn new(name: impl Into<String>, delay_ms: u64, should_fail: bool) -> Self {
Self {
name: name.into(),
delay_ms,
should_fail,
}
}
pub fn success(name: impl Into<String>) -> Self {
Self::new(name, 0, false)
}
pub fn failing(name: impl Into<String>) -> Self {
Self::new(name, 0, true)
}
pub fn delayed(name: impl Into<String>, delay_ms: u64) -> Self {
Self::new(name, delay_ms, false)
}
}
#[async_trait]
impl Warmer for NoopWarmer {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
"Noop warmer for testing"
}
async fn warm(&self) -> Result<(), WarmupError> {
if self.delay_ms > 0 {
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
}
if self.should_fail {
return Err(WarmupError::Custom("Simulated failure".to_string()));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_warmup_error_display() {
assert_eq!(
WarmupError::Io("file not found".to_string()).to_string(),
"IO error: file not found"
);
assert_eq!(
WarmupError::Serialize("invalid json".to_string()).to_string(),
"Serialize error: invalid json"
);
assert_eq!(
WarmupError::Cache("write failed".to_string()).to_string(),
"Cache error: write failed"
);
assert_eq!(
WarmupError::Database("connection refused".to_string()).to_string(),
"Database error: connection refused"
);
assert_eq!(
WarmupError::Custom("custom".to_string()).to_string(),
"custom"
);
}
#[test]
fn test_warmup_error_from_io() {
let io_err = std::io::Error::new(std::io::ErrorKind::NotFound, "missing");
let warmup_err: WarmupError = io_err.into();
assert!(matches!(warmup_err, WarmupError::Io(_)));
}
#[test]
fn test_warmup_status_display() {
assert_eq!(WarmupStatus::Success.to_string(), "success");
assert_eq!(WarmupStatus::Failed.to_string(), "failed");
assert_eq!(WarmupStatus::Timeout.to_string(), "timeout");
assert_eq!(WarmupStatus::Skipped.to_string(), "skipped");
}
#[test]
fn test_warmup_item_success() {
let item = WarmupItem::success("config", "Config warmer", 100);
assert_eq!(item.name, "config");
assert_eq!(item.status, WarmupStatus::Success);
assert_eq!(item.duration_ms, 100);
assert!(item.error.is_none());
assert!(item.is_success());
}
#[test]
fn test_warmup_item_failed() {
let item = WarmupItem::failed("route", "Route warmer", 50, "missing file");
assert_eq!(item.status, WarmupStatus::Failed);
assert_eq!(item.error, Some("missing file".to_string()));
assert!(!item.is_success());
}
#[test]
fn test_warmup_item_timeout() {
let item = WarmupItem::timeout("db", "DB warmer", 30000);
assert_eq!(item.status, WarmupStatus::Timeout);
assert!(item.error.unwrap().contains("Timed out"));
}
#[test]
fn test_warmup_report_empty() {
let report = WarmupReport::new();
assert_eq!(report.total_count(), 0);
assert_eq!(report.success_count(), 0);
assert_eq!(report.failed_count(), 0);
assert_eq!(report.timeout_count(), 0);
assert!(!report.all_success());
}
#[test]
fn test_warmup_report_all_success() {
let mut report = WarmupReport::new();
report.add_item(WarmupItem::success("a", "A", 10));
report.add_item(WarmupItem::success("b", "B", 20));
report.total_duration_ms = 30;
assert_eq!(report.total_count(), 2);
assert_eq!(report.success_count(), 2);
assert_eq!(report.failed_count(), 0);
assert!(report.all_success());
}
#[test]
fn test_warmup_report_mixed() {
let mut report = WarmupReport::new();
report.add_item(WarmupItem::success("a", "A", 10));
report.add_item(WarmupItem::failed("b", "B", 20, "err"));
report.add_item(WarmupItem::timeout("c", "C", 30000));
assert_eq!(report.total_count(), 3);
assert_eq!(report.success_count(), 1);
assert_eq!(report.failed_count(), 1);
assert_eq!(report.timeout_count(), 1);
assert!(!report.all_success());
}
#[test]
fn test_warmup_report_summary_contains_status() {
let mut report = WarmupReport::new();
report.add_item(WarmupItem::success("config", "Config", 100));
report.add_item(WarmupItem::failed("route", "Route", 50, "missing"));
report.total_duration_ms = 150;
let summary = report.summary();
assert!(summary.contains("1/2 succeeded"));
assert!(summary.contains("1 failed"));
assert!(summary.contains("config"));
assert!(summary.contains("success"));
assert!(summary.contains("route"));
assert!(summary.contains("failed"));
assert!(summary.contains("missing"));
}
#[tokio::test]
async fn test_pipeline_empty() {
let pipeline = WarmupPipeline::new();
assert_eq!(pipeline.count(), 0);
let report = pipeline.warm_all().await;
assert_eq!(report.total_count(), 0);
assert!(report.items.is_empty());
}
#[tokio::test]
async fn test_pipeline_single_success() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("config")));
let report = pipeline.warm_all().await;
assert_eq!(report.total_count(), 1);
assert_eq!(report.success_count(), 1);
assert!(report.all_success());
assert_eq!(report.items[0].name, "config");
}
#[tokio::test]
async fn test_pipeline_mixed_results() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("success_warmer")));
pipeline.register(Box::new(NoopWarmer::failing("failing_warmer")));
let report = pipeline.warm_all().await;
assert_eq!(report.total_count(), 2);
assert_eq!(report.success_count(), 1);
assert_eq!(report.failed_count(), 1);
assert!(!report.all_success());
assert_eq!(report.items[0].name, "success_warmer");
assert_eq!(report.items[0].status, WarmupStatus::Success);
assert_eq!(report.items[1].name, "failing_warmer");
assert_eq!(report.items[1].status, WarmupStatus::Failed);
assert!(report.items[1]
.error
.as_ref()
.unwrap()
.contains("Simulated failure"));
}
#[tokio::test]
async fn test_pipeline_failure_does_not_stop_others() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::failing("first_fails")));
pipeline.register(Box::new(NoopWarmer::success("second_success")));
let report = pipeline.warm_all().await;
assert_eq!(report.total_count(), 2);
assert_eq!(report.success_count(), 1);
assert_eq!(report.failed_count(), 1);
}
#[tokio::test]
async fn test_pipeline_timeout() {
let mut pipeline = WarmupPipeline::new();
struct SlowWarmer;
#[async_trait]
impl Warmer for SlowWarmer {
fn name(&self) -> &str {
"slow"
}
fn timeout(&self) -> Duration {
Duration::from_millis(100)
}
async fn warm(&self) -> Result<(), WarmupError> {
tokio::time::sleep(Duration::from_millis(500)).await;
Ok(())
}
}
pipeline.register(Box::new(SlowWarmer));
let report = pipeline.warm_all().await;
assert_eq!(report.total_count(), 1);
assert_eq!(report.timeout_count(), 1);
assert_eq!(report.items[0].status, WarmupStatus::Timeout);
}
#[tokio::test]
async fn test_pipeline_parallel() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::delayed("a", 100)));
pipeline.register(Box::new(NoopWarmer::delayed("b", 100)));
pipeline.register(Box::new(NoopWarmer::delayed("c", 100)));
let serial_report = pipeline.warm_all().await;
assert!(serial_report.total_duration_ms >= 250);
let parallel_report = pipeline.warm_all_parallel().await;
assert!(parallel_report.total_duration_ms < 200);
assert_eq!(parallel_report.success_count(), 3);
}
#[tokio::test]
async fn test_pipeline_warm_one_by_name_found() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("config")));
pipeline.register(Box::new(NoopWarmer::success("route")));
let report = pipeline.warm_one_by_name("config").await;
assert_eq!(report.total_count(), 1);
assert_eq!(report.items[0].name, "config");
assert_eq!(report.items[0].status, WarmupStatus::Success);
}
#[tokio::test]
async fn test_pipeline_warm_one_by_name_not_found() {
let pipeline = WarmupPipeline::new();
let report = pipeline.warm_one_by_name("nonexistent").await;
assert_eq!(report.total_count(), 1);
assert_eq!(report.items[0].status, WarmupStatus::Skipped);
assert!(report.items[0]
.error
.as_ref()
.unwrap()
.contains("not registered"));
}
#[tokio::test]
async fn test_pipeline_names() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("a")));
pipeline.register(Box::new(NoopWarmer::success("b")));
let names = pipeline.names();
assert_eq!(names, vec!["a", "b"]);
}
#[tokio::test]
async fn test_pipeline_with_global_timeout() {
let pipeline = WarmupPipeline::new().with_global_timeout(Duration::from_secs(60));
assert_eq!(pipeline.count(), 0);
}
#[tokio::test]
async fn test_deployment_hook_pre_warmup() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("config")));
let hook = DeploymentHook::new(pipeline);
let report = hook.pre_warmup().await;
assert_eq!(report.success_count(), 1);
}
#[tokio::test]
async fn test_deployment_hook_post_deploy() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("config")));
pipeline.register(Box::new(NoopWarmer::success("route")));
let hook = DeploymentHook::new(pipeline);
let report = hook.post_deploy().await;
assert_eq!(report.success_count(), 2);
}
#[tokio::test]
async fn test_deployment_hook_rollback() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("cleanup")));
let hook = DeploymentHook::new(pipeline);
let report = hook.rollback().await;
assert_eq!(report.success_count(), 1);
}
#[tokio::test]
async fn test_deployment_hook_pipeline_handle() {
let pipeline = WarmupPipeline::new();
let hook = DeploymentHook::new(pipeline);
let handle = hook.pipeline();
let p = handle.lock().await;
assert_eq!(p.count(), 0);
}
#[tokio::test]
async fn test_noop_warmer_success() {
let warmer = NoopWarmer::success("test");
assert_eq!(warmer.name(), "test");
assert_eq!(warmer.description(), "Noop warmer for testing");
let result = warmer.warm().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_noop_warmer_failing() {
let warmer = NoopWarmer::failing("test");
let result = warmer.warm().await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Simulated failure"));
}
#[tokio::test]
async fn test_noop_warmer_delayed() {
let warmer = NoopWarmer::delayed("test", 50);
let start = Instant::now();
warmer.warm().await.unwrap();
assert!(start.elapsed().as_millis() >= 40);
}
#[test]
fn test_noop_warmer_default_timeout() {
let warmer = NoopWarmer::success("test");
assert_eq!(warmer.timeout(), Duration::from_secs(30));
}
#[tokio::test]
async fn test_end_to_end_deployment_scenario() {
let mut pipeline = WarmupPipeline::new();
pipeline.register(Box::new(NoopWarmer::success("config")));
pipeline.register(Box::new(NoopWarmer::success("route")));
pipeline.register(Box::new(NoopWarmer::failing("dict")));
let hook = DeploymentHook::new(pipeline);
let report = hook.pre_warmup().await;
assert_eq!(report.total_count(), 3);
assert_eq!(report.success_count(), 2);
assert_eq!(report.failed_count(), 1);
let summary = report.summary();
assert!(summary.contains("2/3 succeeded"));
assert!(summary.contains("1 failed"));
let failed_item = report
.items
.iter()
.find(|i| i.status == WarmupStatus::Failed)
.unwrap();
assert_eq!(failed_item.name, "dict");
}
}