otel_bootstrap/
axum_middleware.rs

1//! Axum tower middleware for W3C trace context propagation.
2//!
3//! Enabled by the `axum` feature flag. Add to any axum [`Router`] via
4//! [`crate::axum_layer()`]:
5//!
6//! ```no_run
7//! use axum::Router;
8//!
9//! let app: Router = Router::new().layer(otel_bootstrap::axum_layer());
10//! ```
11
12use axum::{
13    body::Body,
14    http::{HeaderMap, HeaderName, HeaderValue, Request, Response},
15};
16use opentelemetry::{
17    context::FutureExt,
18    global,
19    propagation::{Extractor, Injector},
20    trace::{SpanKind, Status, TraceContextExt, Tracer},
21};
22use opentelemetry_semantic_conventions::attribute::{
23    HTTP_REQUEST_METHOD, HTTP_RESPONSE_STATUS_CODE,
24};
25use std::{
26    future::Future,
27    pin::Pin,
28    task::{self, Poll},
29};
30use tower::{Layer, Service};
31
32/// Tower [`Layer`] that instruments incoming HTTP requests with OpenTelemetry
33/// trace context propagation.
34///
35/// Attach to an axum router with [`crate::axum_layer()`].
36///
37/// # Example
38/// ```no_run
39/// use axum::Router;
40///
41/// let app: Router = Router::new().layer(otel_bootstrap::axum_layer());
42/// ```
43#[derive(Clone, Debug)]
44pub struct OtelTraceLayer;
45
46impl<S> Layer<S> for OtelTraceLayer {
47    type Service = OtelTraceService<S>;
48
49    fn layer(&self, inner: S) -> Self::Service {
50        OtelTraceService { inner }
51    }
52}
53
54/// Tower [`Service`] produced by [`OtelTraceLayer`].
55///
56/// This type is not constructed directly. It is returned by
57/// [`OtelTraceLayer`] when wrapping an inner [`tower::Service`].
58#[derive(Clone, Debug)]
59pub struct OtelTraceService<S> {
60    inner: S,
61}
62
63impl<S> Service<Request<Body>> for OtelTraceService<S>
64where
65    S: Service<Request<Body>, Response = Response<Body>> + Send + Clone + 'static,
66    S::Future: Send + 'static,
67    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
68{
69    type Response = Response<Body>;
70    type Error = S::Error;
71    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
72
73    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
74        self.inner.poll_ready(cx)
75    }
76
77    fn call(&mut self, req: Request<Body>) -> Self::Future {
78        let method = req.method().to_string();
79        let route = req.uri().path().to_string();
80
81        // Extract parent context from incoming headers.
82        let parent_cx = global::get_text_map_propagator(|propagator| {
83            propagator.extract(&HeaderExtractor(req.headers()))
84        });
85
86        // Create a child span inheriting the remote parent.
87        let tracer = global::tracer("otel-bootstrap");
88        let span = tracer
89            .span_builder(format!("{method} {route}"))
90            .with_kind(SpanKind::Server)
91            .with_attributes([opentelemetry::KeyValue::new(HTTP_REQUEST_METHOD, method)])
92            .start_with_context(&tracer, &parent_cx);
93
94        let cx = parent_cx.with_span(span);
95
96        // `poll_ready` was called on `self.inner` before `call`, per the
97        // tower contract — some inner services track readiness per-handle
98        // (e.g. a `tower::buffer::Buffer`), so firing on a fresh unpolled
99        // clone can panic. Swap the clone into `self` for next time, and
100        // fire this request on the already-ready original handle.
101        let clone = self.inner.clone();
102        let mut inner = std::mem::replace(&mut self.inner, clone);
103
104        Box::pin(async move {
105            // Without this, `Context::current()` inside the handler (and
106            // anything it calls, e.g. api-bones's `propagation::inject_current`)
107            // sees the empty root context, not this request's parent — every
108            // outbound call the handler makes injects a disconnected
109            // traceparent regardless of what was extracted above. `with_context`
110            // (not a bare `cx.attach()` guard) is required because the inner
111            // future can resume on a different tokio worker thread between
112            // polls, and thread-local attach doesn't survive that hop.
113            let mut response = inner.call(req).with_context(cx.clone()).await?;
114
115            // Record HTTP status on the span.
116            let status_code = response.status().as_u16();
117            cx.span().set_attribute(opentelemetry::KeyValue::new(
118                HTTP_RESPONSE_STATUS_CODE,
119                status_code as i64,
120            ));
121            if response.status().is_server_error() {
122                cx.span().set_status(Status::Error {
123                    description: response.status().canonical_reason().unwrap_or("").into(),
124                });
125            }
126
127            // Inject outgoing trace context into response headers.
128            let mut injector = HeaderInjector(response.headers_mut());
129            global::get_text_map_propagator(|propagator| {
130                propagator.inject_context(&cx, &mut injector);
131            });
132
133            Ok(response)
134        })
135    }
136}
137
138/// Tower [`Layer`] that records span attributes from a `T: EnrichSpan`
139/// extension on each incoming request.
140///
141/// Construct via [`crate::span_enricher_layer`]. When no `T` extension is
142/// found in the request (e.g. a platform-scope route), the service is a no-op.
143///
144/// # Example
145/// ```no_run
146/// use axum::{Router, Extension, routing::get};
147/// use otel_bootstrap::span_enrichment::EnrichSpan;
148/// use tracing_opentelemetry::OpenTelemetrySpanExt as _;
149///
150/// #[derive(Clone)]
151/// struct MyCtx { user_id: String }
152///
153/// impl EnrichSpan for MyCtx {
154///     fn enrich_span(&self, span: &tracing::Span) {
155///         span.set_attribute("enduser.id", self.user_id.clone());
156///     }
157/// }
158///
159/// let app: Router = Router::new()
160///     .route("/", get(|| async { "ok" }))
161///     .layer(otel_bootstrap::span_enricher_layer::<MyCtx>())
162///     .layer(Extension(MyCtx { user_id: "u1".into() }))
163///     .layer(otel_bootstrap::axum_layer());
164/// ```
165#[derive(Debug)]
166pub struct SpanEnricherLayer<T>(std::marker::PhantomData<T>);
167
168impl<T> Default for SpanEnricherLayer<T> {
169    fn default() -> Self {
170        Self(std::marker::PhantomData)
171    }
172}
173
174impl<T> Clone for SpanEnricherLayer<T> {
175    fn clone(&self) -> Self {
176        Self(std::marker::PhantomData)
177    }
178}
179
180impl<T, S> Layer<S> for SpanEnricherLayer<T>
181where
182    T: crate::span_enrichment::EnrichSpan + Clone + Send + Sync + 'static,
183{
184    type Service = SpanEnricherService<T, S>;
185
186    fn layer(&self, inner: S) -> Self::Service {
187        SpanEnricherService {
188            inner,
189            _marker: std::marker::PhantomData,
190        }
191    }
192}
193
194/// Tower [`Service`] produced by [`SpanEnricherLayer`].
195#[derive(Clone, Debug)]
196pub struct SpanEnricherService<T, S> {
197    inner: S,
198    _marker: std::marker::PhantomData<T>,
199}
200
201impl<T, S> Service<Request<Body>> for SpanEnricherService<T, S>
202where
203    T: crate::span_enrichment::EnrichSpan + Clone + Send + Sync + 'static,
204    S: Service<Request<Body>, Response = Response<Body>>,
205{
206    type Response = Response<Body>;
207    type Error = S::Error;
208    type Future = S::Future;
209
210    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
211        self.inner.poll_ready(cx)
212    }
213
214    fn call(&mut self, req: Request<Body>) -> Self::Future {
215        if let Some(ctx) = req.extensions().get::<T>() {
216            ctx.enrich_span(&tracing::Span::current());
217        }
218        self.inner.call(req)
219    }
220}
221
222/// [`Extractor`] that reads from [`HeaderMap`].
223struct HeaderExtractor<'a>(&'a HeaderMap);
224
225impl Extractor for HeaderExtractor<'_> {
226    fn get(&self, key: &str) -> Option<&str> {
227        self.0.get(key).and_then(|v| v.to_str().ok())
228    }
229
230    fn keys(&self) -> Vec<&str> {
231        self.0.keys().map(HeaderName::as_str).collect()
232    }
233}
234
235/// [`Injector`] that writes into a mutable [`HeaderMap`].
236struct HeaderInjector<'a>(&'a mut HeaderMap);
237
238impl Injector for HeaderInjector<'_> {
239    fn set(&mut self, key: &str, value: String) {
240        if let (Ok(name), Ok(val)) = (
241            HeaderName::from_bytes(key.as_bytes()),
242            HeaderValue::from_str(&value),
243        ) {
244            self.0.insert(name, val);
245        }
246    }
247}
otel_bootstrap/axum_middleware.rs

otel_bootstrap/
axum_middleware.rs
