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

use remora_contracts::{
    ActivationResponseDto, AttemptActivateRequestDto, AttemptCheckpointRequestDto,
    AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto,
    AttemptResolvedSourceRequestDto, AttemptStagingMultipartRequestDto,
    AttemptStagingObjectRequestDto, AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto,
    ClaimRequestDto, ClaimResponseDto, JobPhaseProjectionDto, LeaseProjectionDto, ProgressValueDto,
    ResolvedSourceDescriptorRefDto,
};
use remora_domain::OpaqueId;
use reqwest::{Client, Response, StatusCode};
use serde::{Serialize, de::DeserializeOwned};

use crate::WorkerError;

#[derive(Clone)]
pub struct WorkerApiClient {
    base_url: String,
    client: Client,
}

impl WorkerApiClient {
    pub fn new(base_url: impl Into<String>, client: Client) -> Result<Self, WorkerError> {
        let base_url = base_url.into().trim_end_matches('/').to_owned();
        let parsed =
            reqwest::Url::parse(&base_url).map_err(|_| WorkerError::InvalidConfiguration)?;
        if parsed.scheme() != "https" || parsed.host_str().is_none() {
            return Err(WorkerError::InvalidConfiguration);
        }
        Ok(Self { base_url, client })
    }

    pub async fn claim(
        &self,
        job_id: &OpaqueId,
        operation_id: OpaqueId,
        worker_incarnation: OpaqueId,
    ) -> Result<ClaimResponseDto, WorkerError> {
        self.post_json(
            format!("/internal/v1/jobs/{job_id}/claim"),
            ClaimRequestDto {
                operation_id,
                worker_incarnation,
            },
        )
        .await
    }

    pub async fn activate(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptActivateRequestDto,
    ) -> Result<ActivationResponseDto, WorkerError> {
        self.post_json(
            format!("/internal/v1/attempts/{attempt_id}/activate"),
            request,
        )
        .await
    }

    pub async fn renew(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptRenewRequestDto,
    ) -> Result<LeaseProjectionDto, WorkerError> {
        self.post_json(format!("/internal/v1/attempts/{attempt_id}/renew"), request)
            .await
    }

    pub async fn pin_resolved_source(
        &self,
        attempt_id: &OpaqueId,
        operation_id: OpaqueId,
        fence: remora_contracts::PositiveIntegerDto,
        worker_incarnation: OpaqueId,
        resolved_source_ref: ResolvedSourceDescriptorRefDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/resolved-source"),
            AttemptResolvedSourceRequestDto {
                operation_id,
                fence,
                worker_incarnation,
                resolved_source_ref,
            },
        )
        .await
    }

    pub async fn persist_checkpoint(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptCheckpointRequestDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/checkpoint"),
            request,
        )
        .await
    }

    pub async fn report_progress(
        &self,
        attempt_id: &OpaqueId,
        operation_id: OpaqueId,
        fence: remora_contracts::PositiveIntegerDto,
        worker_incarnation: OpaqueId,
        progress: ProgressValueDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/progress"),
            AttemptProgressRequestDto {
                operation_id,
                fence,
                worker_incarnation,
                progress,
            },
        )
        .await
    }

    pub async fn report_outcome(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptOutcomeRequestDto,
    ) -> Result<JobPhaseProjectionDto, WorkerError> {
        self.post_json(
            format!("/internal/v1/attempts/{attempt_id}/outcome"),
            request,
        )
        .await
    }

    pub async fn stop(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptStopRequestDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(format!("/internal/v1/attempts/{attempt_id}/stop"), request)
            .await
    }

    pub async fn register_staging_object(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptStagingObjectRequestDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/staging/object"),
            request,
        )
        .await
    }

    pub async fn register_staging_write_intent(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptStagingWriteIntentRequestDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/staging/write-intent"),
            request,
        )
        .await
    }

    pub async fn register_staging_multipart(
        &self,
        attempt_id: &OpaqueId,
        request: AttemptStagingMultipartRequestDto,
    ) -> Result<(), WorkerError> {
        self.post_empty(
            format!("/internal/v1/attempts/{attempt_id}/staging/multipart"),
            request,
        )
        .await
    }

    async fn post_json<T, R>(&self, path: String, body: T) -> Result<R, WorkerError>
    where
        T: Clone + Serialize,
        R: DeserializeOwned,
    {
        let response = self.send_with_retry(path, body).await?;
        response
            .json()
            .await
            .map_err(|_| WorkerError::InvalidResponse)
    }

    async fn post_empty<T>(&self, path: String, body: T) -> Result<(), WorkerError>
    where
        T: Clone + Serialize,
    {
        self.send_with_retry(path, body).await.map(|_| ())
    }

    async fn send_with_retry<T>(&self, path: String, body: T) -> Result<Response, WorkerError>
    where
        T: Clone + Serialize,
    {
        let url = format!("{}{path}", self.base_url);
        let send = || self.client.post(&url).json(&body).send();
        match send().await {
            Ok(response) => classify(response),
            Err(_) => send()
                .await
                .map_err(|_| WorkerError::ApiUnavailable)
                .and_then(classify),
        }
    }
}

fn classify(response: Response) -> Result<Response, WorkerError> {
    match response.status() {
        status if status.is_success() => Ok(response),
        StatusCode::BAD_REQUEST => Err(WorkerError::InvalidRequest),
        StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => Err(WorkerError::Unauthenticated),
        StatusCode::CONFLICT => Err(WorkerError::PreconditionFailed),
        _ => Err(WorkerError::ApiUnavailable),
    }
}

pub(crate) fn operation_id() -> OpaqueId {
    generated_id('o')
}

pub(crate) fn incarnation_id() -> OpaqueId {
    generated_id('w')
}

fn generated_id(prefix: char) -> OpaqueId {
    let mut bytes = [0_u8; 16];
    getrandom::fill(&mut bytes).expect("operating system randomness is required");
    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).expect("generated identifier is valid")
}