otel_bootstrap/
grpc_middleware.rs

1//! Tonic gRPC trace-context propagation, client and server side.
2//!
3//! Enabled by the `tonic-tracing` feature. Mirrors [`crate::axum_middleware`]
4//! but for raw tonic `Channel`/`Server` usage (services that don't go through
5//! an axum router — e.g. a hand-rolled tonic client/server pair).
6//!
7//! # Client side
8//!
9//! ```no_run
10//! # #[cfg(feature = "tonic-tracing")]
11//! # async fn example() -> Result<(), tonic::transport::Error> {
12//! let channel = tonic::transport::Channel::from_static("http://localhost:50051")
13//!     .connect()
14//!     .await?;
15//! let channel = tower::ServiceBuilder::new()
16//!     .layer(otel_bootstrap::grpc_client_layer())
17//!     .service(channel);
18//! # Ok(())
19//! # }
20//! ```
21//!
22//! # Server side
23//!
24//! ```no_run
25//! # #[cfg(feature = "tonic-tracing")]
26//! # fn example<S>(svc: S) {
27//! let _ = tonic::transport::Server::builder()
28//!     .layer(otel_bootstrap::grpc_server_layer());
29//! # }
30//! ```
31
32use opentelemetry::{
33    global,
34    propagation::{Extractor, Injector},
35    trace::{SpanKind, Status, TraceContextExt, Tracer},
36};
37use std::{
38    future::Future,
39    pin::Pin,
40    task::{self, Poll},
41};
42use tonic::body::Body;
43use tower::{Layer, Service};
44
45/// Tower [`Layer`] that injects the current trace context into outgoing gRPC
46/// request metadata. Wrap a tonic [`tonic::transport::Channel`] with this
47/// before constructing the generated client stub.
48///
49/// Construct via [`crate::grpc_client_layer`].
50#[derive(Clone, Debug, Default)]
51pub struct GrpcClientTraceLayer;
52
53impl<S> Layer<S> for GrpcClientTraceLayer {
54    type Service = GrpcClientTraceService<S>;
55
56    fn layer(&self, inner: S) -> Self::Service {
57        GrpcClientTraceService { inner }
58    }
59}
60
61/// Tower [`Service`] produced by [`GrpcClientTraceLayer`].
62#[derive(Clone, Debug)]
63pub struct GrpcClientTraceService<S> {
64    inner: S,
65}
66
67impl<S> Service<http::Request<Body>> for GrpcClientTraceService<S>
68where
69    S: Service<http::Request<Body>> + Clone + Send + 'static,
70    S::Future: Send + 'static,
71{
72    type Response = S::Response;
73    type Error = S::Error;
74    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
75
76    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
77        self.inner.poll_ready(cx)
78    }
79
80    fn call(&mut self, mut req: http::Request<Body>) -> Self::Future {
81        let path = req.uri().path().to_string();
82
83        let tracer = global::tracer("otel-bootstrap");
84        let span = tracer
85            .span_builder(path)
86            .with_kind(SpanKind::Client)
87            .start(&tracer);
88        let cx = opentelemetry::Context::current_with_span(span);
89
90        global::get_text_map_propagator(|propagator| {
91            propagator.inject_context(&cx, &mut MetadataInjector(req.headers_mut()));
92        });
93
94        let mut inner = self.inner.clone();
95        Box::pin(async move { inner.call(req).await })
96    }
97}
98
99/// Tower [`Layer`] that extracts trace context from incoming gRPC request
100/// metadata and opens a child span. Attach to a tonic
101/// [`tonic::transport::Server`] via `.layer(...)`.
102///
103/// Construct via [`crate::grpc_server_layer`].
104#[derive(Clone, Debug, Default)]
105pub struct GrpcServerTraceLayer;
106
107impl<S> Layer<S> for GrpcServerTraceLayer {
108    type Service = GrpcServerTraceService<S>;
109
110    fn layer(&self, inner: S) -> Self::Service {
111        GrpcServerTraceService { inner }
112    }
113}
114
115/// Tower [`Service`] produced by [`GrpcServerTraceLayer`].
116#[derive(Clone, Debug)]
117pub struct GrpcServerTraceService<S> {
118    inner: S,
119}
120
121impl<S> Service<http::Request<Body>> for GrpcServerTraceService<S>
122where
123    S: Service<http::Request<Body>, Response = http::Response<Body>> + Clone + Send + 'static,
124    S::Future: Send + 'static,
125    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
126{
127    type Response = http::Response<Body>;
128    type Error = S::Error;
129    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
130
131    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
132        self.inner.poll_ready(cx)
133    }
134
135    fn call(&mut self, req: http::Request<Body>) -> Self::Future {
136        let path = req.uri().path().to_string();
137
138        let parent_cx = global::get_text_map_propagator(|propagator| {
139            propagator.extract(&MetadataExtractor(req.headers()))
140        });
141
142        let tracer = global::tracer("otel-bootstrap");
143        let span = tracer
144            .span_builder(path)
145            .with_kind(SpanKind::Server)
146            .start_with_context(&tracer, &parent_cx);
147        let cx = parent_cx.with_span(span);
148
149        let mut inner = self.inner.clone();
150        Box::pin(async move {
151            let result = inner.call(req).await;
152
153            match &result {
154                Ok(resp) => {
155                    // gRPC status is carried in the `grpc-status` trailer, not the
156                    // HTTP status — a non-OK RPC still returns HTTP 200. Tonic
157                    // trailers aren't available at this layer (they're written
158                    // after the body stream completes), so only genuine transport
159                    // failures (HTTP-level errors) are recorded here.
160                    if resp.status().is_server_error() {
161                        cx.span().set_status(Status::Error {
162                            description: resp.status().canonical_reason().unwrap_or("").into(),
163                        });
164                    }
165                }
166                Err(_) => {
167                    cx.span().set_status(Status::Error {
168                        description: "transport error".into(),
169                    });
170                }
171            }
172
173            result
174        })
175    }
176}
177
178/// [`Extractor`] that reads from tonic/http [`http::HeaderMap`].
179struct MetadataExtractor<'a>(&'a http::HeaderMap);
180
181impl Extractor for MetadataExtractor<'_> {
182    fn get(&self, key: &str) -> Option<&str> {
183        self.0.get(key).and_then(|v| v.to_str().ok())
184    }
185
186    fn keys(&self) -> Vec<&str> {
187        self.0.keys().map(http::HeaderName::as_str).collect()
188    }
189}
190
191/// [`Injector`] that writes into a mutable [`http::HeaderMap`].
192struct MetadataInjector<'a>(&'a mut http::HeaderMap);
193
194impl Injector for MetadataInjector<'_> {
195    fn set(&mut self, key: &str, value: String) {
196        if let (Ok(name), Ok(val)) = (
197            http::HeaderName::from_bytes(key.as_bytes()),
198            http::HeaderValue::from_str(&value),
199        ) {
200            self.0.insert(name, val);
201        }
202    }
203}
otel_bootstrap/grpc_middleware.rs

otel_bootstrap/
grpc_middleware.rs
