runledger_runtime/

worker.rs

use std::any::Any;
use std::cmp::min;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;

use futures_util::FutureExt;
use runledger_core::jobs::{JobCompletion, JobContext, JobFailure};
use runledger_postgres::QueryErrorKind;
use runledger_postgres::jobs::{self, JobProgressUpdate};
use tokio::sync::{Semaphore, watch};
use tokio::task::JoinSet;
use tokio::time::{Duration, Instant, MissedTickBehavior, sleep_until};
use tracing::{Instrument, error, info, info_span, warn};

mod completion;
mod dead_letter;
mod observers;

use self::completion::{
    CompletionObservation, complete_job_after_handler, complete_job_failure_after_handler,
};
use self::observers::{JobRunningNotification, TerminalJobObserverEvent, TerminalObserverTasks};
use crate::RuntimeLoopExit;
use crate::WorkerError;
use crate::config::JobsConfig;
use crate::observer::{JobLeaseLostEvent, JobLifecycleObservers, ObservedJob};
use crate::registry::JobRegistry;
use crate::shutdown;

const UNKNOWN_WORKER_ID: &str = "unknown-worker";
// Kept stable for clients that already match this code; it also covers leases
// that expired before the worker's lifecycle update reached storage.
const LEASE_OWNER_MISMATCH_CODE: &str = "job.lease_owner_mismatch";
const LEASE_MAINTENANCE_FAILED_CODE: &str = "job.lease_maintenance_failed";
const HANDLER_PANIC_CODE: &str = "job.handler_panic";
const RUNNING_PROGRESS_PERSIST_FAILED_REASON: &str = "RUNNING_PROGRESS_PERSIST_FAILED";
const UNSTARTED_CLAIM_RETRY_DELAY_MS: i32 = 1_000;

enum JobExecutionFailure {
    Handler(JobFailure),
    LeaseMaintenance(JobFailure),
}

pub async fn run_worker_loop(
    pool: runledger_postgres::DbPool,
    registry: JobRegistry,
    config: JobsConfig,
    shutdown: watch::Receiver<bool>,
) -> RuntimeLoopExit {
    run_worker_loop_with_observer(
        pool,
        registry,
        config,
        shutdown,
        JobLifecycleObservers::empty(),
    )
    .await
}

pub async fn run_worker_loop_with_observer(
    pool: runledger_postgres::DbPool,
    registry: JobRegistry,
    config: JobsConfig,
    mut shutdown: watch::Receiver<bool>,
    observers: JobLifecycleObservers,
) -> RuntimeLoopExit {
    if let Err(error) = config.validate_worker_loop() {
        warn!(%error, "invalid jobs config; stopping worker loop");
        return RuntimeLoopExit::InvalidConfig(error);
    }

    let registry = Arc::new(registry);
    let claimable_job_types = registry.registered_types();
    let semaphore = Arc::new(Semaphore::new(config.max_global_concurrency));
    let mut join_set: JoinSet<()> = JoinSet::new();
    let terminal_observer_tasks = TerminalObserverTasks::owned();

    loop {
        drain_finished_tasks(&mut join_set).await;
        terminal_observer_tasks.drain_finished().await;

        if shutdown::is_requested_or_closed(&shutdown) {
            return drain_worker_tasks(
                join_set,
                terminal_observer_tasks,
                RuntimeLoopExit::Shutdown,
            )
            .await;
        }

        if claimable_job_types.is_empty() {
            if shutdown::wait_for_request_or_timeout(&mut shutdown, config.poll_interval).await {
                return drain_worker_tasks(
                    join_set,
                    terminal_observer_tasks,
                    RuntimeLoopExit::Shutdown,
                )
                .await;
            }
            continue;
        }

        let available = semaphore.available_permits();
        if available == 0 {
            if shutdown::wait_for_request_or_timeout(&mut shutdown, config.poll_interval).await {
                return drain_worker_tasks(
                    join_set,
                    terminal_observer_tasks,
                    RuntimeLoopExit::Shutdown,
                )
                .await;
            }
            continue;
        }

        let claim_limit = min(available, config.claim_batch_size as usize);
        let claimed = match jobs::claim_prestart_jobs_for_types(
            &pool,
            &config.worker_id,
            config.lease_ttl_seconds,
            claim_limit as i64,
            &claimable_job_types,
        )
        .await
        {
            Ok(claimed) => claimed,
            Err(error) => {
                let error = WorkerError::ClaimJobs {
                    worker_id: config.worker_id.clone(),
                    source: error,
                };
                warn!(%error, "worker claim failed");
                Vec::new()
            }
        };

        if claimed.is_empty() {
            shutdown::wait_for_request_or_timeout(&mut shutdown, config.poll_interval).await;
            continue;
        }

        let claimed_len = claimed.len();
        for job in claimed {
            let permit = match Arc::clone(&semaphore).acquire_owned().await {
                Ok(permit) => permit,
                Err(_) => {
                    // The worker owns this semaphore and never closes it. If
                    // this defensive branch fires, surface it as an unexpected
                    // loop completion rather than graceful shutdown.
                    warn!("worker semaphore closed; stopping worker loop");
                    return drain_worker_tasks(
                        join_set,
                        terminal_observer_tasks,
                        RuntimeLoopExit::Completed,
                    )
                    .await;
                }
            };
            let pool_clone = pool.clone();
            let registry_clone = Arc::clone(&registry);
            let lease_ttl_seconds = config.lease_ttl_seconds;
            let observers = observers.clone();
            let terminal_observer_tasks = terminal_observer_tasks.clone();
            join_set.spawn(async move {
                let _permit = permit;
                process_claimed_job_with_terminal_observers(
                    pool_clone,
                    registry_clone,
                    job,
                    lease_ttl_seconds,
                    observers,
                    terminal_observer_tasks,
                )
                .await;
            });
        }

        if claimed_len == claim_limit {
            continue;
        }

        if shutdown::wait_for_request_or_timeout(&mut shutdown, config.poll_interval).await {
            return drain_worker_tasks(
                join_set,
                terminal_observer_tasks,
                RuntimeLoopExit::Shutdown,
            )
            .await;
        }
    }
}

async fn drain_worker_tasks(
    mut join_set: JoinSet<()>,
    terminal_observer_tasks: TerminalObserverTasks,
    exit: RuntimeLoopExit,
) -> RuntimeLoopExit {
    if !join_set.is_empty() {
        match exit {
            RuntimeLoopExit::Shutdown => {
                info!("worker shutdown requested; draining in-flight jobs")
            }
            RuntimeLoopExit::InvalidConfig(_) => {
                warn!("worker loop rejected invalid config; draining in-flight jobs");
            }
            RuntimeLoopExit::Completed => {
                warn!("worker loop completed before shutdown; draining in-flight jobs");
            }
        }
    }
    while let Some(result) = join_set.join_next().await {
        if let Err(error) = result {
            error!(%error, "job task crashed while draining in-flight jobs");
        }
    }
    terminal_observer_tasks.drain_for_shutdown().await;
    exit
}

async fn drain_finished_tasks(join_set: &mut JoinSet<()>) {
    while let Some(result) = join_set.try_join_next() {
        if let Err(error) = result {
            error!(%error, "job task crashed");
        }
    }
}

#[cfg(test)]
async fn process_claimed_job(
    pool: runledger_postgres::DbPool,
    registry: Arc<JobRegistry>,
    job: jobs::JobQueueRecord,
    lease_ttl_seconds: i32,
) {
    process_claimed_job_with_observer(
        pool,
        registry,
        job,
        lease_ttl_seconds,
        JobLifecycleObservers::empty(),
    )
    .await;
}

#[cfg(test)]
async fn process_claimed_job_with_observer(
    pool: runledger_postgres::DbPool,
    registry: Arc<JobRegistry>,
    job: jobs::JobQueueRecord,
    lease_ttl_seconds: i32,
    observers: JobLifecycleObservers,
) {
    process_claimed_job_with_terminal_observers(
        pool,
        registry,
        job,
        lease_ttl_seconds,
        observers,
        TerminalObserverTasks::detached(),
    )
    .await;
}

async fn process_claimed_job_with_terminal_observers(
    pool: runledger_postgres::DbPool,
    registry: Arc<JobRegistry>,
    job: jobs::JobQueueRecord,
    lease_ttl_seconds: i32,
    observers: JobLifecycleObservers,
    terminal_observer_tasks: TerminalObserverTasks,
) {
    let worker_id = job
        .worker_id
        .clone()
        .unwrap_or_else(|| UNKNOWN_WORKER_ID.to_owned());

    let job_span = info_span!(
        "job",
        sentry.name = %job.job_type,
        sentry.op = "runledger.job",
        job_id = %job.id,
        job_type = %job.job_type,
        run_number = job.run_number,
        attempt = job.attempt,
        organization_id = ?job.organization_id,
        worker_id = %worker_id,
    );
    async {
        let start = Instant::now();
        let context = JobContext {
            job_id: job.id,
            run_number: job.run_number,
            attempt: job.attempt,
            organization_id: job.organization_id,
            worker_id: worker_id.clone(),
            checkpoint: job.checkpoint.clone(),
        };
        let observed_job = observed_job(&job, &worker_id);

        if !mark_job_running_or_abort(&pool, &context, &job).await {
            return;
        }
        let mut running_notification =
            JobRunningNotification::spawn(observers.clone(), observed_job.clone());

        match execute_job_handler_with_heartbeats(
            pool.clone(),
            Arc::clone(&registry),
            &context,
            &job,
            lease_ttl_seconds,
        )
        .await
        {
            Ok(completion) => {
                complete_job_after_handler(
                    &pool,
                    registry.as_ref(),
                    &context,
                    &job,
                    completion,
                    CompletionObservation::new(
                        &observers,
                        observed_job.clone(),
                        start.elapsed(),
                        &mut running_notification,
                        &terminal_observer_tasks,
                    ),
                )
                .await;
            }
            Err(JobExecutionFailure::Handler(failure)) => {
                complete_job_failure_after_handler(
                    &pool,
                    registry.as_ref(),
                    &context,
                    &job,
                    failure,
                    CompletionObservation::new(
                        &observers,
                        observed_job.clone(),
                        start.elapsed(),
                        &mut running_notification,
                        &terminal_observer_tasks,
                    ),
                )
                .await;
            }
            Err(JobExecutionFailure::LeaseMaintenance(failure)) => {
                log_lease_maintenance_abort(&job, &failure);
                running_notification
                    .spawn_terminal_observer(
                        &terminal_observer_tasks,
                        &job,
                        observers.clone(),
                        TerminalJobObserverEvent::LeaseLost(JobLeaseLostEvent {
                            job: observed_job.clone(),
                            duration: start.elapsed(),
                            failure,
                        }),
                    )
                    .await;
            }
        }

        info!(
            job_id = %job.id,
            attempt = job.attempt,
            run_number = job.run_number,
            elapsed_ms = start.elapsed().as_millis(),
            "job processed"
        );
    }
    .instrument(job_span)
    .await;
}

fn observed_job(job: &jobs::JobQueueRecord, worker_id: &str) -> ObservedJob {
    ObservedJob {
        job_id: job.id,
        job_type: job.job_type.clone(),
        organization_id: job.organization_id,
        run_number: job.run_number,
        attempt: job.attempt,
        max_attempts: job.max_attempts,
        worker_id: worker_id.to_owned(),
    }
}

fn log_lease_maintenance_abort(job: &jobs::JobQueueRecord, failure: &JobFailure) {
    warn!(
        job_id = %job.id,
        attempt = job.attempt,
        failure_code = failure.code,
        "job processing aborted because durable lease maintenance was lost"
    );
}

async fn mark_job_running_or_abort(
    pool: &runledger_postgres::DbPool,
    context: &JobContext,
    job: &jobs::JobQueueRecord,
) -> bool {
    let running_progress = JobProgressUpdate {
        stage: Some(runledger_core::jobs::JobStage::Running),
        progress_done: None,
        progress_total: None,
        checkpoint: None,
    };

    let Err(source) = jobs::update_job_progress(
        pool,
        job.id,
        job.run_number,
        job.attempt,
        &context.worker_id,
        &running_progress,
    )
    .await
    else {
        return true;
    };

    handle_running_progress_persist_failure(pool, context, job, source).await;
    false
}

async fn handle_running_progress_persist_failure(
    pool: &runledger_postgres::DbPool,
    context: &JobContext,
    job: &jobs::JobQueueRecord,
    source: runledger_postgres::Error,
) {
    let lease_owner_mismatch = is_lease_owner_mismatch_error(&source);
    let error = WorkerError::SetRunningProgress {
        job_id: job.id,
        attempt: job.attempt,
        source,
    };

    if lease_owner_mismatch {
        warn!(
            %error,
            job_id = %job.id,
            attempt = job.attempt,
            "aborting job before execution because lease ownership was already lost"
        );
        return;
    }

    match jobs::release_unstarted_job_claim(
        pool,
        job.id,
        job.run_number,
        job.attempt,
        &context.worker_id,
        RUNNING_PROGRESS_PERSIST_FAILED_REASON,
        UNSTARTED_CLAIM_RETRY_DELAY_MS,
    )
    .await
    {
        Ok(()) => {
            warn!(
                %error,
                job_id = %job.id,
                attempt = job.attempt,
                "running progress could not be persisted; released unstarted claim back to pending"
            );
        }
        Err(release_error) => {
            let no_longer_releasable =
                is_unstarted_claim_release_not_applicable_error(&release_error);
            let release_error = WorkerError::ReleaseUnstartedClaim {
                job_id: job.id,
                attempt: job.attempt,
                source: release_error,
            };
            if no_longer_releasable {
                warn!(
                    %error,
                    %release_error,
                    job_id = %job.id,
                    attempt = job.attempt,
                    "running progress could not be persisted; unstarted release no longer applies and the job will continue under the current lease owner"
                );
                return;
            }

            warn!(
                %error,
                %release_error,
                job_id = %job.id,
                attempt = job.attempt,
                "running progress could not be persisted; leaving claim for reaper recovery"
            );
        }
    }
}

async fn execute_job_handler(
    registry: Arc<JobRegistry>,
    context: &JobContext,
    job: &jobs::JobQueueRecord,
) -> Result<JobCompletion, JobFailure> {
    let Some(handler) = registry.get(job.job_type.as_borrowed()) else {
        return Err(JobFailure::terminal(
            "job.handler_not_registered",
            "No handler is registered for this job type.",
        ));
    };

    handler.execute(context.clone(), job.payload.clone()).await
}

async fn execute_job_handler_with_heartbeats(
    pool: runledger_postgres::DbPool,
    registry: Arc<JobRegistry>,
    context: &JobContext,
    job: &jobs::JobQueueRecord,
    lease_ttl_seconds: i32,
) -> Result<JobCompletion, JobExecutionFailure> {
    let mut execution =
        Box::pin(AssertUnwindSafe(execute_job_handler(registry, context, job)).catch_unwind());
    let timeout_deadline = Instant::now() + Duration::from_secs(job.timeout_seconds.max(1) as u64);
    let mut timeout = Box::pin(sleep_until(timeout_deadline));

    let mut ticker = tokio::time::interval(heartbeat_interval(lease_ttl_seconds));
    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
    ticker.tick().await;

    loop {
        tokio::select! {
            result = &mut execution => {
                return match result {
                    Ok(result) => result.map_err(JobExecutionFailure::Handler),
                    Err(panic_payload) => {
                        Err(JobExecutionFailure::Handler(handler_panic_failure(panic_payload)))
                    }
                };
            }
            _ = &mut timeout => {
                return Err(JobExecutionFailure::Handler(JobFailure::timeout(
                    "job.timeout_exceeded",
                    "Job exceeded the configured timeout.",
                )));
            }
            _ = ticker.tick() => {
                if let Err(error) = jobs::heartbeat_job(
                    &pool,
                    job.id,
                    job.run_number,
                    job.attempt,
                    &context.worker_id,
                    lease_ttl_seconds,
                )
                .await
                {
                    let lease_owner_mismatch = is_lease_owner_mismatch_error(&error);
                    let error = WorkerError::Heartbeat {
                        job_id: job.id,
                        attempt: job.attempt,
                        source: error,
                    };

                    if lease_owner_mismatch {
                        warn!(%error, job_id = %job.id, "job heartbeat lost lease ownership");
                        return Err(JobExecutionFailure::LeaseMaintenance(
                            lease_owner_mismatch_failure(),
                        ));
                    }

                    warn!(
                        %error,
                        job_id = %job.id,
                        "aborting job because lease heartbeat could not be persisted"
                    );
                    return Err(JobExecutionFailure::LeaseMaintenance(
                        lease_maintenance_failure(),
                    ));
                }
            }
        }
    }
}

fn lease_owner_mismatch_failure() -> JobFailure {
    JobFailure::lease_expired(
        LEASE_OWNER_MISMATCH_CODE,
        "Job lease ownership was lost during processing.",
    )
}

fn lease_maintenance_failure() -> JobFailure {
    JobFailure::lease_expired(
        LEASE_MAINTENANCE_FAILED_CODE,
        "Job lease could not be durably maintained during processing.",
    )
}

fn handler_panic_failure(panic_payload: Box<dyn Any + Send>) -> JobFailure {
    JobFailure::panicked(
        HANDLER_PANIC_CODE,
        format!(
            "Job handler panicked: {}",
            panic_payload_message(&*panic_payload)
        ),
    )
}

fn panic_payload_message(panic_payload: &(dyn Any + Send)) -> String {
    if let Some(message) = panic_payload.downcast_ref::<String>() {
        return message.clone();
    }

    if let Some(message) = panic_payload.downcast_ref::<&'static str>() {
        return (*message).to_string();
    }

    "non-string panic payload".to_string()
}

fn has_query_error_kind(error: &runledger_postgres::Error, expected_kind: QueryErrorKind) -> bool {
    matches!(
        error,
        runledger_postgres::Error::QueryError(query_error)
            if query_error.kind() == Some(expected_kind)
    )
}

fn is_lease_owner_mismatch_error(error: &runledger_postgres::Error) -> bool {
    has_query_error_kind(error, QueryErrorKind::JobLeaseOwnerMismatch)
}

fn is_unstarted_claim_release_not_applicable_error(error: &runledger_postgres::Error) -> bool {
    has_query_error_kind(error, QueryErrorKind::JobUnstartedClaimReleaseNotApplicable)
}

fn heartbeat_interval(lease_ttl_seconds: i32) -> Duration {
    // Renew at one-third of the lease TTL so a delayed heartbeat still leaves
    // time for subsequent renewals before the lease expires.
    let seconds = (lease_ttl_seconds.max(1) / 3).max(1) as u64;
    Duration::from_secs(seconds)
}

#[cfg(test)]
mod tests;