lambda_runtime_api_client 1.1.0

#![deny(clippy::all, clippy::cargo)]
#![warn(missing_docs, nonstandard_style, rust_2018_idioms)]
#![allow(clippy::multiple_crate_versions)]
#![cfg_attr(docsrs, feature(doc_cfg))]

//! This crate includes a base HTTP client to interact with
//! the AWS Lambda Runtime API.
use futures_util::{future::BoxFuture, FutureExt, TryFutureExt};
use http::{
    uri::{PathAndQuery, Scheme},
    Request, Response, Uri,
};
use hyper::body::Incoming;
use hyper_util::client::legacy::connect::HttpConnector;
use std::{convert::TryInto, fmt::Debug, future, sync::OnceLock};

const USER_AGENT_HEADER: &str = "User-Agent";
const DEFAULT_USER_AGENT: &str = concat!("aws-lambda-rust/", env!("CARGO_PKG_VERSION"));
const CUSTOM_USER_AGENT: Option<&str> = option_env!("LAMBDA_RUNTIME_USER_AGENT");

mod error;
pub use error::*;
pub mod body;

#[cfg(feature = "tracing")]
#[cfg_attr(docsrs, doc(cfg(feature = "tracing")))]
pub mod tracing;

/// The single method the runtime needs from any client: send a request to the
/// Lambda Runtime API (RAPID). Implemented by both [`Client`] and [`PooledClient`].
pub trait RuntimeApiClient {
    /// Send a given request to the Runtime API.
    fn call(&self, req: Request<body::Body>) -> BoxFuture<'static, Result<Response<Incoming>, BoxError>>;
}

/// API client to interact with the AWS Lambda Runtime API.
///
/// **Superseded by [`PooledClient`]**, which additionally supports rebuilding its
/// connection pool after a SnapStart VM restore via [`PooledClient::reset_pool`].
/// `Client` remains fully functional; it simply cannot reset its pool after a
/// restore. Prefer [`PooledClient`] for new code; `Client` is expected to be
/// removed in the next major version (2.0.0).
#[derive(Debug)]
pub struct Client {
    /// The runtime API URI
    pub base: Uri,
    /// The client that manages the API connections
    pub client: hyper_util::client::legacy::Client<HttpConnector, body::Body>,
}

impl Client {
    /// Create a builder struct to configure the client.
    pub fn builder() -> ClientBuilder {
        ClientBuilder::new()
    }

    /// Send a given request to the Runtime API.
    /// Use the client's base URI to ensure the API endpoint is correct.
    pub fn call(&self, req: Request<body::Body>) -> BoxFuture<'static, Result<Response<Incoming>, BoxError>> {
        <Self as RuntimeApiClient>::call(self, req)
    }

    /// Create a new client with a given base URI, HTTP connector, and optional pool size hint.
    fn with(base: Uri, connector: HttpConnector, pool_size: Option<usize>) -> Self {
        Self {
            base,
            client: build_hyper_client(&connector, pool_size),
        }
    }
}

impl RuntimeApiClient for Client {
    fn call(&self, req: Request<body::Body>) -> BoxFuture<'static, Result<Response<Incoming>, BoxError>> {
        // NOTE: This method returns a boxed future such that the future has a static lifetime.
        //       Due to limitations around the Rust async implementation as of Mar 2024, this is
        //       required to minimize constraints on the handler passed to [lambda_runtime::run].
        let req = match set_origin(&self.base, req) {
            Ok(req) => req,
            Err(err) => return future::ready(Err(err)).boxed(),
        };
        self.client.request(req).map_err(Into::into).boxed()
    }
}

/// API client to interact with the AWS Lambda Runtime API, with support for
/// rebuilding its connection pool after a SnapStart VM restore.
#[derive(Debug)]
pub struct PooledClient {
    /// The runtime API URI
    pub base: Uri,
    /// The client that manages the API connections
    client: hyper_util::client::legacy::Client<HttpConnector, body::Body>,
    /// The HTTP connector used to rebuild the client after a SnapStart restore
    connector: HttpConnector,
    /// Optional pool size hint for the hyper client
    pool_size: Option<usize>,
    /// Holds a freshly built client after a SnapStart restore. When set, [`PooledClient::call`]
    /// uses it instead of [`PooledClient::client`]; otherwise the original client is used.
    ///
    /// This is populated exactly once by [`PooledClient::reset_pool`] during the restore
    /// lifecycle, before any concurrent polling starts, so the `OnceLock` write never
    /// races with reads. Using `OnceLock` keeps `call` lock-free on the hot path.
    restored: OnceLock<hyper_util::client::legacy::Client<HttpConnector, body::Body>>,
}

impl PooledClient {
    /// Create a builder struct to configure the client.
    pub fn builder() -> PooledClientBuilder {
        PooledClientBuilder::new()
    }

    /// Send a given request to the Runtime API.
    /// Use the client's base URI to ensure the API endpoint is correct.
    pub fn call(&self, req: Request<body::Body>) -> BoxFuture<'static, Result<Response<Incoming>, BoxError>> {
        <Self as RuntimeApiClient>::call(self, req)
    }

    /// Create a new client with a given base URI, HTTP connector, and optional pool size hint.
    fn with(base: Uri, connector: HttpConnector, pool_size: Option<usize>) -> Self {
        let client = build_hyper_client(&connector, pool_size);
        Self {
            base,
            client,
            connector,
            pool_size,
            restored: OnceLock::new(),
        }
    }

    /// Reset the internal connection pool by building a fresh hyper client that
    /// subsequent calls will use.
    ///
    /// This is useful after a SnapStart VM restore where existing connections
    /// to the Lambda Runtime API (RAPID) are stale and must be discarded. It is
    /// called exactly once during the restore lifecycle, before concurrent
    /// polling starts, so the underlying `OnceLock` write never races.
    pub fn reset_pool(&self) {
        let _ = self.restored.set(build_hyper_client(&self.connector, self.pool_size));
    }
}

impl RuntimeApiClient for PooledClient {
    fn call(&self, req: Request<body::Body>) -> BoxFuture<'static, Result<Response<Incoming>, BoxError>> {
        let req = match set_origin(&self.base, req) {
            Ok(req) => req,
            Err(err) => return future::ready(Err(err)).boxed(),
        };
        // After a SnapStart restore, `restored` holds a fresh client with no stale
        // connections; otherwise fall back to the original client. Lock-free read.
        let hyper_client = self.restored.get().unwrap_or(&self.client);
        hyper_client.request(req).map_err(Into::into).boxed()
    }
}

/// Build a fresh hyper client from the given connector and pool size settings.
fn build_hyper_client(
    connector: &HttpConnector,
    pool_size: Option<usize>,
) -> hyper_util::client::legacy::Client<HttpConnector, body::Body> {
    let mut builder = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new());
    builder.http1_max_buf_size(1024 * 1024);

    if let Some(size) = pool_size {
        builder.pool_max_idle_per_host(size);
    }

    builder.build(connector.clone())
}

/// Rewrite a request's URI against the client's base (scheme/authority/base path).
fn set_origin<B>(base: &Uri, req: Request<B>) -> Result<Request<B>, BoxError> {
    let (mut parts, body) = req.into_parts();
    let (scheme, authority, base_path) = {
        let scheme = base.scheme().unwrap_or(&Scheme::HTTP);
        let authority = base.authority().expect("Authority not found");
        let base_path = base.path().trim_end_matches('/');
        (scheme, authority, base_path)
    };
    let path = parts.uri.path_and_query().expect("PathAndQuery not found");
    let pq: PathAndQuery = format!("{base_path}{path}").parse().expect("PathAndQuery invalid");

    let uri = Uri::builder()
        .scheme(scheme.as_ref())
        .authority(authority.as_ref())
        .path_and_query(pq)
        .build()
        .map_err(Box::new)?;

    parts.uri = uri;
    Ok(Request::from_parts(parts, body))
}

/// Builder implementation to construct any Runtime API clients.
pub struct ClientBuilder {
    connector: HttpConnector,
    uri: Option<http::Uri>,
    pool_size: Option<usize>,
}

impl ClientBuilder {
    fn new() -> ClientBuilder {
        ClientBuilder {
            connector: HttpConnector::new(),
            uri: None,
            pool_size: None,
        }
    }

    /// Create a new builder with a given HTTP connector.
    pub fn with_connector(self, connector: HttpConnector) -> ClientBuilder {
        ClientBuilder {
            connector,
            uri: self.uri,
            pool_size: self.pool_size,
        }
    }

    /// Create a new builder with a given base URI.
    /// Inherits all other attributes from the existent builder.
    pub fn with_endpoint(self, uri: http::Uri) -> Self {
        Self { uri: Some(uri), ..self }
    }

    /// Provide a pool size hint for the underlying Hyper client.
    ///
    /// When using concurrent polling, this should be at least the maximum
    /// concurrency (e.g., `AWS_LAMBDA_MAX_CONCURRENCY`) to avoid connection
    /// starvation.
    pub fn with_pool_size(self, pool_size: usize) -> Self {
        Self {
            pool_size: Some(pool_size),
            ..self
        }
    }

    /// Create the new client to interact with the Runtime API.
    ///
    /// **Superseded by [`PooledClient::builder`]**, which produces a
    /// [`PooledClient`] that supports SnapStart connection-pool reset. Prefer
    /// `PooledClient::builder().build()` for new code; this method is expected to
    /// be removed in the next major version (2.0.0).
    pub fn build(self) -> Result<Client, Error> {
        let uri = resolve_uri(self.uri);
        Ok(Client::with(uri, self.connector, self.pool_size))
    }
}

/// Builder to construct a [`PooledClient`].
pub struct PooledClientBuilder {
    connector: HttpConnector,
    uri: Option<http::Uri>,
    pool_size: Option<usize>,
}

impl PooledClientBuilder {
    fn new() -> PooledClientBuilder {
        PooledClientBuilder {
            connector: HttpConnector::new(),
            uri: None,
            pool_size: None,
        }
    }

    /// Use a given HTTP connector.
    pub fn with_connector(self, connector: HttpConnector) -> PooledClientBuilder {
        PooledClientBuilder { connector, ..self }
    }

    /// Use a given base URI. Inherits all other attributes from the existing builder.
    pub fn with_endpoint(self, uri: http::Uri) -> Self {
        Self { uri: Some(uri), ..self }
    }

    /// Provide a pool size hint for the underlying Hyper client.
    ///
    /// When using concurrent polling, this should be at least the maximum
    /// concurrency (e.g., `AWS_LAMBDA_MAX_CONCURRENCY`) to avoid connection
    /// starvation.
    pub fn with_pool_size(self, pool_size: usize) -> Self {
        Self {
            pool_size: Some(pool_size),
            ..self
        }
    }

    /// Create the new [`PooledClient`] to interact with the Runtime API.
    pub fn build(self) -> PooledClient {
        let uri = resolve_uri(self.uri);
        PooledClient::with(uri, self.connector, self.pool_size)
    }
}

/// Resolve the base URI from an explicit value or the `AWS_LAMBDA_RUNTIME_API` env var.
fn resolve_uri(uri: Option<http::Uri>) -> Uri {
    match uri {
        Some(uri) => uri,
        None => {
            let uri = std::env::var("AWS_LAMBDA_RUNTIME_API").expect("Missing AWS_LAMBDA_RUNTIME_API env var");
            uri.try_into().expect("Unable to convert to URL")
        }
    }
}

/// Create a request builder.
/// This builder uses `aws-lambda-rust/CRATE_VERSION` as
/// the default User-Agent.
/// Configure environment variable `LAMBDA_RUNTIME_USER_AGENT`
/// at compile time to modify User-Agent value.
pub fn build_request() -> http::request::Builder {
    const USER_AGENT: &str = match CUSTOM_USER_AGENT {
        Some(value) => value,
        None => DEFAULT_USER_AGENT,
    };
    http::Request::builder().header(USER_AGENT_HEADER, USER_AGENT)
}

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

    #[test]
    fn test_set_origin() {
        let base: Uri = "http://localhost:9001".parse().unwrap();
        let req = build_request()
            .uri("/2018-06-01/runtime/invocation/next")
            .body(())
            .unwrap();
        let req = set_origin(&base, req).unwrap();
        assert_eq!(
            "http://localhost:9001/2018-06-01/runtime/invocation/next",
            &req.uri().to_string()
        );
    }

    #[test]
    fn test_set_origin_with_base_path() {
        let base: Uri = "http://localhost:9001/foo".parse().unwrap();
        let req = build_request()
            .uri("/2018-06-01/runtime/invocation/next")
            .body(())
            .unwrap();
        let req = set_origin(&base, req).unwrap();
        assert_eq!(
            "http://localhost:9001/foo/2018-06-01/runtime/invocation/next",
            &req.uri().to_string()
        );

        let base: Uri = "http://localhost:9001/foo/".parse().unwrap();
        let req = build_request()
            .uri("/2018-06-01/runtime/invocation/next")
            .body(())
            .unwrap();
        let req = set_origin(&base, req).unwrap();
        assert_eq!(
            "http://localhost:9001/foo/2018-06-01/runtime/invocation/next",
            &req.uri().to_string()
        );
    }

    #[test]
    fn builder_accepts_pool_size() {
        let base = "http://localhost:9001";
        let expected: Uri = base.parse().unwrap();
        let client = PooledClient::builder()
            .with_pool_size(4)
            .with_endpoint(base.parse().unwrap())
            .build();

        assert_eq!(client.base, expected);
    }

    #[test]
    fn test_reset_pool() {
        let base = "http://localhost:9001";
        let client = PooledClient::builder()
            .with_pool_size(4)
            .with_endpoint(base.parse().unwrap())
            .build();
        client.reset_pool();
        assert_eq!(client.base, base.parse::<Uri>().unwrap());
    }
}