bentoml 0.4.0

//! The async HTTP client and its builder.

mod builder;
mod endpoint;
mod headers;
mod response;

pub mod multipart;

use std::borrow::Cow;
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};

use reqwest::Response as ReqwestResponse;
use reqwest_middleware::{
    ClientBuilder as MiddlewareBuilder, ClientWithMiddleware, RequestBuilder,
};
use reqwest_retry::RetryTransientMiddleware;
use reqwest_retry::policies::ExponentialBackoff;
use url::Url;

pub use self::builder::ClientBuilder;
pub use self::endpoint::Endpoint;
pub(crate) use self::headers::Headers;
pub use self::response::EndpointResponse;
use crate::error::{Error, Result};

/// An async client for a single BentoML service.
///
/// Construct one with [`Client::builder`], then invoke service endpoints through an
/// [`Endpoint`] handle from [`Client::endpoint`]. The client is cheap to clone:
/// internally it is an [`Arc`] around shared state, so clones share one connection
/// pool. Requests pass through a [`reqwest-middleware`] stack that applies a
/// per-request timeout and retries transient failures with exponential backoff.
///
/// [`reqwest-middleware`]: https://docs.rs/reqwest-middleware
#[derive(Debug, Clone)]
pub struct Client {
    inner: Arc<ClientImpl>,
}

/// The shared state behind a [`Client`], held by an [`Arc`] so clones are cheap.
#[derive(Debug)]
struct ClientImpl {
    http: ClientWithMiddleware,
    base_url: Url,
    token: Option<String>,
}

impl Client {
    /// Returns a new [`ClientBuilder`].
    pub fn builder() -> ClientBuilder {
        ClientBuilder::default()
    }

    /// The base URL this client targets.
    pub fn base_url(&self) -> &Url {
        &self.inner.base_url
    }

    /// Returns a handle to the service endpoint at `route`.
    ///
    /// `route` is the endpoint name with or without a leading slash. The returned
    /// [`Endpoint`] carries the route, so calls are made on it rather than passing
    /// the route to each method: `client.endpoint("summarize").call(&req)`.
    ///
    /// Accepts a `&'static str` (borrowed without allocating) or an owned `String`.
    pub fn endpoint(&self, route: impl Into<Cow<'static, str>>) -> Endpoint {
        Endpoint::new(self.clone(), route.into())
    }

    /// Returns whether the service reports itself ready, via `/readyz`.
    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), err))]
    pub async fn is_ready(&self) -> Result<bool> {
        self.health("readyz").await
    }

    /// Returns whether the service reports itself alive, via `/livez`.
    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), err))]
    pub async fn is_live(&self) -> Result<bool> {
        self.health("livez").await
    }

    /// Polls [`is_ready`] until it returns `true` or `timeout` elapses, awaiting
    /// `sleep(interval)` between attempts.
    ///
    /// The crate is runtime-agnostic, so the caller supplies the delay: `sleep` is
    /// invoked with `interval` and the returned future is awaited. With Tokio, pass
    /// `tokio::time::sleep`.
    ///
    /// Returns [`Error::Timeout`] if the service does not become ready in time.
    ///
    /// [`is_ready`]: Self::is_ready
    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, sleep), err))]
    pub async fn wait_until_ready<S, F>(
        &self,
        timeout: Duration,
        interval: Duration,
        mut sleep: S,
    ) -> Result<()>
    where
        S: FnMut(Duration) -> F + Send,
        F: Future<Output = ()> + Send,
    {
        let deadline = Instant::now() + timeout;
        loop {
            // Ignore transient errors (e.g. connection refused during startup);
            // only the deadline ends the loop.
            if matches!(self.is_ready().await, Ok(true)) {
                return Ok(());
            }
            if Instant::now() >= deadline {
                return Err(Error::Timeout { timeout });
            }
            sleep(interval).await;
        }
    }

    /// Joins a route onto the base URL, tolerating an optional leading slash.
    pub(crate) fn route_url(&self, route: &str) -> Result<Url> {
        let route = route.trim_start_matches('/');
        Ok(self.inner.base_url.join(route)?)
    }

    /// Like [`route_url`], but appends a single query parameter.
    ///
    /// [`route_url`]: Self::route_url
    pub(crate) fn endpoint_query(&self, route: &str, key: &str, value: &str) -> Result<Url> {
        let mut url = self.route_url(route)?;
        url.query_pairs_mut().append_pair(key, value);
        Ok(url)
    }

    /// Begins a `POST` request to `route`, with the bearer token applied.
    pub(crate) fn post(&self, route: &str) -> Result<RequestBuilder> {
        Ok(self.post_url(self.route_url(route)?))
    }

    /// Begins a `GET` request to `route`, with the bearer token applied.
    pub(crate) fn get(&self, route: &str) -> Result<RequestBuilder> {
        Ok(self.get_url(self.route_url(route)?))
    }

    /// Begins a `POST` request to a pre-built URL, with the bearer token applied.
    pub(crate) fn post_url(&self, url: Url) -> RequestBuilder {
        self.authed(self.inner.http.post(url))
    }

    /// Begins a `GET` request to a pre-built URL, with the bearer token applied.
    pub(crate) fn get_url(&self, url: Url) -> RequestBuilder {
        self.authed(self.inner.http.get(url))
    }

    /// Begins a `PUT` request to a pre-built URL, with the bearer token applied.
    pub(crate) fn put_url(&self, url: Url) -> RequestBuilder {
        self.authed(self.inner.http.put(url))
    }

    /// Applies the configured bearer token to a request builder, if any.
    pub(crate) fn authed(&self, req: RequestBuilder) -> RequestBuilder {
        match &self.inner.token {
            Some(token) => req.bearer_auth(token),
            None => req,
        }
    }

    /// Sends a request, mapping any non-success status to [`Error::Service`].
    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, err))]
    pub(crate) async fn send(&self, req: RequestBuilder) -> Result<ReqwestResponse> {
        let resp = req.send().await?;
        let status = resp.status();
        if !status.is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(Error::service(status.as_u16(), &body));
        }
        Ok(resp)
    }

    /// Issues a `GET` against a health endpoint, returning whether it is healthy.
    pub(crate) async fn health(&self, route: &str) -> Result<bool> {
        let resp = self.get(route)?.send().await?;
        Ok(resp.status().is_success())
    }

    /// Assembles a [`Client`] from resolved configuration, used by [`ClientBuilder`].
    ///
    /// Builds the HTTP middleware stack: a per-request timeout enforced by reqwest,
    /// with retries layered on top via reqwest-middleware so each attempt gets its
    /// own timeout.
    pub(crate) fn assemble(
        base_url: String,
        token: Option<String>,
        timeout: Option<Duration>,
        max_retries: u32,
        headers: reqwest::header::HeaderMap,
    ) -> Result<Self> {
        let mut http = reqwest::Client::builder().default_headers(headers);
        if let Some(timeout) = timeout {
            http = http.timeout(timeout);
        }
        let inner = http.build()?;

        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(max_retries);

        let http = MiddlewareBuilder::new(inner)
            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
            .build();

        let mut base_url = Url::parse(&base_url)?;

        // A base URL must end in `/` for [`Url::join`] to treat it as a directory
        // rather than replacing the final path segment.
        if !base_url.path().ends_with('/') {
            let path = format!("{}/", base_url.path());
            base_url.set_path(&path);
        }

        Ok(Self {
            inner: Arc::new(ClientImpl {
                http,
                base_url,
                token,
            }),
        })
    }
}

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

    fn client(base_url: &str) -> Client {
        Client::builder().with_base_url(base_url).build().unwrap()
    }

    #[test]
    fn base_url_gets_a_trailing_slash() {
        // Without normalization, `Url::join` would drop the final path segment.
        assert_eq!(
            client("http://localhost:3000").base_url().as_str(),
            "http://localhost:3000/"
        );
        assert_eq!(
            client("http://localhost:3000/").base_url().as_str(),
            "http://localhost:3000/"
        );
    }

    #[test]
    fn endpoint_joins_route_tolerating_leading_slash() {
        let c = client("http://localhost:3000");
        assert_eq!(
            c.route_url("summarize").unwrap().as_str(),
            "http://localhost:3000/summarize"
        );
        assert_eq!(
            c.route_url("/summarize").unwrap().as_str(),
            "http://localhost:3000/summarize"
        );
    }

    #[test]
    fn endpoint_preserves_a_path_prefix() {
        // A `/v1/` prefix must survive route joining (BentoML `path_prefix`).
        let c = client("http://localhost:3000/v1");
        assert_eq!(
            c.route_url("summarize").unwrap().as_str(),
            "http://localhost:3000/v1/summarize"
        );
    }
}