1use std::{sync::Arc, time::Duration};
22
23use bytes::Bytes;
24use http::HeaderMap;
25use pingora_core::{
26 connectors::{ConnectorOptions, http::Connector},
27 upstreams::peer::{HttpPeer, Peer as _},
28};
29use thiserror::Error;
30use tokio::sync::{OwnedSemaphorePermit, Semaphore};
31use tracing::{debug, warn};
32
33use crate::circuit::{CircuitBreakerConfig, CircuitBreakerRegistry, CircuitCheck, CircuitToken, PeerKey};
34
35#[derive(Clone, Debug)]
41pub struct SubRequest {
42 pub method: http::Method,
44
45 pub uri: http::Uri,
47
48 pub headers: HeaderMap,
55
56 pub body: Bytes,
58}
59
60#[derive(Clone, Debug)]
62pub struct SubResponse {
63 pub status: u16,
65
66 pub headers: HeaderMap,
68
69 pub body: Bytes,
71}
72
73pub const DEPTH_HEADER: &str = "x-praxis-iterative-depth";
79
80#[derive(Clone, Debug, Default)]
88pub struct FrameworkHeaders {
89 entries: Vec<(http::header::HeaderName, http::HeaderValue)>,
91}
92
93impl FrameworkHeaders {
94 #[must_use]
96 pub fn new() -> Self {
97 Self::default()
98 }
99
100 pub fn insert(&mut self, name: http::header::HeaderName, value: http::HeaderValue) -> Result<(), SubRequestError> {
108 if is_transport_header(&name) {
109 return Err(SubRequestError::InvalidRequest(format!(
110 "transport header `{name}` cannot be injected as framework metadata"
111 )));
112 }
113 if crate::reserved_headers::is_reserved(name.as_str()) {
114 return Err(SubRequestError::InvalidRequest(format!(
115 "reserved header `{name}` cannot be injected as framework metadata"
116 )));
117 }
118 self.entries.push((name, value));
119 Ok(())
120 }
121
122 pub fn set_depth(&mut self, depth: u8) {
128 let value = http::HeaderValue::from(u16::from(depth));
129 self.entries
130 .push((http::header::HeaderName::from_static(DEPTH_HEADER), value));
131 }
132
133 pub fn iter(&self) -> impl Iterator<Item = &(http::header::HeaderName, http::HeaderValue)> {
135 self.entries.iter()
136 }
137
138 #[must_use]
140 pub fn is_empty(&self) -> bool {
141 self.entries.is_empty()
142 }
143}
144
145#[derive(Debug, Error)]
154#[non_exhaustive]
155pub enum SubRequestError {
156 #[error("sub-request construction error: {0}")]
158 InvalidRequest(String),
159
160 #[error("sub-request admission timeout (all {max_connections} slots busy)")]
163 AdmissionTimeout {
164 max_connections: usize,
166 },
167
168 #[error("sub-request connect error: {0}")]
170 Connect(String),
171
172 #[error("sub-request I/O error: {0}")]
174 Io(String),
175
176 #[error("sub-request deadline exceeded")]
178 DeadlineExceeded,
179
180 #[error("sub-request circuit open for peer {peer}")]
182 CircuitOpen {
183 peer: String,
185 },
186
187 #[error(
189 "sub-request response body exceeded limit \
190 ({actual} > {limit} bytes)"
191 )]
192 ResponseTooLarge {
193 actual: usize,
195 limit: usize,
197 },
198}
199
200#[derive(Debug)]
206pub struct SubRequestConnectorOptions {
207 pub keepalive_pool_size: usize,
209
210 pub max_connections: Option<usize>,
212
213 pub circuit_breaker: Option<CircuitBreakerConfig>,
215}
216
217#[derive(Clone)]
236pub struct SubRequestConnector {
237 inner: Arc<Connector<()>>,
239
240 admission: Option<Arc<Semaphore>>,
242
243 configured_max_connections: Option<usize>,
245
246 circuit_breakers: Option<Arc<CircuitBreakerRegistry>>,
248}
249
250impl SubRequestConnector {
251 pub fn new(keepalive_pool_size: usize, max_connections: Option<usize>) -> Self {
261 let options = ConnectorOptions::new(keepalive_pool_size);
262 Self {
263 inner: Arc::new(Connector::new(Some(options))),
264 admission: max_connections.map(|n| Arc::new(Semaphore::new(n))),
265 configured_max_connections: max_connections,
266 circuit_breakers: None,
267 }
268 }
269
270 pub fn with_options(opts: SubRequestConnectorOptions) -> Self {
282 let options = ConnectorOptions::new(opts.keepalive_pool_size);
283 Self {
284 inner: Arc::new(Connector::new(Some(options))),
285 admission: opts.max_connections.map(|n| Arc::new(Semaphore::new(n))),
286 configured_max_connections: opts.max_connections,
287 circuit_breakers: opts
288 .circuit_breaker
289 .map(|cfg| Arc::new(CircuitBreakerRegistry::new(cfg))),
290 }
291 }
292
293 pub fn connector(&self) -> &Connector<()> {
297 &self.inner
298 }
299
300 pub async fn acquire_permit(&self) -> Option<OwnedSemaphorePermit> {
306 let semaphore = self.admission.as_ref()?;
307 Arc::clone(semaphore).acquire_owned().await.ok()
308 }
309
310 pub async fn try_acquire_permit(&self, timeout: Duration) -> Result<Option<OwnedSemaphorePermit>, SubRequestError> {
321 let Some(semaphore) = self.admission.as_ref() else {
322 return Ok(None);
323 };
324 let configured = self.configured_max_connections.unwrap_or(0);
325 match tokio::time::timeout(timeout, Arc::clone(semaphore).acquire_owned()).await {
326 Ok(Ok(permit)) => Ok(Some(permit)),
327 Ok(Err(_closed)) => Err(SubRequestError::AdmissionTimeout {
328 max_connections: configured,
329 }),
330 Err(_elapsed) => Err(SubRequestError::AdmissionTimeout {
331 max_connections: configured,
332 }),
333 }
334 }
335}
336
337impl std::fmt::Debug for SubRequestConnector {
338 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
339 f.debug_struct("SubRequestConnector")
340 .field("pool", &"Connector<()>")
341 .field("max_connections", &self.configured_max_connections)
342 .field("circuit_breakers", &self.circuit_breakers.is_some())
343 .finish()
344 }
345}
346
347#[derive(Clone, Debug)]
377pub struct SubRequestClient {
378 connector: SubRequestConnector,
380
381 max_response_bytes: usize,
384}
385
386impl SubRequestClient {
387 pub fn new(connector: SubRequestConnector) -> Self {
396 Self {
397 connector,
398 max_response_bytes: crate::config::ABSOLUTE_MAX_BODY_BYTES,
399 }
400 }
401
402 pub fn with_max_response_bytes(connector: SubRequestConnector, max_response_bytes: usize) -> Self {
408 Self {
409 connector,
410 max_response_bytes,
411 }
412 }
413
414 pub fn connector(&self) -> &SubRequestConnector {
416 &self.connector
417 }
418
419 pub fn evict_idle_circuits(&self, idle_threshold: Duration) -> usize {
423 self.connector
424 .circuit_breakers
425 .as_ref()
426 .map_or(0, |registry| registry.evict_idle(idle_threshold))
427 }
428
429 #[expect(
451 clippy::too_many_arguments,
452 reason = "framework_headers is the typed metadata injection point"
453 )]
454 #[expect(
455 clippy::too_many_lines,
456 reason = "circuit + admission + deadline logic is sequential"
457 )]
458 pub async fn execute(
459 &self,
460 peer: &HttpPeer,
461 request: &SubRequest,
462 max_response_bytes: usize,
463 timeout: Duration,
464 framework_headers: Option<&FrameworkHeaders>,
465 ) -> Result<SubResponse, SubRequestError> {
466 let deadline = tokio::time::Instant::now() + timeout;
467 let mut bounded_peer = peer.clone();
468 clamp_peer_timeouts(&mut bounded_peer, timeout);
469
470 let peer_key: Option<PeerKey> = bounded_peer.address().as_inet().copied().map(|addr| {
474 let sni = &bounded_peer.sni;
475 PeerKey::new(addr, sni.as_str())
476 });
477
478 if let (Some(registry), Some(key)) = (&self.connector.circuit_breakers, &peer_key)
480 && !registry.precheck(key)
481 {
482 return Err(SubRequestError::CircuitOpen { peer: key.to_string() });
483 }
484
485 let admission_budget = deadline.saturating_duration_since(tokio::time::Instant::now());
486 if admission_budget.is_zero() {
487 return Err(SubRequestError::DeadlineExceeded);
488 }
489 let _permit = self.connector.try_acquire_permit(admission_budget).await?;
490
491 let circuit_guard = match (&self.connector.circuit_breakers, peer_key) {
494 (Some(registry), Some(key)) => match registry.try_acquire(key.clone()) {
495 CircuitCheck::Rejected => {
496 return Err(SubRequestError::CircuitOpen { peer: key.to_string() });
497 },
498 CircuitCheck::Allowed(token) => Some(CircuitGuard::new(registry, key, token)),
499 },
500 _ => None,
501 };
502
503 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
504 if remaining.is_zero() {
505 return Err(SubRequestError::DeadlineExceeded);
507 }
508 let effective_limit = max_response_bytes.min(self.max_response_bytes);
509 let result: Result<SubResponse, SubRequestError> = tokio::time::timeout(
510 remaining,
511 Box::pin(execute_inner(
512 &self.connector,
513 &bounded_peer,
514 request,
515 effective_limit,
516 remaining,
517 framework_headers,
518 )),
519 )
520 .await
521 .unwrap_or_else(|_elapsed| Err(SubRequestError::DeadlineExceeded));
522
523 if let Some(guard) = circuit_guard {
526 guard.finalize(&result);
527 }
528
529 result
530 }
531}
532
533#[expect(clippy::large_stack_frames, reason = "Pingora session types are large")]
540#[expect(clippy::too_many_lines, reason = "sequential HTTP exchange steps")]
541#[expect(clippy::too_many_arguments, reason = "internal function, all parameters required")]
542async fn execute_inner(
543 connector: &SubRequestConnector,
544 peer: &HttpPeer,
545 request: &SubRequest,
546 max_response_bytes: usize,
547 timeout: Duration,
548 framework_headers: Option<&FrameworkHeaders>,
549) -> Result<SubResponse, SubRequestError> {
550 let (mut session, reused) = Box::pin(connector.connector().get_http_session(peer))
551 .await
552 .map_err(|e| SubRequestError::Connect(e.to_string()))?;
553
554 debug!(
555 peer = %peer.address(),
556 reused,
557 method = %request.method,
558 uri = %request.uri,
559 "sub-request: connected"
560 );
561
562 session.set_read_timeout(Some(min_timeout(peer.options.read_timeout, timeout)));
563 session.set_write_timeout(Some(min_timeout(peer.options.write_timeout, timeout)));
564
565 let path = request
566 .uri
567 .path_and_query()
568 .map_or(b"/".as_slice(), |pq| pq.as_str().as_bytes());
569 let mut req_header = pingora_http::RequestHeader::build(request.method.clone(), path, None)
570 .map_err(|e| SubRequestError::InvalidRequest(e.to_string()))?;
571
572 let mut sanitized = request.headers.clone();
573 strip_hop_by_hop_headers(&mut sanitized);
574 strip_request_framing_headers(&mut sanitized);
575 strip_reserved_headers(&mut sanitized);
576 if let Some(fw) = framework_headers {
577 for (name, value) in fw.iter() {
578 sanitized.insert(name.clone(), value.clone());
579 }
580 }
581
582 for (name, value) in &sanitized {
583 let _append = req_header.append_header(name.clone(), value.clone());
584 }
585
586 ensure_host_header(&mut req_header, peer)?;
587
588 if !request.body.is_empty() || empty_body_needs_framing(&request.method) {
589 let _cl = req_header.insert_header("Content-Length", request.body.len().to_string());
590 }
591
592 session
593 .write_request_header(Box::new(req_header))
594 .await
595 .map_err(|e| SubRequestError::Io(e.to_string()))?;
596
597 if !request.body.is_empty() {
598 session
599 .write_request_body(request.body.clone(), true)
600 .await
601 .map_err(|e| SubRequestError::Io(e.to_string()))?;
602 }
603
604 session
605 .finish_request_body()
606 .await
607 .map_err(|e| SubRequestError::Io(e.to_string()))?;
608
609 session
610 .read_response_header()
611 .await
612 .map_err(|e| SubRequestError::Io(e.to_string()))?;
613
614 let resp_header = session
615 .response_header()
616 .ok_or_else(|| SubRequestError::Io("no response header received".to_owned()))?;
617
618 let status = resp_header.status.as_u16();
619 if !(100..=599).contains(&status) {
620 session.shutdown().await;
621 return Err(SubRequestError::Io(format!(
622 "upstream returned unsupported HTTP status {status}"
623 )));
624 }
625 let mut resp_headers = HeaderMap::new();
626 for (name, value) in &resp_header.headers {
627 if let Ok(v) = http::header::HeaderValue::from_bytes(value.as_bytes()) {
628 resp_headers.append(name.clone(), v);
629 }
630 }
631 strip_hop_by_hop_headers(&mut resp_headers);
632 strip_reserved_headers(&mut resp_headers);
633
634 let mut body_buf = Vec::new();
635 while !session.response_done() {
636 match session.read_response_body().await {
637 Ok(Some(chunk)) => {
638 if body_buf.len() + chunk.len() > max_response_bytes {
639 warn!(
640 current = body_buf.len(),
641 chunk = chunk.len(),
642 limit = max_response_bytes,
643 "sub-request response body exceeded limit"
644 );
645 session.shutdown().await;
646 return Err(SubRequestError::ResponseTooLarge {
647 actual: body_buf.len() + chunk.len(),
648 limit: max_response_bytes,
649 });
650 }
651 body_buf.extend_from_slice(&chunk);
652 },
653 Ok(None) => break,
654 Err(e) => {
655 session.shutdown().await;
656 return Err(SubRequestError::Io(e.to_string()));
657 },
658 }
659 }
660
661 debug!(status, body_bytes = body_buf.len(), "sub-request: response received");
662
663 connector.connector().release_http_session(session, peer, None).await;
664
665 Ok(SubResponse {
666 status,
667 headers: resp_headers,
668 body: Bytes::from(body_buf),
669 })
670}
671
672struct CircuitGuard<'a> {
682 registry: &'a CircuitBreakerRegistry,
684 peer: PeerKey,
686 token: Option<CircuitToken>,
688}
689
690impl<'a> CircuitGuard<'a> {
691 fn new(registry: &'a CircuitBreakerRegistry, peer: PeerKey, token: CircuitToken) -> Self {
693 Self {
694 registry,
695 peer,
696 token: Some(token),
697 }
698 }
699
700 fn finalize(mut self, result: &Result<SubResponse, SubRequestError>) {
707 let Some(token) = self.token.take() else {
708 return;
709 };
710 match result {
711 Err(SubRequestError::Connect(_) | SubRequestError::Io(_) | SubRequestError::DeadlineExceeded) => {
712 self.registry.record_failure(&self.peer, token);
713 },
714 Ok(_) | Err(_) => {
715 self.registry.record_success(&self.peer, token);
716 },
717 }
718 }
719}
720
721impl Drop for CircuitGuard<'_> {
722 fn drop(&mut self) {
723 if let Some(token) = self.token.take() {
724 self.registry.record_failure(&self.peer, token);
725 }
726 }
727}
728
729const HOP_BY_HOP_HEADERS: &[&str] = &[
736 "connection",
737 "keep-alive",
738 "proxy-authenticate",
739 "proxy-authorization",
740 "te",
741 "trailer",
742 "transfer-encoding",
743 "upgrade",
744];
745
746fn strip_hop_by_hop_headers(headers: &mut HeaderMap) {
748 let connection_values: Vec<_> = headers.get_all(http::header::CONNECTION).iter().cloned().collect();
749 for name in HOP_BY_HOP_HEADERS {
750 headers.remove(*name);
751 }
752 for value in connection_values {
753 let Ok(value) = value.to_str() else { continue };
754 for token in value.split(',').map(str::trim).filter(|token| !token.is_empty()) {
755 headers.remove(token);
756 }
757 }
758}
759
760fn strip_request_framing_headers(headers: &mut HeaderMap) {
762 headers.remove(http::header::CONTENT_LENGTH);
763 headers.remove(http::header::TRANSFER_ENCODING);
764}
765
766fn strip_reserved_headers(headers: &mut HeaderMap) {
769 let reserved: Vec<http::header::HeaderName> = headers
770 .keys()
771 .filter(|name| crate::reserved_headers::is_reserved(name.as_str()))
772 .cloned()
773 .collect();
774 for name in reserved {
775 headers.remove(&name);
776 }
777}
778
779fn is_transport_header(name: &http::header::HeaderName) -> bool {
786 HOP_BY_HOP_HEADERS.iter().any(|h| *h == name.as_str()) || name == http::header::CONTENT_LENGTH
787}
788
789fn empty_body_needs_framing(method: &http::Method) -> bool {
791 matches!(*method, http::Method::POST | http::Method::PUT | http::Method::PATCH)
792}
793
794fn ensure_host_header(request: &mut pingora_http::RequestHeader, peer: &HttpPeer) -> Result<(), SubRequestError> {
796 if !request.headers.contains_key(http::header::HOST) {
797 request
798 .insert_header(http::header::HOST, peer.address().to_string())
799 .map_err(|error| SubRequestError::InvalidRequest(error.to_string()))?;
800 }
801 Ok(())
802}
803
804fn clamp_peer_timeouts(peer: &mut HttpPeer, deadline: Duration) {
806 peer.options.connection_timeout = Some(min_timeout(peer.options.connection_timeout, deadline));
807 peer.options.total_connection_timeout = Some(min_timeout(peer.options.total_connection_timeout, deadline));
808}
809
810fn min_timeout(configured: Option<Duration>, deadline: Duration) -> Duration {
812 configured.map_or(deadline, |configured| configured.min(deadline))
813}
814
815#[cfg(test)]
820#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
821#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic, reason = "tests")]
822mod tests {
823 use super::*;
824
825 #[test]
828 fn clone_shares_same_arc() {
829 let a = SubRequestConnector::new(16, None);
830 let b = a.clone();
831 assert!(
832 Arc::ptr_eq(&a.inner, &b.inner),
833 "cloned connectors should share the same Arc"
834 );
835 }
836
837 #[test]
838 fn debug_impl_does_not_panic() {
839 let connector = SubRequestConnector::new(8, None);
840 let debug = format!("{connector:?}");
841 assert!(
842 debug.contains("SubRequestConnector"),
843 "debug output should contain type name"
844 );
845 }
846
847 #[test]
848 fn unbounded_connector_has_no_admission() {
849 let connector = SubRequestConnector::new(8, None);
850 assert!(
851 connector.admission.is_none(),
852 "no max_connections should mean no semaphore"
853 );
854 }
855
856 #[test]
857 fn bounded_connector_has_admission_semaphore() {
858 let connector = SubRequestConnector::new(8, Some(16));
859 let semaphore = connector
860 .admission
861 .as_ref()
862 .expect("max_connections should create semaphore");
863 assert_eq!(
864 semaphore.available_permits(),
865 16,
866 "semaphore should have the configured permits"
867 );
868 }
869
870 #[tokio::test]
871 async fn acquire_permit_returns_none_without_limit() {
872 let connector = SubRequestConnector::new(4, None);
873 assert!(
874 connector.acquire_permit().await.is_none(),
875 "unbounded connector should return None"
876 );
877 }
878
879 #[tokio::test]
880 async fn acquire_permit_returns_some_with_limit() {
881 let connector = SubRequestConnector::new(4, Some(2));
882 assert!(
883 connector.acquire_permit().await.is_some(),
884 "bounded connector should return a permit"
885 );
886 }
887
888 #[tokio::test]
889 async fn dropping_permit_restores_capacity() {
890 let connector = SubRequestConnector::new(4, Some(1));
891 let permit = connector.acquire_permit().await.unwrap();
892 assert_eq!(
893 connector.admission.as_ref().unwrap().available_permits(),
894 0,
895 "all permits should be taken"
896 );
897 drop(permit);
898 assert_eq!(
899 connector.admission.as_ref().unwrap().available_permits(),
900 1,
901 "dropping permit should restore capacity"
902 );
903 }
904
905 #[test]
906 fn clone_shares_admission_semaphore() {
907 let a = SubRequestConnector::new(4, Some(8));
908 let b = a.clone();
909 assert!(
910 Arc::ptr_eq(a.admission.as_ref().unwrap(), b.admission.as_ref().unwrap()),
911 "cloned connectors should share the semaphore"
912 );
913 }
914
915 #[test]
918 fn subrequest_clone_preserves_fields() {
919 let req = SubRequest {
920 method: http::Method::POST,
921 uri: "/v1/chat".parse().unwrap(),
922 headers: HeaderMap::new(),
923 body: Bytes::from_static(b"hello"),
924 };
925 let cloned = req.clone();
926 assert_eq!(cloned.method, http::Method::POST);
927 assert_eq!(cloned.body, Bytes::from_static(b"hello"));
928 }
929
930 #[test]
931 fn subresponse_clone_preserves_fields() {
932 let resp = SubResponse {
933 status: 200,
934 headers: HeaderMap::new(),
935 body: Bytes::from_static(b"world"),
936 };
937 let cloned = resp.clone();
938 assert_eq!(cloned.status, 200);
939 assert_eq!(cloned.body, Bytes::from_static(b"world"));
940 }
941
942 #[test]
945 fn client_wraps_connector() {
946 let connector = SubRequestConnector::new(8, None);
947 let client = SubRequestClient::new(connector);
948 let debug = format!("{client:?}");
949 assert!(
950 debug.contains("SubRequestClient"),
951 "debug output should contain type name"
952 );
953 }
954
955 #[test]
956 fn client_clone_shares_connector() {
957 let connector = SubRequestConnector::new(8, Some(4));
958 let a = SubRequestClient::new(connector);
959 let b = a.clone();
960 assert!(
961 Arc::ptr_eq(&a.connector.inner, &b.connector.inner,),
962 "cloned clients should share the same connector"
963 );
964 }
965
966 #[test]
969 fn subrequest_error_invalid_request_display() {
970 let err = SubRequestError::InvalidRequest("bad header".to_owned());
971 assert!(
972 err.to_string().contains("bad header"),
973 "InvalidRequest error should include reason: {err}"
974 );
975 }
976
977 #[test]
978 fn subrequest_error_admission_timeout_display() {
979 let err = SubRequestError::AdmissionTimeout { max_connections: 64 };
980 let msg = err.to_string();
981 assert!(msg.contains("64"), "should include max_connections: {msg}");
982 assert!(msg.contains("admission"), "should mention admission: {msg}");
983 }
984
985 #[test]
986 fn subrequest_error_connect_display() {
987 let err = SubRequestError::Connect("connection refused".to_owned());
988 assert!(
989 err.to_string().contains("connection refused"),
990 "Connect error should include reason: {err}"
991 );
992 }
993
994 #[test]
995 fn subrequest_error_io_display() {
996 let err = SubRequestError::Io("broken pipe".to_owned());
997 assert!(
998 err.to_string().contains("broken pipe"),
999 "Io error should include reason: {err}"
1000 );
1001 }
1002
1003 #[test]
1004 fn subrequest_error_response_too_large_display() {
1005 let err = SubRequestError::ResponseTooLarge {
1006 actual: 20_000,
1007 limit: 10_000,
1008 };
1009 let msg = err.to_string();
1010 assert!(msg.contains("20000"), "should include actual: {msg}");
1011 assert!(msg.contains("10000"), "should include limit: {msg}");
1012 }
1013
1014 #[test]
1015 fn subrequest_error_deadline_exceeded_display() {
1016 let err = SubRequestError::DeadlineExceeded;
1017 assert!(
1018 !err.to_string().is_empty(),
1019 "DeadlineExceeded should have a display message"
1020 );
1021 }
1022
1023 #[test]
1026 fn strip_hop_by_hop_removes_static_and_connection_nominated() {
1027 let mut headers = HeaderMap::new();
1028 headers.insert("connection", "x-custom, keep-alive".parse().unwrap());
1029 headers.insert("keep-alive", "timeout=5".parse().unwrap());
1030 headers.insert("x-custom", "value".parse().unwrap());
1031 headers.insert("x-safe", "kept".parse().unwrap());
1032 headers.insert("transfer-encoding", "chunked".parse().unwrap());
1033
1034 strip_hop_by_hop_headers(&mut headers);
1035
1036 assert!(!headers.contains_key("connection"));
1037 assert!(!headers.contains_key("keep-alive"));
1038 assert!(!headers.contains_key("x-custom"));
1039 assert!(!headers.contains_key("transfer-encoding"));
1040 assert_eq!(headers.get("x-safe").unwrap(), "kept");
1041 }
1042
1043 #[test]
1044 fn strip_request_framing_removes_content_length_and_transfer_encoding() {
1045 let mut headers = HeaderMap::new();
1046 headers.insert(http::header::CONTENT_LENGTH, "42".parse().unwrap());
1047 headers.insert(http::header::TRANSFER_ENCODING, "chunked".parse().unwrap());
1048 headers.insert("x-safe", "kept".parse().unwrap());
1049
1050 strip_request_framing_headers(&mut headers);
1051
1052 assert!(!headers.contains_key(http::header::CONTENT_LENGTH));
1053 assert!(!headers.contains_key(http::header::TRANSFER_ENCODING));
1054 assert_eq!(headers.get("x-safe").unwrap(), "kept");
1055 }
1056
1057 #[test]
1060 fn empty_entity_methods_get_explicit_framing() {
1061 assert!(empty_body_needs_framing(&http::Method::POST));
1062 assert!(empty_body_needs_framing(&http::Method::PUT));
1063 assert!(empty_body_needs_framing(&http::Method::PATCH));
1064 assert!(!empty_body_needs_framing(&http::Method::GET));
1065 assert!(!empty_body_needs_framing(&http::Method::HEAD));
1066 }
1067
1068 #[test]
1069 fn min_timeout_preserves_stricter_cluster_limit() {
1070 assert_eq!(
1071 min_timeout(Some(Duration::from_secs(1)), Duration::from_secs(10)),
1072 Duration::from_secs(1)
1073 );
1074 assert_eq!(
1075 min_timeout(Some(Duration::from_secs(20)), Duration::from_secs(10)),
1076 Duration::from_secs(10)
1077 );
1078 assert_eq!(min_timeout(None, Duration::from_secs(10)), Duration::from_secs(10));
1079 }
1080
1081 #[test]
1082 fn clamp_peer_timeouts_bounds_connection_setup() {
1083 let mut peer = HttpPeer::new("127.0.0.1:8080", false, String::new());
1084 peer.options.connection_timeout = Some(Duration::from_secs(1));
1085 peer.options.total_connection_timeout = Some(Duration::from_secs(20));
1086
1087 clamp_peer_timeouts(&mut peer, Duration::from_secs(10));
1088
1089 assert_eq!(peer.options.connection_timeout, Some(Duration::from_secs(1)));
1090 assert_eq!(peer.options.total_connection_timeout, Some(Duration::from_secs(10)));
1091 }
1092
1093 #[test]
1094 fn ensure_host_header_uses_peer_address_without_overwriting_explicit_host() {
1095 let peer = HttpPeer::new("127.0.0.1:8443", false, String::new());
1096 let mut generated = pingora_http::RequestHeader::build("GET", b"/", None).unwrap();
1097 ensure_host_header(&mut generated, &peer).unwrap();
1098 assert_eq!(generated.headers.get(http::header::HOST).unwrap(), "127.0.0.1:8443");
1099
1100 let mut explicit = pingora_http::RequestHeader::build("GET", b"/", None).unwrap();
1101 explicit.insert_header(http::header::HOST, "model.example").unwrap();
1102 ensure_host_header(&mut explicit, &peer).unwrap();
1103 assert_eq!(explicit.headers.get(http::header::HOST).unwrap(), "model.example");
1104 }
1105
1106 #[tokio::test]
1109 async fn deadline_bounds_the_complete_exchange() {
1110 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1111 let address = listener.local_addr().unwrap();
1112 let backend = tokio::spawn(async move {
1113 let (_socket, _) = listener.accept().await.unwrap();
1114 tokio::time::sleep(Duration::from_secs(1)).await;
1115 });
1116 let connector = SubRequestConnector::new(1, None);
1117 let client = SubRequestClient::new(connector);
1118 let peer = HttpPeer::new(address.to_string(), false, String::new());
1119 let request = SubRequest {
1120 method: http::Method::GET,
1121 uri: "/".parse().unwrap(),
1122 headers: HeaderMap::new(),
1123 body: Bytes::new(),
1124 };
1125
1126 let started = std::time::Instant::now();
1127 let result = Box::pin(client.execute(&peer, &request, 1024, Duration::from_millis(10), None)).await;
1128 let elapsed = started.elapsed();
1129 backend.abort();
1130
1131 assert!(result.is_err(), "a backend that never responds must time out");
1132 assert!(
1133 elapsed < Duration::from_millis(500),
1134 "exchange exceeded its deadline: {elapsed:?}"
1135 );
1136 }
1137
1138 #[tokio::test]
1139 async fn admission_timeout_returns_typed_error() {
1140 let connector = SubRequestConnector::new(4, Some(1));
1141 let permit = connector.acquire_permit().await.unwrap();
1142
1143 let result = connector.try_acquire_permit(Duration::from_millis(10)).await;
1144
1145 assert!(
1146 matches!(result, Err(SubRequestError::AdmissionTimeout { .. })),
1147 "should return AdmissionTimeout when slots are full: {result:?}"
1148 );
1149 drop(result);
1150 drop(permit);
1151 }
1152
1153 #[tokio::test]
1154 async fn admission_timeout_reports_configured_max() {
1155 let configured_limit = 4;
1156 let connector = SubRequestConnector::new(4, Some(configured_limit));
1157 let mut permits = Vec::new();
1158 for _ in 0..configured_limit {
1159 permits.push(connector.acquire_permit().await.unwrap());
1160 }
1161
1162 let result = connector.try_acquire_permit(Duration::from_millis(10)).await;
1163 match &result {
1164 Err(SubRequestError::AdmissionTimeout { max_connections }) => {
1165 assert_eq!(
1166 *max_connections, configured_limit,
1167 "should report configured limit, not available permits"
1168 );
1169 },
1170 other => panic!("expected AdmissionTimeout, got: {other:?}"),
1171 }
1172 drop(result);
1173 drop(permits);
1174 }
1175
1176 #[tokio::test]
1177 async fn try_acquire_permit_returns_none_without_limit() {
1178 let connector = SubRequestConnector::new(4, None);
1179 let result = connector.try_acquire_permit(Duration::from_millis(10)).await;
1180 assert!(
1181 matches!(result, Ok(None)),
1182 "unbounded connector should return Ok(None): {result:?}"
1183 );
1184 drop(result);
1185 }
1186
1187 #[test]
1190 fn client_with_custom_ceiling() {
1191 let connector = SubRequestConnector::new(8, None);
1192 let client = SubRequestClient::with_max_response_bytes(connector, 4096);
1193 assert_eq!(client.max_response_bytes, 4096);
1194 }
1195
1196 #[test]
1197 fn client_default_ceiling_is_absolute_max() {
1198 let connector = SubRequestConnector::new(8, None);
1199 let client = SubRequestClient::new(connector);
1200 assert_eq!(
1201 client.max_response_bytes,
1202 crate::config::ABSOLUTE_MAX_BODY_BYTES,
1203 "default ceiling should be ABSOLUTE_MAX_BODY_BYTES (64 MiB)"
1204 );
1205 }
1206
1207 #[test]
1210 fn response_hop_by_hop_headers_are_stripped() {
1211 let mut headers = HeaderMap::new();
1212 headers.insert("connection", "x-nominated".parse().unwrap());
1213 headers.insert("transfer-encoding", "chunked".parse().unwrap());
1214 headers.insert("keep-alive", "timeout=5".parse().unwrap());
1215 headers.insert("x-nominated", "internal".parse().unwrap());
1216 headers.insert("content-type", "application/json".parse().unwrap());
1217
1218 strip_hop_by_hop_headers(&mut headers);
1219
1220 assert!(!headers.contains_key("connection"));
1221 assert!(!headers.contains_key("transfer-encoding"));
1222 assert!(!headers.contains_key("keep-alive"));
1223 assert!(!headers.contains_key("x-nominated"));
1224 assert_eq!(headers.get("content-type").unwrap(), "application/json");
1225 }
1226
1227 #[test]
1230 fn strip_reserved_removes_internal_prefixes() {
1231 let mut headers = HeaderMap::new();
1232 headers.insert("x-praxis-route", "internal".parse().unwrap());
1233 headers.insert("x-ext-protocol-model", "gpt-4".parse().unwrap());
1234 headers.insert("x-ext-agent-task", "classify".parse().unwrap());
1235 headers.insert("x-custom", "kept".parse().unwrap());
1236 headers.insert("authorization", "Bearer tok".parse().unwrap());
1237
1238 strip_reserved_headers(&mut headers);
1239
1240 assert!(!headers.contains_key("x-praxis-route"));
1241 assert!(!headers.contains_key("x-ext-protocol-model"));
1242 assert!(!headers.contains_key("x-ext-agent-task"));
1243 assert_eq!(headers.get("x-custom").unwrap(), "kept");
1244 assert_eq!(headers.get("authorization").unwrap(), "Bearer tok");
1245 }
1246
1247 #[test]
1248 fn strip_reserved_is_no_op_for_safe_headers() {
1249 let mut headers = HeaderMap::new();
1250 headers.insert("content-type", "application/json".parse().unwrap());
1251 headers.insert("x-request-id", "abc".parse().unwrap());
1252
1253 strip_reserved_headers(&mut headers);
1254
1255 assert_eq!(headers.len(), 2);
1256 }
1257
1258 #[test]
1261 fn connector_stores_configured_max_connections() {
1262 let connector = SubRequestConnector::new(4, Some(256));
1263 assert_eq!(connector.configured_max_connections, Some(256));
1264
1265 let unbounded = SubRequestConnector::new(4, None);
1266 assert_eq!(unbounded.configured_max_connections, None);
1267 }
1268
1269 #[test]
1272 fn with_options_creates_connector() {
1273 let connector = SubRequestConnector::with_options(SubRequestConnectorOptions {
1274 keepalive_pool_size: 32,
1275 max_connections: Some(64),
1276 circuit_breaker: None,
1277 });
1278 assert_eq!(
1279 connector.configured_max_connections,
1280 Some(64),
1281 "max_connections should be forwarded"
1282 );
1283 assert!(
1284 connector.circuit_breakers.is_none(),
1285 "no circuit breaker config should mean no registry"
1286 );
1287 }
1288
1289 #[test]
1290 fn with_options_circuit_breaker_enabled() {
1291 let connector = SubRequestConnector::with_options(SubRequestConnectorOptions {
1292 keepalive_pool_size: 16,
1293 max_connections: None,
1294 circuit_breaker: Some(CircuitBreakerConfig {
1295 threshold: 3,
1296 recovery_window: Duration::from_secs(30),
1297 half_open_timeout: Duration::from_secs(30),
1298 }),
1299 });
1300 assert!(
1301 connector.circuit_breakers.is_some(),
1302 "circuit breaker config should create a registry"
1303 );
1304 }
1305
1306 fn test_registry(threshold: u32) -> CircuitBreakerRegistry {
1309 CircuitBreakerRegistry::new(CircuitBreakerConfig {
1310 threshold,
1311 recovery_window: Duration::from_secs(9999),
1312 half_open_timeout: Duration::from_secs(9999),
1313 })
1314 }
1315
1316 fn test_peer(addr: &str) -> PeerKey {
1317 PeerKey::new(addr.parse().unwrap(), "")
1318 }
1319
1320 fn acquire_guard(registry: &CircuitBreakerRegistry, key: PeerKey) -> CircuitGuard<'_> {
1321 let CircuitCheck::Allowed(token) = registry.try_acquire(key.clone()) else {
1322 panic!("should be allowed");
1323 };
1324 CircuitGuard::new(registry, key, token)
1325 }
1326
1327 #[test]
1328 fn circuit_guard_success_records_success() {
1329 let registry = test_registry(3);
1330 let key = test_peer("127.0.0.1:8080");
1331 let guard = acquire_guard(®istry, key.clone());
1332 let result: Result<SubResponse, SubRequestError> = Ok(SubResponse {
1333 status: 200,
1334 headers: HeaderMap::new(),
1335 body: Bytes::new(),
1336 });
1337 guard.finalize(&result);
1338 assert!(registry.precheck(&key), "peer should remain healthy after success");
1339 }
1340
1341 #[test]
1342 fn circuit_guard_connect_error_records_failure() {
1343 let registry = test_registry(1);
1344 let key = test_peer("127.0.0.1:8080");
1345 let guard = acquire_guard(®istry, key.clone());
1346 guard.finalize(&Err(SubRequestError::Connect("refused".to_owned())));
1347 assert!(!registry.precheck(&key), "peer should be open after connect failure");
1348 }
1349
1350 #[test]
1351 fn circuit_guard_io_error_records_failure() {
1352 let registry = test_registry(1);
1353 let key = test_peer("127.0.0.1:8080");
1354 let guard = acquire_guard(®istry, key.clone());
1355 guard.finalize(&Err(SubRequestError::Io("broken pipe".to_owned())));
1356 assert!(!registry.precheck(&key), "peer should be open after I/O failure");
1357 }
1358
1359 #[test]
1360 fn circuit_guard_deadline_exceeded_records_failure() {
1361 let registry = test_registry(1);
1362 let key = test_peer("127.0.0.1:8080");
1363 let guard = acquire_guard(®istry, key.clone());
1364 guard.finalize(&Err(SubRequestError::DeadlineExceeded));
1365 assert!(!registry.precheck(&key), "peer should be open after deadline exceeded");
1366 }
1367
1368 #[test]
1369 fn circuit_guard_response_too_large_counts_as_success() {
1370 let registry = test_registry(1);
1371 let key = test_peer("127.0.0.1:8080");
1372 let guard = acquire_guard(®istry, key.clone());
1373 guard.finalize(&Err(SubRequestError::ResponseTooLarge {
1374 actual: 20_000,
1375 limit: 10_000,
1376 }));
1377 assert!(
1378 registry.precheck(&key),
1379 "response too large is not a peer fault — should remain healthy"
1380 );
1381 }
1382
1383 #[test]
1384 fn circuit_guard_admission_timeout_not_peer_fault() {
1385 let registry = test_registry(1);
1386 let key = test_peer("127.0.0.1:8080");
1387 let guard = acquire_guard(®istry, key.clone());
1388 guard.finalize(&Err(SubRequestError::AdmissionTimeout { max_connections: 64 }));
1389 assert!(
1390 registry.precheck(&key),
1391 "admission timeout is not a peer fault — should remain healthy"
1392 );
1393 }
1394
1395 #[test]
1396 fn circuit_guard_drop_without_finalize_records_failure() {
1397 let registry = test_registry(1);
1398 let key = test_peer("127.0.0.1:8080");
1399 let guard = acquire_guard(®istry, key.clone());
1400 drop(guard);
1401 assert!(
1402 !registry.precheck(&key),
1403 "dropped guard should record failure (deadline/panic path)"
1404 );
1405 }
1406
1407 #[test]
1410 fn subrequest_error_circuit_open_display() {
1411 let err = SubRequestError::CircuitOpen {
1412 peer: "127.0.0.1:8080".to_owned(),
1413 };
1414 let msg = err.to_string();
1415 assert!(msg.contains("circuit open"), "should mention circuit open: {msg}");
1416 assert!(msg.contains("127.0.0.1:8080"), "should include peer address: {msg}");
1417 }
1418
1419 #[test]
1422 fn is_transport_header_rejects_hop_by_hop_and_framing() {
1423 let transport_names = [
1424 "connection",
1425 "keep-alive",
1426 "transfer-encoding",
1427 "upgrade",
1428 "content-length",
1429 ];
1430 for name in transport_names {
1431 let hdr: http::header::HeaderName = name.parse().unwrap();
1432 assert!(is_transport_header(&hdr), "{name} should be classified as transport");
1433 }
1434
1435 let safe_names = ["authorization", "x-request-id", "x-custom-header"];
1436 for name in safe_names {
1437 let hdr: http::header::HeaderName = name.parse().unwrap();
1438 assert!(
1439 !is_transport_header(&hdr),
1440 "{name} should not be classified as transport"
1441 );
1442 }
1443 }
1444
1445 #[test]
1446 fn framework_headers_rejects_transport_headers() {
1447 let mut fw = FrameworkHeaders::new();
1448 let val = http::HeaderValue::from_static("1");
1449 let result = fw.insert(http::header::CONTENT_LENGTH, val);
1450 assert!(result.is_err(), "transport header should be rejected");
1451 assert!(fw.is_empty());
1452 }
1453
1454 #[test]
1455 fn framework_headers_rejects_reserved_headers() {
1456 let mut fw = FrameworkHeaders::new();
1457 let val = http::HeaderValue::from_static("1");
1458 let name: http::header::HeaderName = "x-praxis-depth".parse().unwrap();
1459 let result = fw.insert(name, val);
1460 assert!(result.is_err(), "reserved header should be rejected");
1461 assert!(fw.is_empty());
1462 }
1463
1464 #[test]
1465 fn framework_headers_accepts_non_reserved_non_transport() {
1466 let mut fw = FrameworkHeaders::new();
1467 let val = http::HeaderValue::from_static("3");
1468 let name: http::header::HeaderName = "x-request-id".parse().unwrap();
1469 fw.insert(name, val).unwrap();
1470 assert!(!fw.is_empty());
1471 assert_eq!(fw.iter().count(), 1);
1472 }
1473
1474 #[test]
1475 fn framework_headers_set_depth_injects_reserved_header() {
1476 let mut fw = FrameworkHeaders::new();
1477 fw.set_depth(2);
1478 assert_eq!(fw.iter().count(), 1);
1479 let (name, value) = fw.iter().next().unwrap();
1480 assert_eq!(name.as_str(), DEPTH_HEADER);
1481 assert_eq!(value, "2");
1482 }
1483
1484 #[test]
1485 fn framework_headers_set_depth_zero() {
1486 let mut fw = FrameworkHeaders::new();
1487 fw.set_depth(0);
1488 let (_, value) = fw.iter().next().unwrap();
1489 assert_eq!(value, "0");
1490 }
1491}