Skip to main content

kerniq_observability_axum/
layer.rs

1//! The Tower layer.
2
3use std::fmt;
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::{Arc, OnceLock};
7use std::task::{Context, Poll, ready};
8use std::time::Instant;
9
10use axum::extract::MatchedPath;
11use bytes::Buf;
12use http::{Method, Request, Response, StatusCode};
13use kerniq_routing::Middleware;
14use kerniq_telemetry::attributes::{http_method, protocol_version};
15use opentelemetry::{StringValue, Value};
16use pin_project_lite::pin_project;
17use tower::{Layer, Service};
18
19#[cfg(feature = "metrics")]
20use crate::HttpMetrics;
21use crate::body::{Completion, ObservedBody, Outcome};
22
23#[cfg(feature = "tracing")]
24type Propagator = Arc<dyn opentelemetry::propagation::TextMapPropagator + Send + Sync>;
25
26/// Configuration for HTTP instrumentation. Convert it with
27/// [`into_layer`](Self::into_layer), or pass it to `App::middleware`.
28#[derive(Clone)]
29#[must_use]
30pub struct HttpInstrumentation {
31    #[cfg(feature = "metrics")]
32    metrics: Option<HttpMetrics>,
33    #[cfg(feature = "tracing")]
34    tracing: bool,
35    #[cfg(feature = "tracing")]
36    propagator: Option<Propagator>,
37    exclude: Vec<String>,
38    scheme: &'static str,
39    request_id_header: http::HeaderName,
40}
41
42impl fmt::Debug for HttpInstrumentation {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        let mut d = f.debug_struct("HttpInstrumentation");
45        #[cfg(feature = "metrics")]
46        d.field("metrics", &self.metrics.is_some());
47        #[cfg(feature = "tracing")]
48        d.field("tracing", &self.tracing)
49            .field("propagator", &self.propagator.is_some());
50        d.field("exclude", &self.exclude)
51            .field("scheme", &self.scheme)
52            .finish()
53    }
54}
55
56impl Default for HttpInstrumentation {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl HttpInstrumentation {
63    /// Tracing on (when the `tracing` feature is enabled), metrics off until
64    /// a meter is supplied, no propagator, no exclusions.
65    pub fn new() -> Self {
66        Self {
67            #[cfg(feature = "metrics")]
68            metrics: None,
69            #[cfg(feature = "tracing")]
70            tracing: true,
71            #[cfg(feature = "tracing")]
72            propagator: None,
73            exclude: Vec::new(),
74            scheme: "http",
75            request_id_header: http::HeaderName::from_static("x-request-id"),
76        }
77    }
78
79    /// Records HTTP metrics with instruments created on `meter`.
80    #[cfg(feature = "metrics")]
81    pub fn meter(self, meter: &opentelemetry::metrics::Meter) -> Self {
82        self.metrics(HttpMetrics::new(meter))
83    }
84
85    /// Records HTTP metrics with existing instruments.
86    #[cfg(feature = "metrics")]
87    pub fn metrics(mut self, metrics: HttpMetrics) -> Self {
88        self.metrics = Some(metrics);
89        self
90    }
91
92    /// Enables or disables the per-request span.
93    #[cfg(feature = "tracing")]
94    pub fn tracing(mut self, enabled: bool) -> Self {
95        self.tracing = enabled;
96        self
97    }
98
99    /// Extracts the parent context from request headers with `propagator`
100    /// (typically W3C TraceContext + Baggage).
101    #[cfg(feature = "tracing")]
102    pub fn propagator<P>(mut self, propagator: P) -> Self
103    where
104        P: opentelemetry::propagation::TextMapPropagator + Send + Sync + 'static,
105    {
106        self.propagator = Some(Arc::new(propagator));
107        self
108    }
109
110    /// Like [`propagator`](Self::propagator), from a shared instance.
111    #[cfg(feature = "tracing")]
112    pub fn shared_propagator(mut self, propagator: Propagator) -> Self {
113        self.propagator = Some(propagator);
114        self
115    }
116
117    /// Excludes a route template (`/health`, `/users/{id}`) — or, for
118    /// requests that match no route, an exact path — from metrics and
119    /// tracing.
120    pub fn exclude(mut self, route: impl Into<String>) -> Self {
121        self.exclude.push(route.into());
122        self
123    }
124
125    /// The `url.scheme` reported when the request URI has none (HTTP/1).
126    /// Set to `"https"` when TLS terminates in front of the application.
127    /// Default `"http"`.
128    pub fn scheme(mut self, scheme: &'static str) -> Self {
129        self.scheme = scheme;
130        self
131    }
132
133    /// The header whose value is recorded as `kerniq.request_id` on the
134    /// request span. Default `x-request-id` (set by `RequestId`).
135    ///
136    /// # Panics
137    ///
138    /// Panics if `name` is not a valid header name.
139    pub fn request_id_header(mut self, name: &str) -> Self {
140        self.request_id_header =
141            http::HeaderName::try_from(name).expect("invalid request id header name");
142        self
143    }
144
145    /// Converts into a Tower layer.
146    pub fn into_layer(self) -> HttpLayer {
147        HttpLayer {
148            config: Arc::new(Config {
149                #[cfg(feature = "metrics")]
150                metrics: self.metrics.map(Arc::new),
151                #[cfg(feature = "tracing")]
152                tracing: self.tracing,
153                #[cfg(feature = "tracing")]
154                propagator: self.propagator,
155                exclude: self.exclude.into_boxed_slice(),
156                scheme: self.scheme,
157                request_id_header: self.request_id_header,
158            }),
159        }
160    }
161}
162
163impl Middleware for HttpInstrumentation {
164    fn apply<S>(self, router: axum::Router<S>) -> axum::Router<S>
165    where
166        S: Clone + Send + Sync + 'static,
167    {
168        router.layer(self.into_layer())
169    }
170}
171
172struct Config {
173    #[cfg(feature = "metrics")]
174    metrics: Option<Arc<HttpMetrics>>,
175    #[cfg(feature = "tracing")]
176    tracing: bool,
177    #[cfg(feature = "tracing")]
178    propagator: Option<Propagator>,
179    exclude: Box<[String]>,
180    scheme: &'static str,
181    request_id_header: http::HeaderName,
182}
183
184impl Config {
185    fn excluded(&self, route: &str) -> bool {
186        self.exclude.iter().any(|e| e == route)
187    }
188}
189
190/// Tower layer produced by [`HttpInstrumentation::into_layer`].
191#[derive(Clone)]
192pub struct HttpLayer {
193    config: Arc<Config>,
194}
195
196impl fmt::Debug for HttpLayer {
197    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
198        f.debug_struct("HttpLayer").finish_non_exhaustive()
199    }
200}
201
202impl<S> Layer<S> for HttpLayer {
203    type Service = HttpService<S>;
204
205    fn layer(&self, inner: S) -> Self::Service {
206        HttpService {
207            inner,
208            config: self.config.clone(),
209            route: Arc::new(OnceLock::new()),
210        }
211    }
212}
213
214/// The route template of the route this service instance wraps, converted
215/// once into an OpenTelemetry value.
216struct RouteEntry {
217    template: Arc<str>,
218    value: Value,
219    excluded: bool,
220}
221
222/// Service produced by [`HttpLayer`].
223///
224/// Axum's `Router::layer` wraps each route separately, so each instance
225/// usually sees a single route template. The template is converted into an
226/// attribute value on the first request and reused afterwards; requests
227/// with a different template (e.g. the shared 404 fallback) fall back to a
228/// per-request conversion.
229#[derive(Clone)]
230pub struct HttpService<S> {
231    inner: S,
232    config: Arc<Config>,
233    route: Arc<OnceLock<RouteEntry>>,
234}
235
236impl<S> fmt::Debug for HttpService<S> {
237    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
238        f.debug_struct("HttpService").finish_non_exhaustive()
239    }
240}
241
242impl<S> HttpService<S> {
243    /// Returns `(route value, excluded)`.
244    fn route<B>(&self, request: &Request<B>) -> (Option<Value>, bool) {
245        let Some(matched) = request.extensions().get::<MatchedPath>() else {
246            return (None, self.config.excluded(request.uri().path()));
247        };
248        let template = matched.as_str();
249        if let Some(entry) = self.route.get()
250            && &*entry.template == template
251        {
252            return (Some(entry.value.clone()), entry.excluded);
253        }
254        let shared: Arc<str> = Arc::from(template);
255        let entry = RouteEntry {
256            value: Value::String(StringValue::from(shared.clone())),
257            template: shared,
258            excluded: self.config.excluded(template),
259        };
260        let result = (Some(entry.value.clone()), entry.excluded);
261        let _ = self.route.set(entry);
262        result
263    }
264}
265
266impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for HttpService<S>
267where
268    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
269    ResBody: http_body::Body,
270    ResBody::Data: Buf,
271{
272    type Response = Response<ObservedBody<ResBody>>;
273    type Error = S::Error;
274    type Future = ResponseFuture<S::Future>;
275
276    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
277        self.inner.poll_ready(cx)
278    }
279
280    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
281        let start = Instant::now();
282        let (route, excluded) = self.route(&request);
283        if excluded {
284            return ResponseFuture {
285                inner: self.inner.call(request),
286                completion: None,
287                head: false,
288                request_id_header: None,
289            };
290        }
291
292        let config = &*self.config;
293        let method = http_method(request.method());
294        let scheme = request.uri().scheme_str().map_or(config.scheme, |s| {
295            if s == "https" { "https" } else { "http" }
296        });
297        let protocol = protocol_version(request.version());
298        let request_size = request
299            .headers()
300            .get(http::header::CONTENT_LENGTH)
301            .and_then(|v| v.to_str().ok())
302            .and_then(|v| v.parse::<u64>().ok());
303
304        #[cfg(feature = "metrics")]
305        if let Some(metrics) = &config.metrics {
306            use kerniq_telemetry::semconv::http as conv;
307            use opentelemetry::KeyValue;
308            metrics.active.add(
309                1,
310                &[
311                    KeyValue::new(conv::REQUEST_METHOD, method),
312                    KeyValue::new(conv::URL_SCHEME, scheme),
313                ],
314            );
315        }
316
317        #[cfg(feature = "tracing")]
318        let span = if config.tracing {
319            make_span(config, &request, method, scheme, protocol, route.as_ref())
320        } else {
321            tracing::Span::none()
322        };
323
324        let completion = Completion {
325            #[cfg(feature = "metrics")]
326            metrics: config.metrics.clone(),
327            method,
328            scheme,
329            protocol,
330            route,
331            request_size,
332            status: None,
333            start,
334            first_byte: false,
335            bytes: 0,
336            done: false,
337            #[cfg(feature = "tracing")]
338            span,
339        };
340
341        let head = request.method() == Method::HEAD;
342        let request_id_header = (!request.headers().contains_key(&config.request_id_header))
343            .then(|| config.request_id_header.clone());
344        #[cfg(feature = "tracing")]
345        let inner = {
346            let _entered = completion.span().enter();
347            self.inner.call(request)
348        };
349        #[cfg(not(feature = "tracing"))]
350        let inner = self.inner.call(request);
351
352        ResponseFuture {
353            inner,
354            completion: Some(completion),
355            head,
356            request_id_header,
357        }
358    }
359}
360
361#[cfg(feature = "tracing")]
362fn make_span<B>(
363    config: &Config,
364    request: &Request<B>,
365    method: &'static str,
366    scheme: &'static str,
367    protocol: Option<&'static str>,
368    route: Option<&Value>,
369) -> tracing::Span {
370    use tracing::field::Empty;
371    use tracing_opentelemetry::OpenTelemetrySpanExt;
372
373    let span = tracing::info_span!(
374        target: "kerniq::http",
375        "HTTP request",
376        otel.name = Empty,
377        otel.kind = "server",
378        otel.status_code = Empty,
379        http.request.method = method,
380        http.route = Empty,
381        http.response.status_code = Empty,
382        http.response.body.size = Empty,
383        url.path = request.uri().path(),
384        url.scheme = scheme,
385        network.protocol.version = Empty,
386        error.type = Empty,
387        user_agent.original = Empty,
388        kerniq.request_id = Empty,
389    );
390    if span.is_disabled() {
391        return span;
392    }
393    match route {
394        Some(route) => {
395            let route = route.as_str();
396            span.record("http.route", &*route);
397            span.record("otel.name", format!("{method} {route}"));
398        }
399        None => {
400            span.record("otel.name", method);
401        }
402    }
403    if let Some(protocol) = protocol {
404        span.record("network.protocol.version", protocol);
405    }
406    let headers = request.headers();
407    if let Some(ua) = headers
408        .get(http::header::USER_AGENT)
409        .and_then(|v| v.to_str().ok())
410    {
411        span.record("user_agent.original", ua);
412    }
413    if let Some(id) = headers
414        .get(&config.request_id_header)
415        .and_then(|v| v.to_str().ok())
416    {
417        span.record("kerniq.request_id", id);
418    }
419    if let Some(propagator) = &config.propagator {
420        let cx = crate::propagation::extract_context(propagator.as_ref(), headers);
421        use opentelemetry::trace::TraceContextExt;
422        if cx.span().span_context().is_valid() {
423            let _ = span.set_parent(cx);
424        }
425    }
426    span
427}
428
429pin_project! {
430    /// Response future for [`HttpService`].
431    ///
432    /// If it is dropped before completing (the client disconnected while the
433    /// handler was running, or a timeout cancelled it), the request is
434    /// recorded with `error.type = kerniq.aborted`.
435    pub struct ResponseFuture<F> {
436        #[pin]
437        inner: F,
438        completion: Option<Completion>,
439        head: bool,
440        // Set when the request carried no id: read it from the response.
441        request_id_header: Option<http::HeaderName>,
442    }
443}
444
445impl<F> fmt::Debug for ResponseFuture<F> {
446    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
447        f.debug_struct("ResponseFuture").finish_non_exhaustive()
448    }
449}
450
451impl<F, ResBody, E> Future for ResponseFuture<F>
452where
453    F: Future<Output = Result<Response<ResBody>, E>>,
454    ResBody: http_body::Body,
455    ResBody::Data: Buf,
456{
457    type Output = Result<Response<ObservedBody<ResBody>>, E>;
458
459    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
460        let this = self.project();
461        let result = {
462            #[cfg(feature = "tracing")]
463            let _entered = this.completion.as_ref().map(|c| c.span().enter());
464            ready!(this.inner.poll(cx))
465        };
466        let Some(mut completion) = this.completion.take() else {
467            return Poll::Ready(result.map(|r| r.map(|b| ObservedBody::new(b, None))));
468        };
469        match result {
470            Ok(response) => {
471                let status = response.status();
472                completion.status = Some(status.as_u16());
473                #[cfg(feature = "tracing")]
474                {
475                    let span = completion.span();
476                    span.record("http.response.status_code", status.as_u16());
477                    // An inner `RequestId` middleware generates ids after this
478                    // layer saw the request; it echoes them on the response.
479                    if !span.is_disabled()
480                        && let Some(id) = this
481                            .request_id_header
482                            .as_ref()
483                            .and_then(|h| response.headers().get(h))
484                            .and_then(|v| v.to_str().ok())
485                    {
486                        span.record("kerniq.request_id", id);
487                    }
488                }
489                // Hyper never polls these bodies, so they end here.
490                if *this.head
491                    || status == StatusCode::NO_CONTENT
492                    || status == StatusCode::NOT_MODIFIED
493                    || status.is_informational()
494                    || response.body().is_end_stream()
495                {
496                    completion.finish(Outcome::Finished);
497                }
498                Poll::Ready(Ok(
499                    response.map(|body| ObservedBody::new(body, Some(completion)))
500                ))
501            }
502            Err(err) => {
503                completion.finish(Outcome::ServiceError);
504                Poll::Ready(Err(err))
505            }
506        }
507    }
508}