ms2k/remora
SHA256
Watch
0
Code Releases Activity

Separate orchestration by protocol action

Browse source
This commit is contained in:
ms2k 2026-07-19 14:27:59 +09:00
commit 6f0f20f5cc
4 changed files with 6974 additions and 6839 deletions
Download patch file Download diff file Expand all files Collapse all files

688
crates/remora-attempt/src/recovery.rs
View file

@ -22,691 +22,9 @@ use crate::{
};
mod cancellation;
impl AttemptService {
pub async fn register_staging_object(
&self,
attempt_id: &OpaqueId,
object: &StagingObject,
) -> Result<(), AttemptError> {
validate_staging_key(&self.pool, attempt_id, &object.key).await?;
insert_cleanup_identity(
&self.pool,
attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await
}
pub async fn register_staging_write_intent_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
intent: StagingWriteIntent,
) -> Result<bool, AttemptError> {
self.register_staging_operation(
authority,
operation_id,
"register_staging_write_intent",
intent.clone(),
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &intent.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "write_intent",
target_store: &intent.target_store,
bucket: &intent.bucket,
key: &intent.key,
version_id: None,
upload_id: None,
},
)
.await
})
},
)
.await
}
pub async fn register_staging_object_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
object: StagingObject,
completed_upload_id: Option<String>,
) -> Result<bool, AttemptError> {
self.register_staging_operation(
authority,
operation_id,
"register_staging_object",
StagingObjectRegistration {
object: object.clone(),
completed_upload_id: completed_upload_id.clone(),
},
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &object.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await?;
complete_pending_identity(
tx,
&attempt_id,
"write_intent",
&object.target_store,
&object.bucket,
&object.key,
None,
)
.await?;
if let Some(upload_id) = completed_upload_id {
complete_pending_identity(
tx,
&attempt_id,
"multipart",
&object.target_store,
&object.bucket,
&object.key,
Some(&upload_id),
)
.await?;
}
Ok(())
})
},
)
.await
}
pub async fn register_staging_multipart(
&self,
attempt_id: &OpaqueId,
multipart: &StagingMultipart,
) -> Result<(), AttemptError> {
validate_staging_key(&self.pool, attempt_id, &multipart.key).await?;
if multipart.upload_id.is_empty() {
return Err(AttemptError::Invalid(
"multipart upload identity cannot be empty".into(),
));
}
insert_cleanup_identity(
&self.pool,
attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await
}
pub async fn register_staging_multipart_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
multipart: StagingMultipart,
) -> Result<bool, AttemptError> {
if multipart.upload_id.is_empty() {
return Err(AttemptError::Invalid(
"multipart upload identity cannot be empty".into(),
));
}
self.register_staging_operation(
authority,
operation_id,
"register_staging_multipart",
multipart.clone(),
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &multipart.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await
})
},
)
.await
}
async fn register_staging_operation<T, F>(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
kind: &'static str,
request: T,
effect: F,
) -> Result<bool, AttemptError>
where
T: Serialize,
F: for<'c> FnOnce(
&'c mut remora_durable::PgTx<'_>,
OpaqueId,
) -> remora_durable::BoxFuture<'c, Result<(), DurableError>>
+ Send
+ 'static,
{
let request_fingerprint = fingerprint(&RegistrationRequest {
job_id: &authority.job_id,
attempt_id: &authority.attempt_id,
fence: authority.fence,
worker_incarnation: &authority.worker_incarnation,
value: request,
})
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
let authority = authority.clone();
let outcome = execute_mutation(
&self.pool,
MutationScope {
subject: &self.operation_subject,
operation: "attempt.mutation",
key: operation_id.as_str(),
},
kind,
&[],
request_fingerprint,
move |tx| {
Box::pin(async move {
let eligible = lock_registration_authority(tx, &authority).await?;
if eligible {
effect(tx, authority.attempt_id).await?;
}
Ok((
if eligible {
MutationStatus::Accepted
} else {
MutationStatus::Rejected
},
serde_json::Value::Bool(eligible),
))
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
outcome
.response
.as_bool()
.ok_or_else(|| AttemptError::Invalid("invalid staging operation result".into()))
}
pub async fn reconcile_once(
&self,
config: &RecoveryConfig,
) -> Result<ReconcileReport, AttemptError> {
validate_recovery_config(config)?;
let job_ids = sqlx::query_scalar::<_, String>(
"SELECT job_id FROM jobs
WHERE phase NOT IN ('succeeded', 'failed', 'cancelled')
OR coordination_cleaned_at IS NULL
ORDER BY recovery_checked_at NULLS FIRST, recovery_checked_at, job_id LIMIT $1",
)
.bind(i64::from(config.batch_size))
.fetch_all(&self.pool)
.await?;
let mut report = ReconcileReport::default();
for value in job_ids {
let job_id =
OpaqueId::parse(value).map_err(|error| AttemptError::Invalid(error.to_string()))?;
report.merge(self.reconcile_job(&job_id, config).await?);
sqlx::query("UPDATE jobs SET recovery_checked_at = $2 WHERE job_id = $1")
.bind(job_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
}
report.cleanup_items_materialized = self.materialize_cleanup_work().await?;
Ok(report)
}
pub async fn claim_cleanup(
&self,
config: &RecoveryConfig,
) -> Result<Option<CleanupClaim>, AttemptError> {
validate_recovery_config(config)?;
self.materialize_cleanup_work().await?;
let claim_token = new_id('c')?;
let claimed_until = OffsetDateTime::now_utc()
+ time_duration(config.cleanup_claim_ttl, "cleanup_claim_ttl")?;
let row = sqlx::query_as::<_, CleanupRow>(
"WITH candidate AS (
SELECT c.cleanup_id
FROM attempt_staging_cleanup c
JOIN attempts a ON a.attempt_id = c.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE (c.state = 'pending'
OR (c.state = 'claimed' AND c.claimed_until <= $1))
AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ORDER BY c.cleanup_id
FOR UPDATE OF c SKIP LOCKED LIMIT 1
)
UPDATE attempt_staging_cleanup c
SET state = 'claimed', claim_token = $2, claimed_until = $3,
failure_count = CASE WHEN c.state = 'claimed'
THEN c.failure_count + 1 ELSE c.failure_count END
FROM candidate
WHERE c.cleanup_id = candidate.cleanup_id
RETURNING c.cleanup_id, c.attempt_id, c.kind, c.target_store, c.bucket,
c.object_key, c.version_id, c.upload_id",
)
.bind(OffsetDateTime::now_utc())
.bind(claim_token.as_str())
.bind(claimed_until)
.fetch_optional(&self.pool)
.await?;
row.map(|row| row.try_into_claim(claim_token)).transpose()
}
pub async fn complete_cleanup(
&self,
claim: &CleanupClaim,
) -> Result<CleanupCompletion, AttemptError> {
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'complete', claim_token = NULL, claimed_until = NULL,
completed_at = $3, last_error = NULL
WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
)
.bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
.bind(claim.claim_token.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
Ok(if changed.rows_affected() == 1 {
CleanupCompletion::Completed
} else {
CleanupCompletion::ClaimLost
})
}
pub async fn reconcile_write_intent(
&self,
claim: &CleanupClaim,
objects: Vec<StagingObject>,
multiparts: Vec<StagingMultipart>,
) -> Result<CleanupCompletion, AttemptError> {
let CleanupIdentity::WriteIntent(intent) = &claim.identity else {
return Err(AttemptError::Invalid(
"cleanup claim is not a staging write intent".into(),
));
};
for object in &objects {
if object.target_store != intent.target_store
|| object.bucket != intent.bucket
|| object.key != intent.key
{
return Err(AttemptError::Invalid(
"discovered object does not match the write intent".into(),
));
}
}
for multipart in &multiparts {
if multipart.target_store != intent.target_store
|| multipart.bucket != intent.bucket
|| multipart.key != intent.key
|| multipart.upload_id.is_empty()
{
return Err(AttemptError::Invalid(
"discovered multipart does not match the write intent".into(),
));
}
}
let cleanup_id = i64::try_from(claim.cleanup_id).map_err(invalid_number)?;
let claim_token = claim.claim_token.clone();
let attempt_id = claim.attempt_id.clone();
with_transaction(&self.pool, move |tx| {
Box::pin(async move {
let retained = sqlx::query_scalar::<_, i64>(
"SELECT cleanup_id FROM attempt_staging_cleanup
WHERE cleanup_id = $1 AND attempt_id = $2 AND kind = 'write_intent'
AND state = 'claimed' AND claim_token = $3
FOR UPDATE",
)
.bind(cleanup_id)
.bind(attempt_id.as_str())
.bind(claim_token.as_str())
.fetch_optional(tx.as_mut())
.await?;
if retained.is_none() {
return Ok(CleanupCompletion::ClaimLost);
}
for object in &objects {
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await?;
}
for multipart in &multiparts {
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await?;
}
sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'complete', claim_token = NULL, claimed_until = NULL,
completed_at = $2, last_error = NULL
WHERE cleanup_id = $1",
)
.bind(cleanup_id)
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
Ok(CleanupCompletion::Completed)
})
})
.await
.map_err(Into::into)
}
pub async fn retry_cleanup(
&self,
claim: &CleanupClaim,
failure_code: &str,
) -> Result<CleanupCompletion, AttemptError> {
if failure_code.is_empty()
|| failure_code.len() > 64
|| !failure_code
.bytes()
.all(|byte| byte.is_ascii_lowercase() || byte == b'_')
{
return Err(AttemptError::Invalid(
"cleanup failure code must be lowercase snake case".into(),
));
}
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'pending', claim_token = NULL, claimed_until = NULL,
failure_count = failure_count + 1, last_error = $3
WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
)
.bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
.bind(claim.claim_token.as_str())
.bind(failure_code)
.execute(&self.pool)
.await?;
Ok(if changed.rows_affected() == 1 {
CleanupCompletion::RetryScheduled
} else {
CleanupCompletion::ClaimLost
})
}
async fn reconcile_job(
&self,
job_id: &OpaqueId,
config: &RecoveryConfig,
) -> Result<ReconcileReport, AttemptError> {
let Some(mut row) = load_recovery_job(&self.pool, job_id).await? else {
return Ok(ReconcileReport::default());
};
let mut report = ReconcileReport::default();
if row.phase == "retry_wait" {
if row.retry_not_before.is_none() {
let due =
OffsetDateTime::now_utc() + time_duration(config.retry_delay, "retry_delay")?;
sqlx::query(
"UPDATE jobs SET retry_not_before = $2, updated_at = $3
WHERE job_id = $1 AND phase = 'retry_wait' AND retry_not_before IS NULL",
)
.bind(job_id.as_str())
.bind(due)
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
row.retry_not_before = Some(due);
}
if row
.retry_not_before
.is_some_and(|due| due <= OffsetDateTime::now_utc())
{
let mut tx = self.pool.begin().await?;
let now = OffsetDateTime::now_utc();
let changed = sqlx::query(
"UPDATE jobs SET phase = 'queued', dispatch_needed = TRUE,
retry_not_before = NULL, revision = revision + 1, updated_at = $2
WHERE job_id = $1 AND phase = 'retry_wait' AND cancel_requested = FALSE
AND retry_not_before <= $2",
)
.bind(job_id.as_str())
.bind(now)
.execute(tx.as_mut())
.await?;
if changed.rows_affected() == 1 {
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
report.retries_scheduled += 1;
row.phase = "queued".into();
row.dispatch_needed = true;
}
tx.commit().await?;
}
}
if row.phase == "queued"
&& row.dispatch_needed
&& self
.coordination
.project_enqueue(job_id.as_str())
.await?
.ok()
{
report.enqueued += 1;
}
if matches!(row.phase.as_str(), "running" | "cancelling") {
if row.cancel_requested {
if self
.coordination
.project_cancel(job_id.as_str())
.await?
.ok()
{
report.cancellations_published += 1;
}
if self.finalize_cancel(job_id).await? {
report.cancellations_finalized += 1;
row.phase = "cancelled".into();
}
} else if let Some(attempt_id) = row.authority_attempt_id.as_deref()
&& self
.reap_if_inactive(job_id, attempt_id, self.config.max_attempts)
.await?
{
report.attempts_reaped += 1;
}
if self
.coordination
.clear_stale_progress(job_id.as_str())
.await?
.code
== OpCode::Ok
{
report.stale_progress_cleared += 1;
}
}
if parse_phase(&row.phase)?.is_terminal() {
let expired = self.coordination.expire_lease(job_id.as_str()).await?;
self.coordination.clean_terminal(job_id.as_str()).await?;
if expired.code.is_success() {
sqlx::query(
"UPDATE jobs SET coordination_cleaned_at = COALESCE(coordination_cleaned_at, $2)
WHERE job_id = $1 AND phase IN ('succeeded', 'failed', 'cancelled')",
)
.bind(job_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
report.terminal_coordination_cleaned += 1;
}
}
Ok(report)
}
async fn reap_if_inactive(
&self,
job_id: &OpaqueId,
attempt_id: &str,
max_attempts: u64,
) -> Result<bool, AttemptError> {
let attempt_id = attempt_id.to_owned();
let durable_job_id = job_id.clone();
let coordination = self.coordination.clone();
let faults = self.faults.clone();
let reaped = with_transaction(&self.pool, move |tx| {
Box::pin(async move {
reap_lost_attempt(
tx,
&coordination,
faults.as_ref(),
&durable_job_id,
&attempt_id,
max_attempts,
)
.await
})
})
.await?;
let _ = self.coordination.expire_lease(job_id.as_str()).await?;
Ok(reaped.changed)
}
async fn finalize_cancel(&self, job_id: &OpaqueId) -> Result<bool, AttemptError> {
let now = OffsetDateTime::now_utc();
let mut tx = self.pool.begin().await?;
let attempt_id = sqlx::query_scalar::<_, Option<String>>(
"SELECT authority_attempt_id FROM jobs
WHERE job_id = $1 AND phase = 'cancelling' FOR UPDATE",
)
.bind(job_id.as_str())
.fetch_optional(tx.as_mut())
.await?
.flatten();
let Some(attempt_id) = attempt_id else {
tx.rollback().await?;
return Ok(false);
};
sqlx::query(
"UPDATE attempts SET state = 'stopped', stopped_at = $2
WHERE attempt_id = $1 AND state = 'authorized'",
)
.bind(&attempt_id)
.bind(now)
.execute(tx.as_mut())
.await?;
sqlx::query(
"UPDATE jobs SET phase = 'cancelled', authority_attempt_id = NULL,
dispatch_needed = FALSE, terminal_at = $2, coordination_cleaned_at = NULL,
revision = revision + 1, updated_at = $2
WHERE job_id = $1 AND phase = 'cancelling'",
)
.bind(job_id.as_str())
.bind(now)
.execute(tx.as_mut())
.await?;
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
tx.commit().await?;
Ok(true)
}
async fn materialize_cleanup_work(&self) -> Result<u64, AttemptError> {
let now = OffsetDateTime::now_utc();
let members = sqlx::query(
"INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
)
SELECT m.attempt_id, 'object', m.object_store, m.object_bucket,
m.object_key, m.object_version_id, $1
FROM result_inventory_members m
JOIN attempts a ON a.attempt_id = m.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ON CONFLICT DO NOTHING",
)
.bind(now)
.execute(&self.pool)
.await?;
let manifests = sqlx::query(
"INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
)
SELECT r.attempt_id, 'object', r.manifest_store, r.manifest_bucket,
r.manifest_key, r.manifest_version_id, $1
FROM result_inventories r
JOIN attempts a ON a.attempt_id = r.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE r.state = 'sealed'
AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ON CONFLICT DO NOTHING",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(members.rows_affected() + manifests.rows_affected())
}
}
mod cleanup;
mod reconciliation;
mod staging;
async fn cooperative_stop_transaction(
tx: &mut remora_durable::PgTx<'_>,

144
crates/remora-attempt/src/recovery/cleanup.rs Normal file
View file

@ -0,0 +1,144 @@
use super::*;
impl AttemptService {
pub async fn claim_cleanup(
&self,
config: &RecoveryConfig,
) -> Result<Option<CleanupClaim>, AttemptError> {
validate_recovery_config(config)?;
self.materialize_cleanup_work().await?;
let claim_token = new_id('c')?;
let claimed_until = OffsetDateTime::now_utc()
+ time_duration(config.cleanup_claim_ttl, "cleanup_claim_ttl")?;
let row = sqlx::query_as::<_, CleanupRow>(
"WITH candidate AS (
SELECT c.cleanup_id
FROM attempt_staging_cleanup c
JOIN attempts a ON a.attempt_id = c.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE (c.state = 'pending'
OR (c.state = 'claimed' AND c.claimed_until <= $1))
AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ORDER BY c.cleanup_id
FOR UPDATE OF c SKIP LOCKED LIMIT 1
)
UPDATE attempt_staging_cleanup c
SET state = 'claimed', claim_token = $2, claimed_until = $3,
failure_count = CASE WHEN c.state = 'claimed'
THEN c.failure_count + 1 ELSE c.failure_count END
FROM candidate
WHERE c.cleanup_id = candidate.cleanup_id
RETURNING c.cleanup_id, c.attempt_id, c.kind, c.target_store, c.bucket,
c.object_key, c.version_id, c.upload_id",
)
.bind(OffsetDateTime::now_utc())
.bind(claim_token.as_str())
.bind(claimed_until)
.fetch_optional(&self.pool)
.await?;
row.map(|row| row.try_into_claim(claim_token)).transpose()
}
pub async fn complete_cleanup(
&self,
claim: &CleanupClaim,
) -> Result<CleanupCompletion, AttemptError> {
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'complete', claim_token = NULL, claimed_until = NULL,
completed_at = $3, last_error = NULL
WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
)
.bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
.bind(claim.claim_token.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
Ok(if changed.rows_affected() == 1 {
CleanupCompletion::Completed
} else {
CleanupCompletion::ClaimLost
})
}
pub async fn retry_cleanup(
&self,
claim: &CleanupClaim,
failure_code: &str,
) -> Result<CleanupCompletion, AttemptError> {
if failure_code.is_empty()
|| failure_code.len() > 64
|| !failure_code
.bytes()
.all(|byte| byte.is_ascii_lowercase() || byte == b'_')
{
return Err(AttemptError::Invalid(
"cleanup failure code must be lowercase snake case".into(),
));
}
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'pending', claim_token = NULL, claimed_until = NULL,
failure_count = failure_count + 1, last_error = $3
WHERE cleanup_id = $1 AND state = 'claimed' AND claim_token = $2",
)
.bind(i64::try_from(claim.cleanup_id).map_err(invalid_number)?)
.bind(claim.claim_token.as_str())
.bind(failure_code)
.execute(&self.pool)
.await?;
Ok(if changed.rows_affected() == 1 {
CleanupCompletion::RetryScheduled
} else {
CleanupCompletion::ClaimLost
})
}
pub(super) async fn materialize_cleanup_work(&self) -> Result<u64, AttemptError> {
let now = OffsetDateTime::now_utc();
let members = sqlx::query(
"INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
)
SELECT m.attempt_id, 'object', m.object_store, m.object_bucket,
m.object_key, m.object_version_id, $1
FROM result_inventory_members m
JOIN attempts a ON a.attempt_id = m.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ON CONFLICT DO NOTHING",
)
.bind(now)
.execute(&self.pool)
.await?;
let manifests = sqlx::query(
"INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
)
SELECT r.attempt_id, 'object', r.manifest_store, r.manifest_bucket,
r.manifest_key, r.manifest_version_id, $1
FROM result_inventories r
JOIN attempts a ON a.attempt_id = r.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE r.state = 'sealed'
AND j.authority_attempt_id IS DISTINCT FROM a.attempt_id
AND (j.phase IN ('succeeded', 'failed', 'cancelled')
OR j.checkpoint_owner_attempt_id IS DISTINCT FROM a.attempt_id)
AND COALESCE(j.committed_result ->> 'attempt_id', '') <> a.attempt_id
AND (a.state IN ('stopped', 'reaped') OR a.worker_down_at IS NOT NULL)
ON CONFLICT DO NOTHING",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(members.rows_affected() + manifests.rows_affected())
}
}

304
crates/remora-attempt/src/recovery/reconciliation.rs Normal file
View file

@ -0,0 +1,304 @@
use super::*;
impl AttemptService {
pub async fn reconcile_once(
&self,
config: &RecoveryConfig,
) -> Result<ReconcileReport, AttemptError> {
validate_recovery_config(config)?;
let job_ids = sqlx::query_scalar::<_, String>(
"SELECT job_id FROM jobs
WHERE phase NOT IN ('succeeded', 'failed', 'cancelled')
OR coordination_cleaned_at IS NULL
ORDER BY recovery_checked_at NULLS FIRST, recovery_checked_at, job_id LIMIT $1",
)
.bind(i64::from(config.batch_size))
.fetch_all(&self.pool)
.await?;
let mut report = ReconcileReport::default();
for value in job_ids {
let job_id =
OpaqueId::parse(value).map_err(|error| AttemptError::Invalid(error.to_string()))?;
report.merge(self.reconcile_job(&job_id, config).await?);
sqlx::query("UPDATE jobs SET recovery_checked_at = $2 WHERE job_id = $1")
.bind(job_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
}
report.cleanup_items_materialized = self.materialize_cleanup_work().await?;
Ok(report)
}
pub async fn reconcile_write_intent(
&self,
claim: &CleanupClaim,
objects: Vec<StagingObject>,
multiparts: Vec<StagingMultipart>,
) -> Result<CleanupCompletion, AttemptError> {
let CleanupIdentity::WriteIntent(intent) = &claim.identity else {
return Err(AttemptError::Invalid(
"cleanup claim is not a staging write intent".into(),
));
};
for object in &objects {
if object.target_store != intent.target_store
|| object.bucket != intent.bucket
|| object.key != intent.key
{
return Err(AttemptError::Invalid(
"discovered object does not match the write intent".into(),
));
}
}
for multipart in &multiparts {
if multipart.target_store != intent.target_store
|| multipart.bucket != intent.bucket
|| multipart.key != intent.key
|| multipart.upload_id.is_empty()
{
return Err(AttemptError::Invalid(
"discovered multipart does not match the write intent".into(),
));
}
}
let cleanup_id = i64::try_from(claim.cleanup_id).map_err(invalid_number)?;
let claim_token = claim.claim_token.clone();
let attempt_id = claim.attempt_id.clone();
with_transaction(&self.pool, move |tx| {
Box::pin(async move {
let retained = sqlx::query_scalar::<_, i64>(
"SELECT cleanup_id FROM attempt_staging_cleanup
WHERE cleanup_id = $1 AND attempt_id = $2 AND kind = 'write_intent'
AND state = 'claimed' AND claim_token = $3
FOR UPDATE",
)
.bind(cleanup_id)
.bind(attempt_id.as_str())
.bind(claim_token.as_str())
.fetch_optional(tx.as_mut())
.await?;
if retained.is_none() {
return Ok(CleanupCompletion::ClaimLost);
}
for object in &objects {
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await?;
}
for multipart in &multiparts {
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await?;
}
sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'complete', claim_token = NULL, claimed_until = NULL,
completed_at = $2, last_error = NULL
WHERE cleanup_id = $1",
)
.bind(cleanup_id)
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
Ok(CleanupCompletion::Completed)
})
})
.await
.map_err(Into::into)
}
pub(super) async fn reconcile_job(
&self,
job_id: &OpaqueId,
config: &RecoveryConfig,
) -> Result<ReconcileReport, AttemptError> {
let Some(mut row) = load_recovery_job(&self.pool, job_id).await? else {
return Ok(ReconcileReport::default());
};
let mut report = ReconcileReport::default();
if row.phase == "retry_wait" {
if row.retry_not_before.is_none() {
let due =
OffsetDateTime::now_utc() + time_duration(config.retry_delay, "retry_delay")?;
sqlx::query(
"UPDATE jobs SET retry_not_before = $2, updated_at = $3
WHERE job_id = $1 AND phase = 'retry_wait' AND retry_not_before IS NULL",
)
.bind(job_id.as_str())
.bind(due)
.bind(OffsetDateTime::now_utc())
.execute(&self.pool)
.await?;
row.retry_not_before = Some(due);
}
if row
.retry_not_before
.is_some_and(|due| due <= OffsetDateTime::now_utc())
{
let mut tx = self.pool.begin().await?;
let now = OffsetDateTime::now_utc();
let changed = sqlx::query(
"UPDATE jobs SET phase = 'queued', dispatch_needed = TRUE,
retry_not_before = NULL, revision = revision + 1, updated_at = $2
WHERE job_id = $1 AND phase = 'retry_wait' AND cancel_requested = FALSE
AND retry_not_before <= $2",
)
.bind(job_id.as_str())
.bind(now)
.execute(tx.as_mut())
.await?;
if changed.rows_affected() == 1 {
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
report.retries_scheduled += 1;
row.phase = "queued".into();
row.dispatch_needed = true;
}
tx.commit().await?;
}
}
if row.phase == "queued"
&& row.dispatch_needed
&& self
.coordination
.project_enqueue(job_id.as_str())
.await?
.ok()
{
report.enqueued += 1;
}
if matches!(row.phase.as_str(), "running" | "cancelling") {
if row.cancel_requested {
if self
.coordination
.project_cancel(job_id.as_str())
.await?
.ok()
{
report.cancellations_published += 1;
}
if self.finalize_cancel(job_id).await? {
report.cancellations_finalized += 1;
row.phase = "cancelled".into();
}
} else if let Some(attempt_id) = row.authority_attempt_id.as_deref()
&& self
.reap_if_inactive(job_id, attempt_id, self.config.max_attempts)
.await?
{
report.attempts_reaped += 1;
}
if self
.coordination
.clear_stale_progress(job_id.as_str())
.await?
.code
== OpCode::Ok
{
report.stale_progress_cleared += 1;
}
}
if parse_phase(&row.phase)?.is_terminal() {
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)
}
pub(super) async fn reap_if_inactive(
&self,
job_id: &OpaqueId,
attempt_id: &str,
max_attempts: u64,
) -> Result<bool, AttemptError> {
let attempt_id = attempt_id.to_owned();
let durable_job_id = job_id.clone();
let coordination = self.coordination.clone();
let faults = self.faults.clone();
let reaped = with_transaction(&self.pool, move |tx| {
Box::pin(async move {
reap_lost_attempt(
tx,
&coordination,
faults.as_ref(),
&durable_job_id,
&attempt_id,
max_attempts,
)
.await
})
})
.await?;
let _ = self.coordination.expire_lease(job_id.as_str()).await?;
Ok(reaped.changed)
}
pub(super) async fn finalize_cancel(&self, job_id: &OpaqueId) -> Result<bool, AttemptError> {
let now = OffsetDateTime::now_utc();
let mut tx = self.pool.begin().await?;
let attempt_id = sqlx::query_scalar::<_, Option<String>>(
"SELECT authority_attempt_id FROM jobs
WHERE job_id = $1 AND phase = 'cancelling' FOR UPDATE",
)
.bind(job_id.as_str())
.fetch_optional(tx.as_mut())
.await?
.flatten();
let Some(attempt_id) = attempt_id else {
tx.rollback().await?;
return Ok(false);
};
sqlx::query(
"UPDATE attempts SET state = 'stopped', stopped_at = $2
WHERE attempt_id = $1 AND state = 'authorized'",
)
.bind(&attempt_id)
.bind(now)
.execute(tx.as_mut())
.await?;
sqlx::query(
"UPDATE jobs SET phase = 'cancelled', authority_attempt_id = NULL,
dispatch_needed = FALSE, terminal_at = $2, coordination_cleaned_at = NULL,
revision = revision + 1, updated_at = $2
WHERE job_id = $1 AND phase = 'cancelling'",
)
.bind(job_id.as_str())
.bind(now)
.execute(tx.as_mut())
.await?;
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
tx.commit().await?;
Ok(true)
}
}

244
crates/remora-attempt/src/recovery/staging.rs Normal file
View file

@ -0,0 +1,244 @@
use super::*;
impl AttemptService {
pub async fn register_staging_object(
&self,
attempt_id: &OpaqueId,
object: &StagingObject,
) -> Result<(), AttemptError> {
validate_staging_key(&self.pool, attempt_id, &object.key).await?;
insert_cleanup_identity(
&self.pool,
attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await
}
pub async fn register_staging_write_intent_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
intent: StagingWriteIntent,
) -> Result<bool, AttemptError> {
self.register_staging_operation(
authority,
operation_id,
"register_staging_write_intent",
intent.clone(),
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &intent.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "write_intent",
target_store: &intent.target_store,
bucket: &intent.bucket,
key: &intent.key,
version_id: None,
upload_id: None,
},
)
.await
})
},
)
.await
}
pub async fn register_staging_object_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
object: StagingObject,
completed_upload_id: Option<String>,
) -> Result<bool, AttemptError> {
self.register_staging_operation(
authority,
operation_id,
"register_staging_object",
StagingObjectRegistration {
object: object.clone(),
completed_upload_id: completed_upload_id.clone(),
},
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &object.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "object",
target_store: &object.target_store,
bucket: &object.bucket,
key: &object.key,
version_id: object.version_id.as_deref(),
upload_id: None,
},
)
.await?;
complete_pending_identity(
tx,
&attempt_id,
"write_intent",
&object.target_store,
&object.bucket,
&object.key,
None,
)
.await?;
if let Some(upload_id) = completed_upload_id {
complete_pending_identity(
tx,
&attempt_id,
"multipart",
&object.target_store,
&object.bucket,
&object.key,
Some(&upload_id),
)
.await?;
}
Ok(())
})
},
)
.await
}
pub async fn register_staging_multipart(
&self,
attempt_id: &OpaqueId,
multipart: &StagingMultipart,
) -> Result<(), AttemptError> {
validate_staging_key(&self.pool, attempt_id, &multipart.key).await?;
if multipart.upload_id.is_empty() {
return Err(AttemptError::Invalid(
"multipart upload identity cannot be empty".into(),
));
}
insert_cleanup_identity(
&self.pool,
attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await
}
pub async fn register_staging_multipart_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
multipart: StagingMultipart,
) -> Result<bool, AttemptError> {
if multipart.upload_id.is_empty() {
return Err(AttemptError::Invalid(
"multipart upload identity cannot be empty".into(),
));
}
self.register_staging_operation(
authority,
operation_id,
"register_staging_multipart",
multipart.clone(),
move |tx, attempt_id| {
Box::pin(async move {
validate_staging_key_transaction(tx, &attempt_id, &multipart.key).await?;
insert_cleanup_identity_transaction(
tx,
&attempt_id,
NewCleanupIdentity {
kind: "multipart",
target_store: &multipart.target_store,
bucket: &multipart.bucket,
key: &multipart.key,
version_id: None,
upload_id: Some(&multipart.upload_id),
},
)
.await
})
},
)
.await
}
pub(super) async fn register_staging_operation<T, F>(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
kind: &'static str,
request: T,
effect: F,
) -> Result<bool, AttemptError>
where
T: Serialize,
F: for<'c> FnOnce(
&'c mut remora_durable::PgTx<'_>,
OpaqueId,
) -> remora_durable::BoxFuture<'c, Result<(), DurableError>>
+ Send
+ 'static,
{
let request_fingerprint = fingerprint(&RegistrationRequest {
job_id: &authority.job_id,
attempt_id: &authority.attempt_id,
fence: authority.fence,
worker_incarnation: &authority.worker_incarnation,
value: request,
})
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
let authority = authority.clone();
let outcome = execute_mutation(
&self.pool,
MutationScope {
subject: &self.operation_subject,
operation: "attempt.mutation",
key: operation_id.as_str(),
},
kind,
&[],
request_fingerprint,
move |tx| {
Box::pin(async move {
let eligible = lock_registration_authority(tx, &authority).await?;
if eligible {
effect(tx, authority.attempt_id).await?;
}
Ok((
if eligible {
MutationStatus::Accepted
} else {
MutationStatus::Rejected
},
serde_json::Value::Bool(eligible),
))
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
outcome
.response
.as_bool()
.ok_or_else(|| AttemptError::Invalid("invalid staging operation result".into()))
}
}