Implement crash-safe streaming worker runtime

crates/remora-worker/src/client.rs

@@ -0,0 +1,241 @@
+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")
+}