sail-rs 0.1.0

//! Client for the imagebuilder dispatcher service (custom image builds).
//!
//! The image spec is built by the language wrapper, so requests and responses
//! cross this boundary as encoded protobuf bytes: the core decodes the typed
//! message, runs the RPC with transient-failure retry, and re-encodes the
//! response. Build orchestration (the status poll loop) stays in the wrapper.

use std::time::Instant;

use prost::Message;
use tonic::metadata::AsciiMetadataValue;
use tonic::transport::Channel;
use tonic::{Code, Request, Status};

use crate::channels::ChannelCache;
use crate::error::SailError;
use crate::pb::imagebuilder::v1 as pb;
use crate::worker::{
    retry_deadline, sleep_before_retry, EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS,
};
use pb::image_builder_service_client::ImageBuilderServiceClient;

/// Transient imagebuilder transport races worth retrying: a mid-flight
/// connection drop while the dispatcher rolls during a deploy. A bare connect
/// failure is intentionally excluded: it fails fast as a connection error
/// rather than retrying for the whole build budget.
fn is_transient_imagebuilder_transport(status: &Status) -> bool {
    status.code() == Code::Unavailable
        && crate::worker::is_transient_transport_message(status.message())
}

/// Client for the imagebuilder dispatcher gRPC service, bound to one endpoint
/// and bearer key.
pub struct ImageBuilder {
    endpoint: String,
    channels: ChannelCache,
    authorization: AsciiMetadataValue,
}

fn decode<T: Message + Default>(bytes: &[u8]) -> Result<T, SailError> {
    T::decode(bytes).map_err(|err| SailError::Internal {
        message: format!("failed to decode imagebuilder request: {err}"),
    })
}

impl ImageBuilder {
    /// Build a client targeting `endpoint`, authenticating with `api_key`.
    /// Errors if `api_key` contains characters invalid in a gRPC metadata value.
    pub fn new(endpoint: &str, api_key: &str) -> Result<ImageBuilder, SailError> {
        let authorization = format!("Bearer {api_key}")
            .parse()
            .map_err(|_| SailError::Config {
                message: "SAIL_API_KEY contains characters invalid in a gRPC metadata value"
                    .to_string(),
            })?;
        Ok(ImageBuilder {
            endpoint: endpoint.to_string(),
            channels: ChannelCache::new(),
            authorization,
        })
    }

    fn client(&self) -> Result<ImageBuilderServiceClient<Channel>, SailError> {
        Ok(ImageBuilderServiceClient::new(
            self.channels.get(&self.endpoint)?,
        ))
    }

    fn request<T>(&self, message: T) -> Request<T> {
        let mut request = Request::new(message);
        request
            .metadata_mut()
            .insert("authorization", self.authorization.clone());
        request
    }

    /// Retry a transient failure within the deadline (redialing a fresh
    /// channel), or return the classified error.
    async fn handle_retry(
        &self,
        status: &Status,
        deadline: Instant,
        delay: &mut f64,
    ) -> Result<(), SailError> {
        if Instant::now() >= deadline || !is_transient_imagebuilder_transport(status) {
            return Err(SailError::from_rpc_status(status));
        }
        // Every transient the imagebuilder retries is connection-suspect (a
        // lameducking dispatcher ("endpoint closing") or a dead socket), so
        // always redial rather than reuse a channel pinned to the closing pod.
        self.channels.invalidate(&self.endpoint);
        *delay = sleep_before_retry(*delay, deadline).await;
        Ok(())
    }

    /// Submit or resume a build; returns the encoded `BuildImageResponse`.
    pub async fn build_image(
        &self,
        request: &[u8],
        retry_timeout: f64,
    ) -> Result<Vec<u8>, SailError> {
        let message = decode::<pb::BuildImageRequest>(request)?;
        let deadline = retry_deadline(retry_timeout);
        let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
        loop {
            // Re-fetch the client each attempt so a drain-triggered channel
            // invalidation in handle_retry redials on the next iteration.
            match self
                .client()?
                .build_image(self.request(message.clone()))
                .await
            {
                Ok(resp) => return Ok(resp.into_inner().encode_to_vec()),
                Err(status) => self.handle_retry(&status, deadline, &mut delay).await?,
            }
        }
    }

    /// Poll one build's status; returns the encoded `GetImageBuildStatusResponse`.
    pub async fn get_image_build_status(
        &self,
        request: &[u8],
        retry_timeout: f64,
    ) -> Result<Vec<u8>, SailError> {
        let message = decode::<pb::GetImageBuildStatusRequest>(request)?;
        let deadline = retry_deadline(retry_timeout);
        let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
        loop {
            // Re-fetch the client each attempt so a drain-triggered channel
            // invalidation in handle_retry redials on the next iteration.
            match self
                .client()?
                .get_image_build_status(self.request(message.clone()))
                .await
            {
                Ok(resp) => return Ok(resp.into_inner().encode_to_vec()),
                Err(status) => self.handle_retry(&status, deadline, &mut delay).await?,
            }
        }
    }

    /// Prepare a content-addressed local-file upload; single attempt. Returns
    /// the encoded `PrepareLocalFileUploadResponse`.
    pub async fn prepare_local_file_upload(&self, request: &[u8]) -> Result<Vec<u8>, SailError> {
        let message = decode::<pb::PrepareLocalFileUploadRequest>(request)?;
        let resp = self
            .client()?
            .prepare_local_file_upload(self.request(message))
            .await
            .map_err(|status| SailError::from_rpc_status(&status))?;
        Ok(resp.into_inner().encode_to_vec())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn drain_race_is_transient_but_connect_failure_is_not() {
        assert!(is_transient_imagebuilder_transport(&Status::unavailable(
            "endpoint closing"
        )));
        assert!(is_transient_imagebuilder_transport(&Status::unavailable(
            "error reading server preface: connection reset"
        )));
        // A bare connect failure must fail fast, not retry for the build budget.
        assert!(!is_transient_imagebuilder_transport(&Status::unavailable(
            "tcp connect error"
        )));
        assert!(!is_transient_imagebuilder_transport(&Status::internal(
            "boom"
        )));
    }
}