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,
    /// The static handler function name (from `#[durable::workflow]`).
    /// Used to look up the registered resume function on recovery.
    /// Falls back to `name` when `None` (backwards compat).
    pub handler: Option<String>,
}

/// Parse a row from a `RETURNING id, parent_id, name, kind, handler, handler` 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()))?;
    let handler: Option<String> = row.try_get_by_index(4).ok().flatten();
    Ok(RecoveredTask {
        id,
        parent_id,
        name,
        kind,
        handler,
    })
}

/// 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}");
                }
            }
        })
    }

    /// Claim stale tasks from dead workers in a single atomic operation.
    /// Uses `FOR UPDATE SKIP LOCKED` so multiple executors can run this
    /// concurrently without blocking each other — each executor claims a
    /// disjoint set of tasks.
    ///
    /// 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
    /// 3. Orphaned tasks with NULL `executor_id` that have been RUNNING too long
    ///
    /// Claimed tasks go directly to RUNNING (no intermediate PENDING state)
    /// with `started_at` reset and `executor_id` set to this executor.
    pub async fn recover_stale_tasks(
        &self,
        staleness_threshold: Duration,
    ) -> Result<Vec<RecoveredTask>, DurableError> {
        let threshold_secs = staleness_threshold.as_secs();
        let sql = format!(
            "WITH claimable AS ( \
                 SELECT id FROM durable.task \
                 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 != '{eid}') \
                       OR \
                       (executor_id IS NOT NULL \
                        AND executor_id != '{eid}' \
                        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') \
                   ) \
                 FOR UPDATE SKIP LOCKED \
             ) \
             UPDATE durable.task \
             SET status = 'RUNNING', started_at = now(), executor_id = '{eid}' \
             WHERE id IN (SELECT id FROM claimable) \
             RETURNING id, parent_id, name, kind, handler",
            eid = self.executor_id
        );
        let rows = self
            .db
            .query_all(Statement::from_string(DbBackend::Postgres, sql))
            .await?;

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

    /// Claim timed-out tasks in a single atomic operation.
    /// Uses `FOR UPDATE SKIP LOCKED` for safe concurrent claiming.
    ///
    /// A task is considered timed-out if:
    /// - `status = 'RUNNING'` and `started_at IS NOT NULL`
    /// - `timeout_ms` has elapsed, or `deadline_epoch_ms` has passed
    ///
    /// Tasks without `timeout_ms` set are never considered timed-out.
    pub async fn recover(&self) -> Result<Vec<RecoveredTask>, DurableError> {
        let sql = format!(
            "WITH claimable AS ( \
                 SELECT id FROM durable.task \
                 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) \
                   ) \
                 FOR UPDATE SKIP LOCKED \
             ) \
             UPDATE durable.task \
             SET status = 'RUNNING', started_at = now(), deadline_epoch_ms = NULL, \
                 executor_id = '{eid}' \
             WHERE id IN (SELECT id FROM claimable) \
             RETURNING id, parent_id, name, kind, handler",
            eid = self.executor_id
        );

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

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

    /// Reset RUNNING child steps of recovered workflows back to PENDING.
    ///
    /// After a crash, TX1 in `step()` may have committed (setting the step to
    /// RUNNING) but the closure never completed. These orphaned steps block
    /// re-execution because `find_or_create_task` sees RUNNING and returns
    /// `StepLocked`. This method resets them so the recovered workflow can
    /// re-execute its steps from scratch.
    pub async fn reset_orphaned_steps(
        &self,
        recovered_ids: &[Uuid],
    ) -> Result<u64, DurableError> {
        if recovered_ids.is_empty() {
            return Ok(0);
        }
        let id_list: String = recovered_ids
            .iter()
            .map(|id| format!("'{id}'"))
            .collect::<Vec<_>>()
            .join(",");
        let sql = format!(
            "UPDATE durable.task \
             SET status = 'PENDING', started_at = NULL \
             WHERE parent_id IN ({id_list}) \
               AND status = 'RUNNING' \
               AND kind = 'STEP'"
        );
        let result = self
            .db
            .execute(Statement::from_string(DbBackend::Postgres, sql))
            .await?;
        Ok(result.rows_affected())
    }
}

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