use crate::config::UpdateConfig;
use anyhow::{Context, Result};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tracing::{debug, error, info, warn};
#[derive(Debug, Clone)]
pub enum UpdateNotification {
UpdateAvailable {
current_version: String,
latest_version: String,
},
CheckFailed { error: String },
Stopped,
}
pub struct UpdateScheduler {
config: Arc<UpdateConfig>,
update_check_fn: Arc<dyn Fn() -> Result<UpdateCheckResult> + Send + Sync>,
handle: Option<JoinHandle<()>>,
notification_sender: Option<mpsc::UnboundedSender<UpdateNotification>>,
is_running: bool,
}
#[derive(Debug, Clone)]
pub enum UpdateCheckResult {
UpToDate,
UpdateAvailable {
current_version: String,
latest_version: String,
},
Failed { error: String },
}
impl UpdateScheduler {
pub fn new(
config: Arc<UpdateConfig>,
update_check_fn: Arc<dyn Fn() -> Result<UpdateCheckResult> + Send + Sync>,
) -> Self {
Self {
config,
update_check_fn,
handle: None,
notification_sender: None,
is_running: false,
}
}
pub fn create_notification_channel(
&mut self,
) -> Result<mpsc::UnboundedReceiver<UpdateNotification>> {
if self.is_running {
anyhow::bail!("Cannot create channel after scheduler is started");
}
let (sender, receiver) = mpsc::unbounded_channel();
self.notification_sender = Some(sender);
Ok(receiver)
}
pub async fn start(&mut self) -> Result<()> {
if self.is_running {
warn!("Scheduler is already running");
return Ok(());
}
if !self.config.auto_update_enabled {
info!("Auto-update is disabled, scheduler not starting");
return Ok(());
}
let sender = self
.notification_sender
.clone()
.context("No notification channel created. Call create_notification_channel() first")?;
let config = Arc::clone(&self.config);
let check_fn = Arc::clone(&self.update_check_fn);
info!(
"Starting update scheduler with interval: {:?}",
config.auto_update_check_interval
);
let handle = tokio::spawn(async move {
let mut interval = tokio::time::interval(config.auto_update_check_interval);
let _ = Self::perform_check(&check_fn, &sender).await;
loop {
interval.tick().await;
debug!("Performing scheduled update check");
if let Err(e) = Self::perform_check(&check_fn, &sender).await {
error!("Error in scheduled check: {}", e);
}
}
});
self.handle = Some(handle);
self.is_running = true;
info!("Update scheduler started");
Ok(())
}
pub async fn stop(&mut self) -> Result<()> {
if !self.is_running {
debug!("Scheduler is not running");
return Ok(());
}
info!("Stopping update scheduler");
if let Some(handle) = self.handle.take() {
handle.abort();
let _ = tokio::time::timeout(Duration::from_secs(5), handle).await;
}
self.is_running = false;
if let Some(sender) = &self.notification_sender {
let _ = sender.send(UpdateNotification::Stopped);
}
info!("Update scheduler stopped");
Ok(())
}
pub fn is_running(&self) -> bool {
self.is_running
}
pub fn should_check(&self, last_check_time: Option<std::time::Instant>) -> bool {
if !self.config.auto_update_enabled {
return false;
}
match last_check_time {
Some(last) => {
let elapsed = last.elapsed();
elapsed >= self.config.auto_update_check_interval
}
None => true,
}
}
async fn perform_check(
check_fn: &Arc<dyn Fn() -> Result<UpdateCheckResult> + Send + Sync>,
sender: &mpsc::UnboundedSender<UpdateNotification>,
) -> Result<()> {
let result = tokio::task::spawn_blocking({
let check_fn = Arc::clone(check_fn);
move || check_fn()
})
.await
.context("Failed to spawn blocking task for update check")?;
let check_result = result.context("Update check function failed")?;
match check_result {
UpdateCheckResult::UpdateAvailable {
current_version,
latest_version,
} => {
info!(
"Update available: {} -> {}",
current_version, latest_version
);
sender
.send(UpdateNotification::UpdateAvailable {
current_version,
latest_version,
})
.context("Failed to send update notification")?;
}
UpdateCheckResult::UpToDate => {
debug!("Up to date, no notification needed");
}
UpdateCheckResult::Failed { error } => {
warn!("Update check failed: {}", error);
sender
.send(UpdateNotification::CheckFailed { error })
.context("Failed to send error notification")?;
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
#[test]
fn test_scheduler_creation() {
let config = UpdateConfig::default();
let check_count = Arc::new(AtomicUsize::new(0));
let check_count_clone = Arc::clone(&check_count);
let check_fn = Arc::new(move || {
check_count_clone.fetch_add(1, Ordering::SeqCst);
Ok(UpdateCheckResult::UpToDate)
});
let scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
assert!(!scheduler.is_running());
}
#[test]
fn test_should_check_with_interval() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_secs(60),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
assert!(scheduler.should_check(None));
let recent = Instant::now() - Duration::from_secs(30);
assert!(!scheduler.should_check(Some(recent)));
let old = Instant::now() - Duration::from_secs(120);
assert!(scheduler.should_check(Some(old)));
}
#[test]
fn test_should_check_when_disabled() {
let config = UpdateConfig {
auto_update_enabled: false,
auto_update_check_interval: Duration::from_secs(60),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
assert!(!scheduler.should_check(None));
}
#[tokio::test]
async fn test_scheduler_starts_and_stops() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_millis(100),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let mut receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
assert!(scheduler.is_running());
tokio::time::sleep(Duration::from_millis(200)).await;
scheduler.stop().await.unwrap();
assert!(!scheduler.is_running());
let notification = receiver.try_recv().unwrap();
matches!(notification, UpdateNotification::Stopped);
}
#[tokio::test]
async fn test_scheduler_with_short_interval() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_millis(50),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let _receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(150)).await;
scheduler.stop().await.unwrap();
}
#[tokio::test]
async fn test_scheduler_notification_update_available() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_secs(1000),
};
let check_fn = Arc::new(|| {
Ok(UpdateCheckResult::UpdateAvailable {
current_version: "1.0.0".to_string(),
latest_version: "1.1.0".to_string(),
})
});
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let mut receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
let notification = receiver.try_recv().unwrap();
match notification {
UpdateNotification::UpdateAvailable {
current_version,
latest_version,
} => {
assert_eq!(current_version, "1.0.0");
assert_eq!(latest_version, "1.1.0");
}
_ => panic!("Expected UpdateAvailable notification"),
}
scheduler.stop().await.unwrap();
}
#[tokio::test]
async fn test_scheduler_notification_check_failed() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_secs(1000),
};
let check_fn = Arc::new(|| {
Ok(UpdateCheckResult::Failed {
error: "Network error".to_string(),
})
});
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let mut receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
let notification = receiver.try_recv().unwrap();
match notification {
UpdateNotification::CheckFailed { error } => {
assert_eq!(error, "Network error");
}
_ => panic!("Expected CheckFailed notification"),
}
scheduler.stop().await.unwrap();
}
#[tokio::test]
async fn test_scheduler_disabled_does_not_start() {
let config = UpdateConfig {
auto_update_enabled: false,
auto_update_check_interval: Duration::from_millis(50),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let _receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
assert!(!scheduler.is_running());
scheduler.stop().await.unwrap();
}
#[tokio::test]
async fn test_scheduler_already_running() {
let config = UpdateConfig {
auto_update_enabled: true,
auto_update_check_interval: Duration::from_millis(50),
};
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let _receiver = scheduler.create_notification_channel().unwrap();
scheduler.start().await.unwrap();
let result = scheduler.start().await;
assert!(result.is_ok());
assert!(scheduler.is_running());
scheduler.stop().await.unwrap();
}
#[tokio::test]
async fn test_scheduler_stop_when_not_running() {
let config = UpdateConfig::default();
let check_fn = Arc::new(|| Ok(UpdateCheckResult::UpToDate));
let mut scheduler = UpdateScheduler::new(Arc::new(config), check_fn);
let result = scheduler.stop().await;
assert!(result.is_ok());
assert!(!scheduler.is_running());
}
}