ms2k/remora
SHA256
Watch
0
Code Releases Activity

Add authenticated internal worker API

Browse source
This commit is contained in:
ms2k 2026-07-17 09:16:46 +09:00
commit ebf818d23c
4 changed files with 2745 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -20,6 +20,7 @@ pub use recovery_types::{
};
pub use service::{AttemptConfig, AttemptService};
pub use types::{
ActivationOutcome, AttemptAuthority, AttemptFaultInjector, AttemptFaultPoint, AttemptState,
ClaimOutcome, ClaimRejection, RecoveryOutcome, RenewalOutcome,
ActivationMutationOutcome, ActivationOutcome, AttemptAuthority, AttemptFaultInjector,
AttemptFaultPoint, AttemptState, ClaimOutcome, ClaimRejection, ExecutionEnvelopeRecord,
LeaseProjection, OutcomePublication, RecoveryOutcome, RenewalMutationOutcome, RenewalOutcome,
};

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

@ -14,7 +14,7 @@ use time::OffsetDateTime;
use crate::{
AttemptAuthority, AttemptError, AttemptMutationOutcome, AttemptService, CheckpointProjection,
CheckpointRecord, CommittedResultRecord, DurableDescriptor, FailureRecord,
JobMutationProjection, MutationRejection, ProgressRecord,
JobMutationProjection, MutationRejection, OutcomePublication, ProgressRecord,
};
impl AttemptService {
@ -64,6 +64,18 @@ impl AttemptService {
.await
}
pub async fn publish_success_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
result: PublicationInput,
) -> Result<OutcomePublication, AttemptError> {
let outcome = self
.publish_success(authority, operation_id, result)
.await?;
self.publication_outcome(operation_id, outcome).await
}
pub async fn persist_checkpoint(
&self,
authority: &AttemptAuthority,
@ -124,6 +136,16 @@ impl AttemptService {
.await
}
pub async fn publish_failure_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
class: RuntimeFailureClass,
) -> Result<OutcomePublication, AttemptError> {
let outcome = self.publish_failure(authority, operation_id, class).await?;
self.publication_outcome(operation_id, outcome).await
}
pub async fn report_progress(
&self,
authority: &AttemptAuthority,
@ -227,7 +249,7 @@ impl AttemptService {
.execute(tx.as_mut())
.await?;
enqueue(tx, "job.updated", json!({ "job_id": authority.job_id })).await?;
accepted()
accepted_phase(JobPhase::Succeeded)
})
},
)
@ -307,7 +329,11 @@ impl AttemptService {
json!({ "job_id": authority.job_id }),
)
.await?;
accepted()
accepted_phase(if retry {
JobPhase::RetryWait
} else {
JobPhase::Failed
})
})
},
)
@ -329,6 +355,43 @@ impl AttemptService {
row.map(ProjectionRow::try_into_projection).transpose()
}
async fn publication_outcome(
&self,
operation_id: &OpaqueId,
outcome: AttemptMutationOutcome,
) -> Result<OutcomePublication, AttemptError> {
match outcome {
AttemptMutationOutcome::Rejected {
reason,
disposition,
} => Ok(OutcomePublication::Rejected {
reason,
disposition,
}),
AttemptMutationOutcome::Accepted(disposition) => {
let response = sqlx::query_scalar::<_, Json<Value>>(
"SELECT response_body FROM mutation_records
WHERE scope_subject = $1 AND scope_operation = 'attempt.mutation'
AND operation_key = $2",
)
.bind(self.operation_subject.as_str())
.bind(operation_id.as_str())
.fetch_one(&self.pool)
.await
.map_err(DurableError::from)?;
let phase = serde_json::from_value(
response
.0
.get("phase")
.cloned()
.ok_or_else(|| AttemptError::Invalid("missing outcome phase".into()))?,
)
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
Ok(OutcomePublication::Accepted { disposition, phase })
}
}
}
async fn execute_durable_mutation<T, D>(
&self,
authority: &AttemptAuthority,
@ -423,6 +486,13 @@ fn accepted() -> Result<(MutationStatus, Value), DurableError> {
Ok((MutationStatus::Accepted, json!({ "accepted": true })))
}
fn accepted_phase(phase: JobPhase) -> Result<(MutationStatus, Value), DurableError> {
Ok((
MutationStatus::Accepted,
json!({ "accepted": true, "phase": phase }),
))
}
fn rejected(reason: MutationRejection) -> Result<(MutationStatus, Value), DurableError> {
Ok((MutationStatus::Rejected, json!({ "reason": reason })))
}

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

@ -14,9 +14,10 @@ use sqlx::{PgPool, types::Json};
use time::OffsetDateTime;
use crate::{
ActivationOutcome, AttemptAuthority, AttemptError, AttemptFaultInjector, AttemptFaultPoint,
AttemptState, CheckpointProjection, ClaimOutcome, ClaimRejection, RecoveryOutcome,
RenewalOutcome,
ActivationMutationOutcome, ActivationOutcome, AttemptAuthority, AttemptError,
AttemptFaultInjector, AttemptFaultPoint, AttemptState, CheckpointProjection, ClaimOutcome,
ClaimRejection, ExecutionEnvelopeRecord, LeaseProjection, RecoveryOutcome,
RenewalMutationOutcome, RenewalOutcome,
};
struct NoFaults;
@ -179,6 +180,23 @@ impl AttemptService {
authority: &AttemptAuthority,
operation_id: &OpaqueId,
) -> Result<ActivationOutcome, AttemptError> {
Ok(
match self.activate_operation(authority, operation_id).await? {
ActivationMutationOutcome::Activated { disposition, .. } => {
ActivationOutcome::Activated(disposition)
}
ActivationMutationOutcome::PreconditionFailed => {
ActivationOutcome::PreconditionFailed
}
},
)
}
pub async fn activate_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
) -> Result<ActivationMutationOutcome, AttemptError> {
validate_authority(authority)?;
let request_fingerprint = fingerprint(&ActivationRequest {
job_id: &authority.job_id,
@ -203,16 +221,17 @@ impl AttemptService {
request_fingerprint,
move |tx| {
Box::pin(async move {
let accepted =
let lease =
activate_attempt(tx, &coordination, &authority_for_effect, faults.as_ref())
.await?;
Ok((
if accepted {
if lease.is_some() {
MutationStatus::Accepted
} else {
MutationStatus::Rejected
},
json!({ "activated": accepted }),
serde_json::to_value(ActivationStored { lease })
.map_err(DurableError::effect)?,
))
})
},
@ -225,7 +244,14 @@ impl AttemptService {
.exact_authority_is_executing_and_live(authority)
.await?
{
return Ok(ActivationOutcome::Activated(MutationDisposition::First));
let lease = self
.live_lease(authority)
.await?
.ok_or_else(|| AttemptError::Invalid("live lease disappeared".into()))?;
return Ok(ActivationMutationOutcome::Activated {
disposition: MutationDisposition::First,
lease,
});
}
self.release_exact(authority).await;
return Err(error.into());
@ -237,16 +263,24 @@ impl AttemptService {
}
if outcome.status == MutationStatus::Rejected {
self.release_exact(authority).await;
return Ok(ActivationOutcome::PreconditionFailed);
return Ok(ActivationMutationOutcome::PreconditionFailed);
}
if !self
.exact_authority_is_executing_and_live(authority)
.await?
{
self.release_exact(authority).await;
return Ok(ActivationOutcome::PreconditionFailed);
return Ok(ActivationMutationOutcome::PreconditionFailed);
}
Ok(ActivationOutcome::Activated(outcome.disposition))
let stored: ActivationStored = serde_json::from_value(outcome.response)
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
let lease = stored
.lease
.ok_or_else(|| AttemptError::Invalid("accepted activation omitted its lease".into()))?;
Ok(ActivationMutationOutcome::Activated {
disposition: outcome.disposition,
lease,
})
}
pub async fn renew(
@ -269,6 +303,103 @@ impl AttemptService {
Ok(RenewalOutcome::Renewed)
}
pub async fn renew_operation(
&self,
authority: &AttemptAuthority,
operation_id: &OpaqueId,
) -> Result<RenewalMutationOutcome, AttemptError> {
validate_authority(authority)?;
let request_fingerprint = fingerprint(&ActivationRequest {
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 eligibility = authority_binding(authority);
let authority = authority.clone();
let coordination = self.coordination.clone();
let faults = self.faults.clone();
let outcome = execute_mutation(
&self.pool,
MutationScope {
subject: &self.operation_subject,
operation: "attempt.mutation",
key: operation_id.as_str(),
},
"renew_attempt",
&eligibility,
request_fingerprint,
move |tx| {
Box::pin(async move {
let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
return Ok((
MutationStatus::Rejected,
serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
disposition: None,
})
.map_err(DurableError::effect)?,
));
};
if row.state != AttemptState::Executing
|| !row.matches(
&authority.job_id,
authority.fence,
&authority.worker_incarnation,
)
{
return Ok((
MutationStatus::Rejected,
serde_json::to_value(RenewalMutationOutcome::AuthorityLost {
disposition: None,
})
.map_err(DurableError::effect)?,
));
}
faults.hit(AttemptFaultPoint::RenewalBeforeLease);
let expected = lease_owner(&authority).map_err(DurableError::effect)?;
let renewed = coordination
.renew(authority.job_id.as_str(), &expected)
.await
.map_err(DurableError::effect)?;
faults.hit(AttemptFaultPoint::RenewalAfterLease);
let live = coordination
.get_live_lease(authority.job_id.as_str())
.await
.map_err(DurableError::effect)?;
let response = match live
.filter(|lease| renewed.ok() && lease.record.owner() == expected)
{
Some(lease) => RenewalMutationOutcome::Renewed {
lease: LeaseProjection {
expires_at_ms: lease.record.expires_at,
observed_at_ms: lease.now_ms,
},
disposition: None,
},
None => RenewalMutationOutcome::AuthorityLost { disposition: None },
};
let status = if matches!(response, RenewalMutationOutcome::Renewed { .. }) {
MutationStatus::Accepted
} else {
MutationStatus::Rejected
};
Ok((
status,
serde_json::to_value(response).map_err(DurableError::effect)?,
))
})
},
)
.await?;
if outcome.disposition == MutationDisposition::Conflict {
return Err(AttemptError::OperationConflict);
}
let response: RenewalMutationOutcome = serde_json::from_value(outcome.response)
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
Ok(response.with_disposition(outcome.disposition))
}
pub async fn recover_inactive(
&self,
authority: &AttemptAuthority,
@ -372,6 +503,45 @@ impl AttemptService {
.map_err(Into::into)
}
pub async fn historical_authority(
&self,
attempt_id: &OpaqueId,
) -> Result<Option<AttemptAuthority>, AttemptError> {
load_historical_authority(&self.pool, attempt_id)
.await
.map_err(Into::into)
}
pub async fn execution_envelope(
&self,
attempt_id: &OpaqueId,
) -> Result<Option<ExecutionEnvelopeRecord>, AttemptError> {
let value = sqlx::query_scalar::<_, Option<Json<ExecutionEnvelopeRecord>>>(
"SELECT execution_envelope FROM attempts WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_optional(&self.pool)
.await?
.flatten();
Ok(value.map(|value| value.0))
}
pub async fn live_lease(
&self,
authority: &AttemptAuthority,
) -> Result<Option<LeaseProjection>, AttemptError> {
let expected = lease_owner(authority)?;
Ok(self
.coordination
.get_live_lease(authority.job_id.as_str())
.await?
.filter(|lease| lease.record.owner() == expected)
.map(|lease| LeaseProjection {
expires_at_ms: lease.record.expires_at,
observed_at_ms: lease.now_ms,
}))
}
pub fn release_local(&self, job_id: &OpaqueId) -> bool {
self.local_permits.release_all(job_id.as_str())
}
@ -451,6 +621,11 @@ struct ActivationRequest<'a> {
worker_incarnation: &'a OpaqueId,
}
#[derive(serde::Deserialize, Serialize)]
struct ActivationStored {
lease: Option<LeaseProjection>,
}
struct AuthorityRow {
job_id: OpaqueId,
fence: u64,
@ -534,11 +709,14 @@ async fn authorize_attempt(
_ => 0,
};
let now = OffsetDateTime::now_utc();
let envelope =
load_execution_envelope_record(tx, input.job_id, input.attempt_id, fence, input.worker)
.await?;
sqlx::query(
"INSERT INTO attempts (
attempt_id, job_id, fence, worker_incarnation, state,
initial_progress, authorized_at, executing_at, reaped_at
) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL)",
initial_progress, authorized_at, executing_at, reaped_at, execution_envelope
) VALUES ($1, $2, $3, $4, 'authorized', $5, $6, NULL, NULL, $7)",
)
.bind(input.attempt_id.as_str())
.bind(input.job_id.as_str())
@ -546,6 +724,7 @@ async fn authorize_attempt(
.bind(input.worker.as_str())
.bind(i64::try_from(initial_progress).map_err(DurableError::effect)?)
.bind(now)
.bind(Json(&envelope))
.execute(tx.as_mut())
.await?;
sqlx::query(
@ -572,21 +751,79 @@ async fn authorize_attempt(
})
}
async fn load_execution_envelope_record(
tx: &mut remora_durable::PgTx<'_>,
job_id: &OpaqueId,
attempt_id: &OpaqueId,
fence: u64,
worker: &OpaqueId,
) -> Result<ExecutionEnvelopeRecord, DurableError> {
type Row = (
Json<serde_json::Value>,
String,
String,
String,
Json<Vec<String>>,
String,
Option<String>,
Option<String>,
Option<Json<crate::DurableDescriptor>>,
Option<Json<CheckpointProjection>>,
);
let row = sqlx::query_as::<_, Row>(
"SELECT j.source_projection, j.sensitive_source_ref, j.target_store,
j.target_bucket, j.target_prefix, j.execution_grant_id,
g.credential_id, g.credential_version_id, j.resolved_source, j.checkpoint
FROM jobs j
JOIN execution_grants g ON g.grant_id = j.execution_grant_id
WHERE j.job_id = $1",
)
.bind(job_id.as_str())
.fetch_one(tx.as_mut())
.await?;
let grant_id = OpaqueId::parse(row.5).map_err(DurableError::effect)?;
Ok(ExecutionEnvelopeRecord {
id: grant_id.clone(),
job_id: job_id.clone(),
attempt_id: attempt_id.clone(),
fence,
worker_incarnation: worker.clone(),
source: row.0.0,
sensitive_source_ref: OpaqueId::parse(row.1).map_err(DurableError::effect)?,
credential_id: row
.6
.map(OpaqueId::parse)
.transpose()
.map_err(DurableError::effect)?,
credential_version_id: row
.7
.map(OpaqueId::parse)
.transpose()
.map_err(DurableError::effect)?,
resolved_source: row.8.map(|value| value.0),
target_store: row.2,
target_bucket: row.3,
target_prefix: row.4.0,
execution_grant_id: grant_id,
checkpoint: row.9.map(|value| value.0),
})
}
async fn activate_attempt(
tx: &mut remora_durable::PgTx<'_>,
coordination: &CoordinationStore,
authority: &AttemptAuthority,
faults: &dyn AttemptFaultInjector,
) -> Result<bool, DurableError> {
) -> Result<Option<LeaseProjection>, DurableError> {
let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
return Ok(false);
return Ok(None);
};
if !row.matches(
&authority.job_id,
authority.fence,
&authority.worker_incarnation,
) {
return Ok(false);
return Ok(None);
}
if row.state == AttemptState::Executing {
let live = coordination
@ -595,16 +832,28 @@ async fn activate_attempt(
.map_err(DurableError::effect)?;
let expected = lease_owner(authority).map_err(DurableError::effect)?;
if !live.is_some_and(|lease| lease.record.owner() == expected) {
return Ok(false);
return Ok(None);
}
return Ok(coordination
if !coordination
.renew(authority.job_id.as_str(), &expected)
.await
.map_err(DurableError::effect)?
.ok());
.ok()
{
return Ok(None);
}
return Ok(coordination
.get_live_lease(authority.job_id.as_str())
.await
.map_err(DurableError::effect)?
.filter(|lease| lease.record.owner() == expected)
.map(|lease| LeaseProjection {
expires_at_ms: lease.record.expires_at,
observed_at_ms: lease.now_ms,
}));
}
if row.state != AttemptState::Authorized {
return Ok(false);
return Ok(None);
}
faults.hit(AttemptFaultPoint::ActivationBeforeLease);
let owner = lease_owner(authority).map_err(DurableError::effect)?;
@ -620,7 +869,7 @@ async fn activate_attempt(
.await
.map_err(DurableError::effect)?;
if !coordinated.ok() {
return Ok(false);
return Ok(None);
}
faults.hit(AttemptFaultPoint::ActivationAfterLease);
let live = coordination
@ -628,9 +877,9 @@ async fn activate_attempt(
.await
.map_err(DurableError::effect)?;
let expected = lease_owner(authority).map_err(DurableError::effect)?;
if !live.is_some_and(|lease| lease.record.owner() == expected) {
return Ok(false);
}
let Some(live) = live.filter(|lease| lease.record.owner() == expected) else {
return Ok(None);
};
let now = OffsetDateTime::now_utc();
sqlx::query(
"UPDATE attempts SET state = 'executing', executing_at = $2
@ -640,7 +889,10 @@ async fn activate_attempt(
.bind(now)
.execute(tx.as_mut())
.await?;
Ok(true)
Ok(Some(LeaseProjection {
expires_at_ms: live.record.expires_at,
observed_at_ms: live.now_ms,
}))
}
async fn lock_authority(
@ -692,6 +944,30 @@ async fn load_authority(
.transpose()
}
async fn load_historical_authority(
pool: &PgPool,
attempt_id: &OpaqueId,
) -> Result<Option<AttemptAuthority>, DurableError> {
let row = sqlx::query_as::<_, (String, String, i64, String, String)>(
"SELECT job_id, attempt_id, fence, worker_incarnation, state
FROM attempts WHERE attempt_id = $1",
)
.bind(attempt_id.as_str())
.fetch_optional(pool)
.await?;
row.map(|(job, attempt, fence, worker, state)| {
Ok(AttemptAuthority {
job_id: OpaqueId::parse(job).map_err(DurableError::effect)?,
attempt_id: OpaqueId::parse(attempt).map_err(DurableError::effect)?,
fence: u64::try_from(fence).map_err(DurableError::effect)?,
worker_incarnation: OpaqueId::parse(worker).map_err(DurableError::effect)?,
state: AttemptState::parse(&state)
.ok_or_else(|| DurableError::Config(format!("unknown attempt state '{state}'")))?,
})
})
.transpose()
}
fn authority_row(row: (String, i64, String, String)) -> Result<AuthorityRow, DurableError> {
Ok(AuthorityRow {
job_id: OpaqueId::parse(row.0).map_err(DurableError::effect)?,

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

@ -1,4 +1,4 @@
use remora_domain::OpaqueId;
use remora_domain::{JobPhase, OpaqueId};
use remora_durable::MutationDisposition;
use serde::{Deserialize, Serialize};
@ -95,12 +95,83 @@ pub enum ActivationOutcome {
PreconditionFailed,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ActivationMutationOutcome {
Activated {
disposition: MutationDisposition,
lease: LeaseProjection,
},
PreconditionFailed,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RenewalOutcome {
Renewed,
AuthorityLost,
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct LeaseProjection {
pub expires_at_ms: u64,
pub observed_at_ms: u64,
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum RenewalMutationOutcome {
Renewed {
lease: LeaseProjection,
#[serde(skip)]
disposition: Option<MutationDisposition>,
},
AuthorityLost {
#[serde(skip)]
disposition: Option<MutationDisposition>,
},
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OutcomePublication {
Accepted {
disposition: MutationDisposition,
phase: JobPhase,
},
Rejected {
reason: crate::MutationRejection,
disposition: MutationDisposition,
},
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct ExecutionEnvelopeRecord {
pub id: OpaqueId,
pub job_id: OpaqueId,
pub attempt_id: OpaqueId,
pub fence: u64,
pub worker_incarnation: OpaqueId,
pub source: serde_json::Value,
pub sensitive_source_ref: OpaqueId,
pub credential_id: Option<OpaqueId>,
pub credential_version_id: Option<OpaqueId>,
pub resolved_source: Option<crate::DurableDescriptor>,
pub target_store: String,
pub target_bucket: String,
pub target_prefix: Vec<String>,
pub execution_grant_id: OpaqueId,
pub checkpoint: Option<crate::CheckpointProjection>,
}
impl RenewalMutationOutcome {
pub(crate) fn with_disposition(mut self, value: MutationDisposition) -> Self {
match &mut self {
Self::Renewed { disposition, .. } | Self::AuthorityLost { disposition } => {
*disposition = Some(value);
}
}
self
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RecoveryOutcome {
StillOwned,