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

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")
}