use crate::context::Context;
use crate::dal::unified::{RecoveryClaimResult, RecoveryHeartbeatResult};
use crate::dal::DAL;
use crate::database::UniversalUuid;
use crate::executor::{WorkflowExecutionError, WorkflowExecutor};
use crate::models::schedule::ScheduleExecution;
use chrono::Utc;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::watch;
use tracing::{debug, error, info, warn};
#[derive(Debug, Clone)]
pub struct CronRecoveryConfig {
pub check_interval: Duration,
pub lost_threshold_minutes: i32,
pub max_recovery_age: Duration,
pub max_recovery_attempts: usize,
pub recover_disabled_schedules: bool,
pub claim_heartbeat_interval: Duration,
pub claim_stale_after: Duration,
}
impl Default for CronRecoveryConfig {
fn default() -> Self {
Self {
check_interval: Duration::from_secs(300), lost_threshold_minutes: 10,
max_recovery_age: Duration::from_secs(86400), max_recovery_attempts: 3,
recover_disabled_schedules: false,
claim_heartbeat_interval: Duration::from_secs(30),
claim_stale_after: Duration::from_secs(120),
}
}
}
#[derive(Clone)]
pub struct CronRecoveryService {
dal: Arc<DAL>,
executor: Arc<dyn WorkflowExecutor>,
config: CronRecoveryConfig,
shutdown: watch::Receiver<bool>,
owner_id: UniversalUuid,
}
impl CronRecoveryService {
pub fn new(
dal: Arc<DAL>,
executor: Arc<dyn WorkflowExecutor>,
config: CronRecoveryConfig,
shutdown: watch::Receiver<bool>,
) -> Self {
Self {
dal,
executor,
config,
shutdown,
owner_id: UniversalUuid::new_v4(),
}
}
pub fn with_defaults(
dal: Arc<DAL>,
executor: Arc<dyn WorkflowExecutor>,
shutdown: watch::Receiver<bool>,
) -> Self {
Self::new(dal, executor, CronRecoveryConfig::default(), shutdown)
}
pub async fn run_recovery_loop(&mut self) -> Result<(), WorkflowExecutionError> {
info!(
"Starting cron recovery service (interval: {:?}, threshold: {} minutes)",
self.config.check_interval, self.config.lost_threshold_minutes
);
let mut interval = tokio::time::interval(self.config.check_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = interval.tick() => {
if let Err(e) = self.check_and_recover_lost_executions().await {
error!("Error in cron recovery service: {}", e);
}
}
_ = self.shutdown.changed() => {
if *self.shutdown.borrow() {
info!("Cron recovery service received shutdown signal");
break;
}
}
}
}
info!("Cron recovery service stopped");
Ok(())
}
pub async fn recover_lost_executions_once(&self) -> Result<(), WorkflowExecutionError> {
self.check_and_recover_lost_executions().await
}
async fn check_and_recover_lost_executions(&self) -> Result<(), WorkflowExecutionError> {
debug!("Checking for lost cron executions");
let lost_executions = self
.dal
.schedule_execution()
.find_lost_executions(self.config.lost_threshold_minutes)
.await
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Failed to find lost executions: {}", e),
})?;
if lost_executions.is_empty() {
debug!("No lost executions found");
return Ok(());
}
info!("Found {} lost cron execution(s)", lost_executions.len());
for execution in lost_executions {
if let Err(e) = self.recover_execution(&execution).await {
error!(
"Failed to recover execution {} for schedule {}: {}",
execution.id, execution.schedule_id, e
);
}
}
Ok(())
}
async fn recover_execution(
&self,
execution: &ScheduleExecution,
) -> Result<(), WorkflowExecutionError> {
if let Some(workflow_execution_id) = execution.workflow_execution_id {
match self
.dal
.workflow_execution()
.get_by_id(workflow_execution_id)
.await
{
Ok(workflow_execution) => match workflow_execution.status.as_str() {
"Completed" | "Failed" | "Cancelled" => {
info!(
"Execution {} is linked to terminal workflow execution {} (status: {}); backfilling completion accounting instead of re-firing",
execution.id, workflow_execution_id, workflow_execution.status
);
if let Err(e) = self
.dal
.schedule_execution()
.complete(execution.id, Utc::now())
.await
{
warn!(
"Failed to backfill completion for execution {}: {}",
execution.id, e
);
}
return Ok(());
}
_ => {
debug!(
"Execution {} is linked to live workflow execution {} (status: {}); not lost, skipping recovery",
execution.id, workflow_execution_id, workflow_execution.status
);
return Ok(());
}
},
Err(e) => {
warn!(
"Execution {} is linked to workflow execution {} but the row could not be read; skipping recovery this pass: {}",
execution.id, workflow_execution_id, e
);
return Ok(());
}
}
}
let scheduled_time = execution
.scheduled_time
.as_ref()
.map(|t| t.0)
.unwrap_or(execution.created_at.0);
let execution_age = Utc::now() - scheduled_time;
let max_recovery_age = chrono::Duration::from_std(self.config.max_recovery_age)
.unwrap_or_else(|e| {
warn!(
"max_recovery_age out of chrono::Duration range ({:?}): {}; treating as MAX",
self.config.max_recovery_age, e
);
chrono::Duration::MAX
});
if execution_age > max_recovery_age {
warn!(
"Execution {} is too old to recover (age: {:?}), abandoning",
execution.id, execution_age
);
return Ok(());
}
match self
.dal
.schedule_execution()
.claim_for_recovery(execution.id, self.owner_id, self.config.claim_stale_after)
.await
{
Ok(RecoveryClaimResult::Claimed) => {}
Ok(RecoveryClaimResult::NotClaimed) => {
debug!(
"Execution {} is owned by another recovery service (or already completed); skipping",
execution.id
);
return Ok(());
}
Err(e) => {
warn!(
"Failed to claim execution {} for recovery; skipping this pass: {}",
execution.id, e
);
return Ok(());
}
}
let outcome = self.recover_claimed_execution(execution).await;
if let Err(e) = self
.dal
.schedule_execution()
.release_recovery_claim(execution.id, self.owner_id)
.await
{
warn!(
"Failed to release recovery claim on execution {}: {}",
execution.id, e
);
}
outcome
}
async fn recover_claimed_execution(
&self,
execution: &ScheduleExecution,
) -> Result<(), WorkflowExecutionError> {
match self.dal.schedule_execution().get_by_id(execution.id).await {
Ok(fresh) => {
if fresh.workflow_execution_id.is_some() || fresh.completed_at.is_some() {
debug!(
"Execution {} was linked/completed by another recovery pass; skipping",
execution.id
);
return Ok(());
}
}
Err(e) => {
warn!(
"Failed to re-read execution {} under its recovery claim; skipping: {}",
execution.id, e
);
return Ok(());
}
}
let attempt_count = match self
.dal
.schedule_execution()
.increment_recovery_attempts(execution.id)
.await
{
Ok(count) => count,
Err(e) => {
warn!(
"Failed to record a recovery attempt for execution {}; skipping this pass: {}",
execution.id, e
);
return Ok(());
}
};
if attempt_count as usize > self.config.max_recovery_attempts {
error!(
"Execution {} has exceeded max recovery attempts ({}), abandoning",
execution.id, self.config.max_recovery_attempts
);
return Ok(());
}
let scheduled_time = execution
.scheduled_time
.as_ref()
.map(|t| t.0)
.unwrap_or(execution.created_at.0);
info!(
"Attempting recovery of execution {} (schedule: {}, attempt: {}/{})",
execution.id, execution.schedule_id, attempt_count, self.config.max_recovery_attempts
);
let schedule = match self.dal.schedule().get_by_id(execution.schedule_id).await {
Ok(sched) => sched,
Err(e) => {
warn!(
"Schedule {} not found for execution {}, skipping recovery: {}",
execution.schedule_id, execution.id, e
);
return Ok(());
}
};
if !self.config.recover_disabled_schedules && !schedule.enabled.is_true() {
info!(
"Schedule {} is disabled, skipping recovery of execution {}",
schedule.id, execution.id
);
return Ok(());
}
let mut context = Context::new();
context
.insert("is_recovery", serde_json::json!(true))
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert("recovery_attempt", serde_json::json!(attempt_count))
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert(
"original_execution_id",
serde_json::json!(execution.id.to_string()),
)
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert(
"scheduled_time",
serde_json::json!(scheduled_time.to_rfc3339()),
)
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert("schedule_id", serde_json::json!(schedule.id.to_string()))
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert(
"schedule_timezone",
serde_json::json!(schedule.timezone.as_deref().unwrap_or("UTC")),
)
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
context
.insert(
"schedule_expression",
serde_json::json!(schedule.cron_expression.as_deref().unwrap_or("")),
)
.map_err(|e| WorkflowExecutionError::ExecutionFailed {
message: format!("Context error: {}", e),
})?;
info!(
"Executing recovery for workflow '{}' (execution: {}, schedule: {})",
schedule.workflow_name, execution.id, schedule.id
);
let _heartbeat = self.spawn_claim_heartbeat(execution.id);
match self
.executor
.execute(&schedule.workflow_name, context)
.await
{
Ok(workflow_result) => {
if let Err(e) = self
.dal
.schedule_execution()
.update_workflow_execution_id(
execution.id,
crate::database::UniversalUuid(workflow_result.execution_id),
)
.await
{
error!(
"Failed to update audit record for recovered execution {}: {}",
execution.id, e
);
}
if let Err(e) = self
.dal
.schedule_execution()
.complete(execution.id, Utc::now())
.await
{
warn!(
"Failed to mark recovered execution {} complete: {}",
execution.id, e
);
}
info!(
"Successfully recovered execution {} (new workflow execution: {})",
execution.id, workflow_result.execution_id
);
if let Err(e) = self
.dal
.schedule_execution()
.reset_recovery_attempts(execution.id)
.await
{
warn!(
"Failed to reset recovery attempts for execution {}: {}",
execution.id, e
);
}
Ok(())
}
Err(e) => {
error!(
"Failed to recover execution {} for workflow '{}': {}",
execution.id, schedule.workflow_name, e
);
Err(e)
}
}
}
fn spawn_claim_heartbeat(&self, execution_id: UniversalUuid) -> ClaimHeartbeat {
let dal = self.dal.clone();
let owner_id = self.owner_id;
let interval = self.config.claim_heartbeat_interval;
ClaimHeartbeat(tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
ticker.tick().await;
loop {
ticker.tick().await;
match dal
.schedule_execution()
.recovery_heartbeat(execution_id, owner_id)
.await
{
Ok(RecoveryHeartbeatResult::Ok) => {}
Ok(RecoveryHeartbeatResult::ClaimLost) => {
warn!(
"Recovery claim on execution {} was taken over while the re-fire was still running",
execution_id
);
break;
}
Err(e) => {
warn!(
"Failed to beat the recovery claim on execution {}: {}",
execution_id, e
);
}
}
}
}))
}
pub async fn clear_recovery_attempts(
&self,
execution_id: UniversalUuid,
) -> Result<(), crate::error::ValidationError> {
self.dal
.schedule_execution()
.reset_recovery_attempts(execution_id)
.await?;
info!("Cleared recovery attempts for execution {}", execution_id);
Ok(())
}
pub async fn get_recovery_attempts(&self, execution_id: UniversalUuid) -> usize {
match self.dal.schedule_execution().get_by_id(execution_id).await {
Ok(row) => row.recovery_attempts.max(0) as usize,
Err(e) => {
warn!(
"Failed to read recovery attempts for execution {}: {}",
execution_id, e
);
0
}
}
}
}
struct ClaimHeartbeat(tokio::task::JoinHandle<()>);
impl Drop for ClaimHeartbeat {
fn drop(&mut self) {
self.0.abort();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_recovery_config_default() {
let config = CronRecoveryConfig::default();
assert_eq!(config.check_interval, Duration::from_secs(300));
assert_eq!(config.lost_threshold_minutes, 10);
assert_eq!(config.max_recovery_age, Duration::from_secs(86400));
assert_eq!(config.max_recovery_attempts, 3);
assert!(!config.recover_disabled_schedules);
assert!(config.claim_heartbeat_interval < config.claim_stale_after);
assert_eq!(config.claim_heartbeat_interval, Duration::from_secs(30));
assert_eq!(config.claim_stale_after, Duration::from_secs(120));
}
#[test]
fn test_recovery_config_custom() {
let config = CronRecoveryConfig {
check_interval: Duration::from_secs(60),
lost_threshold_minutes: 5,
max_recovery_age: Duration::from_secs(3600),
max_recovery_attempts: 5,
recover_disabled_schedules: true,
claim_heartbeat_interval: Duration::from_secs(5),
claim_stale_after: Duration::from_secs(20),
};
assert_eq!(config.check_interval, Duration::from_secs(60));
assert_eq!(config.lost_threshold_minutes, 5);
assert_eq!(config.max_recovery_age, Duration::from_secs(3600));
assert_eq!(config.max_recovery_attempts, 5);
assert!(config.recover_disabled_schedules);
}
#[test]
fn test_recovery_config_clone() {
let config = CronRecoveryConfig::default();
let cloned = config.clone();
assert_eq!(config.check_interval, cloned.check_interval);
assert_eq!(config.lost_threshold_minutes, cloned.lost_threshold_minutes);
assert_eq!(config.max_recovery_attempts, cloned.max_recovery_attempts);
}
#[test]
fn test_recovery_config_default_recovery_window() {
let config = CronRecoveryConfig::default();
assert_eq!(config.max_recovery_age.as_secs(), 86400);
assert_eq!(config.check_interval.as_secs(), 300);
}
}