1366 lines
48 KiB
Rust
1366 lines
48 KiB
Rust
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
|
|
use remora_coordination::{
|
|
ActivateOptions, CoordinationStore, LeaseOwner, LocalPermitGate, MAX_SAFE_INTEGER, OpCode,
|
|
};
|
|
use remora_domain::{OpaqueId, fingerprint, ports::MutationScope};
|
|
use remora_durable::{
|
|
DurableError, MutationDisposition, MutationStatus, enqueue, execute_mutation, with_transaction,
|
|
};
|
|
use serde::Serialize;
|
|
use serde_json::json;
|
|
use sqlx::{PgPool, types::Json};
|
|
use time::OffsetDateTime;
|
|
|
|
use crate::{
|
|
ActivationMutationOutcome, ActivationOutcome, AttemptAuthority, AttemptError,
|
|
AttemptFaultInjector, AttemptFaultPoint, AttemptState, CheckpointProjection, ClaimOutcome,
|
|
ClaimRejection, ExecutionEnvelopeRecord, LeaseProjection, ProgressRecord, ProgressSnapshot,
|
|
RecoveryOutcome, RenewalMutationOutcome, RenewalOutcome,
|
|
};
|
|
|
|
struct NoFaults;
|
|
|
|
impl AttemptFaultInjector for NoFaults {
|
|
fn hit(&self, _point: AttemptFaultPoint) {}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct AttemptConfig {
|
|
pub max_attempts: u64,
|
|
pub local_limit: u64,
|
|
}
|
|
|
|
impl Default for AttemptConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_attempts: 3,
|
|
local_limit: 4,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AttemptService {
|
|
pub(crate) pool: PgPool,
|
|
pub(crate) coordination: CoordinationStore,
|
|
local_permits: LocalPermitGate,
|
|
pub(crate) operation_subject: OpaqueId,
|
|
pub(crate) faults: Arc<dyn AttemptFaultInjector>,
|
|
pub(crate) config: AttemptConfig,
|
|
}
|
|
|
|
impl AttemptService {
|
|
pub fn new(
|
|
pool: PgPool,
|
|
coordination: CoordinationStore,
|
|
config: AttemptConfig,
|
|
) -> Result<Self, AttemptError> {
|
|
Self::new_with_fault_injector(pool, coordination, config, Arc::new(NoFaults))
|
|
}
|
|
|
|
pub fn new_with_fault_injector(
|
|
pool: PgPool,
|
|
coordination: CoordinationStore,
|
|
config: AttemptConfig,
|
|
faults: Arc<dyn AttemptFaultInjector>,
|
|
) -> Result<Self, AttemptError> {
|
|
if config.max_attempts == 0 || config.max_attempts > MAX_SAFE_INTEGER {
|
|
return Err(AttemptError::Invalid(
|
|
"max_attempts must fit the coordination safe-integer range".into(),
|
|
));
|
|
}
|
|
if config.local_limit == 0 {
|
|
return Err(AttemptError::Invalid(
|
|
"local_limit must be greater than zero".into(),
|
|
));
|
|
}
|
|
Ok(Self {
|
|
pool,
|
|
coordination,
|
|
local_permits: LocalPermitGate::new(config.local_limit),
|
|
operation_subject: OpaqueId::parse("remora_attempt_service")
|
|
.map_err(|error| AttemptError::Invalid(error.to_string()))?,
|
|
faults,
|
|
config,
|
|
})
|
|
}
|
|
|
|
pub async fn progress_snapshot(
|
|
&self,
|
|
job_id: &OpaqueId,
|
|
) -> Result<Option<ProgressSnapshot>, AttemptError> {
|
|
let Some(hint) = self.coordination.get_progress(job_id.as_str()).await? else {
|
|
return Ok(None);
|
|
};
|
|
Ok(Some(ProgressSnapshot {
|
|
attempt_id: OpaqueId::parse(hint.attempt_id)
|
|
.map_err(|error| AttemptError::Invalid(error.to_string()))?,
|
|
fence: hint.fence,
|
|
progress: ProgressRecord {
|
|
completed_bytes: hint.bytes,
|
|
total_bytes: hint.total_bytes,
|
|
completed_members: hint.completed_members,
|
|
total_members: hint.total_members,
|
|
},
|
|
observed_at: OffsetDateTime::now_utc(),
|
|
}))
|
|
}
|
|
|
|
pub async fn claim(
|
|
&self,
|
|
job_id: &OpaqueId,
|
|
worker_incarnation: &OpaqueId,
|
|
operation_id: &OpaqueId,
|
|
) -> Result<ClaimOutcome, AttemptError> {
|
|
if self.local_permits.acquire(job_id.as_str()).is_err() {
|
|
return Ok(ClaimOutcome::Rejected {
|
|
reason: ClaimRejection::LocalCapacity,
|
|
disposition: None,
|
|
});
|
|
}
|
|
let request = ClaimRequest {
|
|
job_id,
|
|
worker_incarnation,
|
|
};
|
|
let request_fingerprint =
|
|
fingerprint(&request).map_err(|error| AttemptError::Invalid(error.to_string()))?;
|
|
let attempt_id = new_id('a')?;
|
|
let job = job_id.clone();
|
|
let worker = worker_incarnation.clone();
|
|
let max_attempts = self.config.max_attempts;
|
|
let coordination = self.coordination.clone();
|
|
let faults = self.faults.clone();
|
|
let reserved = Arc::new(AtomicBool::new(false));
|
|
let effect_reserved = reserved.clone();
|
|
let outcome = execute_mutation(
|
|
&self.pool,
|
|
MutationScope {
|
|
subject: &self.operation_subject,
|
|
operation: "attempt.mutation",
|
|
key: operation_id.as_str(),
|
|
},
|
|
"claim_attempt",
|
|
&[],
|
|
request_fingerprint,
|
|
move |tx| {
|
|
Box::pin(async move {
|
|
let decision = authorize_attempt(
|
|
tx,
|
|
AuthorizeInput {
|
|
coordination: &coordination,
|
|
job_id: &job,
|
|
attempt_id: &attempt_id,
|
|
worker: &worker,
|
|
max_attempts,
|
|
faults: faults.as_ref(),
|
|
reserved: effect_reserved.as_ref(),
|
|
},
|
|
)
|
|
.await?;
|
|
let status = if matches!(decision, ClaimOutcome::Claimed { .. }) {
|
|
MutationStatus::Accepted
|
|
} else {
|
|
MutationStatus::Rejected
|
|
};
|
|
let response = serde_json::to_value(decision)
|
|
.map_err(|error| DurableError::Effect(error.to_string()))?;
|
|
Ok((status, response))
|
|
})
|
|
},
|
|
)
|
|
.await;
|
|
let outcome = match outcome {
|
|
Ok(outcome) => outcome,
|
|
Err(error) => {
|
|
if reserved.load(Ordering::Acquire) {
|
|
self.compensate_reservation(job_id, worker_incarnation)
|
|
.await;
|
|
} else {
|
|
self.local_permits.release(job_id.as_str());
|
|
}
|
|
return Err(error.into());
|
|
}
|
|
};
|
|
self.faults.hit(AttemptFaultPoint::ClaimAfterDurable);
|
|
if outcome.disposition == MutationDisposition::Conflict {
|
|
self.local_permits.release(job_id.as_str());
|
|
return Err(AttemptError::OperationConflict);
|
|
}
|
|
let decision: ClaimOutcome = serde_json::from_value(outcome.response)
|
|
.map_err(|error| AttemptError::Invalid(error.to_string()))?;
|
|
if matches!(decision, ClaimOutcome::Rejected { .. }) {
|
|
self.local_permits.release(job_id.as_str());
|
|
}
|
|
Ok(decision.with_disposition(outcome.disposition))
|
|
}
|
|
|
|
pub async fn activate(
|
|
&self,
|
|
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,
|
|
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_for_effect = 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(),
|
|
},
|
|
"activate_attempt",
|
|
&eligibility,
|
|
request_fingerprint,
|
|
move |tx| {
|
|
Box::pin(async move {
|
|
let lease =
|
|
activate_attempt(tx, &coordination, &authority_for_effect, faults.as_ref())
|
|
.await?;
|
|
Ok((
|
|
if lease.is_some() {
|
|
MutationStatus::Accepted
|
|
} else {
|
|
MutationStatus::Rejected
|
|
},
|
|
serde_json::to_value(ActivationStored { lease })
|
|
.map_err(DurableError::effect)?,
|
|
))
|
|
})
|
|
},
|
|
)
|
|
.await;
|
|
let outcome = match outcome {
|
|
Ok(outcome) => outcome,
|
|
Err(error) => {
|
|
if self
|
|
.exact_authority_is_executing_and_live(authority)
|
|
.await?
|
|
{
|
|
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());
|
|
}
|
|
};
|
|
self.faults.hit(AttemptFaultPoint::ActivationAfterDurable);
|
|
if outcome.disposition == MutationDisposition::Conflict {
|
|
return Err(AttemptError::OperationConflict);
|
|
}
|
|
if outcome.status == MutationStatus::Rejected {
|
|
self.release_exact(authority).await;
|
|
return Ok(ActivationMutationOutcome::PreconditionFailed);
|
|
}
|
|
if !self
|
|
.exact_authority_is_executing_and_live(authority)
|
|
.await?
|
|
{
|
|
self.release_exact(authority).await;
|
|
return Ok(ActivationMutationOutcome::PreconditionFailed);
|
|
}
|
|
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(
|
|
&self,
|
|
authority: &AttemptAuthority,
|
|
) -> Result<RenewalOutcome, AttemptError> {
|
|
validate_authority(authority)?;
|
|
if !self.exact_authority_is_executing(authority).await? {
|
|
return Ok(RenewalOutcome::AuthorityLost);
|
|
}
|
|
self.faults.hit(AttemptFaultPoint::RenewalBeforeLease);
|
|
let result = self
|
|
.coordination
|
|
.renew(authority.job_id.as_str(), &lease_owner(authority)?)
|
|
.await?;
|
|
self.faults.hit(AttemptFaultPoint::RenewalAfterLease);
|
|
if !result.ok() || !self.exact_authority_is_executing(authority).await? {
|
|
return Ok(RenewalOutcome::AuthorityLost);
|
|
}
|
|
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,
|
|
) -> Result<RecoveryOutcome, AttemptError> {
|
|
validate_authority(authority)?;
|
|
let job_id = authority.job_id.clone();
|
|
let attempt_id = authority.attempt_id.clone();
|
|
let worker = authority.worker_incarnation.clone();
|
|
let fence = authority.fence;
|
|
let coordination = self.coordination.clone();
|
|
let faults = self.faults.clone();
|
|
let outcome = with_transaction(&self.pool, move |tx| {
|
|
Box::pin(async move {
|
|
let row = lock_authority(tx, &attempt_id).await?;
|
|
let Some(row) = row else {
|
|
return Ok(RecoveryOutcome::AlreadyInactive);
|
|
};
|
|
if !row.matches(&job_id, fence, &worker) {
|
|
return Ok(RecoveryOutcome::AlreadyInactive);
|
|
}
|
|
if row.state == AttemptState::Executing {
|
|
return Ok(RecoveryOutcome::StillOwned);
|
|
}
|
|
if row.state != AttemptState::Authorized {
|
|
return Ok(RecoveryOutcome::AlreadyInactive);
|
|
}
|
|
faults.hit(AttemptFaultPoint::RecoveryBeforeRevalidation);
|
|
if coordination
|
|
.get_reservation(job_id.as_str())
|
|
.await
|
|
.map_err(DurableError::effect)?
|
|
.is_some()
|
|
{
|
|
return Ok(RecoveryOutcome::StillOwned);
|
|
}
|
|
if let Some(lease) = coordination
|
|
.get_live_lease(job_id.as_str())
|
|
.await
|
|
.map_err(DurableError::effect)?
|
|
{
|
|
if lease.record.owner()
|
|
== LeaseOwner::new(attempt_id.as_str(), fence, worker.as_str())
|
|
.map_err(DurableError::effect)?
|
|
{
|
|
return Ok(RecoveryOutcome::StillOwned);
|
|
}
|
|
return Err(DurableError::Effect(
|
|
"foreign live lease conflicts with durable authority".into(),
|
|
));
|
|
}
|
|
let now = OffsetDateTime::now_utc();
|
|
faults.hit(AttemptFaultPoint::RecoveryBeforeDurable);
|
|
sqlx::query(
|
|
"UPDATE attempts
|
|
SET state = 'reaped', reaped_at = $2
|
|
WHERE attempt_id = $1 AND state = 'authorized'",
|
|
)
|
|
.bind(attempt_id.as_str())
|
|
.bind(now)
|
|
.execute(tx.as_mut())
|
|
.await?;
|
|
let changed = sqlx::query(
|
|
"UPDATE jobs
|
|
SET phase = 'queued', authority_attempt_id = NULL,
|
|
dispatch_needed = TRUE, revision = revision + 1, updated_at = $2
|
|
WHERE job_id = $1 AND authority_attempt_id = $3
|
|
AND current_fence = $4 AND phase = 'running'",
|
|
)
|
|
.bind(job_id.as_str())
|
|
.bind(now)
|
|
.bind(attempt_id.as_str())
|
|
.bind(i64::try_from(fence).map_err(DurableError::effect)?)
|
|
.execute(tx.as_mut())
|
|
.await?;
|
|
if changed.rows_affected() != 1 {
|
|
return Err(DurableError::Effect(
|
|
"attempt authority changed while recovery held its lock".into(),
|
|
));
|
|
}
|
|
enqueue(tx, "job.dispatch", json!({ "job_id": job_id })).await?;
|
|
Ok(RecoveryOutcome::Reaped)
|
|
})
|
|
})
|
|
.await?;
|
|
self.faults.hit(AttemptFaultPoint::RecoveryAfterDurable);
|
|
if matches!(
|
|
outcome,
|
|
RecoveryOutcome::Reaped | RecoveryOutcome::AlreadyInactive
|
|
) {
|
|
self.release_exact(authority).await;
|
|
}
|
|
Ok(outcome)
|
|
}
|
|
|
|
pub async fn authority(
|
|
&self,
|
|
attempt_id: &OpaqueId,
|
|
) -> Result<Option<AttemptAuthority>, AttemptError> {
|
|
load_authority(&self.pool, attempt_id)
|
|
.await
|
|
.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())
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn local_used(&self) -> u64 {
|
|
self.local_permits.used()
|
|
}
|
|
|
|
async fn compensate_reservation(&self, job_id: &OpaqueId, worker: &OpaqueId) {
|
|
let _ = self
|
|
.coordination
|
|
.free_reserve(job_id.as_str(), Some(worker.as_str()))
|
|
.await;
|
|
self.local_permits.release(job_id.as_str());
|
|
}
|
|
|
|
pub(crate) async fn release_exact(&self, authority: &AttemptAuthority) {
|
|
if let Ok(owner) = lease_owner(authority) {
|
|
let _ = self
|
|
.coordination
|
|
.release(authority.job_id.as_str(), &owner)
|
|
.await;
|
|
}
|
|
}
|
|
|
|
async fn exact_authority_is_executing(
|
|
&self,
|
|
authority: &AttemptAuthority,
|
|
) -> Result<bool, AttemptError> {
|
|
Ok(load_authority(&self.pool, &authority.attempt_id)
|
|
.await?
|
|
.is_some_and(|current| current == authority.clone_with_state(AttemptState::Executing)))
|
|
}
|
|
|
|
async fn exact_authority_is_executing_and_live(
|
|
&self,
|
|
authority: &AttemptAuthority,
|
|
) -> Result<bool, AttemptError> {
|
|
if !self.exact_authority_is_executing(authority).await? {
|
|
return Ok(false);
|
|
}
|
|
let Some(lease) = self
|
|
.coordination
|
|
.get_live_lease(authority.job_id.as_str())
|
|
.await?
|
|
else {
|
|
return Ok(false);
|
|
};
|
|
Ok(lease.record.owner() == lease_owner(authority)?)
|
|
}
|
|
}
|
|
|
|
impl AttemptAuthority {
|
|
fn clone_with_state(&self, state: AttemptState) -> Self {
|
|
Self {
|
|
job_id: self.job_id.clone(),
|
|
attempt_id: self.attempt_id.clone(),
|
|
fence: self.fence,
|
|
worker_incarnation: self.worker_incarnation.clone(),
|
|
state,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ClaimRequest<'a> {
|
|
job_id: &'a OpaqueId,
|
|
worker_incarnation: &'a OpaqueId,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ActivationRequest<'a> {
|
|
job_id: &'a OpaqueId,
|
|
attempt_id: &'a OpaqueId,
|
|
fence: u64,
|
|
worker_incarnation: &'a OpaqueId,
|
|
}
|
|
|
|
#[derive(serde::Deserialize, Serialize)]
|
|
struct ActivationStored {
|
|
lease: Option<LeaseProjection>,
|
|
}
|
|
|
|
struct AuthorityRow {
|
|
job_id: OpaqueId,
|
|
fence: u64,
|
|
worker: OpaqueId,
|
|
state: AttemptState,
|
|
}
|
|
|
|
impl AuthorityRow {
|
|
fn matches(&self, job_id: &OpaqueId, fence: u64, worker: &OpaqueId) -> bool {
|
|
self.job_id == *job_id && self.fence == fence && self.worker == *worker
|
|
}
|
|
}
|
|
|
|
struct AuthorizeInput<'a> {
|
|
coordination: &'a CoordinationStore,
|
|
job_id: &'a OpaqueId,
|
|
attempt_id: &'a OpaqueId,
|
|
worker: &'a OpaqueId,
|
|
max_attempts: u64,
|
|
faults: &'a dyn AttemptFaultInjector,
|
|
reserved: &'a AtomicBool,
|
|
}
|
|
|
|
async fn authorize_attempt(
|
|
tx: &mut remora_durable::PgTx<'_>,
|
|
input: AuthorizeInput<'_>,
|
|
) -> Result<ClaimOutcome, DurableError> {
|
|
let row = sqlx::query_as::<
|
|
_,
|
|
(
|
|
String,
|
|
bool,
|
|
i64,
|
|
Option<String>,
|
|
Option<Json<crate::DurableDescriptor>>,
|
|
Option<Json<CheckpointProjection>>,
|
|
Option<String>,
|
|
),
|
|
>(
|
|
"SELECT phase, dispatch_needed, current_fence, authority_attempt_id,
|
|
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,
|
|
checkpoint_owner,
|
|
)) = row
|
|
else {
|
|
return Ok(claim_rejected(ClaimRejection::JobUnavailable));
|
|
};
|
|
if phase != "queued" || !dispatch_needed || authority.is_some() {
|
|
return Ok(claim_rejected(ClaimRejection::JobUnavailable));
|
|
}
|
|
let current_fence = u64::try_from(current_fence).map_err(DurableError::effect)?;
|
|
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
|
|
.reserve(input.job_id.as_str(), input.worker.as_str())
|
|
.await
|
|
.map_err(DurableError::effect)?;
|
|
if !reservation.ok() {
|
|
return Ok(claim_rejected(match reservation.code {
|
|
OpCode::FailFull => ClaimRejection::GlobalCapacity,
|
|
_ => ClaimRejection::ReservationHeld,
|
|
}));
|
|
}
|
|
input.reserved.store(true, Ordering::Release);
|
|
input.faults.hit(AttemptFaultPoint::ClaimAfterReserve);
|
|
let fence = current_fence + 1;
|
|
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?
|
|
}
|
|
_ => 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,
|
|
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,
|
|
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())
|
|
.bind(i64::try_from(fence).map_err(DurableError::effect)?)
|
|
.bind(input.worker.as_str())
|
|
.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,
|
|
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 {
|
|
authority: AttemptAuthority {
|
|
job_id: input.job_id.clone(),
|
|
attempt_id: input.attempt_id.clone(),
|
|
fence,
|
|
worker_incarnation: input.worker.clone(),
|
|
state: AttemptState::Authorized,
|
|
},
|
|
disposition: None,
|
|
})
|
|
}
|
|
|
|
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>,
|
|
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,
|
|
staging_origin_attempt_id: staging_origin_attempt_id.clone(),
|
|
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<Option<LeaseProjection>, DurableError> {
|
|
let Some(row) = lock_authority(tx, &authority.attempt_id).await? else {
|
|
return Ok(None);
|
|
};
|
|
if !row.matches(
|
|
&authority.job_id,
|
|
authority.fence,
|
|
&authority.worker_incarnation,
|
|
) {
|
|
return Ok(None);
|
|
}
|
|
if row.state == AttemptState::Executing {
|
|
let live = coordination
|
|
.get_live_lease(authority.job_id.as_str())
|
|
.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(None);
|
|
}
|
|
if !coordination
|
|
.renew(authority.job_id.as_str(), &expected)
|
|
.await
|
|
.map_err(DurableError::effect)?
|
|
.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(None);
|
|
}
|
|
faults.hit(AttemptFaultPoint::ActivationBeforeLease);
|
|
let owner = lease_owner(authority).map_err(DurableError::effect)?;
|
|
let coordinated = coordination
|
|
.activate(
|
|
authority.job_id.as_str(),
|
|
&owner,
|
|
ActivateOptions {
|
|
seed_progress: true,
|
|
require_reserve: true,
|
|
},
|
|
)
|
|
.await
|
|
.map_err(DurableError::effect)?;
|
|
if !coordinated.ok() {
|
|
return Ok(None);
|
|
}
|
|
faults.hit(AttemptFaultPoint::ActivationAfterLease);
|
|
let live = coordination
|
|
.get_live_lease(authority.job_id.as_str())
|
|
.await
|
|
.map_err(DurableError::effect)?;
|
|
let expected = lease_owner(authority).map_err(DurableError::effect)?;
|
|
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
|
|
WHERE attempt_id = $1 AND state = 'authorized'",
|
|
)
|
|
.bind(authority.attempt_id.as_str())
|
|
.bind(now)
|
|
.execute(tx.as_mut())
|
|
.await?;
|
|
Ok(Some(LeaseProjection {
|
|
expires_at_ms: live.record.expires_at,
|
|
observed_at_ms: live.now_ms,
|
|
}))
|
|
}
|
|
|
|
async fn lock_authority(
|
|
tx: &mut remora_durable::PgTx<'_>,
|
|
attempt_id: &OpaqueId,
|
|
) -> Result<Option<AuthorityRow>, DurableError> {
|
|
let row = sqlx::query_as::<_, (String, i64, String, String)>(
|
|
"SELECT a.job_id, a.fence, a.worker_incarnation, a.state
|
|
FROM attempts a
|
|
JOIN jobs j ON j.job_id = a.job_id
|
|
WHERE a.attempt_id = $1
|
|
AND j.authority_attempt_id = a.attempt_id
|
|
AND j.current_fence = a.fence
|
|
AND j.phase = 'running'
|
|
FOR UPDATE OF a, j",
|
|
)
|
|
.bind(attempt_id.as_str())
|
|
.fetch_optional(tx.as_mut())
|
|
.await?;
|
|
row.map(authority_row).transpose()
|
|
}
|
|
|
|
async fn load_authority(
|
|
pool: &PgPool,
|
|
attempt_id: &OpaqueId,
|
|
) -> Result<Option<AttemptAuthority>, DurableError> {
|
|
let row = sqlx::query_as::<_, (String, String, i64, String, String)>(
|
|
"SELECT a.job_id, a.attempt_id, a.fence, a.worker_incarnation, a.state
|
|
FROM attempts a
|
|
JOIN jobs j ON j.job_id = a.job_id
|
|
WHERE a.attempt_id = $1
|
|
AND j.authority_attempt_id = a.attempt_id
|
|
AND j.current_fence = a.fence
|
|
AND j.phase = 'running'",
|
|
)
|
|
.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()
|
|
}
|
|
|
|
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)?,
|
|
fence: u64::try_from(row.1).map_err(DurableError::effect)?,
|
|
worker: OpaqueId::parse(row.2).map_err(DurableError::effect)?,
|
|
state: AttemptState::parse(&row.3)
|
|
.ok_or_else(|| DurableError::Config(format!("unknown attempt state '{}'", row.3)))?,
|
|
})
|
|
}
|
|
|
|
fn validate_authority(authority: &AttemptAuthority) -> Result<(), AttemptError> {
|
|
if authority.fence == 0 || authority.fence > MAX_SAFE_INTEGER {
|
|
return Err(AttemptError::Invalid("invalid fencing token".into()));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn lease_owner(authority: &AttemptAuthority) -> Result<LeaseOwner, AttemptError> {
|
|
LeaseOwner::new(
|
|
authority.attempt_id.as_str(),
|
|
authority.fence,
|
|
authority.worker_incarnation.as_str(),
|
|
)
|
|
.map_err(Into::into)
|
|
}
|
|
|
|
fn authority_binding(authority: &AttemptAuthority) -> Vec<u8> {
|
|
format!(
|
|
"job={}|attempt={}|fence={}|worker={}",
|
|
authority.job_id, authority.attempt_id, authority.fence, authority.worker_incarnation
|
|
)
|
|
.into_bytes()
|
|
}
|
|
|
|
fn claim_rejected(reason: ClaimRejection) -> ClaimOutcome {
|
|
ClaimOutcome::Rejected {
|
|
reason,
|
|
disposition: None,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn new_id(prefix: char) -> Result<OpaqueId, AttemptError> {
|
|
let mut bytes = [0u8; 16];
|
|
getrandom::fill(&mut bytes)
|
|
.map_err(|error| AttemptError::Invalid(format!("identifier entropy failed: {error}")))?;
|
|
let mut value = String::with_capacity(33);
|
|
value.push(prefix);
|
|
for byte in bytes {
|
|
use std::fmt::Write as _;
|
|
write!(value, "{byte:02x}").expect("writing to String cannot fail");
|
|
}
|
|
OpaqueId::parse(value).map_err(|error| AttemptError::Invalid(error.to_string()))
|
|
}
|