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        // `poll_ready` was called on `self.inner` (the tower contract:
95        // callers must poll_ready before call on the exact same handle).
96        // Some inner services — notably tonic's `Channel`, which wraps a
97        // `tower::buffer::Buffer` — track readiness per-handle: calling on a
98        // fresh clone that was never polled panics ("send_item called
99        // without first calling poll_reserve"). Swap the clone into `self`
100        // for the *next* call, and fire this request on the already-ready
101        // original handle.
102        let clone = self.inner.clone();
103        let mut inner = std::mem::replace(&mut self.inner, clone);
104        Box::pin(async move { inner.call(req).await })
105    }
106}
107
108/// Tower [`Layer`] that extracts trace context from incoming gRPC request
109/// metadata and opens a child span. Attach to a tonic
110/// [`tonic::transport::Server`] via `.layer(...)`.
111///
112/// Construct via [`crate::grpc_server_layer`].
113#[derive(Clone, Debug, Default)]
114pub struct GrpcServerTraceLayer;
115
116impl<S> Layer<S> for GrpcServerTraceLayer {
117    type Service = GrpcServerTraceService<S>;
118
119    fn layer(&self, inner: S) -> Self::Service {
120        GrpcServerTraceService { inner }
121    }
122}
123
124/// Tower [`Service`] produced by [`GrpcServerTraceLayer`].
125#[derive(Clone, Debug)]
126pub struct GrpcServerTraceService<S> {
127    inner: S,
128}
129
130impl<S> Service<http::Request<Body>> for GrpcServerTraceService<S>
131where
132    S: Service<http::Request<Body>, Response = http::Response<Body>> + Clone + Send + 'static,
133    S::Future: Send + 'static,
134    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
135{
136    type Response = http::Response<Body>;
137    type Error = S::Error;
138    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
139
140    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
141        self.inner.poll_ready(cx)
142    }
143
144    fn call(&mut self, req: http::Request<Body>) -> Self::Future {
145        let path = req.uri().path().to_string();
146
147        let parent_cx = global::get_text_map_propagator(|propagator| {
148            propagator.extract(&MetadataExtractor(req.headers()))
149        });
150
151        let tracer = global::tracer("otel-bootstrap");
152        let span = tracer
153            .span_builder(path)
154            .with_kind(SpanKind::Server)
155            .start_with_context(&tracer, &parent_cx);
156        let cx = parent_cx.with_span(span);
157
158        // See the matching comment in `GrpcClientTraceService::call` — fire
159        // on the already-polled handle, not a fresh unpolled clone.
160        let clone = self.inner.clone();
161        let mut inner = std::mem::replace(&mut self.inner, clone);
162        Box::pin(async move {
163            let result = inner.call(req).await;
164
165            match &result {
166                Ok(resp) => {
167                    // gRPC status is carried in the `grpc-status` trailer, not the
168                    // HTTP status — a non-OK RPC still returns HTTP 200. Tonic
169                    // trailers aren't available at this layer (they're written
170                    // after the body stream completes), so only genuine transport
171                    // failures (HTTP-level errors) are recorded here.
172                    if resp.status().is_server_error() {
173                        cx.span().set_status(Status::Error {
174                            description: resp.status().canonical_reason().unwrap_or("").into(),
175                        });
176                    }
177                }
178                Err(_) => {
179                    cx.span().set_status(Status::Error {
180                        description: "transport error".into(),
181                    });
182                }
183            }
184
185            result
186        })
187    }
188}
189
190/// [`Extractor`] that reads from tonic/http [`http::HeaderMap`].
191struct MetadataExtractor<'a>(&'a http::HeaderMap);
192
193impl Extractor for MetadataExtractor<'_> {
194    fn get(&self, key: &str) -> Option<&str> {
195        self.0.get(key).and_then(|v| v.to_str().ok())
196    }
197
198    fn keys(&self) -> Vec<&str> {
199        self.0.keys().map(http::HeaderName::as_str).collect()
200    }
201}
202
203/// [`Injector`] that writes into a mutable [`http::HeaderMap`].
204struct MetadataInjector<'a>(&'a mut http::HeaderMap);
205
206impl Injector for MetadataInjector<'_> {
207    fn set(&mut self, key: &str, value: String) {
208        if let (Ok(name), Ok(val)) = (
209            http::HeaderName::from_bytes(key.as_bytes()),
210            http::HeaderValue::from_str(&value),
211        ) {
212            self.0.insert(name, val);
213        }
214    }
215}
otel_bootstrap/grpc_middleware.rs

otel_bootstrap/
grpc_middleware.rs
