remora/crates/remora-api/src/worker.rs

613 lines
21 KiB
Rust

use std::{
any::Any,
future::{Ready, ready},
};
use actix_web::{
FromRequest, HttpRequest, HttpResponse, ResponseError, dev::Payload, http::StatusCode, web,
};
use remora_attempt::{
ActivationMutationOutcome, AttemptAuthority, AttemptError, AttemptMutationOutcome,
AttemptService, CheckpointProjection, CheckpointRecord, ClaimOutcome, DurableDescriptor,
LeaseProjection, MutationRejection, OutcomePublication, ProgressRecord, RenewalMutationOutcome,
};
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,
};
use remora_domain::OpaqueId;
use remora_result::PublicationInput;
use remora_target::{ContentSha256, ObjectIdentity, RESERVED_NAMESPACE_SEGMENT, StoredObject};
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AuthenticatedWorker {
pub identity: OpaqueId,
}
pub fn authenticate_rustls_worker(
connection: &dyn Any,
extensions: &mut actix_web::dev::Extensions,
) {
type Stream = actix_tls::accept::rustls_0_23::TlsStream<actix_web::rt::net::TcpStream>;
let Some(stream) = connection.downcast_ref::<Stream>() else {
return;
};
let Some(certificate) = stream
.get_ref()
.1
.peer_certificates()
.and_then(|chain| chain.first())
else {
return;
};
let digest = Sha256::digest(certificate.as_ref());
let mut value = String::from("w");
for byte in digest {
use std::fmt::Write as _;
write!(value, "{byte:02x}").expect("writing to String cannot fail");
}
extensions.insert(AuthenticatedWorker {
identity: OpaqueId::parse(value).expect("certificate digest is a valid identity"),
});
}
impl FromRequest for AuthenticatedWorker {
type Error = WorkerApiError;
type Future = Ready<Result<Self, Self::Error>>;
fn from_request(request: &HttpRequest, _payload: &mut Payload) -> Self::Future {
ready(
request
.conn_data::<Self>()
.cloned()
.ok_or_else(WorkerApiError::unauthenticated),
)
}
}
#[derive(Clone)]
pub struct WorkerApiState {
attempts: AttemptService,
}
impl WorkerApiState {
#[must_use]
pub fn new(attempts: AttemptService) -> Self {
Self { attempts }
}
}
pub fn configure_worker_api(config: &mut web::ServiceConfig) {
config.service(
web::scope("/internal/v1")
.app_data(
web::JsonConfig::default()
.limit(1024 * 1024)
.error_handler(|_, _| WorkerApiError::invalid_request().into()),
)
.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}/resolved-source",
web::post().to(pin_resolved_source),
)
.route(
"/attempts/{attempt_id}/checkpoint",
web::post().to(persist_checkpoint),
)
.route(
"/attempts/{attempt_id}/progress",
web::post().to(report_progress),
)
.route(
"/attempts/{attempt_id}/outcome",
web::post().to(report_outcome),
),
);
}
async fn claim(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<ClaimRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let job_id = parse_id(path.into_inner())?;
let outcome = state
.attempts
.claim(&job_id, &request.worker_incarnation, &request.operation_id)
.await
.map_err(WorkerApiError::from_attempt)?;
let ClaimOutcome::Claimed { authority, .. } = outcome else {
return Err(WorkerApiError::precondition_failed());
};
let execution = load_execution_envelope(&state.attempts, &authority).await?;
Ok(HttpResponse::Ok().json(ClaimResponseDto {
authority: authority_dto(&authority)?,
execution,
}))
}
async fn renew(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptRenewRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
match state
.attempts
.renew_operation(&authority, &request.operation_id)
.await
.map_err(WorkerApiError::from_attempt)?
{
RenewalMutationOutcome::Renewed { lease, .. } => {
Ok(HttpResponse::Ok().json(lease_dto(lease)?))
}
RenewalMutationOutcome::AuthorityLost { .. } => Err(WorkerApiError::precondition_failed()),
}
}
async fn activate(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptActivateRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
let lease = match state
.attempts
.activate_operation(&authority, &request.operation_id)
.await
.map_err(WorkerApiError::from_attempt)?
{
ActivationMutationOutcome::Activated { lease, .. } => lease,
ActivationMutationOutcome::PreconditionFailed => {
return Err(WorkerApiError::precondition_failed());
}
};
let execution = load_execution_envelope(&state.attempts, &authority).await?;
Ok(HttpResponse::Ok().json(ActivationResponseDto {
authority: authority_dto(&authority)?,
execution,
lease: lease_dto(lease)?,
}))
}
async fn pin_resolved_source(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptResolvedSourceRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
accepted(
state
.attempts
.pin_resolved_source(
&authority,
&request.operation_id,
descriptor_from_resolved(&request.resolved_source_ref)?,
)
.await
.map_err(WorkerApiError::from_attempt)?,
)
}
async fn persist_checkpoint(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptCheckpointRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
accepted(
state
.attempts
.persist_checkpoint(
&authority,
&request.operation_id,
CheckpointRecord {
descriptor: descriptor_from_checkpoint(&request.checkpoint_ref)?,
resolved_source_id: request.resolved_source_id.clone(),
validated_bytes: request.validated_bytes.get(),
validated_members: request.validated_members.get(),
},
)
.await
.map_err(WorkerApiError::from_attempt)?,
)
}
async fn report_progress(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptProgressRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
accepted(
state
.attempts
.report_progress(
&authority,
&request.operation_id,
ProgressRecord {
completed_bytes: request.progress.completed_bytes.get(),
total_bytes: request.progress.total_bytes.map(SafeIntegerDto::get),
completed_members: request.progress.completed_members.get(),
total_members: request.progress.total_members.map(SafeIntegerDto::get),
},
)
.await
.map_err(WorkerApiError::from_attempt)?,
)
}
async fn report_outcome(
_: AuthenticatedWorker,
state: web::Data<WorkerApiState>,
path: web::Path<String>,
request: web::Json<AttemptOutcomeRequestDto>,
) -> Result<HttpResponse, WorkerApiError> {
let authority = request_authority(
&state,
path.into_inner(),
request.fence.get(),
&request.worker_incarnation,
)
.await?;
let outcome = match &request.outcome {
AttemptOutcomeDto::Success(success) => {
let result = &success.result;
state
.attempts
.publish_success_operation(
&authority,
&request.operation_id,
PublicationInput {
attempt_id: authority.attempt_id.clone(),
manifest_id: result.manifest_id.clone(),
schema_version: u32::try_from(result.schema_version.get())
.map_err(|_| WorkerApiError::invalid_request())?,
manifest: StoredObject {
identity: ObjectIdentity {
target_store_id: result.manifest.store.clone(),
bucket: result.manifest.bucket.clone(),
key: result.manifest.key.clone(),
version_id: result.manifest.version_id.clone(),
},
content_sha256: ContentSha256::parse(&result.manifest.checksum)
.map_err(|_| WorkerApiError::invalid_request())?,
size: result.manifest.size.get(),
},
member_count: result.member_count.get(),
total_bytes: result.total_bytes.get(),
},
)
.await
}
AttemptOutcomeDto::Failure(failure) => {
state
.attempts
.publish_failure_operation(&authority, &request.operation_id, failure.failure.class)
.await
}
}
.map_err(WorkerApiError::from_attempt)?;
let phase = match outcome {
OutcomePublication::Accepted { phase, .. } => phase,
OutcomePublication::Rejected { reason, .. } => {
return Err(WorkerApiError::from_rejection(reason));
}
};
Ok(HttpResponse::Ok().json(JobPhaseProjectionDto { phase }))
}
fn accepted(outcome: AttemptMutationOutcome) -> Result<HttpResponse, WorkerApiError> {
match outcome {
AttemptMutationOutcome::Accepted(_) => Ok(HttpResponse::NoContent().finish()),
AttemptMutationOutcome::Rejected { reason, .. } => {
Err(WorkerApiError::from_rejection(reason))
}
}
}
async fn request_authority(
state: &WorkerApiState,
attempt_id: String,
fence: u64,
worker: &OpaqueId,
) -> Result<AttemptAuthority, WorkerApiError> {
let attempt_id = parse_id(attempt_id)?;
state
.attempts
.historical_authority(&attempt_id)
.await
.map_err(WorkerApiError::from_attempt)?
.filter(|authority| authority.fence == fence && authority.worker_incarnation == *worker)
.ok_or_else(WorkerApiError::precondition_failed)
}
async fn load_execution_envelope(
attempts: &AttemptService,
authority: &AttemptAuthority,
) -> Result<ExecutionEnvelopeDto, WorkerApiError> {
let row = attempts
.execution_envelope(&authority.attempt_id)
.await
.map_err(WorkerApiError::from_attempt)?
.ok_or_else(WorkerApiError::precondition_failed)?;
if row.job_id != authority.job_id
|| row.attempt_id != authority.attempt_id
|| row.fence != authority.fence
|| row.worker_incarnation != authority.worker_incarnation
{
return Err(WorkerApiError::precondition_failed());
}
let source: SourceProjectionDto =
serde_json::from_value(row.source).map_err(|_| WorkerApiError::unavailable())?;
let target_prefix = row.target_prefix.join("/");
let resolved_source_ref = row.resolved_source.map(|descriptor| {
NullableDto(Some(descriptor_ref(
descriptor,
DurableDescriptorKindDto::ResolvedSource,
)))
});
let checkpoint_ref = row.checkpoint.and_then(|checkpoint| match checkpoint {
CheckpointProjection::Reusable { checkpoint } => Some(NullableDto(Some(descriptor_ref(
checkpoint.descriptor,
DurableDescriptorKindDto::Checkpoint,
)))),
CheckpointProjection::Complete { .. } => None,
});
Ok(ExecutionEnvelopeDto {
id: row.id,
job_id: authority.job_id.clone(),
attempt_id: authority.attempt_id.clone(),
fence: positive(authority.fence)?,
worker_incarnation: authority.worker_incarnation.clone(),
source,
sensitive_source_ref: row.sensitive_source_ref,
credential_version: NullableDto(match (row.credential_id, row.credential_version_id) {
(Some(credential_id), Some(version_id)) => Some(CredentialGrantVersionDto {
credential_id,
version_id,
}),
(None, None) => None,
_ => return Err(WorkerApiError::unavailable()),
}),
resolved_source_ref: resolved_source_ref.unwrap_or(NullableDto(None)),
target: TargetDto {
store: row.target_store.clone(),
bucket: row.target_bucket.clone(),
prefix: target_prefix,
},
staging_target: StagingTargetDto {
store: row.target_store,
bucket: row.target_bucket,
prefix: format!(
"{RESERVED_NAMESPACE_SEGMENT}/a/{}/",
authority.attempt_id.as_str()
),
},
execution_grant_ref: row.execution_grant_id,
checkpoint_ref: checkpoint_ref.unwrap_or(NullableDto(None)),
})
}
fn descriptor_ref(
descriptor: DurableDescriptor,
kind: DurableDescriptorKindDto,
) -> DurableDescriptorRefDto {
DurableDescriptorRefDto {
id: descriptor.id,
kind,
schema_version: positive(descriptor.schema_version.into()).expect("validated descriptor"),
byte_length: safe(descriptor.byte_length).expect("validated descriptor"),
item_count: descriptor
.item_count
.map(|value| safe(value).expect("validated descriptor")),
checksum: descriptor.checksum,
}
}
fn descriptor_from_resolved(
value: &ResolvedSourceDescriptorRefDto,
) -> Result<DurableDescriptor, WorkerApiError> {
let ResolvedSourceDescriptorKindDto::ResolvedSource = value.kind;
Ok(DurableDescriptor {
id: value.id.clone(),
schema_version: u32::try_from(value.schema_version.get())
.map_err(|_| WorkerApiError::invalid_request())?,
byte_length: value.byte_length.get(),
item_count: value.item_count.map(SafeIntegerDto::get),
checksum: value.checksum.clone(),
})
}
fn descriptor_from_checkpoint(
value: &CheckpointDescriptorRefDto,
) -> Result<DurableDescriptor, WorkerApiError> {
let CheckpointDescriptorKindDto::Checkpoint = value.kind;
Ok(DurableDescriptor {
id: value.id.clone(),
schema_version: u32::try_from(value.schema_version.get())
.map_err(|_| WorkerApiError::invalid_request())?,
byte_length: value.byte_length.get(),
item_count: value.item_count.map(SafeIntegerDto::get),
checksum: value.checksum.clone(),
})
}
fn authority_dto(authority: &AttemptAuthority) -> Result<AttemptAuthorityDto, WorkerApiError> {
Ok(AttemptAuthorityDto {
attempt_id: authority.attempt_id.clone(),
fence: positive(authority.fence)?,
worker_incarnation: authority.worker_incarnation.clone(),
})
}
fn lease_dto(lease: LeaseProjection) -> Result<LeaseProjectionDto, WorkerApiError> {
let observed = timestamp(lease.observed_at_ms)?;
let expires = timestamp(lease.expires_at_ms)?;
let renew_at = observed + (expires - observed) / 2;
Ok(LeaseProjectionDto {
expires_at: datetime(expires)?,
renew_after: datetime(renew_at)?,
})
}
fn timestamp(milliseconds: u64) -> Result<OffsetDateTime, WorkerApiError> {
let nanoseconds = i128::from(milliseconds) * 1_000_000;
OffsetDateTime::from_unix_timestamp_nanos(nanoseconds)
.map_err(|_| WorkerApiError::unavailable())
}
fn datetime(value: OffsetDateTime) -> Result<DateTimeDto, WorkerApiError> {
let value = value
.format(&Rfc3339)
.map_err(|_| WorkerApiError::unavailable())?;
DateTimeDto::parse(value).map_err(|_| WorkerApiError::unavailable())
}
fn positive(value: u64) -> Result<PositiveIntegerDto, WorkerApiError> {
PositiveIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable())
}
fn safe(value: u64) -> Result<SafeIntegerDto, WorkerApiError> {
SafeIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable())
}
fn parse_id(value: impl Into<String>) -> Result<OpaqueId, WorkerApiError> {
OpaqueId::parse(value).map_err(|_| WorkerApiError::invalid_request())
}
#[derive(Debug, thiserror::Error)]
#[error("worker API request failed")]
pub struct WorkerApiError {
status: StatusCode,
code: ApiErrorCodeDto,
message: &'static str,
}
impl WorkerApiError {
fn invalid_request() -> Self {
Self::new(
StatusCode::BAD_REQUEST,
ApiErrorCodeDto::InvalidRequest,
"invalid request",
)
}
fn unauthenticated() -> Self {
Self::new(
StatusCode::UNAUTHORIZED,
ApiErrorCodeDto::Unauthenticated,
"authentication required",
)
}
fn precondition_failed() -> Self {
Self::new(
StatusCode::CONFLICT,
ApiErrorCodeDto::PreconditionFailed,
"precondition failed",
)
}
fn unavailable() -> Self {
Self::new(
StatusCode::SERVICE_UNAVAILABLE,
ApiErrorCodeDto::Unavailable,
"service unavailable",
)
}
fn new(status: StatusCode, code: ApiErrorCodeDto, message: &'static str) -> Self {
Self {
status,
code,
message,
}
}
fn from_attempt(error: AttemptError) -> Self {
match error {
AttemptError::Invalid(_) => Self::invalid_request(),
AttemptError::OperationConflict => Self::precondition_failed(),
AttemptError::Coordination(_) | AttemptError::Durable(_) => Self::unavailable(),
}
}
fn from_rejection(_reason: MutationRejection) -> Self {
Self::precondition_failed()
}
}
impl ResponseError for WorkerApiError {
fn status_code(&self) -> StatusCode {
self.status
}
fn error_response(&self) -> HttpResponse {
HttpResponse::build(self.status)
.content_type("application/problem+json")
.json(ApiErrorDto {
code: self.code,
message: self.message.into(),
request_id: request_id(),
fields: None,
})
}
}
fn request_id() -> OpaqueId {
let mut bytes = [0_u8; 16];
getrandom::fill(&mut bytes).expect("operating system randomness is required");
let mut value = String::from("r");
for byte in bytes {
use std::fmt::Write as _;
write!(value, "{byte:02x}").expect("writing to String cannot fail");
}
OpaqueId::parse(value).expect("generated request identifier is valid")
}