Skip to main content

praxis_core/subrequest/
internals.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2026 Praxis Contributors
3
4use std::{sync::Arc, time::Duration};
5
6use http::HeaderMap;
7use metrics::{counter, histogram};
8use pingora_core::{
9    connectors::{ConnectorOptions, http::Connector},
10    protocols::http::client::HttpSession,
11    upstreams::peer::{HttpPeer, Peer as _},
12};
13use tokio::sync::{OwnedSemaphorePermit, Semaphore};
14use tracing::debug;
15
16use super::types::{SubRequestError, SubResponse};
17use crate::circuit::{CircuitBreakerConfig, CircuitBreakerRegistry, CircuitToken, PeerKey};
18
19// ---------------------------------------------------------------------------
20// Metric names
21// ---------------------------------------------------------------------------
22
23/// Metric name for total streaming sub-request count (with termination label).
24pub(super) const SUBREQUEST_STREAMS_TOTAL: &str = "praxis_subrequest_streams_total";
25/// Metric name for streaming sub-request duration histogram.
26pub(super) const SUBREQUEST_STREAM_DURATION_SECONDS: &str = "praxis_subrequest_stream_duration_seconds";
27/// Metric name for streaming sub-request total bytes counter.
28pub(super) const SUBREQUEST_STREAM_BYTES_TOTAL: &str = "praxis_subrequest_stream_bytes_total";
29/// Metric name for header phase duration (shared buffered/streaming).
30pub(super) const SUBREQUEST_HEADER_DURATION_SECONDS: &str = "praxis_subrequest_header_duration_seconds";
31
32// ---------------------------------------------------------------------------
33// SubRequestConnectorOptions
34// ---------------------------------------------------------------------------
35
36/// Options for constructing a [`SubRequestConnector`].
37#[derive(Debug)]
38pub struct SubRequestConnectorOptions {
39    /// Number of idle connections to keep in the pool.
40    pub keepalive_pool_size: usize,
41
42    /// Maximum number of concurrently active exchanges.
43    pub max_connections: Option<usize>,
44
45    /// Circuit breaker configuration for peer-level failure tracking.
46    pub circuit_breaker: Option<CircuitBreakerConfig>,
47}
48
49// ---------------------------------------------------------------------------
50// SubRequestConnector
51// ---------------------------------------------------------------------------
52
53/// Shared HTTP connector for sub-requests.
54///
55/// Wraps Pingora's [`Connector`] behind an [`Arc`] so that all
56/// filter instances share a single connection pool. Created once at
57/// server startup and passed through unchanged on config reload.
58///
59/// An optional admission semaphore limits the number of concurrently
60/// active sub-request exchanges.
61///
62/// ```
63/// use praxis_core::subrequest::SubRequestConnector;
64///
65/// let connector = SubRequestConnector::new(128, None);
66/// let _clone = connector.clone();
67/// ```
68///
69/// [`Arc`]: std::sync::Arc
70/// [`Connector`]: pingora_core::connectors::http::Connector
71#[derive(Clone)]
72pub struct SubRequestConnector {
73    /// Shared Pingora HTTP connector.
74    pub(super) inner: Arc<Connector<()>>,
75
76    /// Admission semaphore bounding concurrently active exchanges.
77    pub(super) admission: Option<Arc<Semaphore>>,
78
79    /// The configured concurrency limit, retained for error reporting.
80    pub(super) configured_max_connections: Option<usize>,
81
82    /// Per-peer circuit breaker registry.
83    pub(super) circuit_breakers: Option<Arc<CircuitBreakerRegistry>>,
84}
85
86impl SubRequestConnector {
87    /// Create a connector with the given keepalive pool size and
88    /// optional active-connection limit.
89    ///
90    /// ```
91    /// use praxis_core::subrequest::SubRequestConnector;
92    ///
93    /// let connector = SubRequestConnector::new(64, None);
94    /// let bounded = SubRequestConnector::new(64, Some(256));
95    /// ```
96    pub fn new(keepalive_pool_size: usize, max_connections: Option<usize>) -> Self {
97        let options = ConnectorOptions::new(keepalive_pool_size);
98        Self {
99            inner: Arc::new(Connector::new(Some(options))),
100            admission: max_connections.map(|n| Arc::new(Semaphore::new(n))),
101            configured_max_connections: max_connections,
102            circuit_breakers: None,
103        }
104    }
105
106    /// Create a connector from [`SubRequestConnectorOptions`].
107    ///
108    /// ```
109    /// use praxis_core::subrequest::{SubRequestConnector, SubRequestConnectorOptions};
110    ///
111    /// let connector = SubRequestConnector::with_options(SubRequestConnectorOptions {
112    ///     keepalive_pool_size: 64,
113    ///     max_connections: Some(256),
114    ///     circuit_breaker: None,
115    /// });
116    /// ```
117    pub fn with_options(opts: SubRequestConnectorOptions) -> Self {
118        let options = ConnectorOptions::new(opts.keepalive_pool_size);
119        Self {
120            inner: Arc::new(Connector::new(Some(options))),
121            admission: opts.max_connections.map(|n| Arc::new(Semaphore::new(n))),
122            configured_max_connections: opts.max_connections,
123            circuit_breakers: opts
124                .circuit_breaker
125                .map(|cfg| Arc::new(CircuitBreakerRegistry::new(cfg))),
126        }
127    }
128
129    /// Access the underlying Pingora [`Connector`].
130    ///
131    /// [`Connector`]: pingora_core::connectors::http::Connector
132    pub fn connector(&self) -> &Connector<()> {
133        &self.inner
134    }
135
136    /// Acquire an admission permit if a concurrency limit is
137    /// configured. Returns `None` when no limit is set.
138    ///
139    /// The returned permit must be held for the entire sub-request
140    /// exchange. Dropping it releases the slot.
141    pub async fn acquire_permit(&self) -> Option<OwnedSemaphorePermit> {
142        let semaphore = self.admission.as_ref()?;
143        Arc::clone(semaphore).acquire_owned().await.ok()
144    }
145
146    /// Try to acquire an admission permit within the given deadline.
147    ///
148    /// Returns `Ok(Some(permit))` if acquired, `Ok(None)` if no
149    /// concurrency limit is configured, or `Err` if the deadline
150    /// expires before a slot opens.
151    ///
152    /// # Errors
153    ///
154    /// Returns [`SubRequestError::AdmissionTimeout`] when the
155    /// semaphore cannot be acquired within the timeout.
156    pub async fn try_acquire_permit(&self, timeout: Duration) -> Result<Option<OwnedSemaphorePermit>, SubRequestError> {
157        let Some(semaphore) = self.admission.as_ref() else {
158            return Ok(None);
159        };
160        let configured = self.configured_max_connections.unwrap_or(0);
161        match tokio::time::timeout(timeout, Arc::clone(semaphore).acquire_owned()).await {
162            Ok(Ok(permit)) => Ok(Some(permit)),
163            Ok(Err(_closed)) => Err(SubRequestError::AdmissionTimeout {
164                max_connections: configured,
165            }),
166            Err(_elapsed) => Err(SubRequestError::AdmissionTimeout {
167                max_connections: configured,
168            }),
169        }
170    }
171}
172
173impl std::fmt::Debug for SubRequestConnector {
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        f.debug_struct("SubRequestConnector")
176            .field("pool", &"Connector<()>")
177            .field("max_connections", &self.configured_max_connections)
178            .field("circuit_breakers", &self.circuit_breakers.is_some())
179            .finish()
180    }
181}
182
183// ---------------------------------------------------------------------------
184// RawExchange
185// ---------------------------------------------------------------------------
186
187/// Live HTTP exchange after validated response headers.
188///
189/// Returned by `open_exchange()` and consumed by either
190/// `execute()` (buffered collection) or `send_streaming()`
191/// (body ownership handoff).
192pub(super) struct RawExchange<'a> {
193    /// Live Pingora HTTP session.
194    pub(super) session: HttpSession<()>,
195    /// Peer address (timeout-bounded).
196    pub(super) peer: HttpPeer,
197    /// Connector ref for session release.
198    pub(super) connector: &'a SubRequestConnector,
199    /// HTTP status code.
200    pub(super) status: u16,
201    /// Sanitized response headers.
202    pub(super) headers: HeaderMap,
203    /// Optional circuit breaker guard.
204    pub(super) circuit_guard: Option<CircuitGuard<'a>>,
205    /// Admission permit.
206    pub(super) permit: Option<OwnedSemaphorePermit>,
207    /// Absolute deadline for the entire exchange.
208    pub(super) deadline: tokio::time::Instant,
209}
210
211// ---------------------------------------------------------------------------
212// Circuit Breaker Guard
213// ---------------------------------------------------------------------------
214
215/// RAII guard ensuring every acquired circuit token is finalized.
216///
217/// On drop without explicit [`finalize`](Self::finalize), records a
218/// failure — this covers deadline exits, panics, and any early-return
219/// path after token acquisition.
220pub(super) struct CircuitGuard<'a> {
221    /// The registry that issued the token.
222    registry: &'a CircuitBreakerRegistry,
223    /// Logical peer identity the token was acquired for.
224    peer: PeerKey,
225    /// The generation token; `None` after finalization.
226    token: Option<CircuitToken>,
227}
228
229impl<'a> CircuitGuard<'a> {
230    /// Create a guard from an acquired token.
231    pub(super) fn new(registry: &'a CircuitBreakerRegistry, peer: PeerKey, token: CircuitToken) -> Self {
232        Self {
233            registry,
234            peer,
235            token: Some(token),
236        }
237    }
238
239    /// Finalize the guard as a success regardless of later body outcome.
240    pub(super) fn finalize_success(mut self) {
241        if let Some(token) = self.token.take() {
242            self.registry.record_success(&self.peer, token);
243        }
244    }
245
246    /// Finalize the guard with the actual exchange outcome.
247    pub(super) fn finalize(mut self, result: &Result<SubResponse, SubRequestError>) {
248        let Some(token) = self.token.take() else {
249            return;
250        };
251        match result {
252            Err(SubRequestError::Connect(_) | SubRequestError::Io(_) | SubRequestError::DeadlineExceeded) => {
253                self.registry.record_failure(&self.peer, token);
254            },
255            Ok(_) | Err(_) => {
256                self.registry.record_success(&self.peer, token);
257            },
258        }
259    }
260}
261
262impl Drop for CircuitGuard<'_> {
263    fn drop(&mut self) {
264        if let Some(token) = self.token.take() {
265            self.registry.record_failure(&self.peer, token);
266        }
267    }
268}
269
270// ---------------------------------------------------------------------------
271// Protocol-aware completion check
272// ---------------------------------------------------------------------------
273
274/// Protocol-aware clean completion check.
275///
276/// Returns `Ok(true)` for clean EOF, `Ok(false)` for incomplete,
277/// and `Err` for H2 error-terminated streams.
278pub(super) fn check_clean_completion(session: &mut HttpSession<()>) -> Result<bool, SubRequestError> {
279    use pingora_core::protocols::http::custom::client::Session as _;
280    match session {
281        HttpSession::H1(h1) => Ok(h1.is_body_done()),
282        HttpSession::H2(h2) => h2
283            .check_response_end_or_error()
284            .map_err(|e| SubRequestError::Io(e.to_string())),
285        HttpSession::Custom(c) => Ok(c.response_finished()),
286    }
287}
288
289/// Record metrics and trace for header-phase stream termination.
290///
291/// Used by `send_streaming()` when the response completes at header
292/// time (HEAD, 204, 304, zero-length, or H2 error).
293pub(super) fn record_header_termination(termination: &str) {
294    counter!(SUBREQUEST_STREAMS_TOTAL, "termination" => termination.to_owned()).increment(1);
295    histogram!(SUBREQUEST_STREAM_DURATION_SECONDS).record(0.0);
296    debug!(termination, "sub-request: stream terminated at header phase");
297}
298
299// ---------------------------------------------------------------------------
300// Header sanitization
301// ---------------------------------------------------------------------------
302
303/// Headers that apply only to one HTTP connection and must not be
304/// forwarded across a sub-request boundary.
305pub(super) const HOP_BY_HOP_HEADERS: &[&str] = &[
306    "connection",
307    "keep-alive",
308    "proxy-authenticate",
309    "proxy-authorization",
310    "te",
311    "trailer",
312    "transfer-encoding",
313    "upgrade",
314];
315
316/// Remove hop-by-hop headers and headers nominated by `Connection`.
317pub(super) fn strip_hop_by_hop_headers(headers: &mut HeaderMap) {
318    let connection_values: Vec<_> = headers.get_all(http::header::CONNECTION).iter().cloned().collect();
319    for name in HOP_BY_HOP_HEADERS {
320        headers.remove(*name);
321    }
322    for value in connection_values {
323        let Ok(value) = value.to_str() else { continue };
324        for token in value.split(',').map(str::trim).filter(|token| !token.is_empty()) {
325            headers.remove(token);
326        }
327    }
328}
329
330/// Remove request framing headers that the executor re-computes.
331pub(super) fn strip_request_framing_headers(headers: &mut HeaderMap) {
332    headers.remove(http::header::CONTENT_LENGTH);
333    headers.remove(http::header::TRANSFER_ENCODING);
334}
335
336/// Remove headers matching reserved internal prefixes (`x-praxis-*`,
337/// `x-ext-protocol-*`, `x-ext-agent-*`).
338pub(super) fn strip_reserved_headers(headers: &mut HeaderMap) {
339    let reserved: Vec<http::header::HeaderName> = headers
340        .keys()
341        .filter(|name| crate::reserved_headers::is_reserved(name.as_str()))
342        .cloned()
343        .collect();
344    for name in reserved {
345        headers.remove(&name);
346    }
347}
348
349// ---------------------------------------------------------------------------
350// Helpers
351// ---------------------------------------------------------------------------
352
353/// Whether a header is a transport-level header that must not be
354/// injected via framework metadata.
355pub(super) fn is_transport_header(name: &http::header::HeaderName) -> bool {
356    HOP_BY_HOP_HEADERS.iter().any(|h| *h == name.as_str()) || name == http::header::CONTENT_LENGTH
357}
358
359/// Methods whose empty payload is commonly rejected without explicit framing.
360pub(super) fn empty_body_needs_framing(method: &http::Method) -> bool {
361    matches!(*method, http::Method::POST | http::Method::PUT | http::Method::PATCH)
362}
363
364/// Ensure HTTP/1.1 virtual hosting and HTTP/2 `:authority` are valid.
365pub(super) fn ensure_host_header(
366    request: &mut pingora_http::RequestHeader,
367    peer: &HttpPeer,
368) -> Result<(), SubRequestError> {
369    if !request.headers.contains_key(http::header::HOST) {
370        request
371            .insert_header(http::header::HOST, peer.address().to_string())
372            .map_err(|error| SubRequestError::InvalidRequest(error.to_string()))?;
373    }
374    Ok(())
375}
376
377/// Clamp connect timeouts to the remaining overall deadline.
378pub(super) fn clamp_peer_timeouts(peer: &mut HttpPeer, deadline: Duration) {
379    peer.options.connection_timeout = Some(min_timeout(peer.options.connection_timeout, deadline));
380    peer.options.total_connection_timeout = Some(min_timeout(peer.options.total_connection_timeout, deadline));
381}
382
383/// Keep an operator-configured timeout when it is stricter than the deadline.
384pub(super) fn min_timeout(configured: Option<Duration>, deadline: Duration) -> Duration {
385    configured.map_or(deadline, |configured| configured.min(deadline))
386}
387
388/// Classify a timeout expiry as either `DeadlineExceeded` (the overall
389/// request deadline fired) or `Io` (a shorter operator-configured
390/// read/write timeout fired).  Call this inside the `Err(_elapsed)`
391/// arm of `tokio::time::timeout` to preserve the 502-vs-504 distinction.
392///
393/// Uses the pre-computed budget and configured timeout rather than a
394/// post-hoc `Instant::now() >= deadline` check, which races with
395/// scheduler jitter when the configured timeout equals the remaining
396/// budget.
397pub(super) fn classify_timeout(
398    remaining_budget: Duration,
399    configured_timeout: Option<Duration>,
400    phase: &str,
401) -> SubRequestError {
402    if configured_timeout.is_none_or(|t| t >= remaining_budget) {
403        SubRequestError::DeadlineExceeded
404    } else {
405        SubRequestError::Io(format!("upstream {phase} timeout"))
406    }
407}