use super::{LayerErrorFn, LayerErrorStatic, MakeLayerError};
use crate::{Context, Service};
use rama_utils::macros::define_inner_service_accessors;
use std::{fmt, time::Duration};
mod error;
#[doc(inline)]
pub use error::Elapsed;
mod layer;
#[doc(inline)]
pub use layer::TimeoutLayer;
pub struct Timeout<S, F> {
inner: S,
into_error: F,
timeout: Duration,
}
impl<S, F> Timeout<S, F> {
define_inner_service_accessors!();
}
impl<S: fmt::Debug, F: fmt::Debug> fmt::Debug for Timeout<S, F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Timeout")
.field("inner", &self.inner)
.field("into_error", &self.into_error)
.field("timeout", &self.timeout)
.finish()
}
}
impl<S, F> Clone for Timeout<S, F>
where
S: Clone,
F: Clone,
{
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
into_error: self.into_error.clone(),
timeout: self.timeout,
}
}
}
impl<S> Timeout<S, LayerErrorStatic<Elapsed>> {
pub fn new(inner: S, timeout: Duration) -> Self {
Self::with_error(inner, timeout, error::Elapsed::new(timeout))
}
}
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,
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: FnOnce() -> E + Clone + Send + Sync + 'static,
E: Send + 'static,
{
Self {
inner,
timeout,
into_error: LayerErrorFn::new(error_fn),
}
}
}
impl<S, F> Timeout<S, F>
where
F: MakeLayerError,
{
pub(crate) fn with(inner: S, timeout: Duration, into_error: F) -> Self {
Self {
inner,
timeout,
into_error,
}
}
}
impl<T, F, S, Request, E> Service<S, Request> for Timeout<T, F>
where
Request: Send + 'static,
S: Clone + Send + Sync + 'static,
F: MakeLayerError<Error = E>,
E: Into<T::Error> + Send + 'static,
T: Service<S, Request>,
{
type Response = T::Response;
type Error = T::Error;
async fn serve(
&self,
ctx: Context<S>,
request: Request,
) -> Result<Self::Response, Self::Error> {
tokio::select! {
res = self.inner.serve(ctx, request) => res,
_ = tokio::time::sleep(self.timeout) => Err(self.into_error.make_layer_error().into()),
}
}
}