1use std::path::PathBuf;
34use std::sync::Arc;
35
36use clap::Subcommand;
37use serde::{Deserialize, Serialize};
38
39use sz_rust_core::orm::scheduler::Scheduler;
40
41use crate::error::CliError;
42
43const DEFAULT_CONFIG_PATH: &str = "scheduler.toml";
45
46#[derive(Subcommand, Debug)]
48pub enum SchedulerCommand {
49 #[command(name = "list")]
51 List {
52 #[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
54 config: PathBuf,
55 },
56
57 #[command(name = "run")]
59 Run {
60 #[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
62 config: PathBuf,
63 },
64
65 #[command(name = "start")]
67 Start {
68 #[arg(short = 't', long, default_value = "1000")]
70 tick_ms: u64,
71
72 #[arg(short = 'c', long, default_value = DEFAULT_CONFIG_PATH)]
74 config: PathBuf,
75 },
76}
77
78#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
80pub struct TaskConfig {
81 pub id: String,
83 pub name: String,
85 pub cron: String,
87 #[serde(default)]
89 pub callback: String,
90 #[serde(default = "default_enabled")]
92 pub enabled: bool,
93}
94
95fn default_enabled() -> bool {
96 true
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
101pub struct SchedulerConfig {
102 #[serde(default)]
104 pub tasks: Vec<TaskConfig>,
105}
106
107impl SchedulerConfig {
108 pub fn from_toml_str(s: &str) -> Result<Self, CliError> {
115 let config: SchedulerConfig = toml::from_str(s)
116 .map_err(|e| CliError::Scheduler(format!("config parse error: {e}")))?;
117 config.validate()?;
118 Ok(config)
119 }
120
121 pub fn from_file(path: &std::path::Path) -> Result<Self, CliError> {
129 let content = std::fs::read_to_string(path)?;
130 Self::from_toml_str(&content)
131 }
132
133 fn validate(&self) -> Result<(), CliError> {
135 let mut seen = std::collections::HashSet::new();
136 for task in &self.tasks {
137 if !seen.insert(&task.id) {
138 return Err(CliError::Scheduler(format!(
139 "duplicate task id: {}",
140 task.id
141 )));
142 }
143 }
144 Ok(())
145 }
146
147 pub fn register_to(
151 &self,
152 scheduler: &sz_rust_core::orm::scheduler::CronScheduler,
153 ) -> Result<usize, CliError> {
154 let handler: Arc<dyn sz_rust_core::orm::scheduler::JobHandler> = Arc::new(PrintJobHandler);
155 let mut count = 0;
156 for task in &self.tasks {
157 let scheduled =
158 sz_rust_core::orm::scheduler::ScheduledTask::new(&task.id, &task.name, &task.cron)
159 .with_callback(&task.callback);
160 let scheduled = if task.enabled {
161 scheduled
162 } else {
163 scheduled.disable()
164 };
165 scheduler.schedule(scheduled).map_err(|e| {
166 CliError::Scheduler(format!("schedule task {} failed: {e}", task.id))
167 })?;
168 scheduler.register_handler(&task.id, handler.clone());
169 count += 1;
170 }
171 Ok(count)
172 }
173}
174
175pub fn execute(cmd: &SchedulerCommand) -> Result<(), CliError> {
177 match cmd {
178 SchedulerCommand::List { config } => execute_list(config),
179 SchedulerCommand::Run { config } => execute_run(config),
180 SchedulerCommand::Start { tick_ms, config } => execute_start(*tick_ms, config),
181 }
182}
183
184fn load_config_or_empty(path: &std::path::Path) -> Result<SchedulerConfig, CliError> {
186 if !path.exists() {
187 return Ok(SchedulerConfig::default());
188 }
189 SchedulerConfig::from_file(path)
190}
191
192fn build_scheduler(
194 config: &SchedulerConfig,
195) -> Result<sz_rust_core::orm::scheduler::CronScheduler, CliError> {
196 let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
197 config.register_to(&scheduler)?;
198 Ok(scheduler)
199}
200
201fn execute_list(config_path: &std::path::Path) -> Result<(), CliError> {
203 let config = load_config_or_empty(config_path)?;
204
205 if config.tasks.is_empty() {
206 println!("No scheduled tasks registered.");
207 println!();
208 println!("To register tasks, create a scheduler configuration file at:");
209 println!(" {}", config_path.display());
210 println!();
211 println!("Example scheduler.toml:");
212 println!();
213 println!(" [[tasks]]");
214 println!(" id = \"cleanup-logs\"");
215 println!(" name = \"清理日志\"");
216 println!(" cron = \"0 3 * * *\"");
217 println!(" callback = \"demo::cleanup_logs\"");
218 println!(" enabled = true");
219 return Ok(());
220 }
221
222 let scheduler = build_scheduler(&config)?;
223 let tasks = scheduler.list_tasks();
224
225 println!(
226 "{:<20} {:<25} {:<25} {:<8}",
227 "ID", "Name", "Cron", "Enabled"
228 );
229 println!("{}", "-".repeat(80));
230
231 for task in &tasks {
232 println!(
233 "{:<20} {:<25} {:<25} {:<8}",
234 task.id, task.name, task.cron_expr, task.enabled
235 );
236 }
237
238 println!();
239 println!("Total: {} task(s)", tasks.len());
240 println!("Config: {}", config_path.display());
241 Ok(())
242}
243
244fn execute_run(config_path: &std::path::Path) -> Result<(), CliError> {
248 let config = load_config_or_empty(config_path)?;
249
250 if config.tasks.is_empty() {
251 println!("No scheduled tasks registered. Nothing to run.");
252 println!("Config: {}", config_path.display());
253 return Ok(());
254 }
255
256 let scheduler = build_scheduler(&config)?;
257 let now = chrono::Utc::now();
258
259 println!("Scheduler run at {} (UTC)", now);
261 println!();
262 for task in &config.tasks {
263 match scheduler.next_run_time(&task.cron, now) {
264 Ok(next) => {
265 println!(
266 " {:<20} {:<25} cron={:<20} next={}",
267 task.id, task.name, task.cron, next
268 );
269 }
270 Err(e) => {
271 println!(
272 " {:<20} {:<25} cron={:<20} next=N/A ({})",
273 task.id, task.name, task.cron, e
274 );
275 }
276 }
277 }
278 println!();
279
280 let fired = scheduler.try_fire_due(now);
281 println!("Fired {} task(s) matching current time.", fired);
282 println!("Config: {}", config_path.display());
283 Ok(())
284}
285
286fn execute_start(tick_ms: u64, config_path: &std::path::Path) -> Result<(), CliError> {
290 let config = load_config_or_empty(config_path)?;
291
292 if config.tasks.is_empty() {
293 return Err(CliError::Scheduler(format!(
294 "no tasks to schedule. Please create config file at {}",
295 config_path.display()
296 )));
297 }
298
299 let scheduler = build_scheduler(&config)?;
300 let task_count = config.tasks.len();
301
302 println!("Starting scheduler with tick interval: {}ms", tick_ms);
303 println!(
304 "Loaded {} task(s) from {}",
305 task_count,
306 config_path.display()
307 );
308 println!("Press Ctrl+C to stop.");
309 println!();
310
311 scheduler
313 .start(tick_ms)
314 .map_err(|e| CliError::Scheduler(e.to_string()))?;
315
316 println!("Scheduler started. Waiting for tasks to fire...");
317
318 loop {
321 std::thread::sleep(std::time::Duration::from_secs(1));
322 }
323}
324
325struct PrintJobHandler;
329
330impl sz_rust_core::orm::scheduler::JobHandler for PrintJobHandler {
331 fn handle(&self, task: &sz_rust_core::orm::scheduler::ScheduledTask) -> Result<(), String> {
332 println!(
333 "[{}] Task fired: {} ({}) callback={}",
334 chrono::Utc::now(),
335 task.name,
336 task.id,
337 if task.callback.is_empty() {
338 "(none)"
339 } else {
340 &task.callback
341 }
342 );
343 Ok(())
344 }
345}
346
347#[cfg(test)]
348mod tests {
349 use super::*;
350 use sz_rust_core::orm::scheduler::JobHandler;
351 use tempfile::NamedTempFile;
352
353 const SAMPLE_TOML: &str = r#"
354[[tasks]]
355id = "cleanup-logs"
356name = "清理日志"
357cron = "0 3 * * *"
358callback = "demo::cleanup_logs"
359enabled = true
360
361[[tasks]]
362id = "sync-data"
363name = "数据同步"
364cron = "*/30 * * * *"
365callback = "demo::sync_data"
366
367[[tasks]]
368id = "disabled-task"
369name = "已禁用任务"
370cron = "0 0 * * *"
371callback = "demo::disabled"
372enabled = false
373"#;
374
375 #[test]
376 fn test_config_from_toml_str() {
377 let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
378 assert_eq!(config.tasks.len(), 3);
379
380 assert_eq!(config.tasks[0].id, "cleanup-logs");
381 assert_eq!(config.tasks[0].name, "清理日志");
382 assert_eq!(config.tasks[0].cron, "0 3 * * *");
383 assert_eq!(config.tasks[0].callback, "demo::cleanup_logs");
384 assert!(config.tasks[0].enabled);
385
386 assert!(config.tasks[1].enabled);
388
389 assert!(!config.tasks[2].enabled);
391 }
392
393 #[test]
394 fn test_config_from_toml_str_empty() {
395 let config = SchedulerConfig::from_toml_str("").unwrap();
396 assert!(config.tasks.is_empty());
397 }
398
399 #[test]
400 fn test_config_duplicate_id_rejected() {
401 let toml_str = r#"
402[[tasks]]
403id = "dup"
404name = "任务1"
405cron = "0 * * * *"
406
407[[tasks]]
408id = "dup"
409name = "任务2"
410cron = "0 * * * *"
411"#;
412 let result = SchedulerConfig::from_toml_str(toml_str);
413 assert!(result.is_err());
414 let err = result.unwrap_err().to_string();
415 assert!(err.contains("duplicate task id"), "got: {err}");
416 }
417
418 #[test]
419 fn test_config_invalid_toml_rejected() {
420 let result = SchedulerConfig::from_toml_str("not valid toml [[[[");
421 assert!(result.is_err());
422 }
423
424 #[test]
425 fn test_config_from_file() {
426 let tmp = NamedTempFile::new().unwrap();
427 std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
428
429 let config = SchedulerConfig::from_file(tmp.path()).unwrap();
430 assert_eq!(config.tasks.len(), 3);
431 }
432
433 #[test]
434 fn test_config_from_file_not_found() {
435 let result = SchedulerConfig::from_file(std::path::Path::new("/nonexistent/path.toml"));
436 assert!(result.is_err());
437 }
438
439 #[test]
440 fn test_register_to_scheduler() {
441 let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
442 let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
443
444 let count = config.register_to(&scheduler).unwrap();
445 assert_eq!(count, 3);
446
447 let tasks = scheduler.list_tasks();
448 assert_eq!(tasks.len(), 3);
449
450 let disabled = tasks.iter().find(|t| t.id == "disabled-task").unwrap();
452 assert!(!disabled.enabled);
453 }
454
455 #[test]
456 fn test_register_to_invalid_cron_rejected() {
457 let toml_str = r#"
458[[tasks]]
459id = "bad"
460name = "坏任务"
461cron = "not a cron"
462"#;
463 let config = SchedulerConfig::from_toml_str(toml_str).unwrap();
464 let scheduler = sz_rust_core::orm::scheduler::CronScheduler::new();
465
466 let result = config.register_to(&scheduler);
467 assert!(result.is_err());
468 let err = result.unwrap_err().to_string();
469 assert!(err.contains("schedule task"), "got: {err}");
470 }
471
472 #[test]
473 fn test_load_config_or_empty_missing_file() {
474 let config =
475 load_config_or_empty(std::path::Path::new("/nonexistent/scheduler.toml")).unwrap();
476 assert!(config.tasks.is_empty());
477 }
478
479 #[test]
480 fn test_load_config_or_empty_existing_file() {
481 let tmp = NamedTempFile::new().unwrap();
482 std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
483
484 let config = load_config_or_empty(tmp.path()).unwrap();
485 assert_eq!(config.tasks.len(), 3);
486 }
487
488 #[test]
489 fn test_build_scheduler() {
490 let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
491 let scheduler = build_scheduler(&config).unwrap();
492 let tasks = scheduler.list_tasks();
493 assert_eq!(tasks.len(), 3);
494 }
495
496 #[test]
497 fn test_execute_list_empty_config() {
498 let result = execute_list(std::path::Path::new("/nonexistent/scheduler.toml"));
500 assert!(result.is_ok());
501 }
502
503 #[test]
504 fn test_execute_list_with_config() {
505 let tmp = NamedTempFile::new().unwrap();
506 std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
507
508 let result = execute_list(tmp.path());
509 assert!(result.is_ok());
510 }
511
512 #[test]
513 fn test_execute_run_empty_config() {
514 let result = execute_run(std::path::Path::new("/nonexistent/scheduler.toml"));
515 assert!(result.is_ok());
516 }
517
518 #[test]
519 fn test_execute_run_with_config() {
520 let tmp = NamedTempFile::new().unwrap();
521 std::fs::write(tmp.path(), SAMPLE_TOML).unwrap();
522
523 let result = execute_run(tmp.path());
524 assert!(result.is_ok());
525 }
526
527 #[test]
528 fn test_execute_start_empty_config_errors() {
529 let result = execute_start(1000, std::path::Path::new("/nonexistent/scheduler.toml"));
530 assert!(result.is_err());
531 let err = result.unwrap_err().to_string();
532 assert!(err.contains("no tasks to schedule"), "got: {err}");
533 }
534
535 #[test]
536 fn test_print_job_handler() {
537 let handler = PrintJobHandler;
538 let task = sz_rust_core::orm::scheduler::ScheduledTask::new("test", "测试", "0 * * * *")
539 .with_callback("demo::test");
540 let result = handler.handle(&task);
541 assert!(result.is_ok());
542 }
543
544 #[test]
545 fn test_print_job_handler_empty_callback() {
546 let handler = PrintJobHandler;
547 let task = sz_rust_core::orm::scheduler::ScheduledTask::new("test", "测试", "0 * * * *");
548 let result = handler.handle(&task);
549 assert!(result.is_ok());
550 }
551
552 #[test]
553 fn test_task_config_default_enabled() {
554 let toml_str = r#"
555[[tasks]]
556id = "t1"
557name = "任务"
558cron = "0 * * * *"
559"#;
560 let config = SchedulerConfig::from_toml_str(toml_str).unwrap();
561 assert!(config.tasks[0].enabled, "enabled should default to true");
562 }
563
564 #[test]
565 fn test_config_serialize_roundtrip() {
566 let config = SchedulerConfig::from_toml_str(SAMPLE_TOML).unwrap();
567 let toml_str = toml::to_string(&config).unwrap();
568 let config2 = SchedulerConfig::from_toml_str(&toml_str).unwrap();
569 assert_eq!(config, config2);
570 }
571}