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

//! Tonic gRPC trace-context propagation, client and server side.
//!
//! Enabled by the `tonic-tracing` feature. Mirrors [`crate::axum_middleware`]
//! but for raw tonic `Channel`/`Server` usage (services that don't go through
//! an axum router — e.g. a hand-rolled tonic client/server pair).
//!
//! # Client side
//!
//! ```no_run
//! # #[cfg(feature = "tonic-tracing")]
//! # async fn example() -> Result<(), tonic::transport::Error> {
//! let channel = tonic::transport::Channel::from_static("http://localhost:50051")
//!     .connect()
//!     .await?;
//! let channel = tower::ServiceBuilder::new()
//!     .layer(otel_bootstrap::grpc_client_layer())
//!     .service(channel);
//! # Ok(())
//! # }
//! ```
//!
//! # Server side
//!
//! ```no_run
//! # #[cfg(feature = "tonic-tracing")]
//! # fn example<S>(svc: S) {
//! let _ = tonic::transport::Server::builder()
//!     .layer(otel_bootstrap::grpc_server_layer());
//! # }
//! ```

use opentelemetry::{
    context::FutureExt,
    global,
    propagation::{Extractor, Injector},
    trace::{SpanKind, Status, TraceContextExt, Tracer},
};
use std::{
    future::Future,
    pin::Pin,
    task::{self, Poll},
};
use tonic::body::Body;
use tower::{Layer, Service};

/// Tower [`Layer`] that injects the current trace context into outgoing gRPC
/// request metadata. Wrap a tonic [`tonic::transport::Channel`] with this
/// before constructing the generated client stub.
///
/// Construct via [`crate::grpc_client_layer`].
#[derive(Clone, Debug, Default)]
pub struct GrpcClientTraceLayer;

impl<S> Layer<S> for GrpcClientTraceLayer {
    type Service = GrpcClientTraceService<S>;

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

/// Tower [`Service`] produced by [`GrpcClientTraceLayer`].
#[derive(Clone, Debug)]
pub struct GrpcClientTraceService<S> {
    inner: S,
}

impl<S> Service<http::Request<Body>> for GrpcClientTraceService<S>
where
    S: Service<http::Request<Body>> + Clone + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    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, mut req: http::Request<Body>) -> Self::Future {
        let path = req.uri().path().to_string();

        let tracer = global::tracer("otel-bootstrap");
        let span = tracer
            .span_builder(path)
            .with_kind(SpanKind::Client)
            .start(&tracer);
        let cx = opentelemetry::Context::current_with_span(span);

        global::get_text_map_propagator(|propagator| {
            propagator.inject_context(&cx, &mut MetadataInjector(req.headers_mut()));
        });

        // `poll_ready` was called on `self.inner` (the tower contract:
        // callers must poll_ready before call on the exact same handle).
        // Some inner services — notably tonic's `Channel`, which wraps a
        // `tower::buffer::Buffer` — track readiness per-handle: calling on a
        // fresh clone that was never polled panics ("send_item called
        // without first calling poll_reserve"). Swap the clone into `self`
        // for the *next* call, 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 { inner.call(req).await })
    }
}

/// Tower [`Layer`] that extracts trace context from incoming gRPC request
/// metadata and opens a child span. Attach to a tonic
/// [`tonic::transport::Server`] via `.layer(...)`.
///
/// Construct via [`crate::grpc_server_layer`].
#[derive(Clone, Debug, Default)]
pub struct GrpcServerTraceLayer;

impl<S> Layer<S> for GrpcServerTraceLayer {
    type Service = GrpcServerTraceService<S>;

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

/// Tower [`Service`] produced by [`GrpcServerTraceLayer`].
#[derive(Clone, Debug)]
pub struct GrpcServerTraceService<S> {
    inner: S,
}

impl<S> Service<http::Request<Body>> for GrpcServerTraceService<S>
where
    S: Service<http::Request<Body>, Response = http::Response<Body>> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    type Response = http::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: http::Request<Body>) -> Self::Future {
        let path = req.uri().path().to_string();

        let parent_cx = global::get_text_map_propagator(|propagator| {
            propagator.extract(&MetadataExtractor(req.headers()))
        });

        let tracer = global::tracer("otel-bootstrap");
        let span = tracer
            .span_builder(path)
            .with_kind(SpanKind::Server)
            .start_with_context(&tracer, &parent_cx);
        let cx = parent_cx.with_span(span);

        // See the matching comment in `GrpcClientTraceService::call` — fire
        // on the already-polled handle, not a fresh unpolled clone.
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);
        Box::pin(async move {
            // See the matching comment in axum_middleware's OtelTraceService —
            // without attaching this request's context, `Context::current()`
            // inside the RPC handler (and anything using
            // api-bones's `propagation::inject_current`) sees the empty root
            // context regardless of what was extracted above.
            let result = inner.call(req).with_context(cx.clone()).await;

            match &result {
                Ok(resp) => {
                    // gRPC status is carried in the `grpc-status` trailer, not the
                    // HTTP status — a non-OK RPC still returns HTTP 200. Tonic
                    // trailers aren't available at this layer (they're written
                    // after the body stream completes), so only genuine transport
                    // failures (HTTP-level errors) are recorded here.
                    if resp.status().is_server_error() {
                        cx.span().set_status(Status::Error {
                            description: resp.status().canonical_reason().unwrap_or("").into(),
                        });
                    }
                }
                Err(_) => {
                    cx.span().set_status(Status::Error {
                        description: "transport error".into(),
                    });
                }
            }

            result
        })
    }
}

/// [`Extractor`] that reads from tonic/http [`http::HeaderMap`].
struct MetadataExtractor<'a>(&'a http::HeaderMap);

impl Extractor for MetadataExtractor<'_> {
    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(http::HeaderName::as_str).collect()
    }
}

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

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