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, client: Client) -> Result { 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 { 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 { self.post_json( format!("/internal/v1/attempts/{attempt_id}/activate"), request, ) .await } pub async fn renew( &self, attempt_id: &OpaqueId, request: AttemptRenewRequestDto, ) -> Result { 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 { 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(&self, path: String, body: T) -> Result 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(&self, path: String, body: T) -> Result<(), WorkerError> where T: Clone + Serialize, { self.send_with_retry(path, body).await.map(|_| ()) } async fn send_with_retry(&self, path: String, body: T) -> Result 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 { 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") }