apollo-http-client 0.3.0

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;

use apollo_opentelemetry::metrics::UpDownCounterExt;
use hyper_rustls::HttpsConnector;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::client::pool;
use hyper_util::client::pool::singleton::Singleton;
use hyper_util::rt::{TokioExecutor, TokioIo};
use opentelemetry::KeyValue;
use opentelemetry_semantic_conventions::attribute as semconv;
use tower::BoxError;
use tower::Service;
use tower::ServiceExt as _;

use super::h1::H1Connector;
use super::h2::H2Connector;
use super::handshake::{Http1Handshake, Http2Handshake};
use super::{
    HostSender, Origin, ProtocolVersionCell, into_http_client_error, retain_predicate,
    spawn_eviction_task,
};
use crate::error::HttpClientError;
use crate::metrics::{ConnectionMetrics, STATE_CUSTOM_CONNECTING};

// ---- TLS connector ----------------------------------------------------------

/// Establishes a TCP/TLS connection and returns the raw IO stream.
///
/// Plain HTTP connections always fall back to HTTP/1.1 — h2c is not supported here.
/// Use [`Protocol::Http2`](crate::Protocol::Http2) for HTTP/2 cleartext.
#[derive(Clone)]
struct TlsConnector {
    inner: HttpsConnector<HttpConnector>,
    metrics: ConnectionMetrics,
    connect_timeout: Duration,
}

impl TlsConnector {
    fn new(
        inner: HttpsConnector<HttpConnector>,
        metrics: ConnectionMetrics,
        connect_timeout: Duration,
    ) -> Self {
        Self {
            inner,
            metrics,
            connect_timeout,
        }
    }
}

impl Service<Origin> for TlsConnector {
    type Response = <HttpsConnector<HttpConnector> as Service<http::Uri>>::Response;
    type Error = HttpClientError;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner
            .poll_ready(cx)
            .map_err(|source| HttpClientError::Transport { source })
    }

    fn call(&mut self, origin: Origin) -> Self::Future {
        let connector = self.inner.clone();
        let connector = std::mem::replace(&mut self.inner, connector);
        let metrics = self.metrics.clone();
        let connect_timeout = self.connect_timeout;

        Box::pin(async move {
            // Track as connecting while TCP+TLS handshake runs.
            // Dropped when TLS completes and raw IO is returned to the caller.
            let mut attrs = metrics.connection_attrs(None, &origin);
            attrs.push(KeyValue::new(
                semconv::HTTP_CONNECTION_STATE,
                STATE_CUSTOM_CONNECTING,
            ));
            let _connecting = metrics.counter().track(attrs);

            super::connect_with_timeout(connector, origin, connect_timeout).await
        })
    }
}

// ---- Pool -------------------------------------------------------------------

/// Pool and keep-alive timing parameters threaded into the ALPN sender.
pub(crate) struct AlpnPoolConfig {
    pub(crate) h1_max_idle: usize,
    pub(crate) h1_idle_timeout: Duration,
    pub(crate) h2_keep_alive_interval: Duration,
    pub(crate) h2_keep_alive_timeout: Duration,
    pub(crate) h2_keep_alive_while_idle: bool,
    pub(crate) h2_idle_timeout: Duration,
}

/// Creates a [`HostSender`] that negotiates HTTP/1.1 or HTTP/2 via ALPN over a direct
/// TCP/TLS connection. Thin wrapper around [`new_sender_with_connector`] for the
/// common direct-TLS case; see that function for full behaviour.
pub(crate) fn new_sender(
    h1: H1Connector,
    h1_max_idle: usize,
    h1_idle_timeout: Duration,
    h2: H2Connector,
    h2_idle_timeout: Duration,
    origin: Origin,
) -> HostSender {
    let (h1_inner, h1_metrics, h1_connect_timeout) = h1.into_parts();
    new_sender_with_connector(
        TlsConnector::new(h1_inner, h1_metrics.clone(), h1_connect_timeout),
        |io: &hyper_rustls::MaybeHttpsStream<TokioIo<tokio::net::TcpStream>>| match io {
            // Inspect the ALPN protocol negotiated during the TLS handshake.
            // Plain HTTP has no ALPN and falls back to H1.
            hyper_rustls::MaybeHttpsStream::Https(tls_io) => {
                let (_, rustls_conn) = tls_io.inner().get_ref();
                rustls_conn.alpn_protocol() == Some(b"h2")
            }
            hyper_rustls::MaybeHttpsStream::Http(_) => false,
        },
        h1_metrics,
        AlpnPoolConfig {
            h1_max_idle,
            h1_idle_timeout,
            h2_keep_alive_interval: h2.keep_alive_interval(),
            h2_keep_alive_timeout: h2.keep_alive_timeout(),
            h2_keep_alive_while_idle: h2.keep_alive_while_idle(),
            h2_idle_timeout,
        },
        origin,
    )
}

/// Creates a [`HostSender`] that negotiates HTTP/1.1 or HTTP/2 via ALPN.
///
/// Accepts any [`Service`] producing a raw IO stream keyed by [`Origin`], plus an
/// `inspect` predicate that returns `true` when the IO has negotiated HTTP/2 (e.g.
/// by reading the ALPN protocol from the underlying TLS session).
///
/// HTTP/2 connections are multiplexed over a single TCP connection; HTTP/1.1
/// connections use a pool. If the HTTP/2 connection dies, the next request
/// transparently re-negotiates.
///
/// Idle connections are evicted based on the configured timeouts, and both pools
/// are cleaned up when the sender is dropped.
pub(crate) fn new_sender_with_connector<C, IO, F>(
    connector: C,
    inspect: F,
    metrics: ConnectionMetrics,
    config: AlpnPoolConfig,
    origin: Origin,
) -> HostSender
where
    C: Service<Origin, Response = IO> + Clone + Send + Sync + 'static,
    C::Error: Into<BoxError>,
    C::Future: Send + 'static,
    IO: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static,
    F: Fn(&IO) -> bool + Clone + Send + Sync + 'static,
{
    let AlpnPoolConfig {
        h1_max_idle,
        h1_idle_timeout,
        h2_keep_alive_interval,
        h2_keep_alive_timeout,
        h2_keep_alive_while_idle,
        h2_idle_timeout,
    } = config;

    let mut negotiate =
        pool::negotiate::builder()
            .connect(connector)
            .inspect(inspect)
            .fallback(tower::layer::layer_fn({
                let metrics = metrics.clone();
                let attrs = metrics.connection_attrs(Some(http::Version::HTTP_11), &origin);
                move |inspector| {
                    pool::cache::builder().executor(TokioExecutor::new()).build(
                        Http1Handshake::new(inspector, metrics.clone(), attrs.clone()),
                    )
                }
            }))
            .upgrade(tower::layer::layer_fn({
                let metrics = metrics;
                let attrs = metrics.connection_attrs(Some(http::Version::HTTP_2), &origin);
                move |inspected| {
                    Singleton::new(Http2Handshake {
                        inner: inspected,
                        metrics: metrics.clone(),
                        attrs: attrs.clone(),
                        keep_alive_interval: h2_keep_alive_interval,
                        keep_alive_timeout: h2_keep_alive_timeout,
                        keep_alive_while_idle: h2_keep_alive_while_idle,
                    })
                }
            }))
            .build::<Origin>();

    // Clone pool services for eviction tasks. Clones share the same internal Arc
    // state as the services inside `negotiate`, so eviction affects the live pools.
    let h1_evict = Arc::new(Mutex::new(negotiate.fallback_mut().clone()));
    let h2_evict = Arc::new(Mutex::new(negotiate.upgrade_mut().clone()));

    let h1_guard = spawn_eviction_task(&h1_evict, h1_idle_timeout, move |cache, now| {
        let mut count = 0;
        cache.retain(|conn| {
            retain_predicate(
                conn.idle_since(),
                h1_idle_timeout,
                now,
                &mut count,
                h1_max_idle,
            )
        });
    });
    let h2_guard = spawn_eviction_task(&h2_evict, h2_idle_timeout, move |singleton, now| {
        singleton
            .retain(|conn| retain_predicate(conn.idle_since(), h2_idle_timeout, now, &mut 0, 1));
    });

    Arc::new(move |req, version: ProtocolVersionCell| {
        // Reference both guards so they are captured by this move closure and live
        // as long as the Arc. Without this, they would not be captured and would
        // fire cancel() immediately after new_sender returns.
        let _guards = (&h1_guard, &h2_guard);
        let mut pool = negotiate.clone();
        let k = origin.clone();
        Box::pin(async move {
            pool.ready().await.map_err(into_http_client_error)?;
            let mut svc = pool.call(k).await.map_err(into_http_client_error)?;
            if svc.upgraded_ref().is_some() {
                version.set(http::Version::HTTP_2);
            } else {
                version.set(http::Version::HTTP_11);
            }
            svc.ready()
                .await
                .map_err(|source| HttpClientError::Request { source })?
                .call(req)
                .await
                .map_err(|source| HttpClientError::Request { source })
        })
    })
}