axum_tracing_opentelemetry/middleware/
response_injector.rs

1use futures_core::future::BoxFuture;
2use http::{Request, Response};
3use std::task::{Context, Poll};
4use tower::{Layer, Service};
5use tracing_opentelemetry_instrumentation_sdk as otel;
6use tracing_opentelemetry_instrumentation_sdk::http as otel_http;
7
8#[deprecated(
9    since = "0.12.0",
10    note = "keep for transition, replaced by OtelInResponseLayer"
11)]
12#[must_use]
13pub fn response_with_trace_layer() -> OtelInResponseLayer {
14    OtelInResponseLayer {}
15}
16
17#[derive(Default, Debug, Clone)]
18pub struct OtelInResponseLayer;
19
20impl<S> Layer<S> for OtelInResponseLayer {
21    type Service = OtelInResponseService<S>;
22
23    fn layer(&self, inner: S) -> Self::Service {
24        OtelInResponseService { inner }
25    }
26}
27
28#[derive(Default, Debug, Clone)]
29pub struct OtelInResponseService<S> {
30    inner: S,
31}
32
33impl<S, B, B2> Service<Request<B>> for OtelInResponseService<S>
34where
35    S: Service<Request<B>, Response = Response<B2>> + Send + 'static,
36    S::Future: Send + 'static,
37{
38    type Response = S::Response;
39    type Error = S::Error;
40    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
41    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
42
43    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
44        self.inner.poll_ready(cx)
45    }
46
47    #[allow(unused_mut)]
48    fn call(&mut self, mut request: Request<B>) -> Self::Future {
49        let future = self.inner.call(request);
50
51        Box::pin(async move {
52            let mut response = future.await?;
53            // inject the trace context into the response (optional but useful for debugging and client)
54            otel_http::inject_context(&otel::find_current_context(), response.headers_mut());
55            Ok(response)
56        })
57    }
58}