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 name: String,
    pub kind: String,
}

/// 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 (stale heartbeat > staleness_threshold) to PENDING.
    ///
    /// Excludes own executor_id — own tasks are handled by startup reset.
    /// Returns the number of rows affected.
    pub async fn recover_stale_tasks(
        &self,
        staleness_threshold: Duration,
    ) -> Result<u64, DurableError> {
        let threshold_secs = staleness_threshold.as_secs();
        let sql = format!(
            "UPDATE durable.task SET status = 'PENDING', started_at = NULL \
             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 != '{}'",
            self.executor_id
        );
        let result = self
            .db
            .execute(Statement::from_string(DbBackend::Postgres, sql))
            .await
            .map_err(DurableError::from)?;
        Ok(result.rows_affected())
    }

    /// Spawn a background tokio task that calls `recover_stale_tasks()` every
    /// `config.staleness_threshold` interval.
    pub fn start_recovery_loop(&self, config: HeartbeatConfig) -> JoinHandle<()> {
        let db = self.db.clone();
        let executor_id = self.executor_id.clone();
        let threshold = config.staleness_threshold;

        tokio::spawn(async move {
            let executor = Executor::new(db, executor_id);
            let mut ticker = tokio::time::interval(threshold);
            loop {
                ticker.tick().await;
                match executor.recover_stale_tasks(threshold).await {
                    Ok(n) if n > 0 => {
                        tracing::info!("recovered {n} stale tasks from dead workers");
                    }
                    Ok(_) => {}
                    Err(e) => {
                        tracing::warn!("recover_stale_tasks failed: {e}");
                    }
                }
            }
        })
    }

    /// 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 = "
            UPDATE durable.task
            SET status = 'PENDING', started_at = NULL, deadline_epoch_ms = NULL
            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, name, kind
        ";

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

        let mut recovered = Vec::with_capacity(rows.len());
        for row in rows {
            let id: Uuid = row
                .try_get_by_index(0)
                .map_err(|e| DurableError::custom(e.to_string()))?;
            let name: String = row
                .try_get_by_index(1)
                .map_err(|e| DurableError::custom(e.to_string()))?;
            let kind: String = row
                .try_get_by_index(2)
                .map_err(|e| DurableError::custom(e.to_string()))?;
            recovered.push(RecoveredTask { id, name, kind });
        }

        Ok(recovered)
    }
}

#[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));
    }
}