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::{
    ActivationOutcome, AttemptAuthority, AttemptError, AttemptFaultInjector, AttemptFaultPoint,
    AttemptState, CheckpointProjection, ClaimOutcome, ClaimRejection, RecoveryOutcome,
    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,
    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> {
        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 accepted =
                        activate_attempt(tx, &coordination, &authority_for_effect, faults.as_ref())
                            .await?;
                    Ok((
                        if accepted {
                            MutationStatus::Accepted
                        } else {
                            MutationStatus::Rejected
                        },
                        json!({ "activated": accepted }),
                    ))
                })
            },
        )
        .await;
        let outcome = match outcome {
            Ok(outcome) => outcome,
            Err(error) => {
                if self
                    .exact_authority_is_executing_and_live(authority)
                    .await?
                {
                    return Ok(ActivationOutcome::Activated(MutationDisposition::First));
                }
                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(ActivationOutcome::PreconditionFailed);
        }
        if !self
            .exact_authority_is_executing_and_live(authority)
            .await?
        {
            self.release_exact(authority).await;
            return Ok(ActivationOutcome::PreconditionFailed);
        }
        Ok(ActivationOutcome::Activated(outcome.disposition))
    }

    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 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 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());
    }

    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,
}

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>>,
        ),
    >(
        "SELECT phase, dispatch_needed, current_fence, authority_attempt_id,
                resolved_source, checkpoint
         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)) = 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));
    }
    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 initial_progress = match (resolved_source, checkpoint) {
        (Some(source), Some(checkpoint))
            if checkpoint
                .0
                .reusable()
                .is_some_and(|value| value.resolved_source_id == source.0.id) =>
        {
            checkpoint.0.reusable().unwrap().validated_bytes
        }
        _ => 0,
    };
    let now = OffsetDateTime::now_utc();
    sqlx::query(
        "INSERT INTO attempts (
            attempt_id, job_id, fence, worker_incarnation, state,
            initial_progress, authorized_at, executing_at, reaped_at
         ) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL)",
    )
    .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)
    .execute(tx.as_mut())
    .await?;
    sqlx::query(
        "UPDATE jobs
         SET phase = 'running', current_fence = $2, authority_attempt_id = $3,
             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)
    .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,
    })
}

async fn activate_attempt(
    tx: &mut remora_durable::PgTx<'_>,
    coordination: &CoordinationStore,
    authority: &AttemptAuthority,
    faults: &dyn AttemptFaultInjector,
) -> Result<bool, DurableError> {
    let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
        return Ok(false);
    };
    if !row.matches(
        &authority.job_id,
        authority.fence,
        &authority.worker_incarnation,
    ) {
        return Ok(false);
    }
    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(false);
        }
        return Ok(coordination
            .renew(authority.job_id.as_str(), &expected)
            .await
            .map_err(DurableError::effect)?
            .ok());
    }
    if row.state != AttemptState::Authorized {
        return Ok(false);
    }
    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(false);
    }
    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)?;
    if !live.is_some_and(|lease| lease.record.owner() == expected) {
        return Ok(false);
    }
    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(true)
}

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

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,
    }
}

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