mirror of
https://github.com/actix/actix-web.git
synced 2025-01-31 11:22:19 +00:00
Add awc retry middleware
This commit is contained in:
joshbenaron
parent
e0b2246c68
commit
bb1bf5f2be
2 changed files with 474 additions and 0 deletions
|
|
@ -1,6 +1,8 @@
|
|||
mod redirect;
|
||||
mod retry;
|
||||
|
||||
pub use self::redirect::Redirect;
|
||||
pub use self::retry::Retry;
|
||||
|
||||
use std::marker::PhantomData;
|
||||
|
||||
|
|
|
|||
472
awc/src/middleware/retry.rs
Normal file
472
awc/src/middleware/retry.rs
Normal file
|
|
@ -0,0 +1,472 @@
|
|||
use super::Transform;
|
||||
use std::rc::Rc;
|
||||
use actix_http::RequestHeadType;
|
||||
use actix_http::http::{StatusCode, HeaderMap};
|
||||
use std::ops::Deref;
|
||||
use crate::{ConnectRequest, ConnectResponse};
|
||||
use actix_service::Service;
|
||||
use actix_http::client::SendRequestError;
|
||||
use std::task::{Context, Poll};
|
||||
use crate::RequestHead;
|
||||
use futures_core::future::LocalBoxFuture;
|
||||
use actix_http::body::Body;
|
||||
|
||||
pub struct Retry(Inner);
|
||||
|
||||
struct Inner {
|
||||
/// Number of retries. So each request will be tried [max_retry + 1] times
|
||||
max_retry: u8,
|
||||
policies: Vec<RetryPolicy>,
|
||||
}
|
||||
|
||||
impl Retry {
|
||||
pub fn new(retries: u8) -> Self {
|
||||
Retry(Inner {
|
||||
max_retry: retries,
|
||||
policies: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
/// Allows you to add a retry policy to the [`policies`]
|
||||
/// It allows two types of policy:
|
||||
/// - `Vec<StatusCode>` and will retry if one of them is received
|
||||
/// - `Fn(&ResponseHead) -> bool` and will retry when this function resolves to false
|
||||
///
|
||||
/// # example
|
||||
///
|
||||
///```
|
||||
///
|
||||
/// // Creates a policy which will try each request a max of 5 times if any policies resolve to true
|
||||
/// // i.e.
|
||||
/// // if you receive a 401 or 501 status code
|
||||
/// // or
|
||||
/// // the response doesn't have a [`SOME_HEADER`] header
|
||||
/// use awc::http::{StatusCode, HeaderMap};
|
||||
/// use awc::middleware::Retry;
|
||||
///
|
||||
/// let retry_policies = Retry::new(5)
|
||||
/// .policy(vec![StatusCode::INTERNAL_SERVER_ERROR, StatusCode::UNAUTHORIZED])
|
||||
/// .policy(|code: StatusCode, headers: &HeaderMap| {
|
||||
/// return if headers.contains_key("SOME_HEADER") {
|
||||
/// true
|
||||
/// } else {
|
||||
/// false
|
||||
/// };
|
||||
/// });
|
||||
///
|
||||
/// // Creates awc client
|
||||
/// let client = awc::Client::builder()
|
||||
/// .wrap(retry_policies)
|
||||
/// .finish();
|
||||
///```
|
||||
pub fn policy<T>(mut self, p: T) -> Self
|
||||
where T: IntoRetryPolicy
|
||||
{
|
||||
self.0.policies.push(p.into_policy());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[non_exhaustive]
|
||||
pub enum RetryPolicy {
|
||||
Status(Vec<StatusCode>),
|
||||
Custom(Box<dyn Fn(StatusCode, &HeaderMap) -> bool>),
|
||||
}
|
||||
|
||||
pub trait IntoRetryPolicy {
|
||||
fn into_policy(self) -> RetryPolicy;
|
||||
}
|
||||
|
||||
impl<T> IntoRetryPolicy for T
|
||||
where T: for<'a> Fn(StatusCode, &'a HeaderMap) -> bool + 'static
|
||||
{
|
||||
fn into_policy(self) -> RetryPolicy {
|
||||
RetryPolicy::Custom(Box::new(self))
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoRetryPolicy for Vec<StatusCode> {
|
||||
fn into_policy(self) -> RetryPolicy {
|
||||
RetryPolicy::Status(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Transform<S, ConnectRequest> for Retry
|
||||
where
|
||||
S: Service<ConnectRequest, Response=ConnectResponse, Error=SendRequestError> + 'static,
|
||||
{
|
||||
type Transform = RetryService<S>;
|
||||
|
||||
fn new_transform(self, service: S) -> Self::Transform {
|
||||
RetryService {
|
||||
max_retry: self.0.max_retry,
|
||||
policies: self.0.policies.into_boxed_slice().into(),
|
||||
connector: service.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RetryService<S> {
|
||||
policies: Rc<[RetryPolicy]>,
|
||||
max_retry: u8,
|
||||
connector: Rc<S>,
|
||||
}
|
||||
|
||||
impl<S> Service<ConnectRequest> for RetryService<S>
|
||||
where
|
||||
S: Service<ConnectRequest, Response=ConnectResponse, Error=SendRequestError> + 'static,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = S::Error;
|
||||
type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
|
||||
|
||||
fn poll_ready(&self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.connector.poll_ready(ctx)
|
||||
}
|
||||
|
||||
fn call(&self, req: ConnectRequest) -> Self::Future {
|
||||
let connector = self.connector.clone();
|
||||
let policies = self.policies.clone();
|
||||
let max_retry = self.max_retry;
|
||||
|
||||
Box::pin(async move {
|
||||
let mut tries = 0;
|
||||
match req {
|
||||
ConnectRequest::Client(head, body, addr) => {
|
||||
match body {
|
||||
Body::Bytes(b) => {
|
||||
loop {
|
||||
let h = clone_request_head_type(&head);
|
||||
|
||||
match connector.call(ConnectRequest::Client(h, Body::Bytes(b.clone()), addr)).await
|
||||
{
|
||||
Ok(res) => {
|
||||
// ConnectResponse
|
||||
match &res {
|
||||
ConnectResponse::Client(ref r) => {
|
||||
if is_valid_response(policies.as_ref(), r.status(), r.headers()) {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
ConnectResponse::Tunnel(ref head, _) => {
|
||||
if is_valid_response(policies.as_ref(), head.status, head.headers()) {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
// SendRequestError
|
||||
Err(e) => {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Err(e);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Body::Empty => {
|
||||
loop {
|
||||
let h = clone_request_head_type(&head);
|
||||
|
||||
match connector.call(ConnectRequest::Client(h, Body::Empty, addr)).await
|
||||
{
|
||||
Ok(res) => {
|
||||
// ConnectResponse
|
||||
match &res {
|
||||
ConnectResponse::Client(ref r) => {
|
||||
if is_valid_response(policies.as_ref(), r.status(), r.headers()) {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
ConnectResponse::Tunnel(ref head, _) => {
|
||||
if is_valid_response(policies.as_ref(), head.status, head.headers()) {
|
||||
return Ok(res);
|
||||
} else {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// SendRequestError
|
||||
Err(e) => {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Err(e);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
log::debug!("Non cloneable body type given - defaulting to `Body::None`");
|
||||
loop {
|
||||
let h = clone_request_head_type(&head);
|
||||
|
||||
match connector.call(ConnectRequest::Client(h, Body::None, addr)).await
|
||||
{
|
||||
Ok(res) => {
|
||||
// ConnectResponse
|
||||
match &res {
|
||||
ConnectResponse::Client(ref r) => {
|
||||
if is_valid_response(policies.as_ref(), r.status(), r.headers()) {
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
ConnectResponse::Tunnel(ref head, _) => {
|
||||
if is_valid_response(policies.as_ref(), head.status, head.headers()) {
|
||||
return Ok(res);
|
||||
} else {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// SendRequestError
|
||||
Err(e) => {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Err(e);
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ConnectRequest::Tunnel(head, addr) => {
|
||||
loop {
|
||||
let h = clone_request_head(&head);
|
||||
|
||||
match connector.call(ConnectRequest::Tunnel(h, addr)).await {
|
||||
Ok(res) => {
|
||||
match &res {
|
||||
ConnectResponse::Client(r) => {
|
||||
if is_valid_response(&policies, r.status(), r.headers()) {
|
||||
return Ok(res)
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res)
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
ConnectResponse::Tunnel(head, _) => {
|
||||
if is_valid_response(&policies, head.status, head.headers()) {
|
||||
return Ok(res)
|
||||
}
|
||||
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Ok(res)
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
if tries == max_retry {
|
||||
log::debug!("Request max retry reached");
|
||||
return Err(e)
|
||||
}
|
||||
|
||||
tries += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
/// Clones [RequestHeadType] except for the extensions (not required for this middleware)
|
||||
fn clone_request_head_type(head_type: &RequestHeadType) -> RequestHeadType {
|
||||
match head_type {
|
||||
RequestHeadType::Owned(h) => {
|
||||
let mut inner_head = RequestHead::default();
|
||||
inner_head.uri = h.uri.clone();
|
||||
inner_head.method = h.method.clone();
|
||||
inner_head.version = h.version;
|
||||
inner_head.peer_addr = h.peer_addr;
|
||||
inner_head.headers = h.headers.clone();
|
||||
|
||||
RequestHeadType::Owned(inner_head)
|
||||
}
|
||||
RequestHeadType::Rc(h, header_map) => {
|
||||
RequestHeadType::Rc(h.clone(), header_map.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
/// Clones [RequestHeadType] except for the extensions (not required for this middleware)
|
||||
fn clone_request_head(head: &RequestHead) -> RequestHead {
|
||||
let mut new_head = RequestHead::default();
|
||||
new_head.uri = head.uri.clone();
|
||||
new_head.method = head.method.clone();
|
||||
new_head.version = head.version;
|
||||
new_head.headers = head.headers.clone();
|
||||
new_head.peer_addr = head.peer_addr;
|
||||
|
||||
new_head
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
/// Checks whether the response matches the policies
|
||||
fn is_valid_response(policies: &[RetryPolicy], status_code: StatusCode, headers: &HeaderMap) -> bool {
|
||||
policies.iter().all(|policy| {
|
||||
match policy {
|
||||
RetryPolicy::Status(v) => {
|
||||
// is valid if:
|
||||
// - the list of status codes is empty
|
||||
// or
|
||||
// - the list doesn't contain the received status code
|
||||
v.is_empty() || !v.contains(&status_code)
|
||||
}
|
||||
RetryPolicy::Custom(func) => {
|
||||
// custom policy
|
||||
(func.deref())(status_code, headers)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use actix_web::{web, App, Error, HttpResponse};
|
||||
|
||||
use super::*;
|
||||
use crate::ClientBuilder;
|
||||
|
||||
#[actix_rt::test]
|
||||
async fn test_basic_policy() {
|
||||
let client = ClientBuilder::new()
|
||||
.disable_redirects()
|
||||
.wrap(Retry::new(3)
|
||||
.policy(vec![StatusCode::INTERNAL_SERVER_ERROR])
|
||||
)
|
||||
.finish();
|
||||
|
||||
let srv = actix_test::start(|| {
|
||||
App::new()
|
||||
.service(web::resource("/test").route(web::to(|| async {
|
||||
Ok::<_, Error>(
|
||||
HttpResponse::InternalServerError()
|
||||
.finish(),
|
||||
)
|
||||
})))
|
||||
});
|
||||
|
||||
let res = client.get(srv.url("/test")).send().await.unwrap();
|
||||
|
||||
assert_eq!(res.status().as_u16(), 500);
|
||||
}
|
||||
|
||||
#[actix_rt::test]
|
||||
async fn test_header_policy() {
|
||||
std::env::set_var("RUST_LOG", "RUST_LOG=debug,a=debug");
|
||||
env_logger::init();
|
||||
|
||||
let client = ClientBuilder::new()
|
||||
.disable_redirects()
|
||||
.wrap(Retry::new(3)
|
||||
.policy(|code: StatusCode, headers: &HeaderMap| {
|
||||
code.is_success() && headers.contains_key("SOME_HEADER")
|
||||
})
|
||||
)
|
||||
.finish();
|
||||
|
||||
let srv = actix_test::start(|| {
|
||||
App::new()
|
||||
.service(web::resource("/test").route(web::to(|| async {
|
||||
Ok::<_, Error>(
|
||||
HttpResponse::Ok()
|
||||
.insert_header(("SOME_HEADER", "test"))
|
||||
.finish(),
|
||||
)
|
||||
})))
|
||||
});
|
||||
|
||||
let res = client.get(srv.url("/test")).send().await.unwrap();
|
||||
|
||||
assert_eq!(res.status().as_u16(), 200);
|
||||
}
|
||||
|
||||
#[actix_rt::test]
|
||||
async fn test_bad_header_policy() {
|
||||
std::env::set_var("RUST_LOG", "RUST_LOG=debug,a=debug");
|
||||
env_logger::init();
|
||||
|
||||
let client = ClientBuilder::new()
|
||||
.disable_redirects()
|
||||
.wrap(Retry::new(3)
|
||||
.policy(|code: StatusCode, headers: &HeaderMap| {
|
||||
code.is_success() && headers.contains_key("WRONG_HEADER")
|
||||
})
|
||||
)
|
||||
.finish();
|
||||
|
||||
let srv = actix_test::start(|| {
|
||||
App::new()
|
||||
.service(web::resource("/test").route(web::to(|| async {
|
||||
Ok::<_, Error>(
|
||||
HttpResponse::Ok()
|
||||
.insert_header(("SOME_HEADER", "test"))
|
||||
.finish(),
|
||||
)
|
||||
})))
|
||||
});
|
||||
|
||||
let res = client.get(srv.url("/test")).send().await.unwrap();
|
||||
|
||||
assert_eq!(res.status().as_u16(), 200);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue