otel-bootstrap 2.11.0

One-call OpenTelemetry bootstrap — traces + metrics with OTLP export
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247

//! Axum tower middleware for W3C trace context propagation.
//!
//! Enabled by the `axum` feature flag. Add to any axum [`Router`] via
//! [`crate::axum_layer()`]:
//!
//! ```no_run
//! use axum::Router;
//!
//! let app: Router = Router::new().layer(otel_bootstrap::axum_layer());
//! ```

use axum::{
    body::Body,
    http::{HeaderMap, HeaderName, HeaderValue, Request, Response},
};
use opentelemetry::{
    context::FutureExt,
    global,
    propagation::{Extractor, Injector},
    trace::{SpanKind, Status, TraceContextExt, Tracer},
};
use opentelemetry_semantic_conventions::attribute::{
    HTTP_REQUEST_METHOD, HTTP_RESPONSE_STATUS_CODE,
};
use std::{
    future::Future,
    pin::Pin,
    task::{self, Poll},
};
use tower::{Layer, Service};

/// Tower [`Layer`] that instruments incoming HTTP requests with OpenTelemetry
/// trace context propagation.
///
/// Attach to an axum router with [`crate::axum_layer()`].
///
/// # Example
/// ```no_run
/// use axum::Router;
///
/// let app: Router = Router::new().layer(otel_bootstrap::axum_layer());
/// ```
#[derive(Clone, Debug)]
pub struct OtelTraceLayer;

impl<S> Layer<S> for OtelTraceLayer {
    type Service = OtelTraceService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        OtelTraceService { inner }
    }
}

/// Tower [`Service`] produced by [`OtelTraceLayer`].
///
/// This type is not constructed directly. It is returned by
/// [`OtelTraceLayer`] when wrapping an inner [`tower::Service`].
#[derive(Clone, Debug)]
pub struct OtelTraceService<S> {
    inner: S,
}

impl<S> Service<Request<Body>> for OtelTraceService<S>
where
    S: Service<Request<Body>, Response = Response<Body>> + Send + Clone + 'static,
    S::Future: Send + 'static,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    type Response = Response<Body>;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        let method = req.method().to_string();
        let route = req.uri().path().to_string();

        // Extract parent context from incoming headers.
        let parent_cx = global::get_text_map_propagator(|propagator| {
            propagator.extract(&HeaderExtractor(req.headers()))
        });

        // Create a child span inheriting the remote parent.
        let tracer = global::tracer("otel-bootstrap");
        let span = tracer
            .span_builder(format!("{method} {route}"))
            .with_kind(SpanKind::Server)
            .with_attributes([opentelemetry::KeyValue::new(HTTP_REQUEST_METHOD, method)])
            .start_with_context(&tracer, &parent_cx);

        let cx = parent_cx.with_span(span);

        // `poll_ready` was called on `self.inner` before `call`, per the
        // tower contract — some inner services track readiness per-handle
        // (e.g. a `tower::buffer::Buffer`), so firing on a fresh unpolled
        // clone can panic. Swap the clone into `self` for next time, and
        // fire this request on the already-ready original handle.
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);

        Box::pin(async move {
            // Without this, `Context::current()` inside the handler (and
            // anything it calls, e.g. api-bones's `propagation::inject_current`)
            // sees the empty root context, not this request's parent — every
            // outbound call the handler makes injects a disconnected
            // traceparent regardless of what was extracted above. `with_context`
            // (not a bare `cx.attach()` guard) is required because the inner
            // future can resume on a different tokio worker thread between
            // polls, and thread-local attach doesn't survive that hop.
            let mut response = inner.call(req).with_context(cx.clone()).await?;

            // Record HTTP status on the span.
            let status_code = response.status().as_u16();
            cx.span().set_attribute(opentelemetry::KeyValue::new(
                HTTP_RESPONSE_STATUS_CODE,
                status_code as i64,
            ));
            if response.status().is_server_error() {
                cx.span().set_status(Status::Error {
                    description: response.status().canonical_reason().unwrap_or("").into(),
                });
            }

            // Inject outgoing trace context into response headers.
            let mut injector = HeaderInjector(response.headers_mut());
            global::get_text_map_propagator(|propagator| {
                propagator.inject_context(&cx, &mut injector);
            });

            Ok(response)
        })
    }
}

/// Tower [`Layer`] that records span attributes from a `T: EnrichSpan`
/// extension on each incoming request.
///
/// Construct via [`crate::span_enricher_layer`]. When no `T` extension is
/// found in the request (e.g. a platform-scope route), the service is a no-op.
///
/// # Example
/// ```no_run
/// use axum::{Router, Extension, routing::get};
/// use otel_bootstrap::span_enrichment::EnrichSpan;
/// use tracing_opentelemetry::OpenTelemetrySpanExt as _;
///
/// #[derive(Clone)]
/// struct MyCtx { user_id: String }
///
/// impl EnrichSpan for MyCtx {
///     fn enrich_span(&self, span: &tracing::Span) {
///         span.set_attribute("enduser.id", self.user_id.clone());
///     }
/// }
///
/// let app: Router = Router::new()
///     .route("/", get(|| async { "ok" }))
///     .layer(otel_bootstrap::span_enricher_layer::<MyCtx>())
///     .layer(Extension(MyCtx { user_id: "u1".into() }))
///     .layer(otel_bootstrap::axum_layer());
/// ```
#[derive(Debug)]
pub struct SpanEnricherLayer<T>(std::marker::PhantomData<T>);

impl<T> Default for SpanEnricherLayer<T> {
    fn default() -> Self {
        Self(std::marker::PhantomData)
    }
}

impl<T> Clone for SpanEnricherLayer<T> {
    fn clone(&self) -> Self {
        Self(std::marker::PhantomData)
    }
}

impl<T, S> Layer<S> for SpanEnricherLayer<T>
where
    T: crate::span_enrichment::EnrichSpan + Clone + Send + Sync + 'static,
{
    type Service = SpanEnricherService<T, S>;

    fn layer(&self, inner: S) -> Self::Service {
        SpanEnricherService {
            inner,
            _marker: std::marker::PhantomData,
        }
    }
}

/// Tower [`Service`] produced by [`SpanEnricherLayer`].
#[derive(Clone, Debug)]
pub struct SpanEnricherService<T, S> {
    inner: S,
    _marker: std::marker::PhantomData<T>,
}

impl<T, S> Service<Request<Body>> for SpanEnricherService<T, S>
where
    T: crate::span_enrichment::EnrichSpan + Clone + Send + Sync + 'static,
    S: Service<Request<Body>, Response = Response<Body>>,
{
    type Response = Response<Body>;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        if let Some(ctx) = req.extensions().get::<T>() {
            ctx.enrich_span(&tracing::Span::current());
        }
        self.inner.call(req)
    }
}

/// [`Extractor`] that reads from [`HeaderMap`].
struct HeaderExtractor<'a>(&'a HeaderMap);

impl Extractor for HeaderExtractor<'_> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(|v| v.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(HeaderName::as_str).collect()
    }
}

/// [`Injector`] that writes into a mutable [`HeaderMap`].
struct HeaderInjector<'a>(&'a mut HeaderMap);

impl Injector for HeaderInjector<'_> {
    fn set(&mut self, key: &str, value: String) {
        if let (Ok(name), Ok(val)) = (
            HeaderName::from_bytes(key.as_bytes()),
            HeaderValue::from_str(&value),
        ) {
            self.0.insert(name, val);
        }
    }
}