sail-rs 0.2.19

//! Typed client for the imagebuilder dispatcher gRPC service: submit builds,
//! poll build status, and prepare content-addressed local-file uploads. The
//! image pipeline in [`crate::imagebuild`] is the only caller; transient
//! transport failures retry with backoff against a per-call budget.

use std::time::Instant;

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,
}

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 = crate::worker::bearer_metadata(api_key)?;
        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.
    pub async fn build_image(
        &self,
        message: pb::BuildImageRequest,
        retry_timeout: f64,
    ) -> Result<pb::BuildImageResponse, SailError> {
        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()),
                Err(status) => self.handle_retry(&status, deadline, &mut delay).await?,
            }
        }
    }

    /// Poll one build's status.
    pub async fn get_image_build_status(
        &self,
        message: pb::GetImageBuildStatusRequest,
        retry_timeout: f64,
    ) -> Result<pb::GetImageBuildStatusResponse, SailError> {
        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()),
                Err(status) => self.handle_retry(&status, deadline, &mut delay).await?,
            }
        }
    }

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

#[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"
        )));
    }
}