use http::StatusCode;
use std::time::Duration;
use tokio::time;
use super::{BoxFuture, Middleware, Next};
use crate::{Error, Request, Response, Result};
type RespondWithTimeout<State> = fn(&State) -> Result<Response, Error>;
pub struct Timeout<F> {
duration: Duration,
or_else: F,
}
pub fn timeout<State>(duration: Duration) -> Timeout<RespondWithTimeout<State>> {
Timeout::new(duration)
}
fn respond_with_timeout<State>(_: &State) -> Result<Response, Error> {
let mut message = String::with_capacity(65);
let status = StatusCode::GATEWAY_TIMEOUT;
message.push_str("The server is taking too long to respond. ");
message.push_str("Please try again later.");
Ok(Error::new_with_status(message, status).into_response())
}
impl<State> Timeout<RespondWithTimeout<State>> {
pub fn new(duration: Duration) -> Self {
Self {
duration,
or_else: respond_with_timeout,
}
}
}
impl<F> Timeout<F> {
pub fn or_else<State, O>(self, f: O) -> Timeout<O>
where
O: Fn(&State) -> Result<Response, Error> + Copy + Send + Sync + 'static,
{
Timeout {
duration: self.duration,
or_else: f,
}
}
}
impl<State, F> Middleware<State> for Timeout<F>
where
F: Fn(&State) -> Result<Response, Error> + Copy + Send + Sync + 'static,
State: Send + Sync + 'static,
{
fn call(&self, request: Request<State>, next: Next<State>) -> BoxFuture<Result<Response>> {
let duration = self.duration;
let or_else = self.or_else;
let state = request.state().clone();
Box::pin(async move {
match time::timeout(duration, next.call(request)).await {
Ok(result) => result,
Err(_) => or_else(&state),
}
})
}
}