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

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use remora_coordination::{
    ActivateOptions, CoordinationStore, LeaseOwner, LocalPermitGate, MAX_SAFE_INTEGER, OpCode,
};
use remora_domain::{OpaqueId, fingerprint, ports::MutationScope};
use remora_durable::{
    DurableError, MutationDisposition, MutationStatus, enqueue, execute_mutation, with_transaction,
};
use serde::Serialize;
use serde_json::json;
use sqlx::{PgPool, types::Json};
use time::OffsetDateTime;

use crate::{
    ActivationMutationOutcome, ActivationOutcome, AttemptAuthority, AttemptError,
    AttemptFaultInjector, AttemptFaultPoint, AttemptState, CheckpointProjection, ClaimOutcome,
    ClaimRejection, ExecutionEnvelopeRecord, LeaseProjection, RecoveryOutcome,
    RenewalMutationOutcome, RenewalOutcome,
};

struct NoFaults;

impl AttemptFaultInjector for NoFaults {
    fn hit(&self, _point: AttemptFaultPoint) {}
}

#[derive(Clone, Debug)]
pub struct AttemptConfig {
    pub max_attempts: u64,
    pub local_limit: u64,
}

impl Default for AttemptConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            local_limit: 4,
        }
    }
}

#[derive(Clone)]
pub struct AttemptService {
    pub(crate) pool: PgPool,
    pub(crate) coordination: CoordinationStore,
    local_permits: LocalPermitGate,
    pub(crate) operation_subject: OpaqueId,
    pub(crate) faults: Arc<dyn AttemptFaultInjector>,
    pub(crate) config: AttemptConfig,
}

impl AttemptService {
    pub fn new(
        pool: PgPool,
        coordination: CoordinationStore,
        config: AttemptConfig,
    ) -> Result<Self, AttemptError> {
        Self::new_with_fault_injector(pool, coordination, config, Arc::new(NoFaults))
    }

    pub fn new_with_fault_injector(
        pool: PgPool,
        coordination: CoordinationStore,
        config: AttemptConfig,
        faults: Arc<dyn AttemptFaultInjector>,
    ) -> Result<Self, AttemptError> {
        if config.max_attempts == 0 || config.max_attempts > MAX_SAFE_INTEGER {
            return Err(AttemptError::Invalid(
                "max_attempts must fit the coordination safe-integer range".into(),
            ));
        }
        if config.local_limit == 0 {
            return Err(AttemptError::Invalid(
                "local_limit must be greater than zero".into(),
            ));
        }
        Ok(Self {
            pool,
            coordination,
            local_permits: LocalPermitGate::new(config.local_limit),
            operation_subject: OpaqueId::parse("remora_attempt_service")
                .map_err(|error| AttemptError::Invalid(error.to_string()))?,
            faults,
            config,
        })
    }

    pub async fn claim(
        &self,
        job_id: &OpaqueId,
        worker_incarnation: &OpaqueId,
        operation_id: &OpaqueId,
    ) -> Result<ClaimOutcome, AttemptError> {
        if self.local_permits.acquire(job_id.as_str()).is_err() {
            return Ok(ClaimOutcome::Rejected {
                reason: ClaimRejection::LocalCapacity,
                disposition: None,
            });
        }
        let request = ClaimRequest {
            job_id,
            worker_incarnation,
        };
        let request_fingerprint =
            fingerprint(&request).map_err(|error| AttemptError::Invalid(error.to_string()))?;
        let attempt_id = new_id('a')?;
        let job = job_id.clone();
        let worker = worker_incarnation.clone();
        let max_attempts = self.config.max_attempts;
        let coordination = self.coordination.clone();
        let faults = self.faults.clone();
        let reserved = Arc::new(AtomicBool::new(false));
        let effect_reserved = reserved.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &self.operation_subject,
                operation: "attempt.mutation",
                key: operation_id.as_str(),
            },
            "claim_attempt",
            &[],
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let decision = authorize_attempt(
                        tx,
                        AuthorizeInput {
                            coordination: &coordination,
                            job_id: &job,
                            attempt_id: &attempt_id,
                            worker: &worker,
                            max_attempts,
                            faults: faults.as_ref(),
                            reserved: effect_reserved.as_ref(),
                        },
                    )
                    .await?;
                    let status = if matches!(decision, ClaimOutcome::Claimed { .. }) {
                        MutationStatus::Accepted
                    } else {
                        MutationStatus::Rejected
                    };
                    let response = serde_json::to_value(decision)
                        .map_err(|error| DurableError::Effect(error.to_string()))?;
                    Ok((status, response))
                })
            },
        )
        .await;
        let outcome = match outcome {
            Ok(outcome) => outcome,
            Err(error) => {
                if reserved.load(Ordering::Acquire) {
                    self.compensate_reservation(job_id, worker_incarnation)
                        .await;
                } else {
                    self.local_permits.release(job_id.as_str());
                }
                return Err(error.into());
            }
        };
        self.faults.hit(AttemptFaultPoint::ClaimAfterDurable);
        if outcome.disposition == MutationDisposition::Conflict {
            self.local_permits.release(job_id.as_str());
            return Err(AttemptError::OperationConflict);
        }
        let decision: ClaimOutcome = serde_json::from_value(outcome.response)
            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        if matches!(decision, ClaimOutcome::Rejected { .. }) {
            self.local_permits.release(job_id.as_str());
        }
        Ok(decision.with_disposition(outcome.disposition))
    }

    pub async fn activate(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
    ) -> Result<ActivationOutcome, AttemptError> {
        Ok(
            match self.activate_operation(authority, operation_id).await? {
                ActivationMutationOutcome::Activated { disposition, .. } => {
                    ActivationOutcome::Activated(disposition)
                }
                ActivationMutationOutcome::PreconditionFailed => {
                    ActivationOutcome::PreconditionFailed
                }
            },
        )
    }

    pub async fn activate_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
    ) -> Result<ActivationMutationOutcome, AttemptError> {
        validate_authority(authority)?;
        let request_fingerprint = fingerprint(&ActivationRequest {
            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 eligibility = authority_binding(authority);
        let authority_for_effect = authority.clone();
        let coordination = self.coordination.clone();
        let faults = self.faults.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &self.operation_subject,
                operation: "attempt.mutation",
                key: operation_id.as_str(),
            },
            "activate_attempt",
            &eligibility,
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let lease =
                        activate_attempt(tx, &coordination, &authority_for_effect, faults.as_ref())
                            .await?;
                    Ok((
                        if lease.is_some() {
                            MutationStatus::Accepted
                        } else {
                            MutationStatus::Rejected
                        },
                        serde_json::to_value(ActivationStored { lease })
                            .map_err(DurableError::effect)?,
                    ))
                })
            },
        )
        .await;
        let outcome = match outcome {
            Ok(outcome) => outcome,
            Err(error) => {
                if self
                    .exact_authority_is_executing_and_live(authority)
                    .await?
                {
                    let lease = self
                        .live_lease(authority)
                        .await?
                        .ok_or_else(|| AttemptError::Invalid("live lease disappeared".into()))?;
                    return Ok(ActivationMutationOutcome::Activated {
                        disposition: MutationDisposition::First,
                        lease,
                    });
                }
                self.release_exact(authority).await;
                return Err(error.into());
            }
        };
        self.faults.hit(AttemptFaultPoint::ActivationAfterDurable);
        if outcome.disposition == MutationDisposition::Conflict {
            return Err(AttemptError::OperationConflict);
        }
        if outcome.status == MutationStatus::Rejected {
            self.release_exact(authority).await;
            return Ok(ActivationMutationOutcome::PreconditionFailed);
        }
        if !self
            .exact_authority_is_executing_and_live(authority)
            .await?
        {
            self.release_exact(authority).await;
            return Ok(ActivationMutationOutcome::PreconditionFailed);
        }
        let stored: ActivationStored = serde_json::from_value(outcome.response)
            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        let lease = stored
            .lease
            .ok_or_else(|| AttemptError::Invalid("accepted activation omitted its lease".into()))?;
        Ok(ActivationMutationOutcome::Activated {
            disposition: outcome.disposition,
            lease,
        })
    }

    pub async fn renew(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<RenewalOutcome, AttemptError> {
        validate_authority(authority)?;
        if !self.exact_authority_is_executing(authority).await? {
            return Ok(RenewalOutcome::AuthorityLost);
        }
        self.faults.hit(AttemptFaultPoint::RenewalBeforeLease);
        let result = self
            .coordination
            .renew(authority.job_id.as_str(), &lease_owner(authority)?)
            .await?;
        self.faults.hit(AttemptFaultPoint::RenewalAfterLease);
        if !result.ok() || !self.exact_authority_is_executing(authority).await? {
            return Ok(RenewalOutcome::AuthorityLost);
        }
        Ok(RenewalOutcome::Renewed)
    }

    pub async fn renew_operation(
        &self,
        authority: &AttemptAuthority,
        operation_id: &OpaqueId,
    ) -> Result<RenewalMutationOutcome, AttemptError> {
        validate_authority(authority)?;
        let request_fingerprint = fingerprint(&ActivationRequest {
            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 eligibility = authority_binding(authority);
        let authority = authority.clone();
        let coordination = self.coordination.clone();
        let faults = self.faults.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &self.operation_subject,
                operation: "attempt.mutation",
                key: operation_id.as_str(),
            },
            "renew_attempt",
            &eligibility,
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
                        return Ok((
                            MutationStatus::Rejected,
                            serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
                                disposition: None,
                            })
                            .map_err(DurableError::effect)?,
                        ));
                    };
                    if row.state != AttemptState::Executing
                        || !row.matches(
                            &authority.job_id,
                            authority.fence,
                            &authority.worker_incarnation,
                        )
                    {
                        return Ok((
                            MutationStatus::Rejected,
                            serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
                                disposition: None,
                            })
                            .map_err(DurableError::effect)?,
                        ));
                    }
                    faults.hit(AttemptFaultPoint::RenewalBeforeLease);
                    let expected = lease_owner(&authority).map_err(DurableError::effect)?;
                    let renewed = coordination
                        .renew(authority.job_id.as_str(), &expected)
                        .await
                        .map_err(DurableError::effect)?;
                    faults.hit(AttemptFaultPoint::RenewalAfterLease);
                    let live = coordination
                        .get_live_lease(authority.job_id.as_str())
                        .await
                        .map_err(DurableError::effect)?;
                    let response = match live
                        .filter(|lease| renewed.ok() && lease.record.owner() == expected)
                    {
                        Some(lease) => RenewalMutationOutcome::Renewed {
                            lease: LeaseProjection {
                                expires_at_ms: lease.record.expires_at,
                                observed_at_ms: lease.now_ms,
                            },
                            disposition: None,
                        },
                        None => RenewalMutationOutcome::AuthorityLost { disposition: None },
                    };
                    let status = if matches!(response, RenewalMutationOutcome::Renewed { .. }) {
                        MutationStatus::Accepted
                    } else {
                        MutationStatus::Rejected
                    };
                    Ok((
                        status,
                        serde_json::to_value(response).map_err(DurableError::effect)?,
                    ))
                })
            },
        )
        .await?;
        if outcome.disposition == MutationDisposition::Conflict {
            return Err(AttemptError::OperationConflict);
        }
        let response: RenewalMutationOutcome = serde_json::from_value(outcome.response)
            .map_err(|error| AttemptError::Invalid(error.to_string()))?;
        Ok(response.with_disposition(outcome.disposition))
    }

    pub async fn recover_inactive(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<RecoveryOutcome, AttemptError> {
        validate_authority(authority)?;
        let job_id = authority.job_id.clone();
        let attempt_id = authority.attempt_id.clone();
        let worker = authority.worker_incarnation.clone();
        let fence = authority.fence;
        let coordination = self.coordination.clone();
        let faults = self.faults.clone();
        let outcome = with_transaction(&self.pool, move |tx| {
            Box::pin(async move {
                let row = lock_authority(tx, &attempt_id).await?;
                let Some(row) = row else {
                    return Ok(RecoveryOutcome::AlreadyInactive);
                };
                if !row.matches(&job_id, fence, &worker) {
                    return Ok(RecoveryOutcome::AlreadyInactive);
                }
                if row.state == AttemptState::Executing {
                    return Ok(RecoveryOutcome::StillOwned);
                }
                if row.state != AttemptState::Authorized {
                    return Ok(RecoveryOutcome::AlreadyInactive);
                }
                faults.hit(AttemptFaultPoint::RecoveryBeforeRevalidation);
                if coordination
                    .get_reservation(job_id.as_str())
                    .await
                    .map_err(DurableError::effect)?
                    .is_some()
                {
                    return Ok(RecoveryOutcome::StillOwned);
                }
                if let Some(lease) = coordination
                    .get_live_lease(job_id.as_str())
                    .await
                    .map_err(DurableError::effect)?
                {
                    if lease.record.owner()
                        == LeaseOwner::new(attempt_id.as_str(), fence, worker.as_str())
                            .map_err(DurableError::effect)?
                    {
                        return Ok(RecoveryOutcome::StillOwned);
                    }
                    return Err(DurableError::Effect(
                        "foreign live lease conflicts with durable authority".into(),
                    ));
                }
                let now = OffsetDateTime::now_utc();
                faults.hit(AttemptFaultPoint::RecoveryBeforeDurable);
                sqlx::query(
                    "UPDATE attempts
                     SET state = 'reaped', reaped_at = $2
                     WHERE attempt_id = $1 AND state = 'authorized'",
                )
                .bind(attempt_id.as_str())
                .bind(now)
                .execute(tx.as_mut())
                .await?;
                let changed = sqlx::query(
                    "UPDATE jobs
                     SET phase = 'queued', authority_attempt_id = NULL,
                         dispatch_needed = TRUE, revision = revision + 1, updated_at = $2
                     WHERE job_id = $1 AND authority_attempt_id = $3
                       AND current_fence = $4 AND phase = 'running'",
                )
                .bind(job_id.as_str())
                .bind(now)
                .bind(attempt_id.as_str())
                .bind(i64::try_from(fence).map_err(DurableError::effect)?)
                .execute(tx.as_mut())
                .await?;
                if changed.rows_affected() != 1 {
                    return Err(DurableError::Effect(
                        "attempt authority changed while recovery held its lock".into(),
                    ));
                }
                enqueue(tx, "job.dispatch", json!({ "job_id": job_id })).await?;
                Ok(RecoveryOutcome::Reaped)
            })
        })
        .await?;
        self.faults.hit(AttemptFaultPoint::RecoveryAfterDurable);
        if matches!(
            outcome,
            RecoveryOutcome::Reaped | RecoveryOutcome::AlreadyInactive
        ) {
            self.release_exact(authority).await;
        }
        Ok(outcome)
    }

    pub async fn authority(
        &self,
        attempt_id: &OpaqueId,
    ) -> Result<Option<AttemptAuthority>, AttemptError> {
        load_authority(&self.pool, attempt_id)
            .await
            .map_err(Into::into)
    }

    pub async fn historical_authority(
        &self,
        attempt_id: &OpaqueId,
    ) -> Result<Option<AttemptAuthority>, AttemptError> {
        load_historical_authority(&self.pool, attempt_id)
            .await
            .map_err(Into::into)
    }

    pub async fn execution_envelope(
        &self,
        attempt_id: &OpaqueId,
    ) -> Result<Option<ExecutionEnvelopeRecord>, AttemptError> {
        let value = sqlx::query_scalar::<_, Option<Json<ExecutionEnvelopeRecord>>>(
            "SELECT execution_envelope FROM attempts WHERE attempt_id = $1",
        )
        .bind(attempt_id.as_str())
        .fetch_optional(&self.pool)
        .await?
        .flatten();
        Ok(value.map(|value| value.0))
    }

    pub async fn live_lease(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<Option<LeaseProjection>, AttemptError> {
        let expected = lease_owner(authority)?;
        Ok(self
            .coordination
            .get_live_lease(authority.job_id.as_str())
            .await?
            .filter(|lease| lease.record.owner() == expected)
            .map(|lease| LeaseProjection {
                expires_at_ms: lease.record.expires_at,
                observed_at_ms: lease.now_ms,
            }))
    }

    pub fn release_local(&self, job_id: &OpaqueId) -> bool {
        self.local_permits.release_all(job_id.as_str())
    }

    #[must_use]
    pub fn local_used(&self) -> u64 {
        self.local_permits.used()
    }

    async fn compensate_reservation(&self, job_id: &OpaqueId, worker: &OpaqueId) {
        let _ = self
            .coordination
            .free_reserve(job_id.as_str(), Some(worker.as_str()))
            .await;
        self.local_permits.release(job_id.as_str());
    }

    pub(crate) async fn release_exact(&self, authority: &AttemptAuthority) {
        if let Ok(owner) = lease_owner(authority) {
            let _ = self
                .coordination
                .release(authority.job_id.as_str(), &owner)
                .await;
        }
    }

    async fn exact_authority_is_executing(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<bool, AttemptError> {
        Ok(load_authority(&self.pool, &authority.attempt_id)
            .await?
            .is_some_and(|current| current == authority.clone_with_state(AttemptState::Executing)))
    }

    async fn exact_authority_is_executing_and_live(
        &self,
        authority: &AttemptAuthority,
    ) -> Result<bool, AttemptError> {
        if !self.exact_authority_is_executing(authority).await? {
            return Ok(false);
        }
        let Some(lease) = self
            .coordination
            .get_live_lease(authority.job_id.as_str())
            .await?
        else {
            return Ok(false);
        };
        Ok(lease.record.owner() == lease_owner(authority)?)
    }
}

impl AttemptAuthority {
    fn clone_with_state(&self, state: AttemptState) -> Self {
        Self {
            job_id: self.job_id.clone(),
            attempt_id: self.attempt_id.clone(),
            fence: self.fence,
            worker_incarnation: self.worker_incarnation.clone(),
            state,
        }
    }
}

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

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

#[derive(serde::Deserialize, Serialize)]
struct ActivationStored {
    lease: Option<LeaseProjection>,
}

struct AuthorityRow {
    job_id: OpaqueId,
    fence: u64,
    worker: OpaqueId,
    state: AttemptState,
}

impl AuthorityRow {
    fn matches(&self, job_id: &OpaqueId, fence: u64, worker: &OpaqueId) -> bool {
        self.job_id == *job_id && self.fence == fence && self.worker == *worker
    }
}

struct AuthorizeInput<'a> {
    coordination: &'a CoordinationStore,
    job_id: &'a OpaqueId,
    attempt_id: &'a OpaqueId,
    worker: &'a OpaqueId,
    max_attempts: u64,
    faults: &'a dyn AttemptFaultInjector,
    reserved: &'a AtomicBool,
}

async fn authorize_attempt(
    tx: &mut remora_durable::PgTx<'_>,
    input: AuthorizeInput<'_>,
) -> Result<ClaimOutcome, DurableError> {
    let row = sqlx::query_as::<
        _,
        (
            String,
            bool,
            i64,
            Option<String>,
            Option<Json<crate::DurableDescriptor>>,
            Option<Json<CheckpointProjection>>,
            Option<String>,
        ),
    >(
        "SELECT phase, dispatch_needed, current_fence, authority_attempt_id,
                resolved_source, checkpoint, checkpoint_owner_attempt_id
         FROM jobs WHERE job_id = $1 FOR UPDATE",
    )
    .bind(input.job_id.as_str())
    .fetch_optional(tx.as_mut())
    .await?;
    let Some((
        phase,
        dispatch_needed,
        current_fence,
        authority,
        resolved_source,
        checkpoint,
        checkpoint_owner,
    )) = row
    else {
        return Ok(claim_rejected(ClaimRejection::JobUnavailable));
    };
    if phase != "queued" || !dispatch_needed || authority.is_some() {
        return Ok(claim_rejected(ClaimRejection::JobUnavailable));
    }
    let current_fence = u64::try_from(current_fence).map_err(DurableError::effect)?;
    if current_fence >= input.max_attempts {
        return Ok(claim_rejected(ClaimRejection::AttemptsExhausted));
    }
    let reusable = resolved_source
        .as_ref()
        .zip(checkpoint.as_ref())
        .and_then(|(source, checkpoint)| checkpoint.0.reusable().map(|value| (source, value)))
        .filter(|(source, checkpoint)| checkpoint.resolved_source_id == source.0.id);
    if let (Some(_), Some(owner)) = (reusable, checkpoint_owner.as_deref()) {
        let stopped = sqlx::query_scalar::<_, bool>(
            "SELECT state = 'stopped'
             FROM attempts
             WHERE attempt_id = $1 AND job_id = $2",
        )
        .bind(owner)
        .bind(input.job_id.as_str())
        .fetch_optional(tx.as_mut())
        .await?
        .unwrap_or(false);
        if !stopped {
            return Ok(claim_rejected(ClaimRejection::JobUnavailable));
        }
    }
    input.faults.hit(AttemptFaultPoint::ClaimBeforeReserve);
    let reservation = input
        .coordination
        .reserve(input.job_id.as_str(), input.worker.as_str())
        .await
        .map_err(DurableError::effect)?;
    if !reservation.ok() {
        return Ok(claim_rejected(match reservation.code {
            OpCode::FailFull => ClaimRejection::GlobalCapacity,
            _ => ClaimRejection::ReservationHeld,
        }));
    }
    input.reserved.store(true, Ordering::Release);
    input.faults.hit(AttemptFaultPoint::ClaimAfterReserve);
    let fence = current_fence + 1;
    let checkpoint_stage = match (reusable, checkpoint_owner.as_deref()) {
        (Some((_, checkpoint)), Some(owner)) if checkpoint.validated_bytes > 0 => {
            load_checkpoint_stage(tx, input.job_id, owner, checkpoint).await?
        }
        _ => None,
    };
    if checkpoint.is_some() && checkpoint_stage.is_none() {
        sqlx::query(
            "UPDATE jobs
             SET checkpoint = NULL, checkpoint_owner_attempt_id = NULL,
                 revision = revision + 1, updated_at = $2
             WHERE job_id = $1",
        )
        .bind(input.job_id.as_str())
        .bind(OffsetDateTime::now_utc())
        .execute(tx.as_mut())
        .await?;
    }
    let initial_progress = checkpoint_stage
        .as_ref()
        .map_or(0, |stage| stage.validated_bytes);
    let staging_origin = checkpoint_stage
        .as_ref()
        .map_or(input.attempt_id, |stage| &stage.staging_origin_attempt_id);
    let checkpoint_owner_after_claim = checkpoint_stage.as_ref().map(|_| input.attempt_id.as_str());
    let now = OffsetDateTime::now_utc();
    let envelope = load_execution_envelope_record(
        tx,
        input.job_id,
        input.attempt_id,
        fence,
        input.worker,
        staging_origin,
    )
    .await?;
    sqlx::query(
        "INSERT INTO attempts (
            attempt_id, job_id, fence, worker_incarnation, state,
            initial_progress, authorized_at, executing_at, reaped_at, execution_envelope,
            staging_origin_attempt_id
         ) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL, $7, $8)",
    )
    .bind(input.attempt_id.as_str())
    .bind(input.job_id.as_str())
    .bind(i64::try_from(fence).map_err(DurableError::effect)?)
    .bind(input.worker.as_str())
    .bind(i64::try_from(initial_progress).map_err(DurableError::effect)?)
    .bind(now)
    .bind(Json(&envelope))
    .bind(staging_origin.as_str())
    .execute(tx.as_mut())
    .await?;
    if let Some(stage) = checkpoint_stage {
        transfer_checkpoint_stage(tx, &stage, input.attempt_id).await?;
    }
    sqlx::query(
        "UPDATE jobs
         SET phase = 'running', current_fence = $2, authority_attempt_id = $3,
             checkpoint_owner_attempt_id = $5, dispatch_needed = FALSE,
             revision = revision + 1, updated_at = $4
         WHERE job_id = $1",
    )
    .bind(input.job_id.as_str())
    .bind(i64::try_from(fence).map_err(DurableError::effect)?)
    .bind(input.attempt_id.as_str())
    .bind(now)
    .bind(checkpoint_owner_after_claim)
    .execute(tx.as_mut())
    .await?;
    Ok(ClaimOutcome::Claimed {
        authority: AttemptAuthority {
            job_id: input.job_id.clone(),
            attempt_id: input.attempt_id.clone(),
            fence,
            worker_incarnation: input.worker.clone(),
            state: AttemptState::Authorized,
        },
        disposition: None,
    })
}

struct CheckpointStage {
    owner_attempt_id: OpaqueId,
    staging_origin_attempt_id: OpaqueId,
    validated_bytes: u64,
    validated_members: u64,
    inventory_bytes: u64,
    staging_multipart: Option<crate::CheckpointMultipart>,
}

async fn load_checkpoint_stage(
    tx: &mut remora_durable::PgTx<'_>,
    job_id: &OpaqueId,
    owner: &str,
    checkpoint: &crate::CheckpointRecord,
) -> Result<Option<CheckpointStage>, DurableError> {
    let Some(staging_origin) = sqlx::query_scalar::<_, String>(
        "SELECT staging_origin_attempt_id
         FROM attempts
         WHERE attempt_id = $1 AND job_id = $2 AND state = 'stopped'",
    )
    .bind(owner)
    .bind(job_id.as_str())
    .fetch_optional(tx.as_mut())
    .await?
    else {
        return Ok(None);
    };
    let Some((state, member_count, inventory_origin, inventory_prefix, target_prefix)) =
        sqlx::query_as::<
            _,
            (
                String,
                i64,
                Option<String>,
                Option<String>,
                Json<Vec<String>>,
            ),
        >(
            "SELECT i.state, i.member_count, i.staging_origin_attempt_id,
                    i.target_prefix, j.target_prefix
             FROM result_inventories i
             JOIN attempts a ON a.attempt_id = i.attempt_id
             JOIN jobs j ON j.job_id = a.job_id
             WHERE i.attempt_id = $1
             FOR UPDATE OF i",
        )
        .bind(owner)
        .fetch_optional(tx.as_mut())
        .await?
    else {
        return Ok(None);
    };
    if inventory_origin.as_deref() != Some(staging_origin.as_str())
        || inventory_prefix.as_deref() != Some(target_prefix.0.join("/").as_str())
    {
        return Ok(None);
    }
    let validated_members =
        i64::try_from(checkpoint.validated_members).map_err(DurableError::effect)?;
    let validated_bytes =
        i64::try_from(checkpoint.validated_bytes).map_err(DurableError::effect)?;
    if state != "open" || member_count < validated_members {
        return Ok(None);
    }
    let (retained_members, retained_bytes) = sqlx::query_as::<_, (i64, i64)>(
        "SELECT COUNT(*), COALESCE(SUM(size), 0)::BIGINT
         FROM (
             SELECT size
             FROM result_inventory_members
             WHERE attempt_id = $1
             ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
             LIMIT $2
         ) retained",
    )
    .bind(owner)
    .bind(validated_members)
    .fetch_one(tx.as_mut())
    .await?;
    if retained_members != validated_members || retained_bytes > validated_bytes {
        return Ok(None);
    }
    let partial = retained_bytes < validated_bytes;
    if partial != checkpoint.staging_multipart.is_some() {
        return Ok(None);
    }
    let cleanup_started = sqlx::query_scalar::<_, bool>(
        "WITH retained AS (
             SELECT object_store, object_bucket, object_key, object_version_id
             FROM result_inventory_members
             WHERE attempt_id = $1
             ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
             LIMIT $2
         )
         SELECT EXISTS (
             SELECT 1
             FROM attempt_staging_cleanup c
             JOIN retained r
               ON r.object_store = c.target_store
              AND r.object_bucket = c.bucket
              AND r.object_key = c.object_key
              AND r.object_version_id IS NOT DISTINCT FROM c.version_id
             WHERE c.attempt_id = $1
               AND c.kind = 'object'
               AND c.state <> 'pending'
         )",
    )
    .bind(owner)
    .bind(validated_members)
    .fetch_one(tx.as_mut())
    .await?;
    if cleanup_started {
        return Ok(None);
    }
    if let Some(multipart) = &checkpoint.staging_multipart {
        let available = sqlx::query_scalar::<_, bool>(
            "SELECT EXISTS (
                 SELECT 1 FROM attempt_staging_cleanup
                 WHERE attempt_id = $1 AND kind = 'multipart'
                   AND target_store = $2 AND bucket = $3
                   AND object_key = $4 AND upload_id = $5
                   AND state = 'pending'
             )",
        )
        .bind(owner)
        .bind(&multipart.target_store)
        .bind(&multipart.bucket)
        .bind(&multipart.key)
        .bind(&multipart.upload_id)
        .fetch_one(tx.as_mut())
        .await?;
        if !available {
            return Ok(None);
        }
    }
    Ok(Some(CheckpointStage {
        owner_attempt_id: OpaqueId::parse(owner).map_err(DurableError::effect)?,
        staging_origin_attempt_id: OpaqueId::parse(staging_origin).map_err(DurableError::effect)?,
        validated_bytes: checkpoint.validated_bytes,
        validated_members: checkpoint.validated_members,
        inventory_bytes: u64::try_from(retained_bytes).map_err(DurableError::effect)?,
        staging_multipart: checkpoint.staging_multipart.clone(),
    }))
}

async fn transfer_checkpoint_stage(
    tx: &mut remora_durable::PgTx<'_>,
    stage: &CheckpointStage,
    replacement: &OpaqueId,
) -> Result<(), DurableError> {
    let retained_members = i64::try_from(stage.validated_members).map_err(DurableError::effect)?;
    let retained_bytes = i64::try_from(stage.inventory_bytes).map_err(DurableError::effect)?;
    let now = OffsetDateTime::now_utc();
    sqlx::query(
        "WITH ranked AS (
             SELECT *, ROW_NUMBER() OVER (
                 ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
             ) AS ordinal
             FROM result_inventory_members
             WHERE attempt_id = $1
         )
         INSERT INTO attempt_staging_cleanup (
             attempt_id, kind, target_store, bucket, object_key, version_id, created_at
         )
         SELECT $1, 'object', object_store, object_bucket, object_key,
                object_version_id, $3
         FROM ranked
         WHERE ordinal > $2
         ON CONFLICT DO NOTHING",
    )
    .bind(stage.owner_attempt_id.as_str())
    .bind(retained_members)
    .bind(now)
    .execute(tx.as_mut())
    .await?;
    sqlx::query(
        "WITH retained AS (
             SELECT object_store, object_bucket, object_key, object_version_id
             FROM result_inventory_members
             WHERE attempt_id = $1
             ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
             LIMIT $3
         )
         UPDATE attempt_staging_cleanup c
         SET attempt_id = $2
         FROM retained r
         WHERE c.attempt_id = $1
           AND c.kind = 'object'
           AND c.state = 'pending'
           AND r.object_store = c.target_store
           AND r.object_bucket = c.bucket
           AND r.object_key = c.object_key
           AND r.object_version_id IS NOT DISTINCT FROM c.version_id",
    )
    .bind(stage.owner_attempt_id.as_str())
    .bind(replacement.as_str())
    .bind(retained_members)
    .execute(tx.as_mut())
    .await?;
    if let Some(multipart) = &stage.staging_multipart {
        let changed = sqlx::query(
            "UPDATE attempt_staging_cleanup
             SET attempt_id = $2
             WHERE attempt_id = $1 AND state = 'pending'
               AND (
                   (kind = 'multipart' AND target_store = $3 AND bucket = $4
                    AND object_key = $5 AND upload_id = $6)
                   OR (kind = 'write_intent' AND target_store = $3 AND bucket = $4
                       AND object_key = $5)
               )",
        )
        .bind(stage.owner_attempt_id.as_str())
        .bind(replacement.as_str())
        .bind(&multipart.target_store)
        .bind(&multipart.bucket)
        .bind(&multipart.key)
        .bind(&multipart.upload_id)
        .execute(tx.as_mut())
        .await?;
        if changed.rows_affected() == 0 {
            return Err(DurableError::effect(
                "checkpoint multipart ownership transfer failed",
            ));
        }
    }
    sqlx::query(
        "WITH extras AS (
             SELECT attempt_id, member_id
             FROM (
                 SELECT attempt_id, member_id, ROW_NUMBER() OVER (
                     ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
                 ) AS ordinal
                 FROM result_inventory_members
                 WHERE attempt_id = $1
             ) ranked
             WHERE ordinal > $2
         )
         DELETE FROM result_inventory_members m
         USING extras
         WHERE m.attempt_id = extras.attempt_id
           AND m.member_id = extras.member_id",
    )
    .bind(stage.owner_attempt_id.as_str())
    .bind(retained_members)
    .execute(tx.as_mut())
    .await?;
    let changed = sqlx::query(
        "UPDATE result_inventories
         SET attempt_id = $2, member_count = $3, total_bytes = $4
         WHERE attempt_id = $1 AND state = 'open'",
    )
    .bind(stage.owner_attempt_id.as_str())
    .bind(replacement.as_str())
    .bind(retained_members)
    .bind(retained_bytes)
    .execute(tx.as_mut())
    .await?;
    if changed.rows_affected() != 1 {
        return Err(DurableError::effect(
            "checkpoint staging inventory transfer failed",
        ));
    }
    Ok(())
}

async fn load_execution_envelope_record(
    tx: &mut remora_durable::PgTx<'_>,
    job_id: &OpaqueId,
    attempt_id: &OpaqueId,
    fence: u64,
    worker: &OpaqueId,
    staging_origin_attempt_id: &OpaqueId,
) -> Result<ExecutionEnvelopeRecord, DurableError> {
    type Row = (
        Json<serde_json::Value>,
        String,
        String,
        String,
        Json<Vec<String>>,
        String,
        Option<String>,
        Option<String>,
        Option<Json<crate::DurableDescriptor>>,
        Option<Json<CheckpointProjection>>,
    );
    let row = sqlx::query_as::<_, Row>(
        "SELECT j.source_projection, j.sensitive_source_ref, j.target_store,
                j.target_bucket, j.target_prefix, j.execution_grant_id,
                g.credential_id, g.credential_version_id, j.resolved_source, j.checkpoint
         FROM jobs j
         JOIN execution_grants g ON g.grant_id = j.execution_grant_id
         WHERE j.job_id = $1",
    )
    .bind(job_id.as_str())
    .fetch_one(tx.as_mut())
    .await?;
    let grant_id = OpaqueId::parse(row.5).map_err(DurableError::effect)?;
    Ok(ExecutionEnvelopeRecord {
        id: grant_id.clone(),
        job_id: job_id.clone(),
        attempt_id: attempt_id.clone(),
        fence,
        worker_incarnation: worker.clone(),
        source: row.0.0,
        sensitive_source_ref: OpaqueId::parse(row.1).map_err(DurableError::effect)?,
        credential_id: row
            .6
            .map(OpaqueId::parse)
            .transpose()
            .map_err(DurableError::effect)?,
        credential_version_id: row
            .7
            .map(OpaqueId::parse)
            .transpose()
            .map_err(DurableError::effect)?,
        resolved_source: row.8.map(|value| value.0),
        target_store: row.2,
        target_bucket: row.3,
        target_prefix: row.4.0,
        staging_origin_attempt_id: staging_origin_attempt_id.clone(),
        execution_grant_id: grant_id,
        checkpoint: row.9.map(|value| value.0),
    })
}

async fn activate_attempt(
    tx: &mut remora_durable::PgTx<'_>,
    coordination: &CoordinationStore,
    authority: &AttemptAuthority,
    faults: &dyn AttemptFaultInjector,
) -> Result<Option<LeaseProjection>, DurableError> {
    let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
        return Ok(None);
    };
    if !row.matches(
        &authority.job_id,
        authority.fence,
        &authority.worker_incarnation,
    ) {
        return Ok(None);
    }
    if row.state == AttemptState::Executing {
        let live = coordination
            .get_live_lease(authority.job_id.as_str())
            .await
            .map_err(DurableError::effect)?;
        let expected = lease_owner(authority).map_err(DurableError::effect)?;
        if !live.is_some_and(|lease| lease.record.owner() == expected) {
            return Ok(None);
        }
        if !coordination
            .renew(authority.job_id.as_str(), &expected)
            .await
            .map_err(DurableError::effect)?
            .ok()
        {
            return Ok(None);
        }
        return Ok(coordination
            .get_live_lease(authority.job_id.as_str())
            .await
            .map_err(DurableError::effect)?
            .filter(|lease| lease.record.owner() == expected)
            .map(|lease| LeaseProjection {
                expires_at_ms: lease.record.expires_at,
                observed_at_ms: lease.now_ms,
            }));
    }
    if row.state != AttemptState::Authorized {
        return Ok(None);
    }
    faults.hit(AttemptFaultPoint::ActivationBeforeLease);
    let owner = lease_owner(authority).map_err(DurableError::effect)?;
    let coordinated = coordination
        .activate(
            authority.job_id.as_str(),
            &owner,
            ActivateOptions {
                seed_progress: true,
                require_reserve: true,
            },
        )
        .await
        .map_err(DurableError::effect)?;
    if !coordinated.ok() {
        return Ok(None);
    }
    faults.hit(AttemptFaultPoint::ActivationAfterLease);
    let live = coordination
        .get_live_lease(authority.job_id.as_str())
        .await
        .map_err(DurableError::effect)?;
    let expected = lease_owner(authority).map_err(DurableError::effect)?;
    let Some(live) = live.filter(|lease| lease.record.owner() == expected) else {
        return Ok(None);
    };
    let now = OffsetDateTime::now_utc();
    sqlx::query(
        "UPDATE attempts SET state = 'executing', executing_at = $2
         WHERE attempt_id = $1 AND state = 'authorized'",
    )
    .bind(authority.attempt_id.as_str())
    .bind(now)
    .execute(tx.as_mut())
    .await?;
    Ok(Some(LeaseProjection {
        expires_at_ms: live.record.expires_at,
        observed_at_ms: live.now_ms,
    }))
}

async fn lock_authority(
    tx: &mut remora_durable::PgTx<'_>,
    attempt_id: &OpaqueId,
) -> Result<Option<AuthorityRow>, DurableError> {
    let row = sqlx::query_as::<_, (String, i64, String, String)>(
        "SELECT a.job_id, a.fence, a.worker_incarnation, a.state
         FROM attempts a
         JOIN jobs j ON j.job_id = a.job_id
         WHERE a.attempt_id = $1
           AND j.authority_attempt_id = a.attempt_id
           AND j.current_fence = a.fence
           AND j.phase = 'running'
         FOR UPDATE OF a, j",
    )
    .bind(attempt_id.as_str())
    .fetch_optional(tx.as_mut())
    .await?;
    row.map(authority_row).transpose()
}

async fn load_authority(
    pool: &PgPool,
    attempt_id: &OpaqueId,
) -> Result<Option<AttemptAuthority>, DurableError> {
    let row = sqlx::query_as::<_, (String, String, i64, String, String)>(
        "SELECT a.job_id, a.attempt_id, a.fence, a.worker_incarnation, a.state
         FROM attempts a
         JOIN jobs j ON j.job_id = a.job_id
         WHERE a.attempt_id = $1
           AND j.authority_attempt_id = a.attempt_id
           AND j.current_fence = a.fence
           AND j.phase = 'running'",
    )
    .bind(attempt_id.as_str())
    .fetch_optional(pool)
    .await?;
    row.map(|(job, attempt, fence, worker, state)| {
        Ok(AttemptAuthority {
            job_id: OpaqueId::parse(job).map_err(DurableError::effect)?,
            attempt_id: OpaqueId::parse(attempt).map_err(DurableError::effect)?,
            fence: u64::try_from(fence).map_err(DurableError::effect)?,
            worker_incarnation: OpaqueId::parse(worker).map_err(DurableError::effect)?,
            state: AttemptState::parse(&state)
                .ok_or_else(|| DurableError::Config(format!("unknown attempt state '{state}'")))?,
        })
    })
    .transpose()
}

async fn load_historical_authority(
    pool: &PgPool,
    attempt_id: &OpaqueId,
) -> Result<Option<AttemptAuthority>, DurableError> {
    let row = sqlx::query_as::<_, (String, String, i64, String, String)>(
        "SELECT job_id, attempt_id, fence, worker_incarnation, state
         FROM attempts WHERE attempt_id = $1",
    )
    .bind(attempt_id.as_str())
    .fetch_optional(pool)
    .await?;
    row.map(|(job, attempt, fence, worker, state)| {
        Ok(AttemptAuthority {
            job_id: OpaqueId::parse(job).map_err(DurableError::effect)?,
            attempt_id: OpaqueId::parse(attempt).map_err(DurableError::effect)?,
            fence: u64::try_from(fence).map_err(DurableError::effect)?,
            worker_incarnation: OpaqueId::parse(worker).map_err(DurableError::effect)?,
            state: AttemptState::parse(&state)
                .ok_or_else(|| DurableError::Config(format!("unknown attempt state '{state}'")))?,
        })
    })
    .transpose()
}

fn authority_row(row: (String, i64, String, String)) -> Result<AuthorityRow, DurableError> {
    Ok(AuthorityRow {
        job_id: OpaqueId::parse(row.0).map_err(DurableError::effect)?,
        fence: u64::try_from(row.1).map_err(DurableError::effect)?,
        worker: OpaqueId::parse(row.2).map_err(DurableError::effect)?,
        state: AttemptState::parse(&row.3)
            .ok_or_else(|| DurableError::Config(format!("unknown attempt state '{}'", row.3)))?,
    })
}

fn validate_authority(authority: &AttemptAuthority) -> Result<(), AttemptError> {
    if authority.fence == 0 || authority.fence > MAX_SAFE_INTEGER {
        return Err(AttemptError::Invalid("invalid fencing token".into()));
    }
    Ok(())
}

fn lease_owner(authority: &AttemptAuthority) -> Result<LeaseOwner, AttemptError> {
    LeaseOwner::new(
        authority.attempt_id.as_str(),
        authority.fence,
        authority.worker_incarnation.as_str(),
    )
    .map_err(Into::into)
}

fn authority_binding(authority: &AttemptAuthority) -> Vec<u8> {
    format!(
        "job={}|attempt={}|fence={}|worker={}",
        authority.job_id, authority.attempt_id, authority.fence, authority.worker_incarnation
    )
    .into_bytes()
}

fn claim_rejected(reason: ClaimRejection) -> ClaimOutcome {
    ClaimOutcome::Rejected {
        reason,
        disposition: None,
    }
}

pub(crate) fn new_id(prefix: char) -> Result<OpaqueId, AttemptError> {
    let mut bytes = [0u8; 16];
    getrandom::fill(&mut bytes)
        .map_err(|error| AttemptError::Invalid(format!("identifier entropy failed: {error}")))?;
    let mut value = String::with_capacity(33);
    value.push(prefix);
    for byte in bytes {
        use std::fmt::Write as _;
        write!(value, "{byte:02x}").expect("writing to String cannot fail");
    }
    OpaqueId::parse(value).map_err(|error| AttemptError::Invalid(error.to_string()))
}