remora/crates/remora-attempt/src/recovery.rs

use std::time::Duration;

use remora_coordination::{LeaseOwner, OpCode};
use remora_domain::{
    FailureClass, FailureDisposition, JobPhase, OpaqueId, fingerprint, ports::MutationScope,
};
use remora_durable::{
    DurableError, MutationDisposition, MutationStatus, PgTx, enqueue, execute_mutation,
    with_transaction,
};
use remora_target::{attempt_key_attempt_id, attempt_staging_prefix, validate_attempt_key};
use serde::Serialize;
use serde_json::json;
use sqlx::{FromRow, Postgres, Transaction, types::Json};
use time::OffsetDateTime;

use crate::service::new_id;
use crate::{
    AttemptAuthority, AttemptError, AttemptService, CancelOutcome, CleanupClaim, CleanupCompletion,
    CleanupIdentity, CooperativeStopOutcome, FailureRecord, ReconcileReport, RecoveryConfig,
    StagingMultipart, StagingObject, StagingWriteIntent,
};

impl AttemptService {
    pub async fn request_cancel(&self, job_id: &OpaqueId) -> Result<CancelOutcome, AttemptError> {
        with_transaction(&self.pool, |tx| {
            let job_id = job_id.clone();
            Box::pin(async move { Self::request_cancel_in_tx(tx, &job_id).await })
        })
        .await
        .map_err(Into::into)
    }

    pub async fn request_cancel_in_tx(
        tx: &mut PgTx<'_>,
        job_id: &OpaqueId,
    ) -> Result<CancelOutcome, DurableError> {
        let Some(row) = sqlx::query_as::<_, CancelRow>(
            "SELECT phase, cancel_requested FROM jobs WHERE job_id = $1 FOR UPDATE",
        )
        .bind(job_id.as_str())
        .fetch_optional(tx.as_mut())
        .await?
        else {
            return Ok(CancelOutcome::NotFound);
        };
        let phase = parse_phase(&row.phase)?;
        if phase.is_terminal() {
            return Ok(CancelOutcome::AlreadyTerminal);
        }
        if row.cancel_requested {
            return Ok(CancelOutcome::AlreadyRequested);
        }
        let now = OffsetDateTime::now_utc();
        if phase == JobPhase::Running {
            sqlx::query(
                "UPDATE jobs SET cancel_requested = TRUE, phase = 'cancelling',
                    dispatch_needed = FALSE, coordination_cleaned_at = NULL,
                    revision = revision + 1, updated_at = $2
                 WHERE job_id = $1",
            )
            .bind(job_id.as_str())
            .bind(now)
            .execute(tx.as_mut())
            .await?;
        } else {
            sqlx::query(
                "UPDATE jobs SET cancel_requested = TRUE, phase = 'cancelled',
                    authority_attempt_id = NULL, dispatch_needed = FALSE,
                    terminal_at = $2, coordination_cleaned_at = NULL,
                    revision = revision + 1, updated_at = $2
                 WHERE job_id = $1",
            )
            .bind(job_id.as_str())
            .bind(now)
            .execute(tx.as_mut())
            .await?;
        }
        enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
        Ok(CancelOutcome::Requested)
    }

    pub async fn cooperative_stop(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<CooperativeStopOutcome, AttemptError> {
        let authority_for_effect = authority.clone();
        let coordination = self.coordination.clone();
        let outcome = with_transaction(&self.pool, move |tx| {
            Box::pin(async move {
                cooperative_stop_transaction(tx, &coordination, &authority_for_effect).await
            })
        })
        .await?;
        if outcome == CooperativeStopOutcome::Stopped {
            let _ = self
                .coordination
                .release(authority.job_id.as_str(), &lease_owner(authority))
                .await;
            self.release_local(&authority.job_id);
        }
        Ok(outcome)
    }

    pub async fn cooperative_stop_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
    ) -> Result<CooperativeStopOutcome, AttemptError> {
        let request_fingerprint = fingerprint(&StopRequest {
            job_id: &authority.job_id,
            attempt_id: &authority.attempt_id,
            fence: authority.fence,
            worker_incarnation: &authority.worker_incarnation,
        })
        .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        let authority_for_effect = authority.clone();
        let coordination = self.coordination.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &self.operation_subject,
                operation: "attempt.mutation",
                key: operation_id.as_str(),
            },
            "cooperative_stop",
            &[],
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let decision =
                        cooperative_stop_transaction(tx, &coordination, &authority_for_effect)
                            .await?;
                    let status = if decision == CooperativeStopOutcome::AuthorityMismatch {
                        MutationStatus::Rejected
                    } else {
                        MutationStatus::Accepted
                    };
                    Ok((
                        status,
                        serde_json::to_value(decision).map_err(DurableError::effect)?,
                    ))
                })
            },
        )
        .await?;
        if outcome.disposition == MutationDisposition::Conflict {
            return Err(AttemptError::OperationConflict);
        }
        let decision = serde_json::from_value(outcome.response)
            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        if decision == CooperativeStopOutcome::Stopped {
            let _ = self
                .coordination
                .release(authority.job_id.as_str(), &lease_owner(authority))
                .await;
            self.release_local(&authority.job_id);
        }
        Ok(decision)
    }

    pub async fn mark_worker_down(
        &self,
        worker_incarnation: &OpaqueId,
    ) -> Result<u64, AttemptError> {
        let changed = sqlx::query(
            "UPDATE attempts SET worker_down_at = COALESCE(worker_down_at, $2)
             WHERE worker_incarnation = $1 AND state IN ('authorized', 'executing')",
        )
        .bind(worker_incarnation.as_str())
        .bind(OffsetDateTime::now_utc())
        .execute(&self.pool)
        .await?;
        Ok(changed.rows_affected())
    }

    pub async fn register_staging_object(
        &self,
        attempt_id: &OpaqueId,
        object: &StagingObject,
    ) -> Result<(), AttemptError> {
        validate_staging_key(&self.pool, attempt_id, &object.key).await?;
        insert_cleanup_identity(
            &self.pool,
            attempt_id,
            NewCleanupIdentity {
                kind: "object",
                target_store: &object.target_store,
                bucket: &object.bucket,
                key: &object.key,
                version_id: object.version_id.as_deref(),
                upload_id: None,
            },
        )
        .await
    }

    pub async fn register_staging_write_intent_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
        intent: StagingWriteIntent,
    ) -> Result<bool, AttemptError> {
        self.register_staging_operation(
            authority,
            operation_id,
            "register_staging_write_intent",
            intent.clone(),
            move |tx, attempt_id| {
                Box::pin(async move {
                    validate_staging_key_transaction(tx, &attempt_id, &intent.key).await?;
                    insert_cleanup_identity_transaction(
                        tx,
                        &attempt_id,
                        NewCleanupIdentity {
                            kind: "write_intent",
                            target_store: &intent.target_store,
                            bucket: &intent.bucket,
                            key: &intent.key,
                            version_id: None,
                            upload_id: None,
                        },
                    )
                    .await
                })
            },
        )
        .await
    }

    pub async fn register_staging_object_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
        object: StagingObject,
        completed_upload_id: Option<String>,
    ) -> Result<bool, AttemptError> {
        self.register_staging_operation(
            authority,
            operation_id,
            "register_staging_object",
            StagingObjectRegistration {
                object: object.clone(),
                completed_upload_id: completed_upload_id.clone(),
            },
            move |tx, attempt_id| {
                Box::pin(async move {
                    validate_staging_key_transaction(tx, &attempt_id, &object.key).await?;
                    insert_cleanup_identity_transaction(
                        tx,
                        &attempt_id,
                        NewCleanupIdentity {
                            kind: "object",
                            target_store: &object.target_store,
                            bucket: &object.bucket,
                            key: &object.key,
                            version_id: object.version_id.as_deref(),
                            upload_id: None,
                        },
                    )
                    .await?;
                    complete_pending_identity(
                        tx,
                        &attempt_id,
                        "write_intent",
                        &object.target_store,
                        &object.bucket,
                        &object.key,
                        None,
                    )
                    .await?;
                    if let Some(upload_id) = completed_upload_id {
                        complete_pending_identity(
                            tx,
                            &attempt_id,
                            "multipart",
                            &object.target_store,
                            &object.bucket,
                            &object.key,
                            Some(&upload_id),
                        )
                        .await?;
                    }
                    Ok(())
                })
            },
        )
        .await
    }

    pub async fn register_staging_multipart(
        &self,
        attempt_id: &OpaqueId,
        multipart: &StagingMultipart,
    ) -> Result<(), AttemptError> {
        validate_staging_key(&self.pool, attempt_id, &multipart.key).await?;
        if multipart.upload_id.is_empty() {
            return Err(AttemptError::Invalid(
                "multipart upload identity cannot be empty".into(),
            ));
        }
        insert_cleanup_identity(
            &self.pool,
            attempt_id,
            NewCleanupIdentity {
                kind: "multipart",
                target_store: &multipart.target_store,
                bucket: &multipart.bucket,
                key: &multipart.key,
                version_id: None,
                upload_id: Some(&multipart.upload_id),
            },
        )
        .await
    }

    pub async fn register_staging_multipart_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
        multipart: StagingMultipart,
    ) -> Result<bool, AttemptError> {
        if multipart.upload_id.is_empty() {
            return Err(AttemptError::Invalid(
                "multipart upload identity cannot be empty".into(),
            ));
        }
        self.register_staging_operation(
            authority,
            operation_id,
            "register_staging_multipart",
            multipart.clone(),
            move |tx, attempt_id| {
                Box::pin(async move {
                    validate_staging_key_transaction(tx, &attempt_id, &multipart.key).await?;
                    insert_cleanup_identity_transaction(
                        tx,
                        &attempt_id,
                        NewCleanupIdentity {
                            kind: "multipart",
                            target_store: &multipart.target_store,
                            bucket: &multipart.bucket,
                            key: &multipart.key,
                            version_id: None,
                            upload_id: Some(&multipart.upload_id),
                        },
                    )
                    .await
                })
            },
        )
        .await
    }

    async fn register_staging_operation<T, F>(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
        kind: &'static str,
        request: T,
        effect: F,
    ) -> Result<bool, AttemptError>
    where
        T: Serialize,
        F: for<'c> FnOnce(
                &'c mut remora_durable::PgTx<'_>,
                OpaqueId,
            ) -> remora_durable::BoxFuture<'c, Result<(), DurableError>>
            + Send
            + 'static,
    {
        let request_fingerprint = fingerprint(&RegistrationRequest {
            job_id: &authority.job_id,
            attempt_id: &authority.attempt_id,
            fence: authority.fence,
            worker_incarnation: &authority.worker_incarnation,
            value: request,
        })
        .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        let authority = authority.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &self.operation_subject,
                operation: "attempt.mutation",
                key: operation_id.as_str(),
            },
            kind,
            &[],
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let eligible = lock_registration_authority(tx, &authority).await?;
                    if eligible {
                        effect(tx, authority.attempt_id).await?;
                    }
                    Ok((
                        if eligible {
                            MutationStatus::Accepted
                        } else {
                            MutationStatus::Rejected
                        },
                        serde_json::Value::Bool(eligible),
                    ))
                })
            },
        )
        .await?;
        if outcome.disposition == MutationDisposition::Conflict {
            return Err(AttemptError::OperationConflict);
        }
        outcome
            .response
            .as_bool()
            .ok_or_else(|| AttemptError::Invalid("invalid staging operation result".into()))
    }

    pub async fn reconcile_once(
        &self,
        config: &RecoveryConfig,
    ) -> Result<ReconcileReport, AttemptError> {
        validate_recovery_config(config)?;
        let job_ids = sqlx::query_scalar::<_, String>(
            "SELECT job_id FROM jobs
             WHERE phase NOT IN ('succeeded', 'failed', 'cancelled')
                OR coordination_cleaned_at IS NULL
             ORDER BY recovery_checked_at NULLS FIRST, recovery_checked_at, job_id LIMIT $1",
        )
        .bind(i64::from(config.batch_size))
        .fetch_all(&self.pool)
        .await?;
        let mut report = ReconcileReport::default();
        for value in job_ids {
            let job_id =
                OpaqueId::parse(value).map_err(|error| AttemptError::Invalid(error.to_string()))?;
            report.merge(self.reconcile_job(&job_id, config).await?);
            sqlx::query("UPDATE jobs SET recovery_checked_at = $2 WHERE job_id = $1")
                .bind(job_id.as_str())
                .bind(OffsetDateTime::now_utc())
                .execute(&self.pool)
                .await?;
        }
        report.cleanup_items_materialized = self.materialize_cleanup_work().await?;
        Ok(report)
    }

    pub async fn claim_cleanup(
        &self,
        config: &RecoveryConfig,
    ) -> Result<Option<CleanupClaim>, AttemptError> {
        validate_recovery_config(config)?;
        self.materialize_cleanup_work().await?;
        let claim_token = new_id('c')?;
        let claimed_until = OffsetDateTime::now_utc()
            + time_duration(config.cleanup_claim_ttl, "cleanup_claim_ttl")?;
        let row = sqlx::query_as::<_, CleanupRow>(
            "WITH candidate AS (
                SELECT c.cleanup_id
                FROM attempt_staging_cleanup c
                JOIN attempts a ON a.attempt_id = c.attempt_id
                JOIN jobs j ON j.job_id = a.job_id
                WHERE (c.state = 'pending'
                       OR (c.state = 'claimed' AND c.claimed_until <= $1))
                  AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
                  AND (j.phase IN ('succeeded', 'failed', 'cancelled')
                       OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
                  AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
                  AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
                ORDER BY c.cleanup_id
                FOR UPDATE OF c SKIP LOCKED LIMIT 1
             )
             UPDATE attempt_staging_cleanup c
             SET state = 'claimed', claim_token = $2, claimed_until = $3,
                 failure_count = CASE WHEN c.state = 'claimed'
                                      THEN c.failure_count + 1 ELSE c.failure_count END
             FROM candidate
             WHERE c.cleanup_id = candidate.cleanup_id
             RETURNING c.cleanup_id, c.attempt_id, c.kind, c.target_store, c.bucket,
                       c.object_key, c.version_id, c.upload_id",
        )
        .bind(OffsetDateTime::now_utc())
        .bind(claim_token.as_str())
        .bind(claimed_until)
        .fetch_optional(&self.pool)
        .await?;
        row.map(|row| row.try_into_claim(claim_token)).transpose()
    }

    pub async fn complete_cleanup(
        &self,
        claim: &CleanupClaim,
    ) -> Result<CleanupCompletion, AttemptError> {
        let changed = sqlx::query(
            "UPDATE attempt_staging_cleanup
             SET state = 'complete', claim_token = NULL, claimed_until = NULL,
                 completed_at = $3, last_error = NULL
             WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
        )
        .bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
        .bind(claim.claim_token.as_str())
        .bind(OffsetDateTime::now_utc())
        .execute(&self.pool)
        .await?;
        Ok(if changed.rows_affected() == 1 {
            CleanupCompletion::Completed
        } else {
            CleanupCompletion::ClaimLost
        })
    }

    pub async fn reconcile_write_intent(
        &self,
        claim: &CleanupClaim,
        objects: Vec<StagingObject>,
        multiparts: Vec<StagingMultipart>,
    ) -> Result<CleanupCompletion, AttemptError> {
        let CleanupIdentity::WriteIntent(intent) = &claim.identity else {
            return Err(AttemptError::Invalid(
                "cleanup claim is not a staging write intent".into(),
            ));
        };
        for object in &objects {
            if object.target_store != intent.target_store
                || object.bucket != intent.bucket
                || object.key != intent.key
            {
                return Err(AttemptError::Invalid(
                    "discovered object does not match the write intent".into(),
                ));
            }
        }
        for multipart in &multiparts {
            if multipart.target_store != intent.target_store
                || multipart.bucket != intent.bucket
                || multipart.key != intent.key
                || multipart.upload_id.is_empty()
            {
                return Err(AttemptError::Invalid(
                    "discovered multipart does not match the write intent".into(),
                ));
            }
        }
        let cleanup_id = i64::try_from(claim.cleanup_id).map_err(invalid_number)?;
        let claim_token = claim.claim_token.clone();
        let attempt_id = claim.attempt_id.clone();
        with_transaction(&self.pool, move |tx| {
            Box::pin(async move {
                let retained = sqlx::query_scalar::<_, i64>(
                    "SELECT cleanup_id FROM attempt_staging_cleanup
                     WHERE cleanup_id = $1 AND attempt_id = $2 AND kind = 'write_intent'
                       AND state = 'claimed' AND claim_token = $3
                     FOR UPDATE",
                )
                .bind(cleanup_id)
                .bind(attempt_id.as_str())
                .bind(claim_token.as_str())
                .fetch_optional(tx.as_mut())
                .await?;
                if retained.is_none() {
                    return Ok(CleanupCompletion::ClaimLost);
                }
                for object in &objects {
                    insert_cleanup_identity_transaction(
                        tx,
                        &attempt_id,
                        NewCleanupIdentity {
                            kind: "object",
                            target_store: &object.target_store,
                            bucket: &object.bucket,
                            key: &object.key,
                            version_id: object.version_id.as_deref(),
                            upload_id: None,
                        },
                    )
                    .await?;
                }
                for multipart in &multiparts {
                    insert_cleanup_identity_transaction(
                        tx,
                        &attempt_id,
                        NewCleanupIdentity {
                            kind: "multipart",
                            target_store: &multipart.target_store,
                            bucket: &multipart.bucket,
                            key: &multipart.key,
                            version_id: None,
                            upload_id: Some(&multipart.upload_id),
                        },
                    )
                    .await?;
                }
                sqlx::query(
                    "UPDATE attempt_staging_cleanup
                     SET state = 'complete', claim_token = NULL, claimed_until = NULL,
                         completed_at = $2, last_error = NULL
                     WHERE cleanup_id = $1",
                )
                .bind(cleanup_id)
                .bind(OffsetDateTime::now_utc())
                .execute(tx.as_mut())
                .await?;
                Ok(CleanupCompletion::Completed)
            })
        })
        .await
        .map_err(Into::into)
    }

    pub async fn retry_cleanup(
        &self,
        claim: &CleanupClaim,
        failure_code: &str,
    ) -> Result<CleanupCompletion, AttemptError> {
        if failure_code.is_empty()
            || failure_code.len() > 64
            || !failure_code
                .bytes()
                .all(|byte| byte.is_ascii_lowercase() || byte == b'_')
        {
            return Err(AttemptError::Invalid(
                "cleanup failure code must be lowercase snake case".into(),
            ));
        }
        let changed = sqlx::query(
            "UPDATE attempt_staging_cleanup
             SET state = 'pending', claim_token = NULL, claimed_until = NULL,
                 failure_count = failure_count + 1, last_error = $3
             WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
        )
        .bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
        .bind(claim.claim_token.as_str())
        .bind(failure_code)
        .execute(&self.pool)
        .await?;
        Ok(if changed.rows_affected() == 1 {
            CleanupCompletion::RetryScheduled
        } else {
            CleanupCompletion::ClaimLost
        })
    }

    async fn reconcile_job(
        &self,
        job_id: &OpaqueId,
        config: &RecoveryConfig,
    ) -> Result<ReconcileReport, AttemptError> {
        let Some(mut row) = load_recovery_job(&self.pool, job_id).await? else {
            return Ok(ReconcileReport::default());
        };
        let mut report = ReconcileReport::default();
        if row.phase == "retry_wait" {
            self.release_local_if_stopped(job_id).await?;
            if row.retry_not_before.is_none() {
                let due =
                    OffsetDateTime::now_utc() + time_duration(config.retry_delay, "retry_delay")?;
                sqlx::query(
                    "UPDATE jobs SET retry_not_before = $2, updated_at = $3
                     WHERE job_id = $1 AND phase = 'retry_wait' AND retry_not_before IS NULL",
                )
                .bind(job_id.as_str())
                .bind(due)
                .bind(OffsetDateTime::now_utc())
                .execute(&self.pool)
                .await?;
                row.retry_not_before = Some(due);
            }
            if row
                .retry_not_before
                .is_some_and(|due| due <= OffsetDateTime::now_utc())
            {
                let mut tx = self.pool.begin().await?;
                let now = OffsetDateTime::now_utc();
                let changed = sqlx::query(
                    "UPDATE jobs SET phase = 'queued', dispatch_needed = TRUE,
                        retry_not_before = NULL, revision = revision + 1, updated_at = $2
                     WHERE job_id = $1 AND phase = 'retry_wait' AND cancel_requested = FALSE
                       AND retry_not_before <= $2",
                )
                .bind(job_id.as_str())
                .bind(now)
                .execute(tx.as_mut())
                .await?;
                if changed.rows_affected() == 1 {
                    enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
                    report.retries_scheduled += 1;
                    row.phase = "queued".into();
                    row.dispatch_needed = true;
                }
                tx.commit().await?;
            }
        }
        if row.phase == "queued"
            && row.dispatch_needed
            && self
                .coordination
                .project_enqueue(job_id.as_str())
                .await?
                .ok()
        {
            report.enqueued += 1;
        }
        if matches!(row.phase.as_str(), "running" | "cancelling") {
            if row.cancel_requested {
                if self
                    .coordination
                    .project_cancel(job_id.as_str())
                    .await?
                    .ok()
                {
                    report.cancellations_published += 1;
                }
                if self.finalize_cancel(job_id).await? {
                    report.cancellations_finalized += 1;
                    row.phase = "cancelled".into();
                }
            } else if let Some(attempt_id) = row.authority_attempt_id.as_deref()
                && self
                    .reap_if_inactive(job_id, attempt_id, self.config.max_attempts)
                    .await?
            {
                report.attempts_reaped += 1;
            }
            if self
                .coordination
                .clear_stale_progress(job_id.as_str())
                .await?
                .code
                == OpCode::Ok
            {
                report.stale_progress_cleared += 1;
            }
        }
        if parse_phase(&row.phase)?.is_terminal() {
            self.release_local_if_stopped(job_id).await?;
            let expired = self.coordination.expire_lease(job_id.as_str()).await?;
            self.coordination.clean_terminal(job_id.as_str()).await?;
            if expired.code.is_success() {
                sqlx::query(
                    "UPDATE jobs SET coordination_cleaned_at = COALESCE(coordination_cleaned_at, $2)
                     WHERE job_id = $1 AND phase IN ('succeeded', 'failed', 'cancelled')",
                )
                .bind(job_id.as_str())
                .bind(OffsetDateTime::now_utc())
                .execute(&self.pool)
                .await?;
                report.terminal_coordination_cleaned += 1;
            }
        }
        Ok(report)
    }

    async fn reap_if_inactive(
        &self,
        job_id: &OpaqueId,
        attempt_id: &str,
        max_attempts: u64,
    ) -> Result<bool, AttemptError> {
        let attempt_id = attempt_id.to_owned();
        let durable_job_id = job_id.clone();
        let coordination = self.coordination.clone();
        let faults = self.faults.clone();
        let reaped = with_transaction(&self.pool, move |tx| {
            Box::pin(async move {
                reap_lost_attempt(
                    tx,
                    &coordination,
                    faults.as_ref(),
                    &durable_job_id,
                    &attempt_id,
                    max_attempts,
                )
                .await
            })
        })
        .await?;
        let _ = self.coordination.expire_lease(job_id.as_str()).await?;
        if reaped.release_local {
            self.release_local(job_id);
        }
        Ok(reaped.changed)
    }

    async fn finalize_cancel(&self, job_id: &OpaqueId) -> Result<bool, AttemptError> {
        let now = OffsetDateTime::now_utc();
        let mut tx = self.pool.begin().await?;
        let attempt_id = sqlx::query_scalar::<_, Option<String>>(
            "SELECT authority_attempt_id FROM jobs
             WHERE job_id = $1 AND phase = 'cancelling' FOR UPDATE",
        )
        .bind(job_id.as_str())
        .fetch_optional(tx.as_mut())
        .await?
        .flatten();
        let Some(attempt_id) = attempt_id else {
            tx.rollback().await?;
            return Ok(false);
        };
        sqlx::query(
            "UPDATE attempts SET state = 'stopped', stopped_at = $2
             WHERE attempt_id = $1 AND state = 'authorized'",
        )
        .bind(&attempt_id)
        .bind(now)
        .execute(tx.as_mut())
        .await?;
        sqlx::query(
            "UPDATE jobs SET phase = 'cancelled', authority_attempt_id = NULL,
                dispatch_needed = FALSE, terminal_at = $2, coordination_cleaned_at = NULL,
                revision = revision + 1, updated_at = $2
             WHERE job_id = $1 AND phase = 'cancelling'",
        )
        .bind(job_id.as_str())
        .bind(now)
        .execute(tx.as_mut())
        .await?;
        enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
        tx.commit().await?;
        Ok(true)
    }

    async fn materialize_cleanup_work(&self) -> Result<u64, AttemptError> {
        let now = OffsetDateTime::now_utc();
        let members = sqlx::query(
            "INSERT INTO attempt_staging_cleanup (
                attempt_id, kind, target_store, bucket, object_key, version_id, created_at
             )
             SELECT m.attempt_id, 'object', m.object_store, m.object_bucket,
                    m.object_key, m.object_version_id, $1
             FROM result_inventory_members m
             JOIN attempts a ON a.attempt_id = m.attempt_id
             JOIN jobs j ON j.job_id = a.job_id
             WHERE j.authority_attempt_id IS DISTINCT FROM a.attempt_id
               AND (j.phase IN ('succeeded', 'failed', 'cancelled')
                    OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
               AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
               AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
             ON CONFLICT DO NOTHING",
        )
        .bind(now)
        .execute(&self.pool)
        .await?;
        let manifests = sqlx::query(
            "INSERT INTO attempt_staging_cleanup (
                attempt_id, kind, target_store, bucket, object_key, version_id, created_at
             )
             SELECT r.attempt_id, 'object', r.manifest_store, r.manifest_bucket,
                    r.manifest_key, r.manifest_version_id, $1
             FROM result_inventories r
             JOIN attempts a ON a.attempt_id = r.attempt_id
             JOIN jobs j ON j.job_id = a.job_id
             WHERE r.state = 'sealed'
               AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
               AND (j.phase IN ('succeeded', 'failed', 'cancelled')
                    OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
               AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
               AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
             ON CONFLICT DO NOTHING",
        )
        .bind(now)
        .execute(&self.pool)
        .await?;
        Ok(members.rows_affected() + manifests.rows_affected())
    }

    async fn release_local_if_stopped(&self, job_id: &OpaqueId) -> Result<(), AttemptError> {
        let stopped = sqlx::query_scalar::<_, bool>(
            "SELECT NOT EXISTS (
                SELECT 1 FROM attempts
                WHERE job_id = $1 AND state = 'executing' AND worker_down_at IS NULL
             )",
        )
        .bind(job_id.as_str())
        .fetch_one(&self.pool)
        .await?;
        if stopped {
            self.release_local(job_id);
        }
        Ok(())
    }
}

async fn cooperative_stop_transaction(
    tx: &mut remora_durable::PgTx<'_>,
    coordination: &remora_coordination::CoordinationStore,
    authority: &AttemptAuthority,
) -> Result<CooperativeStopOutcome, DurableError> {
    let row = sqlx::query_as::<_, StopRow>(
        "SELECT a.job_id, a.fence, a.worker_incarnation, a.state,
                j.authority_attempt_id, j.cancel_requested
         FROM attempts a JOIN jobs j ON j.job_id = a.job_id
         WHERE a.attempt_id = $1
         FOR UPDATE OF a, j",
    )
    .bind(authority.attempt_id.as_str())
    .fetch_optional(tx.as_mut())
    .await?;
    let Some(row) = row else {
        return Ok(CooperativeStopOutcome::AuthorityMismatch);
    };
    if row.job_id != authority.job_id.as_str()
        || u64::try_from(row.fence).ok() != Some(authority.fence)
        || row.worker_incarnation != authority.worker_incarnation.as_str()
    {
        return Ok(CooperativeStopOutcome::AuthorityMismatch);
    }
    if row.state == "stopped" {
        return Ok(CooperativeStopOutcome::Stopped);
    }
    if row.state != "executing" {
        return Ok(CooperativeStopOutcome::StopNotRequired);
    }
    let current = row.authority_attempt_id.as_deref() == Some(authority.attempt_id.as_str());
    let live = coordination
        .get_live_lease(authority.job_id.as_str())
        .await
        .map_err(DurableError::effect)?
        .is_some_and(|lease| lease.record.owner() == lease_owner(authority));
    if current && live && !row.cancel_requested {
        return Ok(CooperativeStopOutcome::StopNotRequired);
    }
    let changed = sqlx::query(
        "UPDATE attempts SET state = 'stopped', stopped_at = $2
         WHERE attempt_id = $1 AND state IN ('authorized', 'executing')",
    )
    .bind(authority.attempt_id.as_str())
    .bind(OffsetDateTime::now_utc())
    .execute(tx.as_mut())
    .await?;
    Ok(if changed.rows_affected() == 1 {
        CooperativeStopOutcome::Stopped
    } else {
        CooperativeStopOutcome::StopNotRequired
    })
}

async fn lock_registration_authority(
    tx: &mut remora_durable::PgTx<'_>,
    authority: &AttemptAuthority,
) -> Result<bool, DurableError> {
    let row = sqlx::query_as::<_, RegistrationAuthorityRow>(
        "SELECT a.job_id, a.fence, a.worker_incarnation, a.state,
                j.authority_attempt_id, j.cancel_requested
         FROM attempts a JOIN jobs j ON j.job_id = a.job_id
         WHERE a.attempt_id = $1
         FOR UPDATE OF a, j",
    )
    .bind(authority.attempt_id.as_str())
    .fetch_optional(tx.as_mut())
    .await?;
    Ok(row.is_some_and(|row| {
        row.job_id == authority.job_id.as_str()
            && u64::try_from(row.fence).ok() == Some(authority.fence)
            && row.worker_incarnation == authority.worker_incarnation.as_str()
            && row.state == "executing"
            && row.authority_attempt_id.as_deref() == Some(authority.attempt_id.as_str())
            && !row.cancel_requested
    }))
}

async fn validate_staging_key_transaction(
    tx: &mut remora_durable::PgTx<'_>,
    attempt_id: &OpaqueId,
    key: &str,
) -> Result<(), DurableError> {
    validate_attempt_key(key).map_err(DurableError::effect)?;
    let (staging_origin, target_prefix) = sqlx::query_as::<_, (String, Json<Vec<String>>)>(
        "SELECT a.staging_origin_attempt_id, j.target_prefix
             FROM attempts a JOIN jobs j ON j.job_id = a.job_id
             WHERE a.attempt_id = $1",
    )
    .bind(attempt_id.as_str())
    .fetch_one(tx.as_mut())
    .await?;
    let staging_origin = OpaqueId::parse(staging_origin).map_err(DurableError::effect)?;
    if attempt_key_attempt_id(key).map_err(DurableError::effect)? != staging_origin
        || !key.starts_with(&attempt_staging_prefix(
            &target_prefix.0.join("/"),
            &staging_origin,
        ))
    {
        return Err(DurableError::effect(
            "staging key belongs to a different attempt",
        ));
    }
    Ok(())
}

async fn insert_cleanup_identity_transaction(
    tx: &mut remora_durable::PgTx<'_>,
    attempt_id: &OpaqueId,
    identity: NewCleanupIdentity<'_>,
) -> Result<(), DurableError> {
    if identity.target_store.is_empty() || identity.bucket.is_empty() {
        return Err(DurableError::effect(
            "staging target store and bucket cannot be empty",
        ));
    }
    let changed = sqlx::query(
        "INSERT INTO attempt_staging_cleanup (
            attempt_id, kind, target_store, bucket, object_key, version_id, upload_id, created_at
         )
         SELECT $1, $2, $3, $4, $5, $6, $7, $8
         WHERE EXISTS (
             SELECT 1 FROM attempts a JOIN jobs j ON j.job_id = a.job_id
             WHERE a.attempt_id = $1 AND j.target_store = $3 AND j.target_bucket = $4
         )
         ON CONFLICT DO NOTHING",
    )
    .bind(attempt_id.as_str())
    .bind(identity.kind)
    .bind(identity.target_store)
    .bind(identity.bucket)
    .bind(identity.key)
    .bind(identity.version_id)
    .bind(identity.upload_id)
    .bind(OffsetDateTime::now_utc())
    .execute(tx.as_mut())
    .await?;
    if changed.rows_affected() == 0 {
        let binding = sqlx::query_as::<_, (String, String)>(
            "SELECT j.target_store, j.target_bucket
             FROM attempts a JOIN jobs j ON j.job_id = a.job_id
             WHERE a.attempt_id = $1",
        )
        .bind(attempt_id.as_str())
        .fetch_optional(tx.as_mut())
        .await?;
        match binding {
            None => return Err(DurableError::effect("attempt does not exist")),
            Some((store, bucket))
                if store != identity.target_store || bucket != identity.bucket =>
            {
                return Err(DurableError::effect(
                    "staging identity does not match the attempt target",
                ));
            }
            Some(_) => {}
        }
    }
    Ok(())
}

async fn complete_pending_identity(
    tx: &mut remora_durable::PgTx<'_>,
    attempt_id: &OpaqueId,
    kind: &str,
    target_store: &str,
    bucket: &str,
    key: &str,
    upload_id: Option<&str>,
) -> Result<(), DurableError> {
    let changed = sqlx::query(
        "UPDATE attempt_staging_cleanup
         SET state = 'complete', completed_at = $7
         WHERE attempt_id = $1 AND kind = $2 AND target_store = $3 AND bucket = $4
           AND object_key = $5 AND upload_id IS NOT DISTINCT FROM $6 AND state = 'pending'",
    )
    .bind(attempt_id.as_str())
    .bind(kind)
    .bind(target_store)
    .bind(bucket)
    .bind(key)
    .bind(upload_id)
    .bind(OffsetDateTime::now_utc())
    .execute(tx.as_mut())
    .await?;
    if changed.rows_affected() != 1 {
        return Err(DurableError::effect(
            "staging write lifecycle identity is missing or unavailable",
        ));
    }
    Ok(())
}

async fn insert_cleanup_identity(
    pool: &sqlx::PgPool,
    attempt_id: &OpaqueId,
    identity: NewCleanupIdentity<'_>,
) -> Result<(), AttemptError> {
    if identity.target_store.is_empty() || identity.bucket.is_empty() {
        return Err(AttemptError::Invalid(
            "staging target store and bucket cannot be empty".into(),
        ));
    }
    let changed = sqlx::query(
        "INSERT INTO attempt_staging_cleanup (
            attempt_id, kind, target_store, bucket, object_key, version_id, upload_id, created_at
         )
         SELECT $1, $2, $3, $4, $5, $6, $7, $8
         WHERE EXISTS (
             SELECT 1 FROM attempts a JOIN jobs j ON j.job_id = a.job_id
             WHERE a.attempt_id = $1 AND j.target_store = $3 AND j.target_bucket = $4
         )
         ON CONFLICT DO NOTHING",
    )
    .bind(attempt_id.as_str())
    .bind(identity.kind)
    .bind(identity.target_store)
    .bind(identity.bucket)
    .bind(identity.key)
    .bind(identity.version_id)
    .bind(identity.upload_id)
    .bind(OffsetDateTime::now_utc())
    .execute(pool)
    .await?;
    if changed.rows_affected() == 0 {
        let binding = sqlx::query_as::<_, (String, String)>(
            "SELECT j.target_store, j.target_bucket
             FROM attempts a JOIN jobs j ON j.job_id = a.job_id
             WHERE a.attempt_id = $1",
        )
        .bind(attempt_id.as_str())
        .fetch_optional(pool)
        .await?;
        match binding {
            None => return Err(AttemptError::Invalid("attempt does not exist".into())),
            Some((store, bucket))
                if store != identity.target_store || bucket != identity.bucket =>
            {
                return Err(AttemptError::Invalid(
                    "staging identity does not match the attempt target".into(),
                ));
            }
            Some(_) => {}
        }
    }
    Ok(())
}

struct NewCleanupIdentity<'a> {
    kind: &'a str,
    target_store: &'a str,
    bucket: &'a str,
    key: &'a str,
    version_id: Option<&'a str>,
    upload_id: Option<&'a str>,
}

async fn load_recovery_job(
    pool: &sqlx::PgPool,
    job_id: &OpaqueId,
) -> Result<Option<RecoveryJobRow>, AttemptError> {
    sqlx::query_as(
        "SELECT phase, cancel_requested, dispatch_needed, authority_attempt_id,
                retry_not_before
         FROM jobs WHERE job_id = $1",
    )
    .bind(job_id.as_str())
    .fetch_optional(pool)
    .await
    .map_err(Into::into)
}

async fn reap_lost_attempt(
    tx: &mut Transaction<'_, Postgres>,
    coordination: &remora_coordination::CoordinationStore,
    faults: &dyn crate::AttemptFaultInjector,
    job_id: &OpaqueId,
    attempt_id: &str,
    max_attempts: u64,
) -> Result<ReapEffect, DurableError> {
    let row = sqlx::query_as::<_, LockedReapRow>(
        "SELECT j.current_fence, a.state, a.worker_down_at, a.fence,
                a.worker_incarnation
         FROM jobs j JOIN attempts a ON a.attempt_id = j.authority_attempt_id
         WHERE j.job_id = $1 AND a.attempt_id = $2 AND j.phase = 'running'
         FOR UPDATE OF j, a",
    )
    .bind(job_id.as_str())
    .bind(attempt_id)
    .fetch_optional(tx.as_mut())
    .await?;
    let Some(row) = row else {
        return Ok(ReapEffect::default());
    };
    faults.hit(crate::AttemptFaultPoint::RecoveryBeforeRevalidation);
    let reservation = coordination
        .get_reservation(job_id.as_str())
        .await
        .map_err(DurableError::effect)?;
    let live = coordination
        .get_live_lease(job_id.as_str())
        .await
        .map_err(DurableError::effect)?;
    let exact_live = live.is_some_and(|lease| {
        lease.record.attempt_id == attempt_id
            && i64::try_from(lease.record.fence).ok() == Some(row.fence)
            && lease.record.worker_id == row.worker_incarnation
    });
    let inactive = match row.state.as_str() {
        "authorized" => reservation.is_none() && !exact_live,
        "executing" => row.worker_down_at.is_some() || !exact_live,
        "stopped" => !exact_live,
        _ => false,
    };
    if !inactive {
        return Ok(ReapEffect::default());
    }
    let retry = u64::try_from(row.current_fence).map_err(DurableError::effect)? < max_attempts;
    let now = OffsetDateTime::now_utc();
    let failure = FailureRecord {
        class: FailureClass::LostAttempt,
        disposition: FailureDisposition::Retryable,
    };
    if row.state == "authorized" {
        sqlx::query(
            "UPDATE attempts SET state = 'reaped', reaped_at = $2, failure = $3
             WHERE attempt_id = $1 AND state = 'authorized'",
        )
        .bind(attempt_id)
        .bind(now)
        .bind(Json(&failure))
        .execute(tx.as_mut())
        .await?;
    }
    if retry {
        sqlx::query(
            "UPDATE jobs SET phase = 'retry_wait', authority_attempt_id = NULL,
                dispatch_needed = FALSE, retry_not_before = NULL,
                revision = revision + 1, updated_at = $2
             WHERE job_id = $1 AND authority_attempt_id = $3 AND phase = 'running'",
        )
        .bind(job_id.as_str())
        .bind(now)
        .bind(attempt_id)
        .execute(tx.as_mut())
        .await?;
        enqueue(tx, "job.retry_wait", json!({ "job_id": job_id })).await?;
        enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
    } else {
        sqlx::query(
            "UPDATE jobs SET phase = 'failed', authority_attempt_id = NULL,
                dispatch_needed = FALSE, failure = $2, terminal_at = $3,
                coordination_cleaned_at = NULL, revision = revision + 1, updated_at = $3
             WHERE job_id = $1 AND authority_attempt_id = $4 AND phase = 'running'",
        )
        .bind(job_id.as_str())
        .bind(Json(&failure))
        .bind(now)
        .bind(attempt_id)
        .execute(tx.as_mut())
        .await?;
        enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;
    }
    Ok(ReapEffect {
        changed: true,
        release_local: row.state == "authorized" || row.worker_down_at.is_some(),
    })
}

async fn validate_staging_key(
    pool: &sqlx::PgPool,
    attempt_id: &OpaqueId,
    key: &str,
) -> Result<(), AttemptError> {
    validate_attempt_key(key).map_err(|error| AttemptError::Invalid(error.to_string()))?;
    let (staging_origin, target_prefix) = sqlx::query_as::<_, (String, Json<Vec<String>>)>(
        "SELECT a.staging_origin_attempt_id, j.target_prefix
         FROM attempts a JOIN jobs j ON j.job_id = a.job_id
         WHERE a.attempt_id = $1",
    )
    .bind(attempt_id.as_str())
    .fetch_optional(pool)
    .await?
    .ok_or_else(|| AttemptError::Invalid("attempt does not exist".into()))?;
    let staging_origin = OpaqueId::parse(staging_origin)
        .map_err(|error| AttemptError::Invalid(error.to_string()))?;
    if attempt_key_attempt_id(key).map_err(|error| AttemptError::Invalid(error.to_string()))?
        != staging_origin
        || !key.starts_with(&attempt_staging_prefix(
            &target_prefix.0.join("/"),
            &staging_origin,
        ))
    {
        return Err(AttemptError::Invalid(
            "staging key belongs to a different attempt".into(),
        ));
    }
    Ok(())
}

fn validate_recovery_config(config: &RecoveryConfig) -> Result<(), AttemptError> {
    if config.batch_size == 0 {
        return Err(AttemptError::Invalid(
            "recovery limits must be greater than zero".into(),
        ));
    }
    if config.cleanup_claim_ttl.is_zero() {
        return Err(AttemptError::Invalid(
            "cleanup claim ttl must be greater than zero".into(),
        ));
    }
    Ok(())
}

fn time_duration(value: Duration, name: &str) -> Result<time::Duration, AttemptError> {
    time::Duration::try_from(value)
        .map_err(|_| AttemptError::Invalid(format!("{name} is too large")))
}

fn lease_owner(authority: &AttemptAuthority) -> LeaseOwner {
    LeaseOwner {
        attempt_id: authority.attempt_id.to_string(),
        fence: authority.fence,
        worker_id: authority.worker_incarnation.to_string(),
    }
}

fn parse_phase(value: &str) -> Result<JobPhase, DurableError> {
    match value {
        "queued" => Ok(JobPhase::Queued),
        "running" => Ok(JobPhase::Running),
        "retry_wait" => Ok(JobPhase::RetryWait),
        "cancelling" => Ok(JobPhase::Cancelling),
        "succeeded" => Ok(JobPhase::Succeeded),
        "failed" => Ok(JobPhase::Failed),
        "cancelled" => Ok(JobPhase::Cancelled),
        _ => Err(DurableError::Effect(format!("unknown job phase '{value}'"))),
    }
}

fn invalid_number(error: impl std::fmt::Display) -> AttemptError {
    AttemptError::Invalid(error.to_string())
}

#[derive(FromRow)]
struct CancelRow {
    phase: String,
    cancel_requested: bool,
}

#[derive(FromRow)]
struct StopRow {
    job_id: String,
    fence: i64,
    worker_incarnation: String,
    state: String,
    authority_attempt_id: Option<String>,
    cancel_requested: bool,
}

#[derive(FromRow)]
struct RegistrationAuthorityRow {
    job_id: String,
    fence: i64,
    worker_incarnation: String,
    state: String,
    authority_attempt_id: Option<String>,
    cancel_requested: bool,
}

#[derive(Serialize)]
struct StopRequest<'a> {
    job_id: &'a OpaqueId,
    attempt_id: &'a OpaqueId,
    fence: u64,
    worker_incarnation: &'a OpaqueId,
}

#[derive(Serialize)]
struct RegistrationRequest<'a, T> {
    job_id: &'a OpaqueId,
    attempt_id: &'a OpaqueId,
    fence: u64,
    worker_incarnation: &'a OpaqueId,
    value: T,
}

#[derive(Serialize)]
struct StagingObjectRegistration {
    object: StagingObject,
    completed_upload_id: Option<String>,
}

#[derive(FromRow)]
struct RecoveryJobRow {
    phase: String,
    cancel_requested: bool,
    dispatch_needed: bool,
    authority_attempt_id: Option<String>,
    retry_not_before: Option<OffsetDateTime>,
}

#[derive(FromRow)]
struct LockedReapRow {
    current_fence: i64,
    state: String,
    worker_down_at: Option<OffsetDateTime>,
    fence: i64,
    worker_incarnation: String,
}

#[derive(Default)]
struct ReapEffect {
    changed: bool,
    release_local: bool,
}

#[derive(FromRow)]
struct CleanupRow {
    cleanup_id: i64,
    attempt_id: String,
    kind: String,
    target_store: String,
    bucket: String,
    object_key: String,
    version_id: Option<String>,
    upload_id: Option<String>,
}

impl CleanupRow {
    fn try_into_claim(self, claim_token: OpaqueId) -> Result<CleanupClaim, AttemptError> {
        let identity = match self.kind.as_str() {
            "write_intent" => CleanupIdentity::WriteIntent(StagingWriteIntent {
                target_store: self.target_store,
                bucket: self.bucket,
                key: self.object_key,
            }),
            "object" => CleanupIdentity::Object(StagingObject {
                target_store: self.target_store,
                bucket: self.bucket,
                key: self.object_key,
                version_id: self.version_id,
            }),
            "multipart" => CleanupIdentity::Multipart(StagingMultipart {
                target_store: self.target_store,
                bucket: self.bucket,
                key: self.object_key,
                upload_id: self.upload_id.ok_or_else(|| {
                    AttemptError::Invalid("multipart cleanup identity is incomplete".into())
                })?,
            }),
            value => {
                return Err(AttemptError::Invalid(format!(
                    "unknown cleanup kind '{value}'"
                )));
            }
        };
        Ok(CleanupClaim {
            cleanup_id: u64::try_from(self.cleanup_id).map_err(invalid_number)?,
            attempt_id: OpaqueId::parse(self.attempt_id)
                .map_err(|error| AttemptError::Invalid(error.to_string()))?,
            claim_token,
            identity,
        })
    }
}