use std::pin::Pin;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::task::{Context, Poll};
use bytes::Bytes;
use futures_util::Stream;
use super::RequestBuilder;
use super::route::Route;
use crate::core::Result;
use crate::core::body::{Body, BodyKind, BoxedStream};
use crate::core::deadline::Deadline;
use crate::core::error::{Error, Kind};
use crate::core::response::Response;
use crate::core::retry::RetryPlan;
use crate::core::session::execute::Attempt;
use crate::core::transport::ResponseMode;
use crate::util::is_idempotent;
impl RequestBuilder {
pub async fn send(self) -> Result<Response> {
Box::pin(self.run()).await
}
async fn run(mut self) -> Result<Response> {
self.prepare()?;
let retry_policy = self.retry_policy.clone();
let session = self.session.clone();
let deadline = session.deadline(self.timeouts.as_ref());
let tag = self.tag.take();
let status_errors = self.status_errors;
let attempt = self.into_attempt(deadline);
let method = attempt.method.clone();
let url = attempt.url.clone();
let streamed = attempt.response == ResponseMode::Streamed;
let mut exhausted = false;
let response = session
.traced(
method.as_str(),
&url,
streamed,
run_attempts(
&session,
attempt,
&retry_policy,
deadline,
tag,
&mut exhausted,
),
)
.await?;
if !status_errors {
return Ok(response);
}
let policy_wait = retry_policy.server_wait(&response);
response
.error_for_status_with_body(&deadline)
.await
.map_err(|mut error| {
error.set_retry_outcome(policy_wait, exhausted);
error
})
}
fn into_attempt(mut self, deadline: Deadline) -> Attempt {
let headers = if self.headers.is_empty() {
None
} else {
Some(std::mem::take(&mut self.headers))
};
Attempt {
method: std::mem::take(&mut self.method),
url: std::mem::take(&mut self.url),
preset: self.preset,
body: std::mem::take(&mut self.body),
headers,
deadline,
response: ResponseMode::new(
self.stream_response,
self.status_errors
.then(|| self.session.error_budget(&deadline)),
),
proxy: self.proxy.take(),
header_order: self.header_order.take(),
redirect: self.redirect.take(),
digest: self.digest_auth.take(),
initiator: self.initiator.take(),
trusted_origin: self.trusted_origin.take(),
}
}
pub(crate) fn prepare(&mut self) -> Result<()> {
if let Some(err) = self.builder_error.take() {
return Err(err);
}
self.infer_from_content_type();
if !self.query_params.is_empty() {
let mut url = url::Url::parse(&self.url).map_err(crate::core::Error::from_url_parse)?;
{
let mut pairs = url.query_pairs_mut();
for (k, v) in &self.query_params {
pairs.append_pair(k, v);
}
}
self.url = url.to_string();
}
if let Some(encoding) = self.compress {
match &self.body.0 {
BodyKind::Bytes(b) => {
let compressed = encoding.encode(b)?;
self.headers
.set("content-encoding", encoding.header_value())?;
self.body = Body::from(compressed);
}
BodyKind::Empty => {}
BodyKind::Stream { .. } => {
return Err(Error::new(Kind::Body).with_message(
"request-body compression is not supported for streaming bodies; \
buffer the body into bytes before calling `.compress(..)`",
));
}
}
}
Ok(())
}
}
async fn run_attempts(
session: &crate::core::session::Session,
mut attempt: Attempt,
retry_policy: &crate::core::retry::RetryPolicy,
deadline: Deadline,
tag: Option<String>,
exhausted: &mut bool,
) -> (Result<Response>, u32) {
crate::trace::note_tag(tag);
let mut route = Route::new(session, &attempt, retry_policy);
if retry_policy.is_none() {
return (counted(route.send(attempt, None).await, 1), 1);
}
let retryable_method =
retry_policy.allow_non_idempotent || is_idempotent(attempt.method.as_str());
let replay = attempt.body.replay();
let mut rewind = Rewind::wrap(&mut attempt.body);
let base_headers = attempt.headers.clone();
let mut n: u32 = 0;
loop {
if let Err(err) = deadline.check() {
return (Err(err), n);
}
let attempt_body = std::mem::take(&mut attempt.body);
let this = attempt.again(attempt_body, base_headers.clone());
let result = route.send(this, retry_policy.proxy_for_retry(n)).await;
let body_retryable = || replay.is_some() || rewind.as_mut().is_some_and(Rewind::reclaim);
let sleep = match retry_policy.plan(&result, n, retryable_method, body_retryable) {
RetryPlan::Backoff(sleep) if sleep < deadline.remaining() => sleep,
plan => {
*exhausted = matches!(
plan,
RetryPlan::Exhausted | RetryPlan::Declined | RetryPlan::Backoff(_)
);
return (counted(result, n + 1), n + 1);
}
};
drop(result);
let slept = session
.unless_shut_down(async {
deadline.sleep(sleep).await;
Ok(())
})
.await;
if let Err(err) = slept {
return (counted(Err(err), n + 1), n + 1);
}
n += 1;
attempt.body = replay
.as_ref()
.and_then(Body::replay)
.or_else(|| rewind.as_mut().and_then(Rewind::again))
.unwrap_or_default();
}
}
fn counted(result: Result<Response>, attempts: u32) -> Result<Response> {
result
.map(|mut response| {
response.set_attempts(attempts);
response
})
.map_err(|mut error| {
error.set_attempts(attempts);
error
})
}
struct Slot {
stream: Option<BoxedStream>,
polled: bool,
}
type SharedSlot = Arc<Mutex<Slot>>;
fn lock(slot: &SharedSlot) -> MutexGuard<'_, Slot> {
slot.lock().unwrap_or_else(PoisonError::into_inner)
}
struct Rewind {
slot: SharedSlot,
length_hint: Option<u64>,
held: Option<BoxedStream>,
}
impl Rewind {
fn wrap(body: &mut Body) -> Option<Self> {
if !matches!(body.0, BodyKind::Stream { .. }) {
return None;
}
let BodyKind::Stream {
stream,
length_hint,
} = std::mem::take(&mut body.0)
else {
return None;
};
let rewind = Self {
slot: Arc::new(Mutex::new(Slot {
stream: Some(stream),
polled: false,
})),
length_hint,
held: None,
};
*body = rewind.body();
Some(rewind)
}
fn body(&self) -> Body {
Body(BodyKind::Stream {
stream: Box::pin(Reclaim(Arc::clone(&self.slot))),
length_hint: self.length_hint,
})
}
fn reclaim(&mut self) -> bool {
let mut slot = lock(&self.slot);
if slot.polled {
return false;
}
self.held = slot.stream.take();
self.held.is_some()
}
fn again(&mut self) -> Option<Body> {
let stream = self.held.take()?;
self.slot = Arc::new(Mutex::new(Slot {
stream: Some(stream),
polled: false,
}));
Some(self.body())
}
}
struct Reclaim(SharedSlot);
impl Stream for Reclaim {
type Item = std::io::Result<Bytes>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut slot = lock(&self.0);
slot.polled = true;
match slot.stream.as_mut() {
Some(stream) => stream.as_mut().poll_next(cx),
None => Poll::Ready(None),
}
}
}
impl std::future::IntoFuture for RequestBuilder {
type Output = Result<Response>;
type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send>>;
fn into_future(self) -> Self::IntoFuture {
Box::pin(async move { self.send().await })
}
}
#[cfg(test)]
#[path = "send_tests.rs"]
mod tests;