Implement crash-safe streaming worker runtime
This commit is contained in:
parent
16cba3c317
commit
cdeef3cc1a
1 changed files with 7448 additions and 194 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue