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}