ms2k/remora
SHA256
Watch
0
Code Releases Activity

Implement crash-safe streaming worker runtime

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

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

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