apollo-http-client 0.2.0

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::time::Instant;

use apollo_opentelemetry::metrics::{
    FutureExt, HistogramExt, RecordDurationGuard, TrackGuard, UpDownCounterExt,
};
use hyper::client::conn::http2;
use hyper_rustls::HttpsConnector;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::client::pool::singleton::Singleton;
use opentelemetry::KeyValue;
use opentelemetry::metrics::UpDownCounter;
use opentelemetry_semantic_conventions::attribute as semconv;
use tower::BoxError;
use tower::Service;
use tower::ServiceExt as _;

use super::handshake::{FixedOrigin, Http2Handshake, TcpConnect};
use super::{
    HostSender, HttpBody, Origin, ProtocolVersionCell, into_http_client_error, retain_predicate,
    spawn_eviction_task,
};
use crate::error::HttpClientError;
use crate::metrics::{ConnectionMetrics, STATE_ACTIVE, STATE_IDLE};

// ---- Connection state -------------------------------------------------------

/// Tracks the `http.client.open_connections` state for an HTTP/2 connection.
///
/// Starts idle. Transitions to active when the first request starts, and back to idle
/// when the last in-flight request completes. Decrements the current state counter on drop.
struct ConnectionState {
    guard: Mutex<TrackGuard<i64>>,
    inflight: AtomicI32,
    idle_since: Mutex<Option<Instant>>,
}

impl ConnectionState {
    fn new(counter: &UpDownCounter<i64>, mut base_attrs: Vec<KeyValue>) -> Self {
        base_attrs.push(KeyValue::new(semconv::HTTP_CONNECTION_STATE, STATE_IDLE));
        Self {
            guard: Mutex::new(counter.track(base_attrs)),
            inflight: AtomicI32::new(0),
            idle_since: Mutex::new(Some(Instant::now())),
        }
    }

    fn request_start(&self) {
        // AcqRel: the fetch_add synchronises with any concurrent fetch_sub so only the
        // thread that observes the 0→1 transition (first in-flight request) flips to active.
        if self.inflight.fetch_add(1, Ordering::AcqRel) == 0 {
            *self.idle_since.lock().unwrap() = None;
            self.guard
                .lock()
                .unwrap()
                .set(KeyValue::new(semconv::HTTP_CONNECTION_STATE, STATE_ACTIVE));
        }
    }

    fn request_done(&self) {
        // Only the thread that observes the 1→0 transition (last in-flight request)
        // flips back to idle.
        if self.inflight.fetch_sub(1, Ordering::AcqRel) == 1 {
            *self.idle_since.lock().unwrap() = Some(Instant::now());
            self.guard
                .lock()
                .unwrap()
                .set(KeyValue::new(semconv::HTTP_CONNECTION_STATE, STATE_IDLE));
        }
    }

    fn idle_since(&self) -> Option<Instant> {
        *self.idle_since.lock().unwrap()
    }
}

/// Increments the inflight count on construction and decrements it on drop, transitioning
/// the connection active/idle as the first request starts and the last one completes.
struct RequestGuard(Arc<H2ConnectionInner>);

impl RequestGuard {
    fn new(inner: Arc<H2ConnectionInner>) -> Self {
        inner.state.request_start();
        Self(inner)
    }
}

impl Drop for RequestGuard {
    fn drop(&mut self) {
        self.0.state.request_done();
    }
}

/// Shared state for a single HTTP/2 connection.
struct H2ConnectionInner {
    state: ConnectionState,
    _connection_duration: RecordDurationGuard,
}

// ---- Connection -------------------------------------------------------------

/// An HTTP/2 connection that can be cloned for concurrent requests.
///
/// All clones share the same underlying TCP/TLS connection and its OTel metrics.
#[derive(Clone)]
pub(crate) struct H2Connection {
    sender: http2::SendRequest<HttpBody>,
    inner: Arc<H2ConnectionInner>,
}

impl H2Connection {
    pub(crate) fn idle_since(&self) -> Option<Instant> {
        self.inner.state.idle_since()
    }

    pub(crate) fn new(
        sender: http2::SendRequest<HttpBody>,
        metrics: &ConnectionMetrics,
        base_attrs: Vec<KeyValue>,
    ) -> Self {
        Self {
            sender,
            inner: Arc::new(H2ConnectionInner {
                state: ConnectionState::new(metrics.counter(), base_attrs.clone()),
                _connection_duration: metrics
                    .duration_histogram()
                    .record_duration_on_drop(base_attrs),
            }),
        }
    }
}

impl Service<http::Request<HttpBody>> for H2Connection {
    type Response = http::Response<hyper::body::Incoming>;
    type Error = hyper::Error;
    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.sender.poll_ready(cx)
    }

    fn call(&mut self, req: http::Request<HttpBody>) -> Self::Future {
        Box::pin(
            self.sender
                .send_request(req)
                .with_guard(RequestGuard::new(self.inner.clone())),
        )
    }
}

// ---- Connector --------------------------------------------------------------

/// Establishes HTTP/2 connections to a given [`Origin`].
#[derive(Clone)]
pub(crate) struct H2Connector {
    inner: HttpsConnector<HttpConnector>,
    metrics: ConnectionMetrics,
    connect_timeout: Duration,
    keep_alive_interval: Duration,
    keep_alive_timeout: Duration,
    keep_alive_while_idle: bool,
}

impl H2Connector {
    pub(crate) fn new(
        inner: HttpsConnector<HttpConnector>,
        metrics: ConnectionMetrics,
        connect_timeout: Duration,
        keep_alive_interval: Duration,
        keep_alive_timeout: Duration,
        keep_alive_while_idle: bool,
    ) -> Self {
        Self {
            inner,
            metrics,
            connect_timeout,
            keep_alive_interval,
            keep_alive_timeout,
            keep_alive_while_idle,
        }
    }

    pub(crate) fn keep_alive_interval(&self) -> Duration {
        self.keep_alive_interval
    }

    pub(crate) fn keep_alive_timeout(&self) -> Duration {
        self.keep_alive_timeout
    }

    pub(crate) fn keep_alive_while_idle(&self) -> bool {
        self.keep_alive_while_idle
    }
}

impl Service<Origin> for H2Connector {
    type Response = H2Connection;
    type Error = BoxError;
    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(BoxError::from)
    }

    fn call(&mut self, origin: Origin) -> Self::Future {
        let attrs = self
            .metrics
            .connection_attrs(Some(http::Version::HTTP_2), &origin);
        let clone = self.inner.clone();
        let inner = std::mem::replace(&mut self.inner, clone);
        let mut handshake = Http2Handshake {
            inner: FixedOrigin {
                inner: TcpConnect::new(inner, self.connect_timeout),
                origin,
            },
            metrics: self.metrics.clone(),
            attrs,
            keep_alive_interval: self.keep_alive_interval,
            keep_alive_timeout: self.keep_alive_timeout,
            keep_alive_while_idle: self.keep_alive_while_idle,
        };
        handshake.call(())
    }
}

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

/// Creates a [`HostSender`] backed by an HTTP/2 connection pool for the given host.
///
/// Accepts any [`Service`] producing an [`H2Connection`]. The call-key type `K` is
/// generic so callers can use either a [`Origin`] (when the connector dispatches on
/// scheme/authority) or `()` (when the key is pre-bound at construction, as in proxy
/// flows that wrap an inner connector with a fixed-key adapter). `decorate_request`
/// is called on every outgoing request before dispatch, allowing callers to inject
/// per-host headers or rewrite the request URI.
///
/// At most one TCP connection is maintained per host. Idle connections are evicted
/// based on `idle_timeout`, and the pool is cleaned up when the sender is dropped.
pub(crate) fn new_sender<C, K, D>(
    connector: C,
    idle_timeout: Duration,
    key: K,
    decorate_request: D,
) -> HostSender
where
    C: Service<K, Response = H2Connection> + Clone + Send + 'static,
    C::Error: Into<BoxError> + Send + 'static,
    C::Future: Send + 'static,
    K: Clone + Send + Sync + 'static,
    D: Fn(&mut http::Request<HttpBody>) + Send + Sync + 'static,
{
    // Singleton ensures only one TCP connection is ever established per host.
    // All concurrent callers share the same connection via Clone (HTTP/2 multiplexes).
    let singleton = Arc::new(Mutex::new(Singleton::new(connector)));
    let drop_guard = spawn_eviction_task(&singleton, idle_timeout, move |s, now| {
        let mut idle_count = 0;
        s.retain(|conn| retain_predicate(conn.idle_since(), idle_timeout, now, &mut idle_count, 1));
    });
    Arc::new(move |mut req, version: ProtocolVersionCell| {
        // Reference drop_guard so it is captured by this move closure and lives
        // as long as the Arc. Without this, it would not be captured and would
        // fire cancel() immediately after new_sender returns.
        let _guard = &drop_guard;
        let mut s = singleton.lock().unwrap().clone();
        let k = key.clone();
        decorate_request(&mut req);
        Box::pin(async move {
            version.set(http::Version::HTTP_2);
            // Backpressure is not propagated here: the pool selects a connection only at call time,
            // so there is no single connection whose readiness we could poll in advance. A new
            // connection also starts unready until established, which would cause spurious
            // not-ready signals. Meaningful backpressure would require pinning each caller to
            // exactly one target host.
            let mut conn = s
                .ready()
                .await
                .map_err(|e| into_http_client_error(e.into()))?
                .call(k)
                .await
                .map_err(|e| into_http_client_error(e.into()))?;
            let resp = conn
                .ready()
                .await
                .map_err(|source| HttpClientError::Request { source })?
                .call(req)
                .await
                .map_err(|source| HttpClientError::Request { source })?;
            Ok(resp)
        })
    })
}

#[cfg(test)]
mod tests {
    use apollo_opentelemetry_test::{TelemetryContext, assert_metrics_snapshot};
    use opentelemetry::KeyValue;
    use opentelemetry::global;

    use super::ConnectionState;

    fn make_counter() -> opentelemetry::metrics::UpDownCounter<i64> {
        global::meter_provider()
            .meter("test")
            .i64_up_down_counter("test.connections")
            .build()
    }

    fn base_attrs() -> Vec<KeyValue> {
        vec![KeyValue::new("server.address", "test-host")]
    }

    /// A new connection immediately increments the idle counter.
    #[test]
    fn new_connection_is_idle() {
        let ctx = TelemetryContext::new();
        let _state = ConnectionState::new(&make_counter(), base_attrs());
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 1
            is_monotonic: false
            temporality: Cumulative
        ");
    }

    /// The first in-flight request transitions the connection from idle to active.
    #[test]
    fn first_request_transitions_to_active() {
        let ctx = TelemetryContext::new();
        let state = ConnectionState::new(&make_counter(), base_attrs());
        state.request_start();
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: active
                  server.address: test-host
                value: 1
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 0
            is_monotonic: false
            temporality: Cumulative
        ");
    }

    /// A second concurrent request does not double-flip the counters.
    #[test]
    fn concurrent_requests_stay_active() {
        let ctx = TelemetryContext::new();
        let state = ConnectionState::new(&make_counter(), base_attrs());
        state.request_start();
        state.request_start(); // second concurrent request — should not re-flip
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: active
                  server.address: test-host
                value: 1
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 0
            is_monotonic: false
            temporality: Cumulative
        ");
    }

    /// When the last in-flight request finishes, the connection returns to idle.
    #[test]
    fn last_request_done_returns_to_idle() {
        let ctx = TelemetryContext::new();
        let state = ConnectionState::new(&make_counter(), base_attrs());
        state.request_start();
        state.request_start();
        state.request_done(); // one still in flight — stays active
        state.request_done(); // last one done — back to idle
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: active
                  server.address: test-host
                value: 0
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 1
            is_monotonic: false
            temporality: Cumulative
        ");
    }

    /// Dropping an idle connection decrements the idle counter.
    #[test]
    fn drop_idle_decrements_idle() {
        let ctx = TelemetryContext::new();
        let state = ConnectionState::new(&make_counter(), base_attrs());
        drop(state);
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 0
            is_monotonic: false
            temporality: Cumulative
        ");
    }

    /// Dropping a connection while requests are in-flight decrements the active
    /// counter, not the idle counter. This covers the case where a connection
    /// is torn down mid-request (e.g. GOAWAY or TCP RST).
    #[test]
    fn drop_with_inflight_requests_decrements_active() {
        let ctx = TelemetryContext::new();
        let state = ConnectionState::new(&make_counter(), base_attrs());
        state.request_start();
        drop(state);
        assert_metrics_snapshot!(ctx, @r"
        - name: test.connections
          data:
            type: Sum
            data_points:
              - attributes:
                  http.connection.state: active
                  server.address: test-host
                value: 0
              - attributes:
                  http.connection.state: idle
                  server.address: test-host
                value: 0
            is_monotonic: false
            temporality: Cumulative
        ");
    }
}