291 lines
10 KiB
Rust
291 lines
10 KiB
Rust
pub mod types;
|
|
pub mod store;
|
|
pub mod tools;
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use tokio::time;
|
|
|
|
use crate::config::SchedulerConfig;
|
|
use crate::session::session::HandleResult;
|
|
use crate::session::SessionManager;
|
|
|
|
pub use types::{JobRun, Schedule, ScheduledJob};
|
|
|
|
/// Compute the next execution time (Unix ms) for a schedule, given `from` (Unix ms).
|
|
/// Returns `None` if no next time can be determined (e.g., invalid cron expression).
|
|
pub fn next_run_for_schedule(schedule: &Schedule, from: i64) -> Option<i64> {
|
|
use chrono::{TimeZone, Utc};
|
|
use std::str::FromStr;
|
|
|
|
match schedule {
|
|
Schedule::At { at } => Some(*at),
|
|
Schedule::Every { every_ms } => Some(from + *every_ms as i64),
|
|
Schedule::Cron { expr, tz } => {
|
|
let cron_schedule = cron::Schedule::from_str(expr.as_str()).ok()?;
|
|
let from_secs = (from / 1000) as i64;
|
|
let from_nanos = ((from % 1000) * 1_000_000) as u32;
|
|
let from_dt = Utc.timestamp_opt(from_secs, from_nanos).single()?;
|
|
|
|
let next_utc = if let Some(tz_str) = tz {
|
|
let tz: chrono_tz::Tz = tz_str.parse().ok()?;
|
|
let _from_local = from_dt.with_timezone(&tz);
|
|
let next_local = cron_schedule.upcoming(tz).next()?;
|
|
next_local.with_timezone(&Utc)
|
|
} else {
|
|
cron_schedule.upcoming(Utc).next()?
|
|
};
|
|
|
|
Some(next_utc.timestamp_millis())
|
|
}
|
|
}
|
|
}
|
|
|
|
fn now_ms() -> i64 {
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_millis() as i64
|
|
}
|
|
|
|
/// The scheduler runs as a background tokio task, periodically checking for due jobs
|
|
/// and executing them via `SessionManager::handle_cron_message`.
|
|
pub struct Scheduler {
|
|
pool: sqlx::SqlitePool,
|
|
session_manager: Arc<SessionManager>,
|
|
config: SchedulerConfig,
|
|
}
|
|
|
|
impl Scheduler {
|
|
pub fn new(
|
|
pool: sqlx::SqlitePool,
|
|
session_manager: Arc<SessionManager>,
|
|
config: SchedulerConfig,
|
|
) -> Self {
|
|
Self {
|
|
pool,
|
|
session_manager,
|
|
config,
|
|
}
|
|
}
|
|
|
|
/// Run the scheduler loop. This is a long-running async function meant to be
|
|
/// spawned as a tokio background task.
|
|
pub async fn run(self: Arc<Self>) {
|
|
let poll_duration = time::Duration::from_secs(self.config.poll_interval_secs);
|
|
let mut interval = time::interval(poll_duration);
|
|
|
|
interval.tick().await;
|
|
|
|
tracing::info!(
|
|
"Scheduler started (poll interval: {}s, max concurrent: {})",
|
|
self.config.poll_interval_secs,
|
|
self.config.max_concurrent,
|
|
);
|
|
|
|
loop {
|
|
interval.tick().await;
|
|
|
|
let now = now_ms();
|
|
|
|
let due = match store::SchedulerStore::due_jobs(&self.pool, now, self.config.max_concurrent).await {
|
|
Ok(jobs) => jobs,
|
|
Err(e) => {
|
|
tracing::error!("scheduler: failed to query due jobs: {}", e);
|
|
continue;
|
|
}
|
|
};
|
|
|
|
if due.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
tracing::info!("scheduler: found {} due job(s)", due.len());
|
|
|
|
for job in &due {
|
|
let start = Instant::now();
|
|
let started_at = now_ms();
|
|
|
|
if let Err(e) = store::SchedulerStore::touch_last_run(&self.pool, &job.id, started_at).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to touch last_run_at: {}", e);
|
|
continue;
|
|
}
|
|
|
|
tracing::info!(
|
|
job_id = %job.id,
|
|
job_name = %job.name,
|
|
"scheduler: executing cron job"
|
|
);
|
|
|
|
let result = self
|
|
.session_manager
|
|
.handle_cron_message(
|
|
&job.channel,
|
|
&job.chat_id,
|
|
&job.prompt,
|
|
&job.id,
|
|
&job.name,
|
|
)
|
|
.await;
|
|
|
|
let finished_at = now_ms();
|
|
let duration_ms = start.elapsed().as_millis() as i64;
|
|
|
|
match result {
|
|
Ok(HandleResult::AgentResponse(output)) => {
|
|
let output_truncated = if output.len() > 8000 {
|
|
format!("{}...[truncated]", &output[..8000])
|
|
} else {
|
|
output.clone()
|
|
};
|
|
|
|
let run = JobRun {
|
|
id: 0,
|
|
job_id: job.id.clone(),
|
|
started_at,
|
|
finished_at,
|
|
status: "ok".to_string(),
|
|
output: Some(output_truncated),
|
|
error: None,
|
|
duration_ms,
|
|
};
|
|
|
|
if let Err(e) = store::SchedulerStore::record_run(&self.pool, &run).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to record run: {}", e);
|
|
}
|
|
|
|
if let Err(e) = store::SchedulerStore::set_last_status(&self.pool, &job.id, "ok", None).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to set last_status: {}", e);
|
|
}
|
|
|
|
tracing::info!(
|
|
job_id = %job.id,
|
|
duration_ms = %duration_ms,
|
|
"scheduler: job completed successfully"
|
|
);
|
|
}
|
|
Ok(HandleResult::CommandOutput(output)) => {
|
|
let run = JobRun {
|
|
id: 0,
|
|
job_id: job.id.clone(),
|
|
started_at,
|
|
finished_at,
|
|
status: "ok".to_string(),
|
|
output: Some(output),
|
|
error: None,
|
|
duration_ms,
|
|
};
|
|
|
|
let _ = store::SchedulerStore::record_run(&self.pool, &run).await;
|
|
}
|
|
Err(e) => {
|
|
let error_str = e.to_string();
|
|
let run = JobRun {
|
|
id: 0,
|
|
job_id: job.id.clone(),
|
|
started_at,
|
|
finished_at,
|
|
status: "error".to_string(),
|
|
output: None,
|
|
error: Some(error_str.clone()),
|
|
duration_ms,
|
|
};
|
|
|
|
if let Err(e2) = store::SchedulerStore::record_run(&self.pool, &run).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to record error run: {}", e2);
|
|
}
|
|
|
|
if let Err(e2) = store::SchedulerStore::set_last_status(
|
|
&self.pool, &job.id, "error", Some(&error_str),
|
|
).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to set error status: {}", e2);
|
|
}
|
|
|
|
tracing::error!(
|
|
job_id = %job.id,
|
|
duration_ms = %duration_ms,
|
|
error = %error_str,
|
|
"scheduler: job failed"
|
|
);
|
|
}
|
|
}
|
|
|
|
if let Err(e) = self.reschedule_after_run(job).await {
|
|
tracing::error!(job_id = %job.id, "scheduler: failed to reschedule: {}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// After a job runs, compute its next execution time or disable/delete it.
|
|
async fn reschedule_after_run(
|
|
&self,
|
|
job: &ScheduledJob,
|
|
) -> anyhow::Result<()> {
|
|
let now = now_ms();
|
|
|
|
match &job.schedule {
|
|
Schedule::At { .. } => {
|
|
if job.delete_after_run {
|
|
store::SchedulerStore::remove_job(&self.pool, &job.id).await?;
|
|
tracing::info!(job_id = %job.id, "scheduler: one-shot job deleted after run");
|
|
} else {
|
|
store::SchedulerStore::set_enabled(&self.pool, &job.id, false).await?;
|
|
tracing::info!(job_id = %job.id, "scheduler: one-shot job disabled after run");
|
|
}
|
|
}
|
|
Schedule::Every { .. } | Schedule::Cron { .. } => {
|
|
if let Some(next) = next_run_for_schedule(&job.schedule, now) {
|
|
store::SchedulerStore::set_next_run(&self.pool, &job.id, next).await?;
|
|
tracing::info!(job_id = %job.id, next_run_at = %next, "scheduler: job rescheduled");
|
|
} else {
|
|
tracing::error!(job_id = %job.id, "scheduler: could not compute next run -- disabling job");
|
|
store::SchedulerStore::set_enabled(&self.pool, &job.id, false).await?;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_next_run_at_schedule() {
|
|
let now = 1000000;
|
|
let next = next_run_for_schedule(&Schedule::At { at: 2000000 }, now);
|
|
assert_eq!(next, Some(2000000));
|
|
}
|
|
|
|
#[test]
|
|
fn test_next_run_every_schedule() {
|
|
let now = 1000000;
|
|
let next = next_run_for_schedule(&Schedule::Every { every_ms: 5000 }, now);
|
|
assert_eq!(next, Some(1005000));
|
|
}
|
|
|
|
#[test]
|
|
fn test_next_run_cron_every_minute() {
|
|
let expr = "0 * * * * *".to_string();
|
|
let schedule = Schedule::Cron { expr, tz: None };
|
|
let now = 1000000;
|
|
let next = next_run_for_schedule(&schedule, now);
|
|
assert!(next.is_some());
|
|
assert!(next.unwrap() > now);
|
|
}
|
|
|
|
#[test]
|
|
fn test_next_run_cron_every_day_at_9am() {
|
|
let expr = "0 0 9 * * *".to_string();
|
|
let schedule = Schedule::Cron { expr, tz: None };
|
|
let now = 1000000;
|
|
let next = next_run_for_schedule(&schedule, now);
|
|
assert!(next.is_some());
|
|
let next_ms = next.unwrap();
|
|
assert!(next_ms > now);
|
|
}
|
|
}
|