Add convergent recovery and exact staging cleanup

This commit is contained in:
ms2k 2026-07-17 08:36:39 +09:00
commit 9d96c4d563
15 changed files with 2163 additions and 9 deletions

View file

@ -250,6 +250,13 @@ COPY crates crates
COPY --from=kind-attempt-verified /tmp/remora-attempt-tests.tar.zst /tests/remora-attempt-tests.tar.zst
CMD ["bash", "-ec", "test -s /run/remora-fixture/profile.env && test -n \"$DATABASE_URL\" && test -n \"$REMORA_COORDINATION_URL\" && export REMORA_ATTEMPT_REQUIRE=1 && cargo-nextest nextest run --archive-file /tests/remora-attempt-tests.tar.zst --workspace-remap /workspace --test-threads=1 -E 'binary(mutations)'"]
FROM rust-tools AS kind-recovery-runner
WORKDIR /workspace
COPY Cargo.toml ./
COPY crates crates
COPY --from=kind-attempt-verified /tmp/remora-attempt-tests.tar.zst /tests/remora-attempt-tests.tar.zst
CMD ["bash", "-ec", "test -s /run/remora-fixture/profile.env && test -n \"$DATABASE_URL\" && test -n \"$REMORA_COORDINATION_URL\" && test -n \"$REMORA_TARGET_ENDPOINT\" && case $REMORA_TARGET_PROVIDER in garage|versitygw) ;; *) exit 1 ;; esac && export REMORA_RECOVERY_REQUIRE=1 && cargo-nextest nextest run --archive-file /tests/remora-attempt-tests.tar.zst --workspace-remap /workspace --test-threads=1 --no-fail-fast -E 'binary(recovery)'"]
FROM rust-tools AS kind-http-source-build
WORKDIR /workspace
COPY Cargo.toml Cargo.lock rust-toolchain.toml ./

View file

@ -1,4 +1,4 @@
.PHONY: check-dependencies check-fixture-lib check-kind-sample check-durable check-coordination check-identity check-credentials check-admission check-attempt check-attempt-mutations check-stream check-target check-result check-outbound check-http-source check-object-source check-hf-source check-fast check-full check-fast-host check-full-host rust-check web-install web-check web-test-fast web-test-full
.PHONY: check-dependencies check-fixture-lib check-kind-sample check-durable check-coordination check-identity check-credentials check-admission check-attempt check-attempt-mutations check-recovery check-stream check-target check-result check-outbound check-http-source check-object-source check-hf-source check-fast check-full check-fast-host check-full-host rust-check web-install web-check web-test-fast web-test-full
CONTAINER_ENGINE ?= nerdctl
CONTAINER_RUN_FLAGS ?= --memory=16g --memory-swap=24g --cpus=4
@ -77,6 +77,9 @@ check-attempt:
check-attempt-mutations:
./tests/kind/run-attempt-mutations-job.sh
check-recovery:
./tests/kind/run-recovery-job.sh
check-stream:
cargo fmt --all --check
cargo clippy --package remora-stream --all-targets --all-features -- -D warnings

View file

@ -11,6 +11,7 @@ remora-coordination = { path = "../remora-coordination" }
remora-domain.workspace = true
remora-durable.workspace = true
remora-result = { path = "../remora-result" }
remora-target = { path = "../remora-target" }
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
@ -23,4 +24,4 @@ remora-contracts = { path = "../remora-contracts" }
remora-credentials = { path = "../remora-credentials" }
remora-identity = { path = "../remora-identity" }
remora-target = { path = "../remora-target", features = ["test-utils"] }
tokio.workspace = true
tokio = { workspace = true, features = ["io-util", "net"] }

View file

@ -0,0 +1,63 @@
ALTER TABLE jobs
ADD COLUMN retry_not_before TIMESTAMPTZ,
ADD COLUMN coordination_cleaned_at TIMESTAMPTZ,
ADD COLUMN recovery_checked_at TIMESTAMPTZ;
ALTER TABLE attempts
DROP CONSTRAINT attempts_state_timestamps_check,
ADD COLUMN worker_down_at TIMESTAMPTZ,
ADD CONSTRAINT attempts_state_timestamps_check CHECK (
(state = 'authorized' AND executing_at IS NULL AND reaped_at IS NULL
AND stopped_at IS NULL AND committed_at IS NULL)
OR (state = 'executing' AND executing_at IS NOT NULL AND reaped_at IS NULL
AND stopped_at IS NULL AND committed_at IS NULL)
OR (state = 'reaped' AND executing_at IS NULL AND reaped_at IS NOT NULL
AND stopped_at IS NULL AND committed_at IS NULL)
OR (state = 'stopped' AND reaped_at IS NULL
AND stopped_at IS NOT NULL AND committed_at IS NULL)
OR (state = 'committed' AND executing_at IS NOT NULL AND reaped_at IS NULL
AND stopped_at IS NULL AND committed_at IS NOT NULL)
);
CREATE TABLE attempt_staging_cleanup (
cleanup_id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
attempt_id TEXT NOT NULL REFERENCES attempts (attempt_id),
kind TEXT NOT NULL CHECK (kind IN ('object', 'multipart')),
target_store TEXT NOT NULL,
bucket TEXT NOT NULL,
object_key TEXT NOT NULL,
version_id TEXT,
upload_id TEXT,
state TEXT NOT NULL DEFAULT 'pending'
CHECK (state IN ('pending', 'claimed', 'complete')),
claim_token TEXT,
claimed_until TIMESTAMPTZ,
failure_count BIGINT NOT NULL DEFAULT 0 CHECK (failure_count >= 0),
last_error TEXT,
created_at TIMESTAMPTZ NOT NULL,
completed_at TIMESTAMPTZ,
CONSTRAINT attempt_staging_cleanup_identity_shape_check CHECK (
(kind = 'object' AND upload_id IS NULL)
OR (kind = 'multipart' AND version_id IS NULL AND upload_id IS NOT NULL)
),
CONSTRAINT attempt_staging_cleanup_claim_check CHECK (
(state = 'pending' AND claim_token IS NULL AND claimed_until IS NULL
AND completed_at IS NULL)
OR (state = 'claimed' AND claim_token IS NOT NULL AND claimed_until IS NOT NULL
AND completed_at IS NULL)
OR (state = 'complete' AND claim_token IS NULL AND claimed_until IS NULL
AND completed_at IS NOT NULL)
),
UNIQUE NULLS NOT DISTINCT (
attempt_id, kind, target_store, bucket, object_key, version_id, upload_id
)
);
CREATE INDEX attempt_staging_cleanup_claimable_idx
ON attempt_staging_cleanup (cleanup_id)
WHERE state IN ('pending', 'claimed');
CREATE INDEX jobs_recovery_idx
ON jobs (updated_at, job_id)
WHERE phase NOT IN ('succeeded', 'failed', 'cancelled')
OR coordination_cleaned_at IS NULL;

View file

@ -12,3 +12,9 @@ pub enum AttemptError {
#[error("durable operation failed: {0}")]
Durable(#[from] DurableError),
}
impl From<sqlx::Error> for AttemptError {
fn from(error: sqlx::Error) -> Self {
Self::Durable(DurableError::Database(error))
}
}

View file

@ -2,6 +2,8 @@ mod error;
mod migrate;
mod mutation_types;
mod mutations;
mod recovery;
mod recovery_types;
mod service;
mod types;
@ -12,6 +14,10 @@ pub use mutation_types::{
CommittedResultRecord, DurableDescriptor, FailureRecord, JobMutationProjection,
MutationRejection, ProgressRecord,
};
pub use recovery_types::{
CancelOutcome, CleanupClaim, CleanupCompletion, CleanupIdentity, CooperativeStopOutcome,
ReconcileReport, RecoveryConfig, StagingMultipart, StagingObject,
};
pub use service::{AttemptConfig, AttemptService};
pub use types::{
ActivationOutcome, AttemptAuthority, AttemptFaultInjector, AttemptFaultPoint, AttemptState,

View file

@ -1,4 +1,4 @@
use remora_domain::{FailureDisposition, JobPhase, OpaqueId, RuntimeFailureClass};
use remora_domain::{FailureClass, FailureDisposition, JobPhase, OpaqueId};
use remora_durable::MutationDisposition;
use remora_result::PublicationInput;
use serde::{Deserialize, Serialize};
@ -125,6 +125,6 @@ impl From<&PublicationInput> for CommittedResultRecord {
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct FailureRecord {
pub class: RuntimeFailureClass,
pub class: FailureClass,
pub disposition: FailureDisposition,
}

View file

@ -257,7 +257,10 @@ impl AttemptService {
let disposition = class.disposition();
let retry = disposition == FailureDisposition::Retryable
&& authority.fence < max_attempts;
let failure = FailureRecord { class, disposition };
let failure = FailureRecord {
class: class.into(),
disposition,
};
let now = OffsetDateTime::now_utc();
let changed = if retry {
sqlx::query(

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

View file

@ -0,0 +1,97 @@
use std::time::Duration;
use remora_domain::OpaqueId;
#[derive(Clone, Debug)]
pub struct RecoveryConfig {
pub retry_delay: Duration,
pub cleanup_claim_ttl: Duration,
pub batch_size: u32,
}
impl Default for RecoveryConfig {
fn default() -> Self {
Self {
retry_delay: Duration::from_secs(1),
cleanup_claim_ttl: Duration::from_secs(30),
batch_size: 128,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CancelOutcome {
Requested,
AlreadyRequested,
AlreadyTerminal,
NotFound,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CooperativeStopOutcome {
Stopped,
StopNotRequired,
AuthorityMismatch,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct ReconcileReport {
pub enqueued: u64,
pub cancellations_published: u64,
pub cancellations_finalized: u64,
pub attempts_reaped: u64,
pub retries_scheduled: u64,
pub stale_progress_cleared: u64,
pub terminal_coordination_cleaned: u64,
pub cleanup_items_materialized: u64,
}
impl ReconcileReport {
pub(crate) fn merge(&mut self, other: Self) {
self.enqueued += other.enqueued;
self.cancellations_published += other.cancellations_published;
self.cancellations_finalized += other.cancellations_finalized;
self.attempts_reaped += other.attempts_reaped;
self.retries_scheduled += other.retries_scheduled;
self.stale_progress_cleared += other.stale_progress_cleared;
self.terminal_coordination_cleaned += other.terminal_coordination_cleaned;
self.cleanup_items_materialized += other.cleanup_items_materialized;
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StagingObject {
pub target_store: String,
pub bucket: String,
pub key: String,
pub version_id: Option<String>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StagingMultipart {
pub target_store: String,
pub bucket: String,
pub key: String,
pub upload_id: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CleanupIdentity {
Object(StagingObject),
Multipart(StagingMultipart),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CleanupClaim {
pub cleanup_id: u64,
pub attempt_id: OpaqueId,
pub claim_token: OpaqueId,
pub identity: CleanupIdentity,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CleanupCompletion {
Completed,
RetryScheduled,
ClaimLost,
}

View file

@ -46,7 +46,7 @@ pub struct AttemptService {
pub(crate) coordination: CoordinationStore,
local_permits: LocalPermitGate,
pub(crate) operation_subject: OpaqueId,
faults: Arc<dyn AttemptFaultInjector>,
pub(crate) faults: Arc<dyn AttemptFaultInjector>,
pub(crate) config: AttemptConfig,
}
@ -733,7 +733,7 @@ fn claim_rejected(reason: ClaimRejection) -> ClaimOutcome {
}
}
fn new_id(prefix: char) -> Result<OpaqueId, AttemptError> {
pub(crate) fn new_id(prefix: char) -> Result<OpaqueId, AttemptError> {
let mut bytes = [0u8; 16];
getrandom::fill(&mut bytes)
.map_err(|error| AttemptError::Invalid(format!("identifier entropy failed: {error}")))?;

View file

@ -573,7 +573,7 @@ async fn failure_disposition_retries_with_budget_and_becomes_terminal_when_exhau
&harness.pool,
&first.attempt_id,
FailureRecord {
class: RuntimeFailureClass::TransientTarget,
class: RuntimeFailureClass::TransientTarget.into(),
disposition: FailureDisposition::Retryable,
},
)
@ -623,7 +623,7 @@ async fn failure_disposition_retries_with_budget_and_becomes_terminal_when_exhau
assert_eq!(
exhausted.failure,
Some(FailureRecord {
class: RuntimeFailureClass::Internal,
class: RuntimeFailureClass::Internal.into(),
disposition: FailureDisposition::Retryable,
})
);

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,34 @@
apiVersion: batch/v1
kind: Job
metadata:
name: @@JOB_NAME@@
namespace: remora-testbed
spec:
activeDeadlineSeconds: 900
backoffLimit: 0
ttlSecondsAfterFinished: 300
template:
spec:
automountServiceAccountToken: false
enableServiceLinks: false
restartPolicy: Never
containers:
- name: tests
image: @@TEST_IMAGE@@
imagePullPolicy: Never
envFrom:
- configMapRef: {name: @@PROFILE_NAME@@}
env:
- name: REMORA_RECOVERY_REQUIRE
value: "1"
volumeMounts:
- {name: fixture-profile, mountPath: /run/remora-fixture, readOnly: true}
resources:
requests: {cpu: 250m, memory: 256Mi}
limits: {cpu: "2", memory: 1Gi}
volumes:
- name: fixture-profile
configMap:
name: @@PROFILE_NAME@@-file
items:
- {key: profile.env, path: profile.env}

50
tests/kind/run-recovery-job.sh Executable file
View file

@ -0,0 +1,50 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
source "$ROOT/tests/kind/lib.sh"
USER_KUBECONFIG_BEFORE="$(user_kubeconfig_digest)"
SUCCEEDED=false
OWNS_FIXTURE=false
cleanup() {
local status=$?
trap - EXIT INT TERM
if [[ "${REMORA_TESTBED_RETAIN:-0}" = 1 ]] && [[ "$SUCCEEDED" != true ]]; then
printf 'retained kubeconfig: %s\n' "$KUBECONFIG_PATH"
printf 'retained context: %s\n' "$CONTEXT"
exit "$status"
fi
if [[ "$OWNS_FIXTURE" = true ]]; then
destroy_fixture || status=1
fi
if [[ "$(user_kubeconfig_digest)" != "$USER_KUBECONFIG_BEFORE" ]]; then
printf 'user kubeconfig changed unexpectedly\n' >&2
status=1
fi
exit "$status"
}
trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
acquire_fixture_lock
REMORA_REAL_NERDCTL="$(command -v nerdctl)"
export REMORA_REAL_NERDCTL
"$REMORA_REAL_NERDCTL" --namespace default build \
--file "$ROOT/Dockerfile.test" \
--target kind-recovery-runner \
--tag remora-kind-recovery:local \
"$ROOT"
OWNS_FIXTURE=true
setup_fixture durable coordination target
load_local_image remora-kind-recovery:local
run_job_matrix \
recovery \
"$TESTBED/recovery-job.yaml.in" \
remora-kind-recovery:local \
durable coordination target
SUCCEEDED=true
printf 'recovery package test Jobs passed\n'