use motore::{layer::Layer, service::Service};
use volo::context::Context;
use crate::{
context::client::Config,
error::ClientError,
request::{Request, RequestPartsExt},
};
#[derive(Clone, Debug, Default)]
pub struct Timeout;
impl<S> Layer<S> for Timeout {
type Service = TimeoutService<S>;
fn layer(self, inner: S) -> Self::Service {
TimeoutService { inner }
}
}
pub struct TimeoutService<S> {
inner: S,
}
impl<Cx, B, S> Service<Cx, Request<B>> for TimeoutService<S>
where
Cx: Context<Config = Config> + Send,
B: Send,
S: Service<Cx, Request<B>, Error = ClientError> + Send + Sync,
{
type Response = S::Response;
type Error = S::Error;
async fn call(&self, cx: &mut Cx, req: Request<B>) -> Result<Self::Response, Self::Error> {
let timeout = cx.rpc_info().config().timeout().cloned();
if let Some(duration) = timeout {
let url = req.url();
let fut = self.inner.call(cx, req);
let sleep = tokio::time::sleep(duration);
tokio::select! {
res = fut => res,
_ = sleep => {
if let Some(url) = url {
tracing::warn!("[Volo-HTTP] request timeout on `{url}`");
}
Err(crate::error::client::timeout().with_endpoint(cx.rpc_info().callee()))
}
}
} else {
self.inner.call(cx, req).await
}
}
}