fastrace_axum/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::future::Future;
4use std::pin::Pin;
5use std::task::Context;
6use std::task::Poll;
7
8use axum::extract::MatchedPath;
9use axum::response::Response;
10use fastrace::future::InSpan;
11use fastrace::local::LocalSpan;
12use fastrace::prelude::*;
13use opentelemetry_semantic_conventions::trace::HTTP_REQUEST_METHOD;
14use opentelemetry_semantic_conventions::trace::HTTP_RESPONSE_STATUS_CODE;
15use opentelemetry_semantic_conventions::trace::HTTP_ROUTE;
16use opentelemetry_semantic_conventions::trace::URL_PATH;
17use tower_layer::Layer;
18use tower_service::Service;
19
20/// The standard [W3C Trace Context](https://www.w3.org/TR/trace-context/) header name for passing trace information.
21///
22/// This is the header key used to propagate trace context between services according to
23/// the W3C Trace Context specification.
24pub const TRACEPARENT_HEADER: &str = "traceparent";
25
26/// Layer for intercepting and processing trace context in incoming requests.
27///
28/// This layer extracts tracing context from incoming requests and creates a new span
29/// for each request. Add this to your axum server to automatically handle trace context
30/// propagation.
31#[derive(Clone)]
32pub struct FastraceLayer;
33
34impl<S> Layer<S> for FastraceLayer {
35    type Service = FastraceService<S>;
36
37    fn layer(&self, service: S) -> Self::Service {
38        FastraceService { service }
39    }
40}
41
42/// A service that handles trace context propagation.
43///
44/// This service extracts trace context from incoming requests and creates
45/// spans to track the request processing. It wraps the inner service and augments
46/// it with tracing capabilities.
47#[derive(Clone)]
48pub struct FastraceService<S> {
49    service: S,
50}
51
52use axum::extract::Request;
53
54impl<S> Service<Request> for FastraceService<S>
55where
56    S: Service<Request, Response = Response> + Send + 'static,
57    S::Future: Send + 'static,
58{
59    type Response = S::Response;
60    type Error = S::Error;
61    type Future = InSpan<InspectHttpResponse<S::Future>>;
62
63    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
64        self.service.poll_ready(cx)
65    }
66
67    fn call(&mut self, req: Request) -> Self::Future {
68        let headers = req.headers();
69        let parent = headers.get(TRACEPARENT_HEADER).and_then(|traceparent| {
70            SpanContext::decode_w3c_traceparent(traceparent.to_str().ok()?)
71        });
72
73        let span = if let Some(parent) = parent {
74            // https://opentelemetry.io/docs/specs/semconv/http/http-spans/#name
75            let name = if let Some(target) = req.extensions().get::<MatchedPath>() {
76                format!("{} {}", req.method(), target.as_str())
77            } else {
78                req.method().to_string()
79            };
80
81            let root = Span::root(name, parent);
82
83            root.add_properties(|| {
84                [
85                    (HTTP_REQUEST_METHOD, req.method().to_string()),
86                    (URL_PATH, req.uri().path().to_string()),
87                ]
88            });
89
90            if let Some(route) = req.extensions().get::<MatchedPath>() {
91                root.add_property(|| (HTTP_ROUTE, route.as_str().to_string()));
92            }
93
94            root
95        } else {
96            Span::noop()
97        };
98
99        let fut = self.service.call(req);
100        let fut = InspectHttpResponse { inner: fut };
101        fut.in_span(span)
102    }
103}
104
105#[pin_project::pin_project]
106pub struct InspectHttpResponse<F> {
107    #[pin]
108    inner: F,
109}
110
111impl<F, E> Future for InspectHttpResponse<F>
112where F: Future<Output = Result<Response, E>>
113{
114    type Output = F::Output;
115
116    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
117        let this = self.project();
118        let poll = this.inner.poll(cx);
119
120        if let Poll::Ready(Ok(response)) = &poll {
121            LocalSpan::add_property(|| {
122                (
123                    HTTP_RESPONSE_STATUS_CODE,
124                    response.status().as_u16().to_string(),
125                )
126            });
127        }
128
129        poll
130    }
131}
fastrace_axum/lib.rs

fastrace_axum/
lib.rs
