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, 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 { with_transaction(&self.pool, |tx| { let job_id = job_id.clone(); Box::pin(async move { 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) }) }) .await .map_err(Into::into) } pub async fn cooperative_stop( &self, authority: &AttemptAuthority, ) -> Result { 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 { 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 { 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 { 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, ) -> Result { 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 { 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( &self, authority: &AttemptAuthority, operation_id: &OpaqueId, kind: &'static str, request: T, effect: F, ) -> Result 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 { 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, 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 { 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, multiparts: Vec, ) -> Result { 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 { 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 { 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 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(OffsetDateTime::now_utc()) .execute(&self.pool) .await?; if changed.rows_affected() == 1 { report.retries_scheduled += 1; row.phase = "queued".into(); row.dispatch_needed = true; } } } 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 { 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 { let now = OffsetDateTime::now_utc(); let mut tx = self.pool.begin().await?; let attempt_id = sqlx::query_scalar::<_, Option>( "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 { 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 { 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 { 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>)>( "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, 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 { 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?; } 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>)>( "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::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 { 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, cancel_requested: bool, } #[derive(FromRow)] struct RegistrationAuthorityRow { job_id: String, fence: i64, worker_incarnation: String, state: String, authority_attempt_id: Option, 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, } #[derive(FromRow)] struct RecoveryJobRow { phase: String, cancel_requested: bool, dispatch_needed: bool, authority_attempt_id: Option, retry_not_before: Option, } #[derive(FromRow)] struct LockedReapRow { current_fence: i64, state: String, worker_down_at: Option, 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, upload_id: Option, } impl CleanupRow { fn try_into_claim(self, claim_token: OpaqueId) -> Result { 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, }) } }