relay_knowledge/net/http/
qos_client.rs

1use std::{error::Error, fmt};
2
3use serde::de::DeserializeOwned;
4
5use crate::net::{
6    http::qos_request_context_active,
7    qos::{QosPermit, QosPolicy, QosRuntime, RejectReason},
8};
9
10/// Error raised by QoS-gated outbound reqwest calls.
11#[derive(Debug)]
12pub enum QosHttpClientError {
13    QosRejected(RejectReason),
14    Transport(reqwest::Error),
15}
16
17impl QosHttpClientError {
18    /// Returns whether the transport layer reported a timeout.
19    pub fn is_timeout(&self) -> bool {
20        matches!(self, Self::Transport(error) if error.is_timeout())
21    }
22}
23
24impl fmt::Display for QosHttpClientError {
25    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
26        match self {
27            Self::QosRejected(reason) => {
28                write!(formatter, "request rejected by QoS: {}", reason.as_str())
29            }
30            Self::Transport(error) => error.fmt(formatter),
31        }
32    }
33}
34
35impl Error for QosHttpClientError {}
36
37/// Reqwest response that keeps the QoS request permit until the body is consumed.
38pub struct QosHttpResponse {
39    inner: reqwest::Response,
40    qos: Option<QosRuntime>,
41    _permit: Option<QosPermit>,
42}
43
44impl QosHttpResponse {
45    fn with_permit(inner: reqwest::Response, qos: QosRuntime, permit: QosPermit) -> Self {
46        Self {
47            inner,
48            qos: Some(qos),
49            _permit: Some(permit),
50        }
51    }
52
53    fn without_permit(inner: reqwest::Response, qos: QosRuntime) -> Self {
54        Self {
55            inner,
56            qos: Some(qos),
57            _permit: None,
58        }
59    }
60
61    pub fn unmetered(inner: reqwest::Response) -> Self {
62        Self {
63            inner,
64            qos: None,
65            _permit: None,
66        }
67    }
68
69    pub fn status(&self) -> reqwest::StatusCode {
70        self.inner.status()
71    }
72
73    pub fn content_length(&self) -> Option<u64> {
74        self.inner.content_length()
75    }
76
77    pub async fn json<T>(self) -> Result<T, reqwest::Error>
78    where
79        T: DeserializeOwned,
80    {
81        let qos = self.qos.clone();
82        record_body_timeout(qos.as_ref(), self.inner.json::<T>().await)
83    }
84
85    pub async fn text(self) -> Result<String, reqwest::Error> {
86        let qos = self.qos.clone();
87        record_body_timeout(qos.as_ref(), self.inner.text().await)
88    }
89
90    pub async fn bytes(self) -> Result<Vec<u8>, reqwest::Error> {
91        let qos = self.qos.clone();
92        record_body_timeout(
93            qos.as_ref(),
94            self.inner.bytes().await.map(|bytes| bytes.to_vec()),
95        )
96    }
97
98    pub async fn chunk(&mut self) -> Result<Option<Vec<u8>>, reqwest::Error> {
99        record_body_timeout(
100            self.qos.as_ref(),
101            self.inner
102                .chunk()
103                .await
104                .map(|chunk| chunk.map(|bytes| bytes.to_vec())),
105        )
106    }
107}
108
109/// Sends an outbound reqwest request after acquiring a QoS request permit.
110pub async fn send_request_with_qos(
111    qos: &QosRuntime,
112    policy: &QosPolicy,
113    request: reqwest::RequestBuilder,
114) -> Result<QosHttpResponse, QosHttpClientError> {
115    if qos_request_context_active() {
116        return send_request_without_new_permit(qos, request).await;
117    }
118
119    let permit = qos
120        .admit_request(policy)
121        .map_err(QosHttpClientError::QosRejected)?;
122    match request.send().await {
123        Ok(response) => Ok(QosHttpResponse::with_permit(response, qos.clone(), permit)),
124        Err(error) => {
125            if error.is_timeout() {
126                qos.record_timed_out();
127            }
128            Err(QosHttpClientError::Transport(error))
129        }
130    }
131}
132
133async fn send_request_without_new_permit(
134    qos: &QosRuntime,
135    request: reqwest::RequestBuilder,
136) -> Result<QosHttpResponse, QosHttpClientError> {
137    match request.send().await {
138        Ok(response) => Ok(QosHttpResponse::without_permit(response, qos.clone())),
139        Err(error) => {
140            if error.is_timeout() {
141                qos.record_timed_out();
142            }
143            Err(QosHttpClientError::Transport(error))
144        }
145    }
146}
147
148fn record_body_timeout<T>(
149    qos: Option<&QosRuntime>,
150    result: Result<T, reqwest::Error>,
151) -> Result<T, reqwest::Error> {
152    if matches!(&result, Err(error) if error.is_timeout()) {
153        if let Some(qos) = qos {
154            qos.record_timed_out();
155        }
156    }
157    result
158}
relay_knowledge/net/http/qos_client.rs

relay_knowledge/net/http/
qos_client.rs
