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 ──────────────────────────────────────────────

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

/// 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"
                );
                match Ctx::from_id(&db_inner, task_id).await {
                    Ok(ctx) => {
                        if let Err(e) = (resume)(ctx).await {
                            tracing::error!(
                                workflow = %task_name,
                                id = %task_id,
                                error = %e,
                                "recovered workflow failed"
                            );
                        }
                    }
                    Err(e) => {
                        tracing::error!(
                            workflow = %task_name,
                            id = %task_id,
                            error = %e,
                            "failed to attach to recovered workflow"
                        );
                    }
                }
            });
        } 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)
}