ms2k/remora
SHA256
Watch
0
Code Releases Activity

Implement crash-safe streaming worker runtime

Browse source
This commit is contained in:
ms2k 2026-07-17 18:51:44 +09:00
commit cdeef3cc1a
7 changed files with 7448 additions and 194 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/remora-attempt/src/lib.rs
View file

@ -10,13 +10,13 @@ mod types;
pub use error::AttemptError;
pub use migrate::migrate;
pub use mutation_types::{
AttemptMutationOutcome, CheckpointProjection, CheckpointRecord, CommittedManifestRecord,
CommittedResultRecord, DurableDescriptor, FailureRecord, JobMutationProjection,
MutationRejection, ProgressRecord,
AttemptMutationOutcome, CheckpointMultipart, CheckpointProjection, CheckpointRecord,
CommittedManifestRecord, CommittedResultRecord, DurableDescriptor, FailureRecord,
JobMutationProjection, MutationRejection, ProgressRecord,
};
pub use recovery_types::{
CancelOutcome, CleanupClaim, CleanupCompletion, CleanupIdentity, CooperativeStopOutcome,
ReconcileReport, RecoveryConfig, StagingMultipart, StagingObject,
ReconcileReport, RecoveryConfig, StagingMultipart, StagingObject, StagingWriteIntent,
};
pub use service::{AttemptConfig, AttemptService};
pub use types::{

10
crates/remora-attempt/src/mutation_types.rs
View file

@ -18,6 +18,16 @@ pub struct CheckpointRecord {
pub resolved_source_id: OpaqueId,
pub validated_bytes: u64,
pub validated_members: u64,
#[serde(default)]
pub staging_multipart: Option<CheckpointMultipart>,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct CheckpointMultipart {
pub target_store: String,
pub bucket: String,
pub key: String,
pub upload_id: String,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]

40
crates/remora-attempt/src/mutations.rs
View file

@ -73,7 +73,11 @@ impl AttemptService {
let outcome = self
.publish_success(authority, operation_id, result)
.await?;
self.publication_outcome(operation_id, outcome).await
let publication = self.publication_outcome(operation_id, outcome).await?;
if matches!(publication, OutcomePublication::Accepted { .. }) {
self.release_exact(authority).await;
}
Ok(publication)
}
pub async fn persist_checkpoint(
@ -118,13 +122,16 @@ impl AttemptService {
}
let now = OffsetDateTime::now_utc();
sqlx::query(
"UPDATE jobs SET checkpoint = $2, revision = revision + 1,
updated_at = $3 WHERE job_id = $1",
"UPDATE jobs
SET checkpoint = $2, checkpoint_owner_attempt_id = $3,
revision = revision + 1, updated_at = $4
WHERE job_id = $1",
)
.bind(authority.job_id.as_str())
.bind(Json(&CheckpointProjection::Reusable {
checkpoint: checkpoint.clone(),
}))
.bind(authority.attempt_id.as_str())
.bind(now)
.execute(tx.as_mut())
.await?;
@ -143,7 +150,11 @@ impl AttemptService {
class: RuntimeFailureClass,
) -> Result<OutcomePublication, AttemptError> {
let outcome = self.publish_failure(authority, operation_id, class).await?;
self.publication_outcome(operation_id, outcome).await
let publication = self.publication_outcome(operation_id, outcome).await?;
if matches!(publication, OutcomePublication::Accepted { .. }) {
self.release_exact(authority).await;
}
Ok(publication)
}
pub async fn report_progress(
@ -227,12 +238,14 @@ impl AttemptService {
};
let changed = sqlx::query(
"UPDATE jobs SET phase = 'succeeded', authority_attempt_id = NULL,
checkpoint = $2, committed_result = $3, terminal_at = $4,
revision = revision + 1, updated_at = $4
checkpoint = $2, checkpoint_owner_attempt_id = $3,
committed_result = $4, terminal_at = $5,
revision = revision + 1, updated_at = $5
WHERE job_id = $1 AND phase = 'running' AND cancel_requested = FALSE",
)
.bind(authority.job_id.as_str())
.bind(Json(&complete_checkpoint))
.bind(authority.attempt_id.as_str())
.bind(stored)
.bind(now)
.execute(tx.as_mut())
@ -544,11 +557,26 @@ fn validate_checkpoint(checkpoint: &CheckpointRecord) -> Result<(), AttemptError
if (checkpoint.validated_bytes == 0 && checkpoint.validated_members == 0)
|| checkpoint.validated_bytes > MAX_SAFE_INTEGER
|| checkpoint.validated_members > MAX_SAFE_INTEGER
|| (checkpoint.validated_members == 0) != checkpoint.staging_multipart.is_some()
{
return Err(AttemptError::Invalid(
"checkpoint exceeds the safe-integer range".into(),
));
}
if checkpoint
.staging_multipart
.as_ref()
.is_some_and(|identity| {
identity.target_store.is_empty()
|| identity.bucket.is_empty()
|| identity.key.is_empty()
|| identity.upload_id.is_empty()
})
{
return Err(AttemptError::Invalid(
"checkpoint multipart identity is invalid".into(),
));
}
Ok(())
}

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

@ -1,9 +1,14 @@
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 remora_domain::{
FailureClass, FailureDisposition, JobPhase, OpaqueId, fingerprint, ports::MutationScope,
};
use remora_durable::{
DurableError, MutationDisposition, MutationStatus, enqueue, execute_mutation, with_transaction,
};
use remora_target::{attempt_key_attempt_id, attempt_staging_prefix, validate_attempt_key};
use serde::Serialize;
use serde_json::json;
use sqlx::{FromRow, Postgres, Transaction, types::Json};
use time::OffsetDateTime;
@ -12,7 +17,7 @@ use crate::service::new_id;
use crate::{
AttemptAuthority, AttemptError, AttemptService, CancelOutcome, CleanupClaim, CleanupCompletion,
CleanupIdentity, CooperativeStopOutcome, FailureRecord, ReconcileReport, RecoveryConfig,
StagingMultipart, StagingObject,
StagingMultipart, StagingObject, StagingWriteIntent,
};
impl AttemptService {
@ -73,57 +78,79 @@ impl AttemptService {
&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)
let authority_for_effect = authority.clone();
let coordination = self.coordination.clone();
let outcome = with_transaction(&self.pool, move |tx| {
Box::pin(async move {
cooperative_stop_transaction(tx, &coordination, &authority_for_effect).await
})
})
.await?;
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 outcome == CooperativeStopOutcome::Stopped {
let _ = self
.coordination
.release(authority.job_id.as_str(), &lease_owner(authority))
.await;
self.release_local(&authority.job_id);
}
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')",
Ok(outcome)
}
pub async fn cooperative_stop_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
) -> Result<CooperativeStopOutcome, AttemptError> {
let request_fingerprint = fingerprint(&StopRequest {
job_id: &authority.job_id,
attempt_id: &authority.attempt_id,
fence: authority.fence,
worker_incarnation: &authority.worker_incarnation,
})
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
let authority_for_effect = authority.clone();
let coordination = self.coordination.clone();
let outcome = execute_mutation(
&self.pool,
MutationScope {
subject: &self.operation_subject,
operation: "attempt.mutation",
key: operation_id.as_str(),
},
"cooperative_stop",
&[],
request_fingerprint,
move |tx| {
Box::pin(async move {
let decision =
cooperative_stop_transaction(tx, &coordination, &authority_for_effect)
.await?;
let status = if decision == CooperativeStopOutcome::AuthorityMismatch {
MutationStatus::Rejected
} else {
MutationStatus::Accepted
};
Ok((
status,
serde_json::to_value(decision).map_err(DurableError::effect)?,
))
})
},
)
.bind(authority.attempt_id.as_str())
.bind(now)
.execute(&self.pool)
.await?;
if changed.rows_affected() == 0 {
return Ok(CooperativeStopOutcome::StopNotRequired);
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
let _ = self
.coordination
.release(authority.job_id.as_str(), &lease_owner(authority))
.await;
self.release_local(&authority.job_id);
Ok(CooperativeStopOutcome::Stopped)
let decision = serde_json::from_value(outcome.response)
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
if decision == CooperativeStopOutcome::Stopped {
let _ = self
.coordination
.release(authority.job_id.as_str(), &lease_owner(authority))
.await;
self.release_local(&authority.job_id);
}
Ok(decision)
}
pub async fn mark_worker_down(
@ -146,7 +173,7 @@ impl AttemptService {
attempt_id: &OpaqueId,
object: &StagingObject,
) -> Result<(), AttemptError> {
validate_staging_key(attempt_id, &object.key)?;
validate_staging_key(&self.pool, attempt_id, &object.key).await?;
insert_cleanup_identity(
&self.pool,
attempt_id,
@ -162,12 +189,105 @@ impl AttemptService {
.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(attempt_id, &multipart.key)?;
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(),
@ -188,6 +308,107 @@ impl AttemptService {
.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,
@ -235,6 +456,8 @@ impl AttemptService {
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
@ -279,6 +502,104 @@ impl AttemptService {
})
}
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,
@ -499,6 +820,8 @@ impl AttemptService {
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",
@ -517,6 +840,8 @@ impl AttemptService {
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",
@ -544,6 +869,199 @@ impl AttemptService {
}
}
async fn cooperative_stop_transaction(
tx: &mut remora_durable::PgTx<'_>,
coordination: &remora_coordination::CoordinationStore,
authority: &AttemptAuthority,
) -> Result<CooperativeStopOutcome, DurableError> {
let row = sqlx::query_as::<_, StopRow>(
"SELECT a.job_id, a.fence, a.worker_incarnation, a.state,
j.authority_attempt_id, j.cancel_requested
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1
FOR UPDATE OF a, j",
)
.bind(authority.attempt_id.as_str())
.fetch_optional(tx.as_mut())
.await?;
let Some(row) = row else {
return Ok(CooperativeStopOutcome::AuthorityMismatch);
};
if row.job_id != authority.job_id.as_str()
|| u64::try_from(row.fence).ok() != Some(authority.fence)
|| row.worker_incarnation != authority.worker_incarnation.as_str()
{
return Ok(CooperativeStopOutcome::AuthorityMismatch);
}
if row.state == "stopped" {
return Ok(CooperativeStopOutcome::Stopped);
}
if row.state != "executing" {
return Ok(CooperativeStopOutcome::StopNotRequired);
}
let current = row.authority_attempt_id.as_deref() == Some(authority.attempt_id.as_str());
let live = coordination
.get_live_lease(authority.job_id.as_str())
.await
.map_err(DurableError::effect)?
.is_some_and(|lease| lease.record.owner() == lease_owner(authority));
if current && live && !row.cancel_requested {
return Ok(CooperativeStopOutcome::StopNotRequired);
}
let changed = sqlx::query(
"UPDATE attempts SET state = 'stopped', stopped_at = $2
WHERE attempt_id = $1 AND state IN ('authorized', 'executing')",
)
.bind(authority.attempt_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
Ok(if changed.rows_affected() == 1 {
CooperativeStopOutcome::Stopped
} else {
CooperativeStopOutcome::StopNotRequired
})
}
async fn lock_registration_authority(
tx: &mut remora_durable::PgTx<'_>,
authority: &AttemptAuthority,
) -> Result<bool, DurableError> {
let row = sqlx::query_as::<_, RegistrationAuthorityRow>(
"SELECT a.job_id, a.fence, a.worker_incarnation, a.state,
j.authority_attempt_id, j.cancel_requested
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1
FOR UPDATE OF a, j",
)
.bind(authority.attempt_id.as_str())
.fetch_optional(tx.as_mut())
.await?;
Ok(row.is_some_and(|row| {
row.job_id == authority.job_id.as_str()
&& u64::try_from(row.fence).ok() == Some(authority.fence)
&& row.worker_incarnation == authority.worker_incarnation.as_str()
&& row.state == "executing"
&& row.authority_attempt_id.as_deref() == Some(authority.attempt_id.as_str())
&& !row.cancel_requested
}))
}
async fn validate_staging_key_transaction(
tx: &mut remora_durable::PgTx<'_>,
attempt_id: &OpaqueId,
key: &str,
) -> Result<(), DurableError> {
validate_attempt_key(key).map_err(DurableError::effect)?;
let (staging_origin, target_prefix) = sqlx::query_as::<_, (String, Json<Vec<String>>)>(
"SELECT a.staging_origin_attempt_id, j.target_prefix
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_one(tx.as_mut())
.await?;
let staging_origin = OpaqueId::parse(staging_origin).map_err(DurableError::effect)?;
if attempt_key_attempt_id(key).map_err(DurableError::effect)? != staging_origin
|| !key.starts_with(&attempt_staging_prefix(
&target_prefix.0.join("/"),
&staging_origin,
))
{
return Err(DurableError::effect(
"staging key belongs to a different attempt",
));
}
Ok(())
}
async fn insert_cleanup_identity_transaction(
tx: &mut remora_durable::PgTx<'_>,
attempt_id: &OpaqueId,
identity: NewCleanupIdentity<'_>,
) -> Result<(), DurableError> {
if identity.target_store.is_empty() || identity.bucket.is_empty() {
return Err(DurableError::effect(
"staging target store and bucket cannot be empty",
));
}
let changed = sqlx::query(
"INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, upload_id, created_at
)
SELECT $1, $2, $3, $4, $5, $6, $7, $8
WHERE EXISTS (
SELECT 1 FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1 AND j.target_store = $3 AND j.target_bucket = $4
)
ON CONFLICT DO NOTHING",
)
.bind(attempt_id.as_str())
.bind(identity.kind)
.bind(identity.target_store)
.bind(identity.bucket)
.bind(identity.key)
.bind(identity.version_id)
.bind(identity.upload_id)
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
if changed.rows_affected() == 0 {
let binding = sqlx::query_as::<_, (String, String)>(
"SELECT j.target_store, j.target_bucket
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_optional(tx.as_mut())
.await?;
match binding {
None => return Err(DurableError::effect("attempt does not exist")),
Some((store, bucket))
if store != identity.target_store || bucket != identity.bucket =>
{
return Err(DurableError::effect(
"staging identity does not match the attempt target",
));
}
Some(_) => {}
}
}
Ok(())
}
async fn complete_pending_identity(
tx: &mut remora_durable::PgTx<'_>,
attempt_id: &OpaqueId,
kind: &str,
target_store: &str,
bucket: &str,
key: &str,
upload_id: Option<&str>,
) -> Result<(), DurableError> {
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET state = 'complete', completed_at = $7
WHERE attempt_id = $1 AND kind = $2 AND target_store = $3 AND bucket = $4
AND object_key = $5 AND upload_id IS NOT DISTINCT FROM $6 AND state = 'pending'",
)
.bind(attempt_id.as_str())
.bind(kind)
.bind(target_store)
.bind(bucket)
.bind(key)
.bind(upload_id)
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
if changed.rows_affected() != 1 {
return Err(DurableError::effect(
"staging write lifecycle identity is missing or unavailable",
));
}
Ok(())
}
async fn insert_cleanup_identity(
pool: &sqlx::PgPool,
attempt_id: &OpaqueId,
@ -662,6 +1180,7 @@ async fn reap_lost_attempt(
let inactive = match row.state.as_str() {
"authorized" => reservation.is_none() && !exact_live,
"executing" => row.worker_down_at.is_some() || !exact_live,
"stopped" => !exact_live,
_ => false,
};
if !inactive {
@ -718,9 +1237,30 @@ async fn reap_lost_attempt(
})
}
fn validate_staging_key(attempt_id: &OpaqueId, key: &str) -> Result<(), AttemptError> {
async fn validate_staging_key(
pool: &sqlx::PgPool,
attempt_id: &OpaqueId,
key: &str,
) -> Result<(), AttemptError> {
validate_attempt_key(key).map_err(|error| AttemptError::Invalid(error.to_string()))?;
if key.split('/').nth(2) != Some(attempt_id.as_str()) {
let (staging_origin, target_prefix) = sqlx::query_as::<_, (String, Json<Vec<String>>)>(
"SELECT a.staging_origin_attempt_id, j.target_prefix
FROM attempts a JOIN jobs j ON j.job_id = a.job_id
WHERE a.attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_optional(pool)
.await?
.ok_or_else(|| AttemptError::Invalid("attempt does not exist".into()))?;
let staging_origin = OpaqueId::parse(staging_origin)
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
if attempt_key_attempt_id(key).map_err(|error| AttemptError::Invalid(error.to_string()))?
!= staging_origin
|| !key.starts_with(&attempt_staging_prefix(
&target_prefix.0.join("/"),
&staging_origin,
))
{
return Err(AttemptError::Invalid(
"staging key belongs to a different attempt".into(),
));
@ -788,6 +1328,39 @@ struct StopRow {
cancel_requested: bool,
}
#[derive(FromRow)]
struct RegistrationAuthorityRow {
job_id: String,
fence: i64,
worker_incarnation: String,
state: String,
authority_attempt_id: Option<String>,
cancel_requested: bool,
}
#[derive(Serialize)]
struct StopRequest<'a> {
job_id: &'a OpaqueId,
attempt_id: &'a OpaqueId,
fence: u64,
worker_incarnation: &'a OpaqueId,
}
#[derive(Serialize)]
struct RegistrationRequest<'a, T> {
job_id: &'a OpaqueId,
attempt_id: &'a OpaqueId,
fence: u64,
worker_incarnation: &'a OpaqueId,
value: T,
}
#[derive(Serialize)]
struct StagingObjectRegistration {
object: StagingObject,
completed_upload_id: Option<String>,
}
#[derive(FromRow)]
struct RecoveryJobRow {
phase: String,
@ -827,6 +1400,11 @@ struct CleanupRow {
impl CleanupRow {
fn try_into_claim(self, claim_token: OpaqueId) -> Result<CleanupClaim, AttemptError> {
let identity = match self.kind.as_str() {
"write_intent" => CleanupIdentity::WriteIntent(StagingWriteIntent {
target_store: self.target_store,
bucket: self.bucket,
key: self.object_key,
}),
"object" => CleanupIdentity::Object(StagingObject {
target_store: self.target_store,
bucket: self.bucket,

16
crates/remora-attempt/src/recovery_types.rs
View file

@ -1,6 +1,7 @@
use std::time::Duration;
use remora_domain::OpaqueId;
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug)]
pub struct RecoveryConfig {
@ -27,7 +28,8 @@ pub enum CancelOutcome {
NotFound,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CooperativeStopOutcome {
Stopped,
StopNotRequired,
@ -59,7 +61,7 @@ impl ReconcileReport {
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct StagingObject {
pub target_store: String,
pub bucket: String,
@ -67,7 +69,7 @@ pub struct StagingObject {
pub version_id: Option<String>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct StagingMultipart {
pub target_store: String,
pub bucket: String,
@ -75,8 +77,16 @@ pub struct StagingMultipart {
pub upload_id: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct StagingWriteIntent {
pub target_store: String,
pub bucket: String,
pub key: String,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CleanupIdentity {
WriteIntent(StagingWriteIntent),
Object(StagingObject),
Multipart(StagingMultipart),
}

358
crates/remora-attempt/src/service.rs
View file

@ -559,7 +559,7 @@ impl AttemptService {
self.local_permits.release(job_id.as_str());
}
async fn release_exact(&self, authority: &AttemptAuthority) {
pub(crate) async fn release_exact(&self, authority: &AttemptAuthority) {
if let Ok(owner) = lease_owner(authority) {
let _ = self
.coordination
@ -662,16 +662,25 @@ async fn authorize_attempt(
Option<String>,
Option<Json<crate::DurableDescriptor>>,
Option<Json<CheckpointProjection>>,
Option<String>,
),
>(
"SELECT phase, dispatch_needed, current_fence, authority_attempt_id,
resolved_source, checkpoint
resolved_source, checkpoint, checkpoint_owner_attempt_id
FROM jobs WHERE job_id = $1 FOR UPDATE",
)
.bind(input.job_id.as_str())
.fetch_optional(tx.as_mut())
.await?;
let Some((phase, dispatch_needed, current_fence, authority, resolved_source, checkpoint)) = row
let Some((
phase,
dispatch_needed,
current_fence,
authority,
resolved_source,
checkpoint,
checkpoint_owner,
)) = row
else {
return Ok(claim_rejected(ClaimRejection::JobUnavailable));
};
@ -682,6 +691,26 @@ async fn authorize_attempt(
if current_fence >= input.max_attempts {
return Ok(claim_rejected(ClaimRejection::AttemptsExhausted));
}
let reusable = resolved_source
.as_ref()
.zip(checkpoint.as_ref())
.and_then(|(source, checkpoint)| checkpoint.0.reusable().map(|value| (source, value)))
.filter(|(source, checkpoint)| checkpoint.resolved_source_id == source.0.id);
if let (Some(_), Some(owner)) = (reusable, checkpoint_owner.as_deref()) {
let stopped = sqlx::query_scalar::<_, bool>(
"SELECT state = 'stopped'
FROM attempts
WHERE attempt_id = $1 AND job_id = $2",
)
.bind(owner)
.bind(input.job_id.as_str())
.fetch_optional(tx.as_mut())
.await?
.unwrap_or(false);
if !stopped {
return Ok(claim_rejected(ClaimRejection::JobUnavailable));
}
}
input.faults.hit(AttemptFaultPoint::ClaimBeforeReserve);
let reservation = input
.coordination
@ -697,26 +726,47 @@ async fn authorize_attempt(
input.reserved.store(true, Ordering::Release);
input.faults.hit(AttemptFaultPoint::ClaimAfterReserve);
let fence = current_fence + 1;
let initial_progress = match (resolved_source, checkpoint) {
(Some(source), Some(checkpoint))
if checkpoint
.0
.reusable()
.is_some_and(|value| value.resolved_source_id == source.0.id) =>
{
checkpoint.0.reusable().unwrap().validated_bytes
let checkpoint_stage = match (reusable, checkpoint_owner.as_deref()) {
(Some((_, checkpoint)), Some(owner)) if checkpoint.validated_bytes > 0 => {
load_checkpoint_stage(tx, input.job_id, owner, checkpoint).await?
}
_ => 0,
_ => None,
};
if checkpoint.is_some() && checkpoint_stage.is_none() {
sqlx::query(
"UPDATE jobs
SET checkpoint = NULL, checkpoint_owner_attempt_id = NULL,
revision = revision + 1, updated_at = $2
WHERE job_id = $1",
)
.bind(input.job_id.as_str())
.bind(OffsetDateTime::now_utc())
.execute(tx.as_mut())
.await?;
}
let initial_progress = checkpoint_stage
.as_ref()
.map_or(0, |stage| stage.validated_bytes);
let staging_origin = checkpoint_stage
.as_ref()
.map_or(input.attempt_id, |stage| &stage.staging_origin_attempt_id);
let checkpoint_owner_after_claim = checkpoint_stage.as_ref().map(|_| input.attempt_id.as_str());
let now = OffsetDateTime::now_utc();
let envelope =
load_execution_envelope_record(tx, input.job_id, input.attempt_id, fence, input.worker)
.await?;
let envelope = load_execution_envelope_record(
tx,
input.job_id,
input.attempt_id,
fence,
input.worker,
staging_origin,
)
.await?;
sqlx::query(
"INSERT INTO attempts (
attempt_id, job_id, fence, worker_incarnation, state,
initial_progress, authorized_at, executing_at, reaped_at, execution_envelope
) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL, $7)",
initial_progress, authorized_at, executing_at, reaped_at, execution_envelope,
staging_origin_attempt_id
) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL, $7, $8)",
)
.bind(input.attempt_id.as_str())
.bind(input.job_id.as_str())
@ -725,18 +775,24 @@ async fn authorize_attempt(
.bind(i64::try_from(initial_progress).map_err(DurableError::effect)?)
.bind(now)
.bind(Json(&envelope))
.bind(staging_origin.as_str())
.execute(tx.as_mut())
.await?;
if let Some(stage) = checkpoint_stage {
transfer_checkpoint_stage(tx, &stage, input.attempt_id).await?;
}
sqlx::query(
"UPDATE jobs
SET phase = 'running', current_fence = $2, authority_attempt_id = $3,
dispatch_needed = FALSE, revision = revision + 1, updated_at = $4
checkpoint_owner_attempt_id = $5, dispatch_needed = FALSE,
revision = revision + 1, updated_at = $4
WHERE job_id = $1",
)
.bind(input.job_id.as_str())
.bind(i64::try_from(fence).map_err(DurableError::effect)?)
.bind(input.attempt_id.as_str())
.bind(now)
.bind(checkpoint_owner_after_claim)
.execute(tx.as_mut())
.await?;
Ok(ClaimOutcome::Claimed {
@ -751,12 +807,277 @@ async fn authorize_attempt(
})
}
struct CheckpointStage {
owner_attempt_id: OpaqueId,
staging_origin_attempt_id: OpaqueId,
validated_bytes: u64,
validated_members: u64,
inventory_bytes: u64,
staging_multipart: Option<crate::CheckpointMultipart>,
}
async fn load_checkpoint_stage(
tx: &mut remora_durable::PgTx<'_>,
job_id: &OpaqueId,
owner: &str,
checkpoint: &crate::CheckpointRecord,
) -> Result<Option<CheckpointStage>, DurableError> {
let Some(staging_origin) = sqlx::query_scalar::<_, String>(
"SELECT staging_origin_attempt_id
FROM attempts
WHERE attempt_id = $1 AND job_id = $2 AND state = 'stopped'",
)
.bind(owner)
.bind(job_id.as_str())
.fetch_optional(tx.as_mut())
.await?
else {
return Ok(None);
};
let Some((state, member_count, inventory_origin, inventory_prefix, target_prefix)) =
sqlx::query_as::<
_,
(
String,
i64,
Option<String>,
Option<String>,
Json<Vec<String>>,
),
>(
"SELECT i.state, i.member_count, i.staging_origin_attempt_id,
i.target_prefix, j.target_prefix
FROM result_inventories i
JOIN attempts a ON a.attempt_id = i.attempt_id
JOIN jobs j ON j.job_id = a.job_id
WHERE i.attempt_id = $1
FOR UPDATE OF i",
)
.bind(owner)
.fetch_optional(tx.as_mut())
.await?
else {
return Ok(None);
};
if inventory_origin.as_deref() != Some(staging_origin.as_str())
|| inventory_prefix.as_deref() != Some(target_prefix.0.join("/").as_str())
{
return Ok(None);
}
let validated_members =
i64::try_from(checkpoint.validated_members).map_err(DurableError::effect)?;
let validated_bytes =
i64::try_from(checkpoint.validated_bytes).map_err(DurableError::effect)?;
if state != "open" || member_count < validated_members {
return Ok(None);
}
let (retained_members, retained_bytes) = sqlx::query_as::<_, (i64, i64)>(
"SELECT COUNT(*), COALESCE(SUM(size), 0)::BIGINT
FROM (
SELECT size
FROM result_inventory_members
WHERE attempt_id = $1
ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
LIMIT $2
) retained",
)
.bind(owner)
.bind(validated_members)
.fetch_one(tx.as_mut())
.await?;
if retained_members != validated_members || retained_bytes > validated_bytes {
return Ok(None);
}
let partial = retained_bytes < validated_bytes;
if partial != checkpoint.staging_multipart.is_some() {
return Ok(None);
}
let cleanup_started = sqlx::query_scalar::<_, bool>(
"WITH retained AS (
SELECT object_store, object_bucket, object_key, object_version_id
FROM result_inventory_members
WHERE attempt_id = $1
ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
LIMIT $2
)
SELECT EXISTS (
SELECT 1
FROM attempt_staging_cleanup c
JOIN retained r
ON r.object_store = c.target_store
AND r.object_bucket = c.bucket
AND r.object_key = c.object_key
AND r.object_version_id IS NOT DISTINCT FROM c.version_id
WHERE c.attempt_id = $1
AND c.kind = 'object'
AND c.state <> 'pending'
)",
)
.bind(owner)
.bind(validated_members)
.fetch_one(tx.as_mut())
.await?;
if cleanup_started {
return Ok(None);
}
if let Some(multipart) = &checkpoint.staging_multipart {
let available = sqlx::query_scalar::<_, bool>(
"SELECT EXISTS (
SELECT 1 FROM attempt_staging_cleanup
WHERE attempt_id = $1 AND kind = 'multipart'
AND target_store = $2 AND bucket = $3
AND object_key = $4 AND upload_id = $5
AND state = 'pending'
)",
)
.bind(owner)
.bind(&multipart.target_store)
.bind(&multipart.bucket)
.bind(&multipart.key)
.bind(&multipart.upload_id)
.fetch_one(tx.as_mut())
.await?;
if !available {
return Ok(None);
}
}
Ok(Some(CheckpointStage {
owner_attempt_id: OpaqueId::parse(owner).map_err(DurableError::effect)?,
staging_origin_attempt_id: OpaqueId::parse(staging_origin).map_err(DurableError::effect)?,
validated_bytes: checkpoint.validated_bytes,
validated_members: checkpoint.validated_members,
inventory_bytes: u64::try_from(retained_bytes).map_err(DurableError::effect)?,
staging_multipart: checkpoint.staging_multipart.clone(),
}))
}
async fn transfer_checkpoint_stage(
tx: &mut remora_durable::PgTx<'_>,
stage: &CheckpointStage,
replacement: &OpaqueId,
) -> Result<(), DurableError> {
let retained_members = i64::try_from(stage.validated_members).map_err(DurableError::effect)?;
let retained_bytes = i64::try_from(stage.inventory_bytes).map_err(DurableError::effect)?;
let now = OffsetDateTime::now_utc();
sqlx::query(
"WITH ranked AS (
SELECT *, ROW_NUMBER() OVER (
ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
) AS ordinal
FROM result_inventory_members
WHERE attempt_id = $1
)
INSERT INTO attempt_staging_cleanup (
attempt_id, kind, target_store, bucket, object_key, version_id, created_at
)
SELECT $1, 'object', object_store, object_bucket, object_key,
object_version_id, $3
FROM ranked
WHERE ordinal > $2
ON CONFLICT DO NOTHING",
)
.bind(stage.owner_attempt_id.as_str())
.bind(retained_members)
.bind(now)
.execute(tx.as_mut())
.await?;
sqlx::query(
"WITH retained AS (
SELECT object_store, object_bucket, object_key, object_version_id
FROM result_inventory_members
WHERE attempt_id = $1
ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
LIMIT $3
)
UPDATE attempt_staging_cleanup c
SET attempt_id = $2
FROM retained r
WHERE c.attempt_id = $1
AND c.kind = 'object'
AND c.state = 'pending'
AND r.object_store = c.target_store
AND r.object_bucket = c.bucket
AND r.object_key = c.object_key
AND r.object_version_id IS NOT DISTINCT FROM c.version_id",
)
.bind(stage.owner_attempt_id.as_str())
.bind(replacement.as_str())
.bind(retained_members)
.execute(tx.as_mut())
.await?;
if let Some(multipart) = &stage.staging_multipart {
let changed = sqlx::query(
"UPDATE attempt_staging_cleanup
SET attempt_id = $2
WHERE attempt_id = $1 AND state = 'pending'
AND (
(kind = 'multipart' AND target_store = $3 AND bucket = $4
AND object_key = $5 AND upload_id = $6)
OR (kind = 'write_intent' AND target_store = $3 AND bucket = $4
AND object_key = $5)
)",
)
.bind(stage.owner_attempt_id.as_str())
.bind(replacement.as_str())
.bind(&multipart.target_store)
.bind(&multipart.bucket)
.bind(&multipart.key)
.bind(&multipart.upload_id)
.execute(tx.as_mut())
.await?;
if changed.rows_affected() == 0 {
return Err(DurableError::effect(
"checkpoint multipart ownership transfer failed",
));
}
}
sqlx::query(
"WITH extras AS (
SELECT attempt_id, member_id
FROM (
SELECT attempt_id, member_id, ROW_NUMBER() OVER (
ORDER BY path COLLATE \"C\", member_id COLLATE \"C\"
) AS ordinal
FROM result_inventory_members
WHERE attempt_id = $1
) ranked
WHERE ordinal > $2
)
DELETE FROM result_inventory_members m
USING extras
WHERE m.attempt_id = extras.attempt_id
AND m.member_id = extras.member_id",
)
.bind(stage.owner_attempt_id.as_str())
.bind(retained_members)
.execute(tx.as_mut())
.await?;
let changed = sqlx::query(
"UPDATE result_inventories
SET attempt_id = $2, member_count = $3, total_bytes = $4
WHERE attempt_id = $1 AND state = 'open'",
)
.bind(stage.owner_attempt_id.as_str())
.bind(replacement.as_str())
.bind(retained_members)
.bind(retained_bytes)
.execute(tx.as_mut())
.await?;
if changed.rows_affected() != 1 {
return Err(DurableError::effect(
"checkpoint staging inventory transfer failed",
));
}
Ok(())
}
async fn load_execution_envelope_record(
tx: &mut remora_durable::PgTx<'_>,
job_id: &OpaqueId,
attempt_id: &OpaqueId,
fence: u64,
worker: &OpaqueId,
staging_origin_attempt_id: &OpaqueId,
) -> Result<ExecutionEnvelopeRecord, DurableError> {
type Row = (
Json<serde_json::Value>,
@ -804,6 +1125,7 @@ async fn load_execution_envelope_record(
target_store: row.2,
target_bucket: row.3,
target_prefix: row.4.0,
staging_origin_attempt_id: staging_origin_attempt_id.clone(),
execution_grant_id: grant_id,
checkpoint: row.9.map(|value| value.0),
})

1
crates/remora-attempt/src/types.rs
View file

@ -157,6 +157,7 @@ pub struct ExecutionEnvelopeRecord {
pub target_store: String,
pub target_bucket: String,
pub target_prefix: Vec<String>,
pub staging_origin_attempt_id: OpaqueId,
pub execution_grant_id: OpaqueId,
pub checkpoint: Option<crate::CheckpointProjection>,
}