ms2k/remora
SHA256
Watch
0
Code Releases Activity

Split orchestration by durable use case

Browse source
This commit is contained in:
ms2k 2026-07-19 01:33:14 +09:00
commit fa81cb7f21
2 changed files with 609 additions and 513 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -21,157 +21,9 @@ use crate::{
StagingMultipart, StagingObject, StagingWriteIntent,
};
mod cancellation;
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 { Self::request_cancel_in_tx(tx, &job_id).await })
})
.await
.map_err(Into::into)
}
pub async fn request_cancel_in_tx(
tx: &mut PgTx<'_>,
job_id: &OpaqueId,
) -> Result<CancelOutcome, DurableError> {
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)
}
pub async fn cooperative_stop(
&self,
authority: &AttemptAuthority,
) -> Result<CooperativeStopOutcome, AttemptError> {
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?;
if outcome == CooperativeStopOutcome::Stopped {
let _ = self
.coordination
.release(authority.job_id.as_str(), &lease_owner(authority))
.await;
}
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)?,
))
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
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;
}
Ok(decision)
}
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,

153
crates/remora-attempt/src/recovery/cancellation.rs Normal file
View file

@ -0,0 +1,153 @@
use super::*;
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 { Self::request_cancel_in_tx(tx, &job_id).await })
})
.await
.map_err(Into::into)
}
pub async fn request_cancel_in_tx(
tx: &mut PgTx<'_>,
job_id: &OpaqueId,
) -> Result<CancelOutcome, DurableError> {
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)
}
pub async fn cooperative_stop(
&self,
authority: &AttemptAuthority,
) -> Result<CooperativeStopOutcome, AttemptError> {
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?;
if outcome == CooperativeStopOutcome::Stopped {
let _ = self
.coordination
.release(authority.job_id.as_str(), &lease_owner(authority))
.await;
}
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)?,
))
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
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;
}
Ok(decision)
}
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())
}
}