Add convergent recovery and exact staging cleanup
This commit is contained in:
parent
80c631f780
commit
9d96c4d563
1 changed files with 2163 additions and 9 deletions
858
crates/remora-attempt/src/recovery.rs
Normal file
858
crates/remora-attempt/src/recovery.rs
Normal file
|
|
@ -0,0 +1,858 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use remora_coordination::{LeaseOwner, OpCode};
|
||||
use remora_domain::{FailureClass, FailureDisposition, JobPhase, OpaqueId};
|
||||
use remora_durable::{DurableError, enqueue, with_transaction};
|
||||
use remora_target::validate_attempt_key;
|
||||
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,
|
||||
};
|
||||
|
||||
impl AttemptService {
|
||||
pub async fn request_cancel(&self, job_id: &OpaqueId) -> Result<CancelOutcome, AttemptError> {
|
||||
with_transaction(&self.pool, |tx| {
|
||||
let job_id = job_id.clone();
|
||||
Box::pin(async move {
|
||||
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<CooperativeStopOutcome, AttemptError> {
|
||||
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",
|
||||
)
|
||||
.bind(authority.attempt_id.as_str())
|
||||
.fetch_optional(&self.pool)
|
||||
.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 = self
|
||||
.coordination
|
||||
.get_live_lease(authority.job_id.as_str())
|
||||
.await?
|
||||
.is_some_and(|lease| lease.record.owner() == lease_owner(authority));
|
||||
if current && live && !row.cancel_requested {
|
||||
return Ok(CooperativeStopOutcome::StopNotRequired);
|
||||
}
|
||||
let now = OffsetDateTime::now_utc();
|
||||
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(now)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
if changed.rows_affected() == 0 {
|
||||
return Ok(CooperativeStopOutcome::StopNotRequired);
|
||||
}
|
||||
let _ = self
|
||||
.coordination
|
||||
.release(authority.job_id.as_str(), &lease_owner(authority))
|
||||
.await;
|
||||
self.release_local(&authority.job_id);
|
||||
Ok(CooperativeStopOutcome::Stopped)
|
||||
}
|
||||
|
||||
pub async fn mark_worker_down(
|
||||
&self,
|
||||
worker_incarnation: &OpaqueId,
|
||||
) -> Result<u64, AttemptError> {
|
||||
let changed = sqlx::query(
|
||||
"UPDATE attempts SET worker_down_at = COALESCE(worker_down_at, $2)
|
||||
WHERE worker_incarnation = $1 AND state IN ('authorized', 'executing')",
|
||||
)
|
||||
.bind(worker_incarnation.as_str())
|
||||
.bind(OffsetDateTime::now_utc())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(changed.rows_affected())
|
||||
}
|
||||
|
||||
pub async fn register_staging_object(
|
||||
&self,
|
||||
attempt_id: &OpaqueId,
|
||||
object: &StagingObject,
|
||||
) -> Result<(), AttemptError> {
|
||||
validate_staging_key(attempt_id, &object.key)?;
|
||||
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_multipart(
|
||||
&self,
|
||||
attempt_id: &OpaqueId,
|
||||
multipart: &StagingMultipart,
|
||||
) -> Result<(), AttemptError> {
|
||||
validate_staging_key(attempt_id, &multipart.key)?;
|
||||
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 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 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
|
||||
})
|
||||
}
|
||||
|
||||
async fn reconcile_job(
|
||||
&self,
|
||||
job_id: &OpaqueId,
|
||||
config: &RecoveryConfig,
|
||||
) -> Result<ReconcileReport, AttemptError> {
|
||||
let Some(mut row) = load_recovery_job(&self.pool, job_id).await? else {
|
||||
return Ok(ReconcileReport::default());
|
||||
};
|
||||
let mut report = ReconcileReport::default();
|
||||
if row.phase == "retry_wait" {
|
||||
self.release_local_if_stopped(job_id).await?;
|
||||
if row.retry_not_before.is_none() {
|
||||
let due =
|
||||
OffsetDateTime::now_utc() + time_duration(config.retry_delay, "retry_delay")?;
|
||||
sqlx::query(
|
||||
"UPDATE jobs SET retry_not_before = $2, updated_at = $3
|
||||
WHERE job_id = $1 AND phase = 'retry_wait' AND retry_not_before IS NULL",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.bind(due)
|
||||
.bind(OffsetDateTime::now_utc())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
row.retry_not_before = Some(due);
|
||||
}
|
||||
if row
|
||||
.retry_not_before
|
||||
.is_some_and(|due| due <= OffsetDateTime::now_utc())
|
||||
{
|
||||
let 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<bool, AttemptError> {
|
||||
let attempt_id = attempt_id.to_owned();
|
||||
let durable_job_id = job_id.clone();
|
||||
let coordination = self.coordination.clone();
|
||||
let faults = self.faults.clone();
|
||||
let reaped = with_transaction(&self.pool, move |tx| {
|
||||
Box::pin(async move {
|
||||
reap_lost_attempt(
|
||||
tx,
|
||||
&coordination,
|
||||
faults.as_ref(),
|
||||
&durable_job_id,
|
||||
&attempt_id,
|
||||
max_attempts,
|
||||
)
|
||||
.await
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
let _ = self.coordination.expire_lease(job_id.as_str()).await?;
|
||||
if reaped.release_local {
|
||||
self.release_local(job_id);
|
||||
}
|
||||
Ok(reaped.changed)
|
||||
}
|
||||
|
||||
async fn finalize_cancel(&self, job_id: &OpaqueId) -> Result<bool, AttemptError> {
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let mut tx = self.pool.begin().await?;
|
||||
let attempt_id = sqlx::query_scalar::<_, Option<String>>(
|
||||
"SELECT authority_attempt_id FROM jobs
|
||||
WHERE job_id = $1 AND phase = 'cancelling' FOR UPDATE",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.fetch_optional(tx.as_mut())
|
||||
.await?
|
||||
.flatten();
|
||||
let Some(attempt_id) = attempt_id else {
|
||||
tx.rollback().await?;
|
||||
return Ok(false);
|
||||
};
|
||||
sqlx::query(
|
||||
"UPDATE attempts SET state = 'stopped', stopped_at = $2
|
||||
WHERE attempt_id = $1 AND state = 'authorized'",
|
||||
)
|
||||
.bind(&attempt_id)
|
||||
.bind(now)
|
||||
.execute(tx.as_mut())
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"UPDATE jobs SET phase = 'cancelled', authority_attempt_id = NULL,
|
||||
dispatch_needed = FALSE, terminal_at = $2, coordination_cleaned_at = NULL,
|
||||
revision = revision + 1, updated_at = $2
|
||||
WHERE job_id = $1 AND phase = 'cancelling'",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.bind(now)
|
||||
.execute(tx.as_mut())
|
||||
.await?;
|
||||
enqueue(&mut tx, "job.updated", json!({ "job_id": job_id })).await?;
|
||||
tx.commit().await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn materialize_cleanup_work(&self) -> Result<u64, AttemptError> {
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let members = sqlx::query(
|
||||
"INSERT INTO attempt_staging_cleanup (
|
||||
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
|
||||
)
|
||||
SELECT m.attempt_id, 'object', m.object_store, m.object_bucket,
|
||||
m.object_key, m.object_version_id, $1
|
||||
FROM result_inventory_members m
|
||||
JOIN attempts a ON a.attempt_id = m.attempt_id
|
||||
JOIN jobs j ON j.job_id = a.job_id
|
||||
WHERE j.authority_attempt_id IS DISTINCT FROM a.attempt_id
|
||||
AND 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 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 insert_cleanup_identity(
|
||||
pool: &sqlx::PgPool,
|
||||
attempt_id: &OpaqueId,
|
||||
identity: NewCleanupIdentity<'_>,
|
||||
) -> Result<(), AttemptError> {
|
||||
if identity.target_store.is_empty() || identity.bucket.is_empty() {
|
||||
return Err(AttemptError::Invalid(
|
||||
"staging target store and bucket cannot be empty".into(),
|
||||
));
|
||||
}
|
||||
let changed = sqlx::query(
|
||||
"INSERT INTO attempt_staging_cleanup (
|
||||
attempt_id, kind, target_store, bucket, object_key, version_id, upload_id, created_at
|
||||
)
|
||||
SELECT $1, $2, $3, $4, $5, $6, $7, $8
|
||||
WHERE EXISTS (
|
||||
SELECT 1 FROM attempts a JOIN jobs j ON j.job_id = a.job_id
|
||||
WHERE a.attempt_id = $1 AND j.target_store = $3 AND j.target_bucket = $4
|
||||
)
|
||||
ON CONFLICT DO NOTHING",
|
||||
)
|
||||
.bind(attempt_id.as_str())
|
||||
.bind(identity.kind)
|
||||
.bind(identity.target_store)
|
||||
.bind(identity.bucket)
|
||||
.bind(identity.key)
|
||||
.bind(identity.version_id)
|
||||
.bind(identity.upload_id)
|
||||
.bind(OffsetDateTime::now_utc())
|
||||
.execute(pool)
|
||||
.await?;
|
||||
if changed.rows_affected() == 0 {
|
||||
let binding = sqlx::query_as::<_, (String, String)>(
|
||||
"SELECT j.target_store, j.target_bucket
|
||||
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
|
||||
WHERE a.attempt_id = $1",
|
||||
)
|
||||
.bind(attempt_id.as_str())
|
||||
.fetch_optional(pool)
|
||||
.await?;
|
||||
match binding {
|
||||
None => return Err(AttemptError::Invalid("attempt does not exist".into())),
|
||||
Some((store, bucket))
|
||||
if store != identity.target_store || bucket != identity.bucket =>
|
||||
{
|
||||
return Err(AttemptError::Invalid(
|
||||
"staging identity does not match the attempt target".into(),
|
||||
));
|
||||
}
|
||||
Some(_) => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct NewCleanupIdentity<'a> {
|
||||
kind: &'a str,
|
||||
target_store: &'a str,
|
||||
bucket: &'a str,
|
||||
key: &'a str,
|
||||
version_id: Option<&'a str>,
|
||||
upload_id: Option<&'a str>,
|
||||
}
|
||||
|
||||
async fn load_recovery_job(
|
||||
pool: &sqlx::PgPool,
|
||||
job_id: &OpaqueId,
|
||||
) -> Result<Option<RecoveryJobRow>, AttemptError> {
|
||||
sqlx::query_as(
|
||||
"SELECT phase, cancel_requested, dispatch_needed, authority_attempt_id,
|
||||
retry_not_before
|
||||
FROM jobs WHERE job_id = $1",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
async fn reap_lost_attempt(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
coordination: &remora_coordination::CoordinationStore,
|
||||
faults: &dyn crate::AttemptFaultInjector,
|
||||
job_id: &OpaqueId,
|
||||
attempt_id: &str,
|
||||
max_attempts: u64,
|
||||
) -> Result<ReapEffect, DurableError> {
|
||||
let row = sqlx::query_as::<_, LockedReapRow>(
|
||||
"SELECT j.current_fence, a.state, a.worker_down_at, a.fence,
|
||||
a.worker_incarnation
|
||||
FROM jobs j JOIN attempts a ON a.attempt_id = j.authority_attempt_id
|
||||
WHERE j.job_id = $1 AND a.attempt_id = $2 AND j.phase = 'running'
|
||||
FOR UPDATE OF j, a",
|
||||
)
|
||||
.bind(job_id.as_str())
|
||||
.bind(attempt_id)
|
||||
.fetch_optional(tx.as_mut())
|
||||
.await?;
|
||||
let Some(row) = row else {
|
||||
return Ok(ReapEffect::default());
|
||||
};
|
||||
faults.hit(crate::AttemptFaultPoint::RecoveryBeforeRevalidation);
|
||||
let reservation = coordination
|
||||
.get_reservation(job_id.as_str())
|
||||
.await
|
||||
.map_err(DurableError::effect)?;
|
||||
let live = coordination
|
||||
.get_live_lease(job_id.as_str())
|
||||
.await
|
||||
.map_err(DurableError::effect)?;
|
||||
let exact_live = live.is_some_and(|lease| {
|
||||
lease.record.attempt_id == attempt_id
|
||||
&& i64::try_from(lease.record.fence).ok() == Some(row.fence)
|
||||
&& lease.record.worker_id == row.worker_incarnation
|
||||
});
|
||||
let inactive = match row.state.as_str() {
|
||||
"authorized" => reservation.is_none() && !exact_live,
|
||||
"executing" => row.worker_down_at.is_some() || !exact_live,
|
||||
_ => 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(),
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_staging_key(attempt_id: &OpaqueId, key: &str) -> Result<(), AttemptError> {
|
||||
validate_attempt_key(key).map_err(|error| AttemptError::Invalid(error.to_string()))?;
|
||||
if key.split('/').nth(2) != Some(attempt_id.as_str()) {
|
||||
return Err(AttemptError::Invalid(
|
||||
"staging key belongs to a different attempt".into(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_recovery_config(config: &RecoveryConfig) -> Result<(), AttemptError> {
|
||||
if config.batch_size == 0 {
|
||||
return Err(AttemptError::Invalid(
|
||||
"recovery limits must be greater than zero".into(),
|
||||
));
|
||||
}
|
||||
if config.cleanup_claim_ttl.is_zero() {
|
||||
return Err(AttemptError::Invalid(
|
||||
"cleanup claim ttl must be greater than zero".into(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn time_duration(value: Duration, name: &str) -> Result<time::Duration, AttemptError> {
|
||||
time::Duration::try_from(value)
|
||||
.map_err(|_| AttemptError::Invalid(format!("{name} is too large")))
|
||||
}
|
||||
|
||||
fn lease_owner(authority: &AttemptAuthority) -> LeaseOwner {
|
||||
LeaseOwner {
|
||||
attempt_id: authority.attempt_id.to_string(),
|
||||
fence: authority.fence,
|
||||
worker_id: authority.worker_incarnation.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_phase(value: &str) -> Result<JobPhase, DurableError> {
|
||||
match value {
|
||||
"queued" => Ok(JobPhase::Queued),
|
||||
"running" => Ok(JobPhase::Running),
|
||||
"retry_wait" => Ok(JobPhase::RetryWait),
|
||||
"cancelling" => Ok(JobPhase::Cancelling),
|
||||
"succeeded" => Ok(JobPhase::Succeeded),
|
||||
"failed" => Ok(JobPhase::Failed),
|
||||
"cancelled" => Ok(JobPhase::Cancelled),
|
||||
_ => Err(DurableError::Effect(format!("unknown job phase '{value}'"))),
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_number(error: impl std::fmt::Display) -> AttemptError {
|
||||
AttemptError::Invalid(error.to_string())
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct CancelRow {
|
||||
phase: String,
|
||||
cancel_requested: bool,
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct StopRow {
|
||||
job_id: String,
|
||||
fence: i64,
|
||||
worker_incarnation: String,
|
||||
state: String,
|
||||
authority_attempt_id: Option<String>,
|
||||
cancel_requested: bool,
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct RecoveryJobRow {
|
||||
phase: String,
|
||||
cancel_requested: bool,
|
||||
dispatch_needed: bool,
|
||||
authority_attempt_id: Option<String>,
|
||||
retry_not_before: Option<OffsetDateTime>,
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct LockedReapRow {
|
||||
current_fence: i64,
|
||||
state: String,
|
||||
worker_down_at: Option<OffsetDateTime>,
|
||||
fence: i64,
|
||||
worker_incarnation: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReapEffect {
|
||||
changed: bool,
|
||||
release_local: bool,
|
||||
}
|
||||
|
||||
#[derive(FromRow)]
|
||||
struct CleanupRow {
|
||||
cleanup_id: i64,
|
||||
attempt_id: String,
|
||||
kind: String,
|
||||
target_store: String,
|
||||
bucket: String,
|
||||
object_key: String,
|
||||
version_id: Option<String>,
|
||||
upload_id: Option<String>,
|
||||
}
|
||||
|
||||
impl CleanupRow {
|
||||
fn try_into_claim(self, claim_token: OpaqueId) -> Result<CleanupClaim, AttemptError> {
|
||||
let identity = match self.kind.as_str() {
|
||||
"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,
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue