kerniq_observability_axum/
layer.rs1use 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#[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 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 #[cfg(feature = "metrics")]
81 pub fn meter(self, meter: &opentelemetry::metrics::Meter) -> Self {
82 self.metrics(HttpMetrics::new(meter))
83 }
84
85 #[cfg(feature = "metrics")]
87 pub fn metrics(mut self, metrics: HttpMetrics) -> Self {
88 self.metrics = Some(metrics);
89 self
90 }
91
92 #[cfg(feature = "tracing")]
94 pub fn tracing(mut self, enabled: bool) -> Self {
95 self.tracing = enabled;
96 self
97 }
98
99 #[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 #[cfg(feature = "tracing")]
112 pub fn shared_propagator(mut self, propagator: Propagator) -> Self {
113 self.propagator = Some(propagator);
114 self
115 }
116
117 pub fn exclude(mut self, route: impl Into<String>) -> Self {
121 self.exclude.push(route.into());
122 self
123 }
124
125 pub fn scheme(mut self, scheme: &'static str) -> Self {
129 self.scheme = scheme;
130 self
131 }
132
133 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 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#[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
214struct RouteEntry {
217 template: Arc<str>,
218 value: Value,
219 excluded: bool,
220}
221
222#[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 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 pub struct ResponseFuture<F> {
436 #[pin]
437 inner: F,
438 completion: Option<Completion>,
439 head: bool,
440 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 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 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}