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
2 changed files with 7448 additions and 194 deletions
Download patch file Download diff file Expand all files Collapse all files

177
crates/remora-api/src/worker.rs
View file

@ -8,22 +8,25 @@ use actix_web::{
};
use remora_attempt::{
ActivationMutationOutcome, AttemptAuthority, AttemptError, AttemptMutationOutcome,
AttemptService, CheckpointProjection, CheckpointRecord, ClaimOutcome, DurableDescriptor,
LeaseProjection, MutationRejection, OutcomePublication, ProgressRecord, RenewalMutationOutcome,
AttemptService, CheckpointMultipart, CheckpointProjection, CheckpointRecord, ClaimOutcome,
DurableDescriptor, LeaseProjection, MutationRejection, OutcomePublication, ProgressRecord,
RenewalMutationOutcome, StagingMultipart, StagingObject, StagingWriteIntent,
};
use remora_contracts::{
ActivationResponseDto, ApiErrorCodeDto, ApiErrorDto, AttemptActivateRequestDto,
AttemptAuthorityDto, AttemptCheckpointRequestDto, AttemptOutcomeDto, AttemptOutcomeRequestDto,
AttemptProgressRequestDto, AttemptRenewRequestDto, AttemptResolvedSourceRequestDto,
CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto,
CredentialGrantVersionDto, DateTimeDto, DurableDescriptorKindDto, DurableDescriptorRefDto,
ExecutionEnvelopeDto, JobPhaseProjectionDto, LeaseProjectionDto, NullableDto,
PositiveIntegerDto, ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto,
SafeIntegerDto, SourceProjectionDto, StagingTargetDto, TargetDto,
AttemptStagingMultipartRequestDto, AttemptStagingObjectRequestDto,
AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto, CheckpointDescriptorKindDto,
CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto, CredentialGrantVersionDto,
DateTimeDto, DurableDescriptorKindDto, DurableDescriptorRefDto, ExecutionEnvelopeDto,
JobPhaseProjectionDto, LeaseProjectionDto, NullableDto, PositiveIntegerDto,
ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto, SafeIntegerDto,
SourceProjectionDto, StagingTargetDto, TargetDto,
};
use remora_domain::OpaqueId;
use remora_result::PublicationInput;
use remora_target::{ContentSha256, ObjectIdentity, RESERVED_NAMESPACE_SEGMENT, StoredObject};
use remora_target::{ContentSha256, ObjectIdentity, StoredObject, attempt_staging_prefix};
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
@ -96,6 +99,19 @@ pub fn configure_worker_api(config: &mut web::ServiceConfig) {
.route("/jobs/{job_id}/claim", web::post().to(claim))
.route("/attempts/{attempt_id}/renew", web::post().to(renew))
.route("/attempts/{attempt_id}/activate", web::post().to(activate))
.route("/attempts/{attempt_id}/stop", web::post().to(stop))
.route(
"/attempts/{attempt_id}/staging/write-intent",
web::post().to(register_staging_write_intent),
)
.route(
"/attempts/{attempt_id}/staging/object",
web::post().to(register_staging_object),
)
.route(
"/attempts/{attempt_id}/staging/multipart",
web::post().to(register_staging_multipart),
)
.route(
"/attempts/{attempt_id}/resolved-source",
web::post().to(pin_resolved_source),
@ -115,6 +131,137 @@ pub fn configure_worker_api(config: &mut web::ServiceConfig) {
);
}
async fn register_staging_object(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptStagingObjectRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
let accepted = state
.attempts
.register_staging_object_operation(
&authority,
&request.operation_id,
StagingObject {
target_store: request.object.store.clone(),
bucket: request.object.bucket.clone(),
key: request.object.key.clone(),
version_id: request.object.version_id.clone(),
},
request.completed_upload_id.clone(),
)
.await
.map_err(WorkerApiError::from_attempt)?;
if accepted {
Ok(HttpResponse::NoContent().finish())
} else {
Err(WorkerApiError::precondition_failed())
}
}
async fn register_staging_write_intent(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptStagingWriteIntentRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
let accepted = state
.attempts
.register_staging_write_intent_operation(
&authority,
&request.operation_id,
StagingWriteIntent {
target_store: request.intent.store.clone(),
bucket: request.intent.bucket.clone(),
key: request.intent.key.clone(),
},
)
.await
.map_err(WorkerApiError::from_attempt)?;
if accepted {
Ok(HttpResponse::NoContent().finish())
} else {
Err(WorkerApiError::precondition_failed())
}
}
async fn register_staging_multipart(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptStagingMultipartRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
let accepted = state
.attempts
.register_staging_multipart_operation(
&authority,
&request.operation_id,
StagingMultipart {
target_store: request.multipart.store.clone(),
bucket: request.multipart.bucket.clone(),
key: request.multipart.key.clone(),
upload_id: request.multipart.upload_id.clone(),
},
)
.await
.map_err(WorkerApiError::from_attempt)?;
if accepted {
Ok(HttpResponse::NoContent().finish())
} else {
Err(WorkerApiError::precondition_failed())
}
}
async fn stop(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptStopRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
match state
.attempts
.cooperative_stop_operation(&authority, &request.operation_id)
.await
.map_err(WorkerApiError::from_attempt)?
{
remora_attempt::CooperativeStopOutcome::Stopped
| remora_attempt::CooperativeStopOutcome::StopNotRequired => {
Ok(HttpResponse::NoContent().finish())
}
remora_attempt::CooperativeStopOutcome::AuthorityMismatch => {
Err(WorkerApiError::precondition_failed())
}
}
}
async fn claim(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
@ -245,6 +392,14 @@ async fn persist_checkpoint(
resolved_source_id: request.resolved_source_id.clone(),
validated_bytes: request.validated_bytes.get(),
validated_members: request.validated_members.get(),
staging_multipart: request.staging_multipart.as_ref().map(|identity| {
CheckpointMultipart {
target_store: identity.store.clone(),
bucket: identity.bucket.clone(),
key: identity.key.clone(),
upload_id: identity.upload_id.clone(),
}
}),
},
)
.await
@ -387,6 +542,7 @@ async fn load_execution_envelope(
let source: SourceProjectionDto =
serde_json::from_value(row.source).map_err(|_| WorkerApiError::unavailable())?;
let target_prefix = row.target_prefix.join("/");
let staging_prefix = attempt_staging_prefix(&target_prefix, &row.staging_origin_attempt_id);
let resolved_source_ref = row.resolved_source.map(|descriptor| {
NullableDto(Some(descriptor_ref(
descriptor,
@ -425,10 +581,7 @@ async fn load_execution_envelope(
staging_target: StagingTargetDto {
store: row.target_store,
bucket: row.target_bucket,
prefix: format!(
"{RESERVED_NAMESPACE_SEGMENT}/a/{}/",
authority.attempt_id.as_str()
),
prefix: staging_prefix,
},
execution_grant_ref: row.execution_grant_id,
checkpoint_ref: checkpoint_ref.unwrap_or(NullableDto(None)),

127
crates/remora-api/tests/worker_api.rs
View file

@ -15,11 +15,13 @@ use remora_attempt::{AttemptConfig, AttemptService, RecoveryConfig, migrate as a
use remora_contracts::{
ApiErrorDto, AttemptActivateRequestDto, AttemptCheckpointRequestDto, AttemptFailureOutcomeDto,
AttemptOutcomeDto, AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto,
AttemptResolvedSourceRequestDto, AttemptRuntimeFailureDto, CheckpointDescriptorKindDto,
CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto, HttpSourceSubmissionDto,
JobPhaseProjectionDto, JobSubmissionDto, PositiveIntegerDto, ProgressValueDto,
ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto, ResultDescriptorDto,
ResultKindDto, SafeIntegerDto, SourceSubmissionDto,
AttemptResolvedSourceRequestDto, AttemptRuntimeFailureDto, AttemptStagingMultipartRequestDto,
AttemptStagingObjectRequestDto, AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto,
CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto,
HttpSourceSubmissionDto, JobPhaseProjectionDto, JobSubmissionDto, PositiveIntegerDto,
ProgressValueDto, ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto,
ResultDescriptorDto, ResultKindDto, SafeIntegerDto, SourceSubmissionDto,
StagingMultipartIdentityDto, StagingObjectIdentityDto, StagingWriteIntentDto,
};
use remora_coordination::{CoordinationStore, StoreConfig};
use remora_credentials::migrate as credentials_migrate;
@ -29,7 +31,7 @@ use remora_domain::{
use remora_durable::{PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{DefaultTarget, Principal};
use remora_result::migrate as result_migrate;
use remora_target::attempt_manifest_key;
use remora_target::attempt_manifest_key_in;
use reqwest::{Certificate as ClientCertificate, Client, Identity};
use rustls::{
RootCertStore, ServerConfig,
@ -153,14 +155,15 @@ impl Harness {
async fn seal_result(&self, attempt_id: &OpaqueId) -> ResultDescriptorDto {
let manifest_id = unique_id("manifest");
let checksum = format!("sha256:{}", "11".repeat(32));
let key = attempt_manifest_key(attempt_id, &manifest_id);
let key = attempt_manifest_key_in("users/alice", attempt_id, &manifest_id);
sqlx::query(
"INSERT INTO result_inventories (
attempt_id, state, manifest_schema_version, manifest_id, manifest_store,
manifest_bucket, manifest_key, manifest_version_id, manifest_sha256,
manifest_size, member_count, total_bytes
manifest_size, member_count, total_bytes,
staging_origin_attempt_id, target_prefix
) VALUES ($1, 'sealed', 1, $2, 'primary', 'remora-results', $3,
'version-1', $4, 128, 1, 4096)",
'version-1', $4, 128, 1, 4096, $1, 'users/alice')",
)
.bind(attempt_id.as_str())
.bind(manifest_id.as_str())
@ -336,6 +339,92 @@ async fn worker_operations_are_authenticated_fenced_replayable_and_redacted() {
assert_eq!(response.status(), 409);
}
let staging_key = format!(
"{}o/{}",
claimed.execution.staging_target.prefix,
unique_id("object")
);
let staging_intent = AttemptStagingWriteIntentRequestDto {
operation_id: unique_id("operation"),
fence: claimed.authority.fence,
worker_incarnation: worker.clone(),
intent: StagingWriteIntentDto {
store: claimed.execution.target.store.clone(),
bucket: claimed.execution.target.bucket.clone(),
key: staging_key.clone(),
},
};
for _ in 0..2 {
assert_eq!(
client
.post(format!("{attempt_url}/staging/write-intent"))
.json(&staging_intent)
.send()
.await
.unwrap()
.status(),
204
);
}
let staging_multipart = AttemptStagingMultipartRequestDto {
operation_id: unique_id("operation"),
fence: claimed.authority.fence,
worker_incarnation: worker.clone(),
multipart: StagingMultipartIdentityDto {
store: claimed.execution.target.store.clone(),
bucket: claimed.execution.target.bucket.clone(),
key: staging_key.clone(),
upload_id: "upload-identity".into(),
},
};
assert_eq!(
client
.post(format!("{attempt_url}/staging/multipart"))
.json(&staging_multipart)
.send()
.await
.unwrap()
.status(),
204
);
let staging_operation = unique_id("operation");
let staging_object = AttemptStagingObjectRequestDto {
operation_id: staging_operation.clone(),
fence: claimed.authority.fence,
worker_incarnation: worker.clone(),
object: StagingObjectIdentityDto {
store: claimed.execution.target.store.clone(),
bucket: claimed.execution.target.bucket.clone(),
key: staging_key.clone(),
version_id: None,
},
completed_upload_id: Some("upload-identity".into()),
};
for _ in 0..2 {
assert_eq!(
client
.post(format!("{attempt_url}/staging/object"))
.json(&staging_object)
.send()
.await
.unwrap()
.status(),
204
);
}
let mut changed_staging = staging_object;
changed_staging.object.key.push_str("-changed");
assert_eq!(
client
.post(format!("{attempt_url}/staging/object"))
.json(&changed_staging)
.send()
.await
.unwrap()
.status(),
409
);
let resolved_id = unique_id("resolved");
let resolved = AttemptResolvedSourceRequestDto {
operation_id: unique_id("operation"),
@ -375,6 +464,7 @@ async fn worker_operations_are_authenticated_fenced_replayable_and_redacted() {
resolved_source_id: resolved_id,
validated_bytes: safe(1024),
validated_members: safe(1),
staging_multipart: None,
};
assert_eq!(
client
@ -470,7 +560,7 @@ async fn worker_operations_are_authenticated_fenced_replayable_and_redacted() {
.json(&AttemptProgressRequestDto {
operation_id: renewal_operation,
fence: claimed.authority.fence,
worker_incarnation: worker,
worker_incarnation: worker.clone(),
progress: ProgressValueDto {
completed_bytes: safe(2048),
total_bytes: Some(safe(4096)),
@ -488,6 +578,23 @@ async fn worker_operations_are_authenticated_fenced_replayable_and_redacted() {
);
let error: ApiErrorDto = cross_kind.json().await.unwrap();
assert_eq!(error.message, "precondition failed");
let stop = AttemptStopRequestDto {
operation_id: unique_id("operation"),
fence: claimed.authority.fence,
worker_incarnation: worker,
};
for _ in 0..2 {
assert_eq!(
client
.post(format!("{attempt_url}/stop"))
.json(&stop)
.send()
.await
.unwrap()
.status(),
204
);
}
server.stop(true).await;
}