fastrace_poem/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use fastrace::prelude::*;
4use poem::Endpoint;
5use poem::IntoResponse;
6use poem::Middleware;
7use poem::Request;
8use poem::Response;
9use poem::Result;
10
11/// The standard [W3C Trace Context](https://www.w3.org/TR/trace-context/) header name for passing trace information.
12///
13/// This is the header key used to propagate trace context between services according to
14/// the W3C Trace Context specification.
15pub const TRACEPARENT_HEADER: &str = "traceparent";
16
17/// Middleware for integrating fastrace distributed tracing with Poem web framework.
18///
19/// This middleware extracts trace context from incoming HTTP requests and creates
20/// a new root span for each request, properly linking it to any parent context
21/// that might exist from upstream services.
22///
23/// # Example
24///
25/// ```
26/// use fastrace_poem::FastraceMiddleware;
27/// use poem::Route;
28/// use poem::get;
29/// use poem::handler;
30///
31/// let app = Route::new().at("/ping", get(ping)).with(FastraceMiddleware);
32/// ```
33#[derive(Default)]
34pub struct FastraceMiddleware;
35
36impl<E: Endpoint> Middleware<E> for FastraceMiddleware {
37    type Output = FastraceEndpoint<E>;
38
39    fn transform(&self, ep: E) -> Self::Output {
40        FastraceEndpoint { inner: ep }
41    }
42}
43
44/// An endpoint wrapper created by [`FastraceMiddleware`].
45///
46/// This type is created by the `FastraceMiddleware` and handles the extraction
47/// of trace context from requests and the creation of spans around request handlers.
48pub struct FastraceEndpoint<E> {
49    inner: E,
50}
51
52impl<E: Endpoint> Endpoint for FastraceEndpoint<E> {
53    type Output = Response;
54
55    async fn call(&self, req: Request) -> Result<Self::Output> {
56        let headers = req.headers();
57        let parent = headers
58            .get(TRACEPARENT_HEADER)
59            .and_then(|traceparent| SpanContext::decode_w3c_traceparent(traceparent.to_str().ok()?))
60            .unwrap_or(SpanContext::random());
61        let root = Span::root(req.uri().to_string(), parent);
62        self.inner
63            .call(req)
64            .in_span(root)
65            .await
66            .map(|resp| resp.into_response())
67    }
68}