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    context::FutureExt,
34    global,
35    propagation::{Extractor, Injector},
36    trace::{SpanKind, Status, TraceContextExt, Tracer},
37};
38use std::{
39    future::Future,
40    pin::Pin,
41    task::{self, Poll},
42};
43use tonic::body::Body;
44use tower::{Layer, Service};
45
46/// Tower [`Layer`] that injects the current trace context into outgoing gRPC
47/// request metadata. Wrap a tonic [`tonic::transport::Channel`] with this
48/// before constructing the generated client stub.
49///
50/// Construct via [`crate::grpc_client_layer`].
51#[derive(Clone, Debug, Default)]
52pub struct GrpcClientTraceLayer;
53
54impl<S> Layer<S> for GrpcClientTraceLayer {
55    type Service = GrpcClientTraceService<S>;
56
57    fn layer(&self, inner: S) -> Self::Service {
58        GrpcClientTraceService { inner }
59    }
60}
61
62/// Tower [`Service`] produced by [`GrpcClientTraceLayer`].
63#[derive(Clone, Debug)]
64pub struct GrpcClientTraceService<S> {
65    inner: S,
66}
67
68impl<S> Service<http::Request<Body>> for GrpcClientTraceService<S>
69where
70    S: Service<http::Request<Body>> + Clone + Send + 'static,
71    S::Future: Send + 'static,
72{
73    type Response = S::Response;
74    type Error = S::Error;
75    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
76
77    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
78        self.inner.poll_ready(cx)
79    }
80
81    fn call(&mut self, mut req: http::Request<Body>) -> Self::Future {
82        let path = req.uri().path().to_string();
83
84        let tracer = global::tracer("otel-bootstrap");
85        let span = tracer
86            .span_builder(path)
87            .with_kind(SpanKind::Client)
88            .start(&tracer);
89        let cx = opentelemetry::Context::current_with_span(span);
90
91        global::get_text_map_propagator(|propagator| {
92            propagator.inject_context(&cx, &mut MetadataInjector(req.headers_mut()));
93        });
94
95        // `poll_ready` was called on `self.inner` (the tower contract:
96        // callers must poll_ready before call on the exact same handle).
97        // Some inner services — notably tonic's `Channel`, which wraps a
98        // `tower::buffer::Buffer` — track readiness per-handle: calling on a
99        // fresh clone that was never polled panics ("send_item called
100        // without first calling poll_reserve"). Swap the clone into `self`
101        // for the *next* call, and fire this request on the already-ready
102        // original handle.
103        let clone = self.inner.clone();
104        let mut inner = std::mem::replace(&mut self.inner, clone);
105        Box::pin(async move { inner.call(req).await })
106    }
107}
108
109/// Tower [`Layer`] that extracts trace context from incoming gRPC request
110/// metadata and opens a child span. Attach to a tonic
111/// [`tonic::transport::Server`] via `.layer(...)`.
112///
113/// Construct via [`crate::grpc_server_layer`].
114#[derive(Clone, Debug, Default)]
115pub struct GrpcServerTraceLayer;
116
117impl<S> Layer<S> for GrpcServerTraceLayer {
118    type Service = GrpcServerTraceService<S>;
119
120    fn layer(&self, inner: S) -> Self::Service {
121        GrpcServerTraceService { inner }
122    }
123}
124
125/// Tower [`Service`] produced by [`GrpcServerTraceLayer`].
126#[derive(Clone, Debug)]
127pub struct GrpcServerTraceService<S> {
128    inner: S,
129}
130
131impl<S> Service<http::Request<Body>> for GrpcServerTraceService<S>
132where
133    S: Service<http::Request<Body>, Response = http::Response<Body>> + Clone + Send + 'static,
134    S::Future: Send + 'static,
135    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
136{
137    type Response = http::Response<Body>;
138    type Error = S::Error;
139    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
140
141    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
142        self.inner.poll_ready(cx)
143    }
144
145    fn call(&mut self, req: http::Request<Body>) -> Self::Future {
146        let path = req.uri().path().to_string();
147
148        let parent_cx = global::get_text_map_propagator(|propagator| {
149            propagator.extract(&MetadataExtractor(req.headers()))
150        });
151
152        let tracer = global::tracer("otel-bootstrap");
153        let span = tracer
154            .span_builder(path)
155            .with_kind(SpanKind::Server)
156            .start_with_context(&tracer, &parent_cx);
157        let cx = parent_cx.with_span(span);
158
159        // See the matching comment in `GrpcClientTraceService::call` — fire
160        // on the already-polled handle, not a fresh unpolled clone.
161        let clone = self.inner.clone();
162        let mut inner = std::mem::replace(&mut self.inner, clone);
163        Box::pin(async move {
164            // See the matching comment in axum_middleware's OtelTraceService —
165            // without attaching this request's context, `Context::current()`
166            // inside the RPC handler (and anything using
167            // api-bones's `propagation::inject_current`) sees the empty root
168            // context regardless of what was extracted above.
169            let result = inner.call(req).with_context(cx.clone()).await;
170
171            match &result {
172                Ok(resp) => {
173                    // gRPC status is carried in the `grpc-status` trailer, not the
174                    // HTTP status — a non-OK RPC still returns HTTP 200. Tonic
175                    // trailers aren't available at this layer (they're written
176                    // after the body stream completes), so only genuine transport
177                    // failures (HTTP-level errors) are recorded here.
178                    if resp.status().is_server_error() {
179                        cx.span().set_status(Status::Error {
180                            description: resp.status().canonical_reason().unwrap_or("").into(),
181                        });
182                    }
183                }
184                Err(_) => {
185                    cx.span().set_status(Status::Error {
186                        description: "transport error".into(),
187                    });
188                }
189            }
190
191            result
192        })
193    }
194}
195
196/// [`Extractor`] that reads from tonic/http [`http::HeaderMap`].
197struct MetadataExtractor<'a>(&'a http::HeaderMap);
198
199impl Extractor for MetadataExtractor<'_> {
200    fn get(&self, key: &str) -> Option<&str> {
201        self.0.get(key).and_then(|v| v.to_str().ok())
202    }
203
204    fn keys(&self) -> Vec<&str> {
205        self.0.keys().map(http::HeaderName::as_str).collect()
206    }
207}
208
209/// [`Injector`] that writes into a mutable [`http::HeaderMap`].
210struct MetadataInjector<'a>(&'a mut http::HeaderMap);
211
212impl Injector for MetadataInjector<'_> {
213    fn set(&mut self, key: &str, value: String) {
214        if let (Ok(name), Ok(val)) = (
215            http::HeaderName::from_bytes(key.as_bytes()),
216            http::HeaderValue::from_str(&value),
217        ) {
218            self.0.insert(name, val);
219        }
220    }
221}
otel_bootstrap/grpc_middleware.rs

otel_bootstrap/
grpc_middleware.rs
