use super::{LayerErrorFn, LayerErrorStatic, MakeLayerError};
use crate::Service;
use rama_error::{BoxError, ErrorContext, ErrorExt};
use rama_utils::macros::define_inner_service_accessors;
use std::time::Duration;
mod error;
#[doc(inline)]
pub use error::Elapsed;
mod layer;
#[doc(inline)]
pub use layer::TimeoutLayer;
#[derive(Debug, Clone)]
pub struct Timeout<S, F> {
inner: S,
into_error: F,
timeout: Option<Duration>,
}
impl<S, F> Timeout<S, F> {
define_inner_service_accessors!();
}
pub type DefaultTimeout<S> = Timeout<S, LayerErrorStatic<Elapsed>>;
impl<S> DefaultTimeout<S> {
pub fn new(inner: S, timeout: Duration) -> Self {
Self::with_error(inner, timeout, error::Elapsed::new(Some(timeout)))
}
pub fn never(inner: S) -> Self {
Self {
inner,
timeout: None,
into_error: LayerErrorStatic::new(error::Elapsed::new(None)),
}
}
}
impl<S, E> Timeout<S, LayerErrorStatic<E>> {
pub fn with_error(inner: S, timeout: Duration, error: E) -> Self
where
E: Clone + Send + Sync + 'static,
{
Self {
inner,
timeout: Some(timeout),
into_error: LayerErrorStatic::new(error),
}
}
}
impl<S, F> Timeout<S, LayerErrorFn<F>> {
pub fn with_error_fn<E>(inner: S, timeout: Duration, error_fn: F) -> Self
where
F: Fn() -> E + Send + Sync + 'static,
E: Send + 'static,
{
Self {
inner,
timeout: Some(timeout),
into_error: LayerErrorFn::new(error_fn),
}
}
}
impl<S, F> Timeout<S, F>
where
F: MakeLayerError,
{
pub(crate) fn with(inner: S, timeout: Option<Duration>, into_error: F) -> Self {
Self {
inner,
timeout,
into_error,
}
}
}
impl<T, F, Input, E> Service<Input> for Timeout<T, F>
where
Input: Send + 'static,
F: MakeLayerError<Error = E>,
E: Into<BoxError> + Send + 'static,
T: Service<Input, Error: Into<BoxError>>,
{
type Output = T::Output;
type Error = BoxError;
async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
match self.timeout {
Some(duration) => tokio::select! {
res = self.inner.serve(input) => res.into_box_error(),
_ = tokio::time::sleep(duration) => Err(self.into_error.make_layer_error().into_box_error()),
},
None => self.inner.serve(input).await.into_box_error(),
}
}
}