use std::path::PathBuf;
use std::sync::Arc;
use clap::Subcommand;
use serde::{Deserialize, Serialize};
use sz_rust_core::orm::scheduler::Scheduler;
use crate::error::CliError;
const DEFAULT_CONFIG_PATH: &str = "scheduler.toml";
#[derive(Subcommand, Debug)]
pub enum SchedulerCommand {
#[command(name = "list")]
List {
#[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
config: PathBuf,
},
#[command(name = "run")]
Run {
#[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
config: PathBuf,
},
#[command(name = "start")]
Start {
#[arg(short = 't', long, default_value = "1000")]
tick_ms: u64,
#[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
config: PathBuf,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TaskConfig {
pub id: String,
pub name: String,
pub cron: String,
#[serde(default)]
pub callback: String,
#[serde(default = "default_enabled")]
pub enabled: bool,
}
fn default_enabled() -> bool {
true
}
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
pub struct SchedulerConfig {
#[serde(default)]
pub tasks: Vec<TaskConfig>,
}
impl SchedulerConfig {
pub fn from_toml_str(s: &str) -> Result<Self, CliError> {
let config: SchedulerConfig = toml::from_str(s)
.map_err(|e| CliError::Scheduler(format!("config parse error: {e}")))?;
config.validate()?;
Ok(config)
}
pub fn from_file(path: &std::path::Path) -> Result<Self, CliError> {
let content = std::fs::read_to_string(path)?;
Self::from_toml_str(&content)
}
fn validate(&self) -> Result<(), CliError> {
let mut seen = std::collections::HashSet::new();
for task in &self.tasks {
if !seen.insert(&task.id) {
return Err(CliError::Scheduler(format!(
"duplicate task id: {}",
task.id
)));
}
}
Ok(())
}
pub fn register_to(
&self,
scheduler: &sz_rust_core::orm::scheduler::CronScheduler,
) -> Result<usize, CliError> {
let handler: Arc<dyn sz_rust_core::orm::scheduler::JobHandler> = Arc::new(PrintJobHandler);
let mut count = 0;
for task in &self.tasks {
let scheduled =
sz_rust_core::orm::scheduler::ScheduledTask::new(&task.id, &task.name, &task.cron)
.with_callback(&task.callback);
let scheduled = if task.enabled {
scheduled
} else {
scheduled.disable()
};
scheduler.schedule(scheduled).map_err(|e| {
CliError::Scheduler(format!("schedule task {} failed: {e}", task.id))
})?;
scheduler.register_handler(&task.id, handler.clone());
count += 1;
}
Ok(count)
}
}
pub fn execute(cmd: &SchedulerCommand) -> Result<(), CliError> {
match cmd {
SchedulerCommand::List { config } => execute_list(config),
SchedulerCommand::Run { config } => execute_run(config),
SchedulerCommand::Start { tick_ms, config } => execute_start(*tick_ms, config),
}
}
fn load_config_or_empty(path: &std::path::Path) -> Result<SchedulerConfig, CliError> {
if !path.exists() {
return Ok(SchedulerConfig::default());
}
SchedulerConfig::from_file(path)
}
fn build_scheduler(
config: &SchedulerConfig,
) -> Result<sz_rust_core::orm::scheduler::CronScheduler, CliError> {
let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
config.register_to(&scheduler)?;
Ok(scheduler)
}
fn execute_list(config_path: &std::path::Path) -> Result<(), CliError> {
let config = load_config_or_empty(config_path)?;
if config.tasks.is_empty() {
println!("No scheduled tasks registered.");
println!();
println!("To register tasks, create a scheduler configuration file at:");
println!(" {}", config_path.display());
println!();
println!("Example scheduler.toml:");
println!();
println!(" [[tasks]]");
println!(" id = \"cleanup-logs\"");
println!(" name = \"清理日志\"");
println!(" cron = \"0 3 * * *\"");
println!(" callback = \"demo::cleanup_logs\"");
println!(" enabled = true");
return Ok(());
}
let scheduler = build_scheduler(&config)?;
let tasks = scheduler.list_tasks();
println!(
"{:<20} {:<25} {:<25} {:<8}",
"ID", "Name", "Cron", "Enabled"
);
println!("{}", "-".repeat(80));
for task in &tasks {
println!(
"{:<20} {:<25} {:<25} {:<8}",
task.id, task.name, task.cron_expr, task.enabled
);
}
println!();
println!("Total: {} task(s)", tasks.len());
println!("Config: {}", config_path.display());
Ok(())
}
fn execute_run(config_path: &std::path::Path) -> Result<(), CliError> {
let config = load_config_or_empty(config_path)?;
if config.tasks.is_empty() {
println!("No scheduled tasks registered. Nothing to run.");
println!("Config: {}", config_path.display());
return Ok(());
}
let scheduler = build_scheduler(&config)?;
let now = chrono::Utc::now();
println!("Scheduler run at {} (UTC)", now);
println!();
for task in &config.tasks {
match scheduler.next_run_time(&task.cron, now) {
Ok(next) => {
println!(
" {:<20} {:<25} cron={:<20} next={}",
task.id, task.name, task.cron, next
);
}
Err(e) => {
println!(
" {:<20} {:<25} cron={:<20} next=N/A ({})",
task.id, task.name, task.cron, e
);
}
}
}
println!();
let fired = scheduler.try_fire_due(now);
println!("Fired {} task(s) matching current time.", fired);
println!("Config: {}", config_path.display());
Ok(())
}
fn execute_start(tick_ms: u64, config_path: &std::path::Path) -> Result<(), CliError> {
let config = load_config_or_empty(config_path)?;
if config.tasks.is_empty() {
return Err(CliError::Scheduler(format!(
"no tasks to schedule. Please create config file at {}",
config_path.display()
)));
}
let scheduler = build_scheduler(&config)?;
let task_count = config.tasks.len();
println!("Starting scheduler with tick interval: {}ms", tick_ms);
println!(
"Loaded {} task(s) from {}",
task_count,
config_path.display()
);
println!("Press Ctrl+C to stop.");
println!();
scheduler
.start(tick_ms)
.map_err(|e| CliError::Scheduler(e.to_string()))?;
println!("Scheduler started. Waiting for tasks to fire...");
loop {
std::thread::sleep(std::time::Duration::from_secs(1));
}
}
struct PrintJobHandler;
impl sz_rust_core::orm::scheduler::JobHandler for PrintJobHandler {
fn handle(&self, task: &sz_rust_core::orm::scheduler::ScheduledTask) -> Result<(), String> {
println!(
"[{}] Task fired: {} ({}) callback={}",
chrono::Utc::now(),
task.name,
task.id,
if task.callback.is_empty() {
"(none)"
} else {
&task.callback
}
);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use sz_rust_core::orm::scheduler::JobHandler;
use tempfile::NamedTempFile;
const SAMPLE_TOML: &str = r#"
[[tasks]]
id = "cleanup-logs"
name = "清理日志"
cron = "0 3 * * *"
callback = "demo::cleanup_logs"
enabled = true
[[tasks]]
id = "sync-data"
name = "数据同步"
cron = "*/30 * * * *"
callback = "demo::sync_data"
[[tasks]]
id = "disabled-task"
name = "已禁用任务"
cron = "0 0 * * *"
callback = "demo::disabled"
enabled = false
"#;
#[test]
fn test_config_from_toml_str() {
let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
assert_eq!(config.tasks.len(), 3);
assert_eq!(config.tasks[0].id, "cleanup-logs");
assert_eq!(config.tasks[0].name, "清理日志");
assert_eq!(config.tasks[0].cron, "0 3 * * *");
assert_eq!(config.tasks[0].callback, "demo::cleanup_logs");
assert!(config.tasks[0].enabled);
assert!(config.tasks[1].enabled);
assert!(!config.tasks[2].enabled);
}
#[test]
fn test_config_from_toml_str_empty() {
let config = SchedulerConfig::from_toml_str("").unwrap();
assert!(config.tasks.is_empty());
}
#[test]
fn test_config_duplicate_id_rejected() {
let toml_str = r#"
[[tasks]]
id = "dup"
name = "任务1"
cron = "0 * * * *"
[[tasks]]
id = "dup"
name = "任务2"
cron = "0 * * * *"
"#;
let result = SchedulerConfig::from_toml_str(toml_str);
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("duplicate task id"), "got: {err}");
}
#[test]
fn test_config_invalid_toml_rejected() {
let result = SchedulerConfig::from_toml_str("not valid toml [[[[");
assert!(result.is_err());
}
#[test]
fn test_config_from_file() {
let tmp = NamedTempFile::new().unwrap();
std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
let config = SchedulerConfig::from_file(tmp.path()).unwrap();
assert_eq!(config.tasks.len(), 3);
}
#[test]
fn test_config_from_file_not_found() {
let result = SchedulerConfig::from_file(std::path::Path::new("/nonexistent/path.toml"));
assert!(result.is_err());
}
#[test]
fn test_register_to_scheduler() {
let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
let count = config.register_to(&scheduler).unwrap();
assert_eq!(count, 3);
let tasks = scheduler.list_tasks();
assert_eq!(tasks.len(), 3);
let disabled = tasks.iter().find(|t| t.id == "disabled-task").unwrap();
assert!(!disabled.enabled);
}
#[test]
fn test_register_to_invalid_cron_rejected() {
let toml_str = r#"
[[tasks]]
id = "bad"
name = "坏任务"
cron = "not a cron"
"#;
let config = SchedulerConfig::from_toml_str(toml_str).unwrap();
let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
let result = config.register_to(&scheduler);
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("schedule task"), "got: {err}");
}
#[test]
fn test_load_config_or_empty_missing_file() {
let config =
load_config_or_empty(std::path::Path::new("/nonexistent/scheduler.toml")).unwrap();
assert!(config.tasks.is_empty());
}
#[test]
fn test_load_config_or_empty_existing_file() {
let tmp = NamedTempFile::new().unwrap();
std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
let config = load_config_or_empty(tmp.path()).unwrap();
assert_eq!(config.tasks.len(), 3);
}
#[test]
fn test_build_scheduler() {
let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
let scheduler = build_scheduler(&config).unwrap();
let tasks = scheduler.list_tasks();
assert_eq!(tasks.len(), 3);
}
#[test]
fn test_execute_list_empty_config() {
let result = execute_list(std::path::Path::new("/nonexistent/scheduler.toml"));
assert!(result.is_ok());
}
#[test]
fn test_execute_list_with_config() {
let tmp = NamedTempFile::new().unwrap();
std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
let result = execute_list(tmp.path());
assert!(result.is_ok());
}
#[test]
fn test_execute_run_empty_config() {
let result = execute_run(std::path::Path::new("/nonexistent/scheduler.toml"));
assert!(result.is_ok());
}
#[test]
fn test_execute_run_with_config() {
let tmp = NamedTempFile::new().unwrap();
std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
let result = execute_run(tmp.path());
assert!(result.is_ok());
}
#[test]
fn test_execute_start_empty_config_errors() {
let result = execute_start(1000, std::path::Path::new("/nonexistent/scheduler.toml"));
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(err.contains("no tasks to schedule"), "got: {err}");
}
#[test]
fn test_print_job_handler() {
let handler = PrintJobHandler;
let task = sz_rust_core::orm::scheduler::ScheduledTask::new("test", "测试", "0 * * * *")
.with_callback("demo::test");
let result = handler.handle(&task);
assert!(result.is_ok());
}
#[test]
fn test_print_job_handler_empty_callback() {
let handler = PrintJobHandler;
let task = sz_rust_core::orm::scheduler::ScheduledTask::new("test", "测试", "0 * * * *");
let result = handler.handle(&task);
assert!(result.is_ok());
}
#[test]
fn test_task_config_default_enabled() {
let toml_str = r#"
[[tasks]]
id = "t1"
name = "任务"
cron = "0 * * * *"
"#;
let config = SchedulerConfig::from_toml_str(toml_str).unwrap();
assert!(config.tasks[0].enabled, "enabled should default to true");
}
#[test]
fn test_config_serialize_roundtrip() {
let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
let toml_str = toml::to_string(&config).unwrap();
let config2 = SchedulerConfig::from_toml_str(&toml_str).unwrap();
assert_eq!(config, config2);
}
}