durable/

lib.rs

pub mod ctx;
pub mod error;
pub mod executor;

pub use ctx::Ctx;
pub use ctx::RetryPolicy;
pub use ctx::{TaskQuery, TaskSort, TaskSummary};
pub use durable_db::entity::sea_orm_active_enums::TaskStatus;
pub use durable_macros::{step, workflow};
pub use error::DurableError;
pub use executor::{Executor, HeartbeatConfig, RecoveredTask};
pub use sea_orm::DatabaseTransaction;

// Re-export so macro-generated code can reference `durable::inventory::submit!`
pub use inventory;

use sea_orm::{ConnectOptions, Database, DatabaseConnection};
use sea_orm_migration::MigratorTrait;
use std::future::Future;
use std::pin::Pin;
use std::sync::RwLock;
use uuid::Uuid;

/// Global executor ID set by [`init`]. Read by [`Ctx::start`] to tag tasks
/// so heartbeat-based recovery can find them after a crash.
static EXECUTOR_ID: RwLock<Option<String>> = RwLock::new(None);

/// Returns the executor ID set by [`init`], or `None` if `init` was not called.
pub fn executor_id() -> Option<String> {
    EXECUTOR_ID.read().ok().and_then(|g| g.clone())
}

// ── Workflow auto-registration ──────────────────────────────────

/// A compiled-in registration of a workflow function for automatic crash recovery.
///
/// Produced by `#[durable::workflow]` and collected at link time by `inventory`.
/// Only workflows with a single `ctx: Ctx` parameter are registered.
pub struct WorkflowRegistration {
    /// The handler name — matches the function name, used for recovery lookup
    /// via the `handler` column in `durable.task`.
    pub name: &'static str,
    /// Resumes the workflow given a `Ctx`. Return type is erased; during
    /// recovery we only care about driving the workflow to completion.
    pub resume_fn: fn(Ctx) -> Pin<Box<dyn Future<Output = Result<(), DurableError>> + Send>>,
}

inventory::collect!(WorkflowRegistration);

/// Look up a registered workflow by name.
pub fn find_workflow(name: &str) -> Option<&'static WorkflowRegistration> {
    inventory::iter::<WorkflowRegistration>().find(|r| r.name == name)
}

// ── Initialization ──────────────────────────────────────────────

/// Resume a failed workflow from its last completed step (DBOS-style).
///
/// Resets the workflow from FAILED → RUNNING, resets failed child steps to
/// PENDING, and re-dispatches the registered workflow function. Completed
/// steps replay from their checkpointed outputs.
///
/// If no `WorkflowRegistration` is found for the workflow's handler, the
/// workflow is still reset but not auto-dispatched — use `Ctx::from_id()`
/// to drive it manually.
///
/// ```ignore
/// durable::resume_workflow(&db, workflow_id).await?;
/// ```
pub async fn resume_workflow(
    db: &DatabaseConnection,
    task_id: uuid::Uuid,
) -> Result<(), DurableError> {
    let handler = Ctx::resume_failed(db, task_id).await?;

    // Look up registered handler and auto-dispatch with recovery
    let lookup_key = handler.as_deref().unwrap_or("");
    if let Some(reg) = find_workflow(lookup_key) {
        let db_inner = db.clone();
        let resume = reg.resume_fn;
        tokio::spawn(async move {
            tracing::info!(
                id = %task_id,
                "re-dispatching resumed workflow"
            );
            run_workflow_with_recovery(db_inner, task_id, resume).await;
        });
    } else {
        tracing::warn!(
            id = %task_id,
            handler = ?handler,
            "no registered handler for resumed workflow — use Ctx::from_id() to resume manually"
        );
    }

    Ok(())
}

/// Initialize durable: connect to Postgres, run migrations, start heartbeat,
/// recover stale tasks, and auto-resume registered workflows.
///
/// After this call, [`Ctx::start`] automatically tags tasks with the executor
/// ID for crash recovery. Recovered root workflows that have a matching
/// `#[durable::workflow]` registration are automatically spawned as tokio
/// tasks.
///
/// Uses [`HeartbeatConfig::default()`] (60 s heartbeat, 180 s staleness).
/// For custom intervals use [`init_with_config`].
///
/// ```ignore
/// let (db, recovered) = durable::init("postgres://localhost/mydb").await?;
/// ```
pub async fn init(
    database_url: &str,
) -> Result<(DatabaseConnection, Vec<RecoveredTask>), DurableError> {
    init_with_config(database_url, HeartbeatConfig::default()).await
}

/// Like [`init`] but with a custom [`HeartbeatConfig`].
pub async fn init_with_config(
    database_url: &str,
    config: HeartbeatConfig,
) -> Result<(DatabaseConnection, Vec<RecoveredTask>), DurableError> {
    let mut opt = ConnectOptions::new(database_url);
    opt.set_schema_search_path("public,durable");
    let db = Database::connect(opt).await?;

    // Run migrations
    durable_db::Migrator::up(&db, None).await?;

    // Create executor with a unique ID for this process
    let eid = format!("exec-{}-{}", std::process::id(), uuid::Uuid::new_v4());
    if let Ok(mut guard) = EXECUTOR_ID.write() {
        *guard = Some(eid.clone());
    }
    let executor = Executor::new(db.clone(), eid);

    // Write initial heartbeat so other workers know we're alive
    executor.heartbeat().await?;

    let mut all_recovered = Vec::new();

    // Recover stale tasks: timeout/deadline-based
    let recovered = executor.recover().await?;
    if !recovered.is_empty() {
        tracing::info!(
            "recovered {} stale tasks (timeout/deadline)",
            recovered.len()
        );
    }
    all_recovered.extend(recovered);

    // Recover stale tasks: heartbeat-based (dead workers, unknown executors, orphaned NULL)
    let recovered = executor
        .recover_stale_tasks(config.staleness_threshold)
        .await?;
    if !recovered.is_empty() {
        tracing::info!(
            "recovered {} stale tasks from dead/unknown workers",
            recovered.len()
        );
    }
    all_recovered.extend(recovered);

    // Reset orphaned RUNNING steps of recovered workflows back to PENDING
    if !all_recovered.is_empty() {
        let recovered_ids: Vec<Uuid> = all_recovered.iter().map(|r| r.id).collect();
        match executor.reset_orphaned_steps(&recovered_ids).await {
            Ok(n) if n > 0 => {
                tracing::info!("reset {n} orphaned steps to PENDING");
            }
            Err(e) => tracing::warn!("failed to reset orphaned steps: {e}"),
            _ => {}
        }
    }

    // Auto-dispatch registered workflows
    dispatch_recovered(&db, &all_recovered);

    // Start background heartbeat loop
    executor.start_heartbeat(&config);

    // Start recovery loop that also auto-dispatches
    start_recovery_dispatch_loop(
        db.clone(),
        executor.executor_id().to_string(),
        config.staleness_threshold,
    );

    tracing::info!("durable initialized (executor={})", executor.executor_id());
    Ok((db, all_recovered))
}

/// Default backoff between automatic workflow recovery attempts.
const RECOVERY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(5);

/// Run a workflow with automatic recovery on failure (DBOS-style).
///
/// When the workflow function returns `Err`:
/// 1. Mark the workflow FAILED
/// 2. If `recovery_count < max_recovery_attempts`, reset and re-execute
/// 3. If recovery limit exceeded, leave it FAILED
///
/// Completed steps replay from checkpoints on each recovery attempt.
async fn run_workflow_with_recovery(
    db: DatabaseConnection,
    task_id: Uuid,
    resume_fn: fn(Ctx) -> Pin<Box<dyn Future<Output = Result<(), DurableError>> + Send>>,
) {
    loop {
        match Ctx::from_id(&db, task_id).await {
            Ok(ctx) => {
                match (resume_fn)(ctx).await {
                    Ok(()) => return, // workflow completed successfully
                    Err(e) => {
                        // Mark workflow as FAILED
                        if let Err(fail_err) = Ctx::fail_by_id(&db, task_id, &e.to_string()).await
                        {
                            tracing::error!(
                                id = %task_id,
                                error = %fail_err,
                                "failed to mark workflow as FAILED"
                            );
                            return;
                        }

                        tracing::warn!(
                            id = %task_id,
                            error = %e,
                            "workflow failed, attempting automatic recovery"
                        );

                        // Attempt auto-recovery
                        match Ctx::resume_failed(&db, task_id).await {
                            Ok(_) => {
                                tracing::info!(
                                    id = %task_id,
                                    "workflow auto-recovery: reset succeeded, re-executing"
                                );
                                tokio::time::sleep(RECOVERY_BACKOFF).await;
                                continue; // re-execute the workflow
                            }
                            Err(DurableError::MaxRecoveryExceeded(_)) => {
                                tracing::error!(
                                    id = %task_id,
                                    "workflow exceeded max recovery attempts, staying FAILED"
                                );
                                return;
                            }
                            Err(recover_err) => {
                                tracing::error!(
                                    id = %task_id,
                                    error = %recover_err,
                                    "workflow auto-recovery failed"
                                );
                                return;
                            }
                        }
                    }
                }
            }
            Err(e) => {
                tracing::error!(
                    id = %task_id,
                    error = %e,
                    "failed to attach to workflow"
                );
                return;
            }
        }
    }
}

/// Spawn registered workflow functions for recovered root tasks.
/// Tasks are already RUNNING (claimed atomically by the recovery SQL).
fn dispatch_recovered(db: &DatabaseConnection, recovered: &[RecoveredTask]) {
    for task in recovered {
        // Only auto-dispatch root workflows (children are driven by their parent)
        if task.parent_id.is_some() {
            continue;
        }

        let lookup_key = task.handler.as_deref().unwrap_or(&task.name);
        if let Some(reg) = find_workflow(lookup_key) {
            let db_inner = db.clone();
            let task_id = task.id;
            let task_name = task.name.clone();
            let resume = reg.resume_fn;
            tokio::spawn(async move {
                tracing::info!(
                    workflow = %task_name,
                    id = %task_id,
                    "auto-resuming recovered workflow"
                );
                run_workflow_with_recovery(db_inner, task_id, resume).await;
            });
        } else {
            tracing::warn!(
                workflow = %task.name,
                handler = ?task.handler,
                id = %task.id,
                "no registered handler for recovered task — use Ctx::from_id() to resume manually"
            );
        }
    }
}

/// Spawn a background loop that recovers stale tasks AND auto-dispatches them.
fn start_recovery_dispatch_loop(
    db: DatabaseConnection,
    executor_id: String,
    staleness_threshold: std::time::Duration,
) {
    tokio::spawn(async move {
        let executor = Executor::new(db.clone(), executor_id);
        let mut ticker = tokio::time::interval(staleness_threshold);
        loop {
            ticker.tick().await;

            // Collect all recovered workflow IDs so we can reset their orphaned steps.
            let mut all_recovered_ids = Vec::new();

            // Timeout-based recovery
            match executor.recover().await {
                Ok(ref recovered) if !recovered.is_empty() => {
                    tracing::info!(
                        "recovered {} stale tasks (timeout/deadline)",
                        recovered.len()
                    );
                    all_recovered_ids.extend(recovered.iter().map(|r| r.id));
                    dispatch_recovered(&db, recovered);
                }
                Err(e) => tracing::warn!("timeout recovery failed: {e}"),
                _ => {}
            }

            // Heartbeat-based recovery
            match executor.recover_stale_tasks(staleness_threshold).await {
                Ok(ref recovered) if !recovered.is_empty() => {
                    tracing::info!(
                        "recovered {} stale tasks from dead workers",
                        recovered.len()
                    );
                    all_recovered_ids.extend(recovered.iter().map(|r| r.id));
                    dispatch_recovered(&db, recovered);
                }
                Err(e) => tracing::warn!("heartbeat recovery failed: {e}"),
                _ => {}
            }

            // Reset orphaned RUNNING steps of recovered workflows back to
            // PENDING so they can be re-executed. Without this, steps that
            // committed TX1 (set RUNNING) but crashed before TX2 (checkpoint)
            // would be permanently stuck.
            if !all_recovered_ids.is_empty() {
                match executor.reset_orphaned_steps(&all_recovered_ids).await {
                    Ok(n) if n > 0 => {
                        tracing::info!("reset {n} orphaned steps to PENDING");
                    }
                    Err(e) => tracing::warn!("failed to reset orphaned steps: {e}"),
                    _ => {}
                }
            }
        }
    });
}

/// Initialize durable: connect to Postgres and run migrations only.
///
/// Does **not** start heartbeat or recovery loops, and does **not** set the
/// global executor ID. Use this for tests, migrations-only scripts, or when
/// you manage the [`Executor`] yourself.
///
/// ```ignore
/// let db = durable::init_db("postgres://localhost/mydb").await?;
/// ```
pub async fn init_db(database_url: &str) -> Result<DatabaseConnection, DurableError> {
    let mut opt = ConnectOptions::new(database_url);
    opt.set_schema_search_path("public,durable");
    let db = Database::connect(opt).await?;
    durable_db::Migrator::up(&db, None).await?;
    tracing::info!("durable initialized (db only)");
    Ok(db)
}