tracewrapper.rs

#[derive(Clone)]
struct TraceWrapper<T> {
    inner: Trace<T, SharedClassifier<GrpcErrorsAsFailures>>,
}

impl<T> TraceWrapper<T> {
    fn new(inner: T) -> Self {
        let wrapped = ServiceBuilder::new()
            .layer(TraceLayer::new_for_grpc())
            .service(inner);
        Self { inner: wrapped }
    }
}

impl<T> NamedService for TraceWrapper<T>
where
    T: NamedService,
{
    const NAME: &'static str = T::NAME;
}

impl<T, B> tower_service::Service<http::Request<B>> for TraceWrapper<T>
where
    T: tower_service::Service<http::Request<B>, Response = http::Response<BoxBody>> + Clone,
    <T as Service<http::Request<B>>>::Future: std::marker::Send + 'static,
    B: http_body::Body + Send + Sync + 'static,
    B::Error: Into<tonic::codegen::StdError> + Send + 'static,
    <T as Service<http::Request<B>>>::Error: std::fmt::Display + 'static,
{
    type Response = T::Response;
    type Error = T::Error;
    type Future = tonic::codegen::BoxFuture<Self::Response, Self::Error>;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: http::Request<B>) -> Self::Future {
        let mut inner = self.inner.clone();
        inner
            .call(req)
            .and_then(|res| async move { Ok(res.map(BoxBody::new)) })
            .boxed()
    }
}