photos/src/jobs.rs

use std::time::Duration;

use uuid::Uuid;

use crate::models::{JobKind, JobStatus, PhotoStatus};
use crate::state::AppState;

/// Hard ceiling on a single job run; the sole bound for a live-but-hung worker
/// (a hung S3 read or exiftool child), since the heartbeat keeps the reaper away.
const JOB_TIMEOUT: Duration = Duration::from_secs(30 * 60);
/// How often a running job refreshes its lock. Must stay well below the
/// reaper's staleness threshold.
const HEARTBEAT_EVERY: Duration = Duration::from_secs(300);
/// A 'running' job whose lock is older than this had its worker die.
const STALE_AFTER: &str = "15 minutes";

#[derive(Debug, sqlx::FromRow)]
pub struct Job {
    pub id: Uuid,
    pub kind: String,
    pub payload: serde_json::Value,
    pub attempts: i32,
    pub max_attempts: i32,
}

pub async fn enqueue<'e, E>(
    executor: E,
    kind: JobKind,
    payload: serde_json::Value,
) -> Result<(), sqlx::Error>
where
    E: sqlx::PgExecutor<'e>,
{
    sqlx::query("insert into jobs (kind, payload) values ($1, $2)")
        .bind(kind.as_str())
        .bind(payload)
        .execute(executor)
        .await?;
    Ok(())
}

/// Make sure a process_photo job will (re)run for this photo: bump a queued
/// one to run now with fresh attempts; leave a running one alone (resetting
/// its attempts wouldn't reach the in-flight worker, which decides exhaustion
/// from its claim-time copy — it will finish or fail on its own and the photo
/// can be retried again); enqueue fresh otherwise.
pub async fn ensure_process_photo(
    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
    photo_id: Uuid,
) -> Result<(), sqlx::Error> {
    let requeued = sqlx::query(
        "update jobs set run_at = now(), attempts = 0
         where kind = $2 and status = $3 and payload->>'photo_id' = $1",
    )
    .bind(photo_id.to_string())
    .bind(JobKind::ProcessPhoto.as_str())
    .bind(JobStatus::Queued.as_str())
    .execute(&mut **tx)
    .await?;
    if requeued.rows_affected() > 0 {
        return Ok(());
    }
    let running: Option<(Uuid,)> = sqlx::query_as(
        "select id from jobs
         where kind = $2 and status = $3 and payload->>'photo_id' = $1",
    )
    .bind(photo_id.to_string())
    .bind(JobKind::ProcessPhoto.as_str())
    .bind(JobStatus::Running.as_str())
    .fetch_optional(&mut **tx)
    .await?;
    if running.is_some() {
        return Ok(());
    }
    enqueue(
        &mut **tx,
        JobKind::ProcessPhoto,
        serde_json::json!({ "photo_id": photo_id }),
    )
    .await
}

pub async fn run_worker(state: AppState) {
    let concurrency = state.config.worker_concurrency.max(1);
    // Unique per process so locked_by distinguishes workers across replicas.
    let instance = crate::auth::random_token(6);
    tracing::info!("starting worker {instance} with concurrency {concurrency}");
    let mut handles = Vec::new();
    handles.push(tokio::spawn(reaper_loop(state.clone())));
    for i in 0..concurrency {
        let state = state.clone();
        handles.push(tokio::spawn(worker_loop(state, format!("worker-{instance}-{i}"))));
    }
    for handle in handles {
        let _ = handle.await;
    }
}

/// Requeue stale jobs whose worker died mid-run — but only while they have
/// attempts left; exhausted stale jobs are failed outright so a job that
/// crashes its worker (e.g. OOM during decode) cannot crash-loop forever.
async fn reaper_loop(state: AppState) {
    loop {
        let failed: Result<Vec<(String, serde_json::Value)>, sqlx::Error> = sqlx::query_as(&format!(
            "update jobs set status = $1, locked_by = null,
                    last_error = coalesce(last_error, 'worker lost repeatedly (crash loop?)')
             where (status = $2 and locked_at < now() - interval '{STALE_AFTER}'
                    or status = $3)
               and attempts >= max_attempts
             returning kind, payload"
        ))
        .bind(JobStatus::Failed.as_str())
        .bind(JobStatus::Running.as_str())
        .bind(JobStatus::Queued.as_str())
        .fetch_all(&state.db)
        .await;
        match failed {
            Ok(jobs) => {
                for (kind, payload) in jobs {
                    tracing::error!(kind, "reaper failed exhausted job");
                    mark_photo_error(&state, &kind, &payload, "processing failed repeatedly").await;
                }
            }
            Err(e) => tracing::error!("job reaper (fail pass) errored: {e}"),
        }

        let requeued = sqlx::query(&format!(
            "update jobs set status = $1, locked_by = null, locked_at = null
             where status = $2 and locked_at < now() - interval '{STALE_AFTER}'
               and attempts < max_attempts"
        ))
        .bind(JobStatus::Queued.as_str())
        .bind(JobStatus::Running.as_str())
        .execute(&state.db)
        .await;
        match requeued {
            Ok(r) if r.rows_affected() > 0 => {
                tracing::warn!("requeued {} stale running job(s)", r.rows_affected())
            }
            Ok(_) => {}
            Err(e) => tracing::error!("job reaper (requeue pass) errored: {e}"),
        }
        tokio::time::sleep(Duration::from_secs(60)).await;
    }
}

/// Terminal-failure side effect for process_photo jobs: surface the error on
/// the photo, but never overwrite a photo a newer job already finished.
async fn mark_photo_error(state: &AppState, kind: &str, payload: &serde_json::Value, message: &str) {
    if kind != JobKind::ProcessPhoto.as_str() {
        return;
    }
    let Some(photo_id) = payload
        .get("photo_id")
        .and_then(|v| v.as_str())
        .and_then(|s| Uuid::parse_str(s).ok())
    else {
        return;
    };
    // A job can die before its first status write, so rescue photos stuck in
    // 'uploaded' as well as 'processing' — but never overwrite 'ready'.
    let _ = sqlx::query(
        "update photos set status = $3, error = $2
         where id = $1 and status in ($4, $5)",
    )
    .bind(photo_id)
    .bind(message)
    .bind(PhotoStatus::Error.as_str())
    .bind(PhotoStatus::Processing.as_str())
    .bind(PhotoStatus::Uploaded.as_str())
    .execute(&state.db)
    .await;
}

async fn worker_loop(state: AppState, name: String) {
    loop {
        match claim(&state, &name).await {
            Ok(Some(job)) => execute(&state, job, &name).await,
            Ok(None) => tokio::time::sleep(Duration::from_secs(2)).await,
            Err(e) => {
                tracing::error!("failed to claim job: {e}");
                tokio::time::sleep(Duration::from_secs(5)).await;
            }
        }
    }
}

async fn claim(state: &AppState, name: &str) -> Result<Option<Job>, sqlx::Error> {
    sqlx::query_as(
        "update jobs
         set status = $2, locked_by = $1, locked_at = now(), attempts = attempts + 1
         where id = (
             select id from jobs
             where status = $3 and run_at <= now() and attempts < max_attempts
             order by created_at
             limit 1
             for update skip locked
         )
         returning id, kind, payload, attempts, max_attempts",
    )
    .bind(name)
    .bind(JobStatus::Running.as_str())
    .bind(JobStatus::Queued.as_str())
    .fetch_optional(&state.db)
    .await
}

/// Keep locked_at fresh while a job runs so the reaper never requeues a job
/// whose worker is alive. Never completes; raced against the job in select!.
async fn heartbeat(state: &AppState, job_id: Uuid, name: &str) {
    loop {
        tokio::time::sleep(HEARTBEAT_EVERY).await;
        let _ = sqlx::query(
            "update jobs set locked_at = now()
             where id = $1 and locked_by = $2 and status = $3",
        )
        .bind(job_id)
        .bind(name)
        .bind(JobStatus::Running.as_str())
        .execute(&state.db)
        .await;
    }
}

async fn execute(state: &AppState, job: Job, name: &str) {
    tracing::info!(job_id = %job.id, kind = %job.kind, attempt = job.attempts, "job started");
    // Parse the kind here (not at claim decode) so an unknown kind — e.g.
    // enqueued by a newer deploy — fails THIS job normally instead of
    // poisoning the claim loop.
    let run = async {
        match job.kind.parse::<JobKind>() {
            Ok(JobKind::ProcessPhoto) => {
                crate::imaging::process_photo_job(state, &job.payload).await
            }
            Ok(JobKind::DeleteS3Prefix) => delete_s3_prefix_job(state, &job.payload).await,
            Err(e) => Err(anyhow::anyhow!(e)),
        }
    };
    let result = tokio::select! {
        result = tokio::time::timeout(JOB_TIMEOUT, run) => match result {
            Ok(result) => result,
            Err(_) => Err(anyhow::anyhow!(
                "job timed out after {}s",
                JOB_TIMEOUT.as_secs()
            )),
        },
        _ = heartbeat(state, job.id, name) => unreachable!("heartbeat never completes"),
    };
    // Finalization is guarded on locked_by so a worker whose job was reclaimed
    // (reaper) cannot overwrite the state written by the new owner.
    match result {
        Ok(()) => {
            let updated = sqlx::query(
                "update jobs set status = $3, locked_by = null, last_error = null
                 where id = $1 and locked_by = $2 and status = $4",
            )
            .bind(job.id)
            .bind(name)
            .bind(JobStatus::Done.as_str())
            .bind(JobStatus::Running.as_str())
            .execute(&state.db)
            .await;
            match updated {
                Ok(r) if r.rows_affected() == 0 => {
                    tracing::warn!(job_id = %job.id, "job was reclaimed by another worker; result discarded")
                }
                Ok(_) => tracing::info!(job_id = %job.id, kind = %job.kind, "job done"),
                Err(e) => tracing::error!(job_id = %job.id, "failed to finalize job: {e}"),
            }
        }
        Err(e) => {
            let message = format!("{e:#}");
            let exhausted = job.attempts >= job.max_attempts;
            tracing::error!(job_id = %job.id, kind = %job.kind, exhausted, "job failed: {message}");
            // Two self-contained query+bind branches — the placeholder lists
            // and bind chains must never be shared across branches.
            let finalize = if exhausted {
                sqlx::query(
                    "update jobs set status = $4, locked_by = null, last_error = $2
                     where id = $1 and locked_by = $3 and status = $5",
                )
                .bind(job.id)
                .bind(&message)
                .bind(name)
                .bind(JobStatus::Failed.as_str())
                .bind(JobStatus::Running.as_str())
            } else {
                let backoff = 30.0 * f64::from(job.attempts * job.attempts);
                sqlx::query(
                    "update jobs
                     set status = $4, locked_by = null, last_error = $2,
                         run_at = now() + make_interval(secs => $6)
                     where id = $1 and locked_by = $3 and status = $5",
                )
                .bind(job.id)
                .bind(&message)
                .bind(name)
                .bind(JobStatus::Queued.as_str())
                .bind(JobStatus::Running.as_str())
                .bind(backoff)
            };
            let owned = match finalize.execute(&state.db).await {
                Ok(r) => r.rows_affected() > 0,
                Err(e) => {
                    tracing::error!(job_id = %job.id, "failed to finalize job: {e}");
                    false
                }
            };
            if owned && exhausted {
                mark_photo_error(state, &job.kind, &job.payload, &message).await;
            }
        }
    }
}

async fn delete_s3_prefix_job(
    state: &AppState,
    payload: &serde_json::Value,
) -> anyhow::Result<()> {
    let prefix = payload
        .get("prefix")
        .and_then(|v| v.as_str())
        .ok_or_else(|| anyhow::anyhow!("missing prefix in payload"))?;
    anyhow::ensure!(
        prefix.starts_with("photos/") && prefix.ends_with('/'),
        "refusing to delete suspicious prefix {prefix:?}"
    );
    crate::s3::delete_prefix(state, prefix).await
}