durable/

executor.rs

use sea_orm::sea_query::OnConflict;
use sea_orm::{ConnectionTrait, DatabaseConnection, DbBackend, EntityTrait, Set, Statement};
use std::time::Duration;
use tokio::task::JoinHandle;
use uuid::Uuid;

use crate::error::DurableError;
use durable_db::entity::executor_heartbeat;

/// Configuration for the executor heartbeat loop and recovery.
///
/// The staleness_threshold should be strictly greater than heartbeat_interval
/// (recommended: 3x). Default is 60s interval and 180s threshold.
pub struct HeartbeatConfig {
    /// How often the executor writes a heartbeat. Default: 60s.
    pub heartbeat_interval: Duration,
    /// Tasks from executors whose last heartbeat is older than this are reset to PENDING. Default: 180s.
    pub staleness_threshold: Duration,
}

impl Default for HeartbeatConfig {
    fn default() -> Self {
        Self {
            heartbeat_interval: Duration::from_secs(60),
            staleness_threshold: Duration::from_secs(180),
        }
    }
}

/// Information about a task that was recovered from a stale RUNNING state.
pub struct RecoveredTask {
    pub id: Uuid,
    pub parent_id: Option<Uuid>,
    pub name: String,
    pub kind: String,
}

/// Parse a row from a `RETURNING id, parent_id, name, kind` clause.
fn parse_recovered_row(row: &sea_orm::QueryResult) -> Result<RecoveredTask, DurableError> {
    let id: Uuid = row
        .try_get_by_index(0)
        .map_err(|e| DurableError::custom(e.to_string()))?;
    let parent_id: Option<Uuid> = row.try_get_by_index(1).ok().flatten();
    let name: String = row
        .try_get_by_index(2)
        .map_err(|e| DurableError::custom(e.to_string()))?;
    let kind: String = row
        .try_get_by_index(3)
        .map_err(|e| DurableError::custom(e.to_string()))?;
    Ok(RecoveredTask {
        id,
        parent_id,
        name,
        kind,
    })
}

/// Executor configuration for a durable worker.
pub struct Executor {
    db: DatabaseConnection,
    executor_id: String,
}

impl Executor {
    pub fn new(db: DatabaseConnection, executor_id: String) -> Self {
        Self { db, executor_id }
    }

    pub fn db(&self) -> &DatabaseConnection {
        &self.db
    }

    pub fn executor_id(&self) -> &str {
        &self.executor_id
    }

    /// Write (or upsert) a heartbeat row for this executor.
    pub async fn heartbeat(&self) -> Result<(), DurableError> {
        let now = chrono::Utc::now().into();
        let model = executor_heartbeat::ActiveModel {
            executor_id: Set(self.executor_id.clone()),
            last_seen: Set(now),
        };

        executor_heartbeat::Entity::insert(model)
            .on_conflict(
                OnConflict::column(executor_heartbeat::Column::ExecutorId)
                    .update_column(executor_heartbeat::Column::LastSeen)
                    .to_owned(),
            )
            .exec(&self.db)
            .await
            .map_err(DurableError::from)?;

        Ok(())
    }

    /// Spawn a background tokio task that calls `heartbeat()` every `config.heartbeat_interval`.
    ///
    /// Returns a `JoinHandle` so the caller can abort it on graceful shutdown.
    pub fn start_heartbeat(&self, config: &HeartbeatConfig) -> JoinHandle<()> {
        let db = self.db.clone();
        let executor_id = self.executor_id.clone();
        let interval = config.heartbeat_interval;

        tokio::spawn(async move {
            let executor = Executor::new(db, executor_id);
            let mut ticker = tokio::time::interval(interval);
            loop {
                ticker.tick().await;
                if let Err(e) = executor.heartbeat().await {
                    tracing::warn!("heartbeat write failed: {e}");
                }
            }
        })
    }

    /// Reset RUNNING tasks from dead workers to PENDING. Catches three cases:
    ///
    /// 1. Tasks whose `executor_id` has a stale heartbeat (> staleness_threshold)
    /// 2. Tasks whose `executor_id` has no heartbeat row at all (executor never registered or was cleaned up)
    /// 3. Orphaned tasks with NULL `executor_id` that have been RUNNING too long
    ///
    /// Excludes own executor_id — own tasks are handled by the current process.
    /// Returns the recovered tasks.
    pub async fn recover_stale_tasks(
        &self,
        staleness_threshold: Duration,
    ) -> Result<Vec<RecoveredTask>, DurableError> {
        let threshold_secs = staleness_threshold.as_secs();
        let sql = format!(
            "UPDATE durable.task \
             SET status = 'PENDING', started_at = NULL, \
                 executor_id = '{}' \
             WHERE status = 'RUNNING' \
               AND ( \
                   (executor_id IN ( \
                       SELECT executor_id FROM durable.executor_heartbeat \
                       WHERE last_seen < now() - interval '{threshold_secs} seconds' \
                   ) AND executor_id != '{}') \
                   OR \
                   (executor_id IS NOT NULL \
                    AND executor_id != '{}' \
                    AND executor_id NOT IN ( \
                        SELECT executor_id FROM durable.executor_heartbeat \
                    )) \
                   OR \
                   (executor_id IS NULL \
                    AND started_at IS NOT NULL \
                    AND started_at < now() - interval '{threshold_secs} seconds') \
               ) \
             RETURNING id, parent_id, name, kind",
            self.executor_id, self.executor_id, self.executor_id
        );
        let rows = self
            .db
            .query_all(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

        rows.iter().map(parse_recovered_row).collect()
    }

    /// Scan for stale tasks by timeout and reset them to PENDING.
    ///
    /// A task is considered stale if:
    /// - `status = 'RUNNING'`
    /// - `timeout_ms IS NOT NULL` or `deadline_epoch_ms IS NOT NULL`
    /// - `started_at IS NOT NULL`
    /// - The timeout/deadline has elapsed
    ///
    /// Tasks without `timeout_ms` set are never considered stale.
    pub async fn recover(&self) -> Result<Vec<RecoveredTask>, DurableError> {
        let sql = format!(
            "UPDATE durable.task \
             SET status = 'PENDING', started_at = NULL, deadline_epoch_ms = NULL, \
                 executor_id = '{}' \
             WHERE status = 'RUNNING' \
               AND started_at IS NOT NULL \
               AND ( \
                 (timeout_ms IS NOT NULL AND started_at < now() - make_interval(secs => timeout_ms::double precision / 1000.0)) \
                 OR \
                 (deadline_epoch_ms IS NOT NULL AND deadline_epoch_ms < EXTRACT(EPOCH FROM now()) * 1000) \
               ) \
             RETURNING id, parent_id, name, kind",
            self.executor_id
        );

        let rows = self
            .db
            .query_all(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

        rows.iter().map(parse_recovered_row).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_heartbeat_config_default() {
        let config = HeartbeatConfig::default();
        assert_eq!(
            config.heartbeat_interval,
            Duration::from_secs(60),
            "default heartbeat_interval should be 60s"
        );
        assert_eq!(
            config.staleness_threshold,
            Duration::from_secs(180),
            "default staleness_threshold should be 180s"
        );
    }

    #[test]
    fn test_heartbeat_config_custom() {
        let config = HeartbeatConfig {
            heartbeat_interval: Duration::from_secs(30),
            staleness_threshold: Duration::from_secs(90),
        };
        assert_eq!(config.heartbeat_interval, Duration::from_secs(30));
        assert_eq!(config.staleness_threshold, Duration::from_secs(90));
    }
}