Skip to main content

leyline/core/request/
send.rs

1use std::pin::Pin;
2use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
3use std::task::{Context, Poll};
4
5use bytes::Bytes;
6use futures_util::Stream;
7
8use super::RequestBuilder;
9use super::route::Route;
10use crate::core::Result;
11use crate::core::body::{Body, BodyKind, BoxedStream};
12use crate::core::deadline::Deadline;
13use crate::core::error::{Error, Kind};
14use crate::core::response::Response;
15use crate::core::retry::RetryPlan;
16use crate::core::session::execute::Attempt;
17use crate::core::transport::ResponseMode;
18use crate::util::is_idempotent;
19
20impl RequestBuilder {
21    pub async fn send(self) -> Result<Response> {
22        Box::pin(self.run()).await
23    }
24
25    async fn run(mut self) -> Result<Response> {
26        self.prepare()?;
27        let retry_policy = self.retry_policy.clone();
28        let session = self.session.clone();
29        let deadline = session.deadline(self.timeouts.as_ref());
30        let tag = self.tag.take();
31        let status_errors = self.status_errors;
32        let attempt = self.into_attempt(deadline);
33        let method = attempt.method.clone();
34        let url = attempt.url.clone();
35        let streamed = attempt.response == ResponseMode::Streamed;
36        let mut exhausted = false;
37        let response = session
38            .traced(
39                method.as_str(),
40                &url,
41                streamed,
42                run_attempts(
43                    &session,
44                    attempt,
45                    &retry_policy,
46                    deadline,
47                    tag,
48                    &mut exhausted,
49                ),
50            )
51            .await?;
52        if !status_errors {
53            return Ok(response);
54        }
55        let policy_wait = retry_policy.server_wait(&response);
56        response
57            .error_for_status_with_body(&deadline)
58            .await
59            .map_err(|mut error| {
60                error.set_retry_outcome(policy_wait, exhausted);
61                error
62            })
63    }
64
65    fn into_attempt(mut self, deadline: Deadline) -> Attempt {
66        let headers = if self.headers.is_empty() {
67            None
68        } else {
69            Some(std::mem::take(&mut self.headers))
70        };
71        Attempt {
72            method: std::mem::take(&mut self.method),
73            url: std::mem::take(&mut self.url),
74            preset: self.preset,
75            body: std::mem::take(&mut self.body),
76            headers,
77            deadline,
78            response: ResponseMode::new(
79                self.stream_response,
80                self.status_errors
81                    .then(|| self.session.error_budget(&deadline)),
82            ),
83            proxy: self.proxy.take(),
84            header_order: self.header_order.take(),
85            redirect: self.redirect.take(),
86            digest: self.digest_auth.take(),
87            initiator: self.initiator.take(),
88            trusted_origin: self.trusted_origin.take(),
89        }
90    }
91
92    pub(crate) fn prepare(&mut self) -> Result<()> {
93        if let Some(err) = self.builder_error.take() {
94            return Err(err);
95        }
96        self.infer_from_content_type();
97
98        if !self.query_params.is_empty() {
99            let mut url = url::Url::parse(&self.url).map_err(crate::core::Error::from_url_parse)?;
100            {
101                let mut pairs = url.query_pairs_mut();
102                for (k, v) in &self.query_params {
103                    pairs.append_pair(k, v);
104                }
105            }
106            self.url = url.to_string();
107        }
108
109        if let Some(encoding) = self.compress {
110            match &self.body.0 {
111                BodyKind::Bytes(b) => {
112                    let compressed = encoding.encode(b)?;
113                    self.headers
114                        .set("content-encoding", encoding.header_value())?;
115                    self.body = Body::from(compressed);
116                }
117                BodyKind::Empty => {}
118                BodyKind::Stream { .. } => {
119                    return Err(Error::new(Kind::Body).with_message(
120                        "request-body compression is not supported for streaming bodies; \
121                         buffer the body into bytes before calling `.compress(..)`",
122                    ));
123                }
124            }
125        }
126        Ok(())
127    }
128}
129
130async fn run_attempts(
131    session: &crate::core::session::Session,
132    mut attempt: Attempt,
133    retry_policy: &crate::core::retry::RetryPolicy,
134    deadline: Deadline,
135    tag: Option<String>,
136    exhausted: &mut bool,
137) -> (Result<Response>, u32) {
138    crate::trace::note_tag(tag);
139    let mut route = Route::new(session, &attempt, retry_policy);
140    if retry_policy.is_none() {
141        return (counted(route.send(attempt, None).await, 1), 1);
142    }
143    let retryable_method =
144        retry_policy.allow_non_idempotent || is_idempotent(attempt.method.as_str());
145    let replay = attempt.body.replay();
146    let mut rewind = Rewind::wrap(&mut attempt.body);
147    let base_headers = attempt.headers.clone();
148
149    let mut n: u32 = 0;
150    loop {
151        if let Err(err) = deadline.check() {
152            return (Err(err), n);
153        }
154        let attempt_body = std::mem::take(&mut attempt.body);
155        let this = attempt.again(attempt_body, base_headers.clone());
156        let result = route.send(this, retry_policy.proxy_for_retry(n)).await;
157
158        let body_retryable = || replay.is_some() || rewind.as_mut().is_some_and(Rewind::reclaim);
159        let sleep = match retry_policy.plan(&result, n, retryable_method, body_retryable) {
160            RetryPlan::Backoff(sleep) if sleep < deadline.remaining() => sleep,
161            plan => {
162                *exhausted = matches!(
163                    plan,
164                    RetryPlan::Exhausted | RetryPlan::Declined | RetryPlan::Backoff(_)
165                );
166                return (counted(result, n + 1), n + 1);
167            }
168        };
169        drop(result);
170
171        let slept = session
172            .unless_shut_down(async {
173                deadline.sleep(sleep).await;
174                Ok(())
175            })
176            .await;
177        if let Err(err) = slept {
178            return (counted(Err(err), n + 1), n + 1);
179        }
180        n += 1;
181        attempt.body = replay
182            .as_ref()
183            .and_then(Body::replay)
184            .or_else(|| rewind.as_mut().and_then(Rewind::again))
185            .unwrap_or_default();
186    }
187}
188
189fn counted(result: Result<Response>, attempts: u32) -> Result<Response> {
190    result
191        .map(|mut response| {
192            response.set_attempts(attempts);
193            response
194        })
195        .map_err(|mut error| {
196            error.set_attempts(attempts);
197            error
198        })
199}
200
201struct Slot {
202    stream: Option<BoxedStream>,
203    polled: bool,
204}
205
206type SharedSlot = Arc<Mutex<Slot>>;
207
208fn lock(slot: &SharedSlot) -> MutexGuard<'_, Slot> {
209    slot.lock().unwrap_or_else(PoisonError::into_inner)
210}
211
212struct Rewind {
213    slot: SharedSlot,
214    length_hint: Option<u64>,
215    held: Option<BoxedStream>,
216}
217
218impl Rewind {
219    fn wrap(body: &mut Body) -> Option<Self> {
220        if !matches!(body.0, BodyKind::Stream { .. }) {
221            return None;
222        }
223        let BodyKind::Stream {
224            stream,
225            length_hint,
226        } = std::mem::take(&mut body.0)
227        else {
228            return None;
229        };
230        let rewind = Self {
231            slot: Arc::new(Mutex::new(Slot {
232                stream: Some(stream),
233                polled: false,
234            })),
235            length_hint,
236            held: None,
237        };
238        *body = rewind.body();
239        Some(rewind)
240    }
241
242    fn body(&self) -> Body {
243        Body(BodyKind::Stream {
244            stream: Box::pin(Reclaim(Arc::clone(&self.slot))),
245            length_hint: self.length_hint,
246        })
247    }
248
249    fn reclaim(&mut self) -> bool {
250        let mut slot = lock(&self.slot);
251        if slot.polled {
252            return false;
253        }
254        self.held = slot.stream.take();
255        self.held.is_some()
256    }
257
258    fn again(&mut self) -> Option<Body> {
259        let stream = self.held.take()?;
260        self.slot = Arc::new(Mutex::new(Slot {
261            stream: Some(stream),
262            polled: false,
263        }));
264        Some(self.body())
265    }
266}
267
268struct Reclaim(SharedSlot);
269
270impl Stream for Reclaim {
271    type Item = std::io::Result<Bytes>;
272
273    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
274        let mut slot = lock(&self.0);
275        slot.polled = true;
276        match slot.stream.as_mut() {
277            Some(stream) => stream.as_mut().poll_next(cx),
278            None => Poll::Ready(None),
279        }
280    }
281}
282
283impl std::future::IntoFuture for RequestBuilder {
284    type Output = Result<Response>;
285    type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send>>;
286
287    fn into_future(self) -> Self::IntoFuture {
288        Box::pin(async move { self.send().await })
289    }
290}
291
292#[cfg(test)]
293#[path = "send_tests.rs"]
294mod tests;