linera_rpc/
propagation.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! OpenTelemetry context propagation for gRPC.
5//!
6//! This module provides utilities for propagating OpenTelemetry context (trace context and baggage)
7//! across gRPC service boundaries using tonic metadata.
8//!
9//! # Usage
10//!
11//! ## Client-side injection
12//!
13//! ```ignore
14//! use linera_rpc::propagation::inject_context;
15//! use opentelemetry::Context;
16//!
17//! let mut request = tonic::Request::new(payload);
18//! inject_context(&Context::current(), request.metadata_mut());
19//! ```
20//!
21//! ## Server-side extraction
22//!
23//! ```ignore
24//! use linera_rpc::propagation::extract_context;
25//!
26//! let cx = extract_context(request.metadata());
27//! // Use cx.with_baggage() to access baggage values
28//! ```
29
30use std::task::{Context as TaskContext, Poll};
31
32use futures::{future::BoxFuture, FutureExt};
33use opentelemetry::{
34    global,
35    propagation::{Extractor, Injector},
36    Context,
37};
38use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
39use tower::{Layer, Service};
40use tracing::warn;
41
42/// Baggage key for traffic type labeling.
43///
44/// Used to distinguish organic traffic from synthetic (benchmark) traffic.
45/// Valid values are "organic" and "synthetic".
46pub const TRAFFIC_TYPE_KEY: &str = "traffic_type";
47
48/// Traffic type for normal production traffic.
49pub const TRAFFIC_TYPE_ORGANIC: &str = "organic";
50
51/// Traffic type for synthetic benchmark traffic.
52pub const TRAFFIC_TYPE_SYNTHETIC: &str = "synthetic";
53
54/// Traffic type when OpenTelemetry feature is disabled.
55pub const TRAFFIC_TYPE_UNKNOWN: &str = "unknown";
56
57/// Environment variable to override the traffic type.
58///
59/// Set this to "synthetic" to mark all outgoing requests as benchmark traffic.
60/// This is useful for benchmark tools that cannot easily set OpenTelemetry baggage.
61pub const TRAFFIC_TYPE_ENV_VAR: &str = "LINERA_TRAFFIC_TYPE";
62
63/// Tower layer that extracts OpenTelemetry context from incoming gRPC requests.
64///
65/// This layer extracts W3C TraceContext and Baggage headers from the request
66/// metadata and stores the extracted context in the request extensions.
67///
68/// # Usage
69///
70/// ```ignore
71/// use linera_rpc::propagation::OtelContextLayer;
72/// use tower::ServiceBuilder;
73///
74/// let service = ServiceBuilder::new()
75///     .layer(OtelContextLayer)
76///     .service(my_service);
77/// ```
78#[derive(Clone, Copy, Debug, Default)]
79pub struct OtelContextLayer;
80
81/// Service wrapper that extracts OpenTelemetry context from requests.
82#[derive(Clone, Debug)]
83pub struct OtelContextService<S> {
84    inner: S,
85}
86
87/// Extension type to store the extracted OpenTelemetry context.
88#[derive(Clone, Debug)]
89pub struct ExtractedOtelContext(pub Context);
90
91/// Trait for request types that can provide access to extracted OpenTelemetry context.
92///
93/// This trait abstracts over `http::Request` and `tonic::Request` to allow
94/// generic functions that work with either request type.
95pub trait HasOtelContext {
96    /// Returns the extracted OpenTelemetry context, if present.
97    fn get_otel_context(&self) -> Option<&ExtractedOtelContext>;
98}
99
100impl<B> HasOtelContext for http::Request<B> {
101    fn get_otel_context(&self) -> Option<&ExtractedOtelContext> {
102        self.extensions().get::<ExtractedOtelContext>()
103    }
104}
105
106impl<T> HasOtelContext for tonic::Request<T> {
107    fn get_otel_context(&self) -> Option<&ExtractedOtelContext> {
108        self.extensions().get::<ExtractedOtelContext>()
109    }
110}
111
112/// Injects the OpenTelemetry context into tonic metadata.
113///
114/// This injects both W3C TraceContext (`traceparent`, `tracestate`) and
115/// W3C Baggage (`baggage`) headers into the metadata, enabling distributed
116/// tracing and baggage propagation across gRPC service boundaries.
117///
118/// # Arguments
119///
120/// * `cx` - The OpenTelemetry context to inject
121/// * `metadata` - The tonic metadata map to inject into
122pub fn inject_context(cx: &Context, metadata: &mut MetadataMap) {
123    global::get_text_map_propagator(|propagator| {
124        propagator.inject_context(cx, &mut MetadataInjector(metadata));
125    });
126}
127
128/// Extracts the OpenTelemetry context from tonic metadata.
129///
130/// This extracts both W3C TraceContext and W3C Baggage headers from the
131/// metadata, returning a context that can be used as a parent for new spans
132/// or to read baggage values.
133///
134/// # Arguments
135///
136/// * `metadata` - The tonic metadata map to extract from
137///
138/// # Returns
139///
140/// The extracted OpenTelemetry context, or an empty context if no propagation
141/// headers were found.
142pub fn extract_context(metadata: &MetadataMap) -> Context {
143    global::get_text_map_propagator(|propagator| propagator.extract(&MetadataExtractor(metadata)))
144}
145
146/// Returns the current OpenTelemetry context, enriched with traffic type baggage
147/// if the `LINERA_TRAFFIC_TYPE` environment variable is set.
148///
149/// This function provides a workaround for async code that cannot hold a `ContextGuard`
150/// across `.await` points (since `ContextGuard` is `!Send`).
151///
152/// If `LINERA_TRAFFIC_TYPE=synthetic` is set, the returned context will have
153/// synthetic traffic baggage attached.
154pub fn get_context_with_traffic_type() -> Context {
155    use opentelemetry::{baggage::BaggageExt, Key, KeyValue};
156
157    let cx = Context::current();
158
159    if std::env::var(TRAFFIC_TYPE_ENV_VAR).is_ok_and(|v| v == TRAFFIC_TYPE_SYNTHETIC) {
160        cx.with_baggage(vec![KeyValue::new(
161            Key::new(TRAFFIC_TYPE_KEY),
162            TRAFFIC_TYPE_SYNTHETIC,
163        )])
164    } else {
165        cx
166    }
167}
168
169/// Extracts the traffic type from the current context.
170///
171/// Returns "organic" if no traffic type baggage is set.
172pub fn get_traffic_type(cx: &Context) -> &'static str {
173    use opentelemetry::baggage::BaggageExt;
174
175    cx.baggage()
176        .get(TRAFFIC_TYPE_KEY)
177        .map(|v| v.as_str())
178        .and_then(|v| {
179            if v == TRAFFIC_TYPE_SYNTHETIC {
180                Some(TRAFFIC_TYPE_SYNTHETIC)
181            } else {
182                None
183            }
184        })
185        .unwrap_or(TRAFFIC_TYPE_ORGANIC)
186}
187
188/// Gets the traffic type from a request's extensions.
189///
190/// Works with both `http::Request` and `tonic::Request` via the `HasOtelContext` trait.
191/// Returns "organic" if no context was extracted or no traffic type baggage was set.
192pub fn get_traffic_type_from_request<R: HasOtelContext>(request: &R) -> &'static str {
193    request
194        .get_otel_context()
195        .map_or(TRAFFIC_TYPE_ORGANIC, |ext| get_traffic_type(&ext.0))
196}
197
198/// Gets the OpenTelemetry context from a tonic::Request's extensions.
199///
200/// Returns `None` if no context was extracted by OtelContextLayer.
201/// Tower middleware extensions are preserved in tonic::Request extensions.
202pub fn get_otel_context_from_tonic_request<T>(request: &tonic::Request<T>) -> Option<Context> {
203    request
204        .extensions()
205        .get::<ExtractedOtelContext>()
206        .map(|ext| ext.0.clone())
207}
208
209/// Creates a new tonic::Request with the OpenTelemetry context injected into metadata.
210///
211/// This is used to propagate context when forwarding requests to downstream services.
212/// If `cx` is `None`, returns a request without injected context.
213pub fn create_request_with_context<T>(inner: T, cx: Option<&Context>) -> tonic::Request<T> {
214    let mut request = tonic::Request::new(inner);
215    if let Some(cx) = cx {
216        inject_context(cx, request.metadata_mut());
217    }
218    request
219}
220
221/// Creates a new tonic::Request with the current tracing span's context injected.
222///
223/// This gets the OpenTelemetry context from the current tracing span and injects it
224/// into the request metadata. Use this when forwarding requests to downstream services
225/// to ensure proper distributed tracing (child spans linked to parent spans).
226pub fn create_request_with_current_span_context<T>(inner: T) -> tonic::Request<T> {
227    use tracing_opentelemetry::OpenTelemetrySpanExt;
228    let cx = tracing::Span::current().context();
229    create_request_with_context(inner, Some(&cx))
230}
231
232impl<S> Layer<S> for OtelContextLayer {
233    type Service = OtelContextService<S>;
234
235    fn layer(&self, service: S) -> Self::Service {
236        OtelContextService { inner: service }
237    }
238}
239
240impl<S, B> Service<http::Request<B>> for OtelContextService<S>
241where
242    S: Service<http::Request<B>> + Clone + Send + 'static,
243    S::Future: Send,
244    B: Send + 'static,
245{
246    type Response = S::Response;
247    type Error = S::Error;
248    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
249
250    fn poll_ready(&mut self, cx: &mut TaskContext<'_>) -> Poll<Result<(), Self::Error>> {
251        self.inner.poll_ready(cx)
252    }
253
254    fn call(&mut self, mut request: http::Request<B>) -> Self::Future {
255        use tracing::Instrument;
256        use tracing_opentelemetry::OpenTelemetrySpanExt;
257
258        let cx = global::get_text_map_propagator(|propagator| {
259            propagator.extract(&HttpHeaderExtractor(request.headers()))
260        });
261
262        request
263            .extensions_mut()
264            .insert(ExtractedOtelContext(cx.clone()));
265
266        let span = tracing::info_span!("grpc_request");
267        span.set_parent(cx);
268
269        let mut inner = self.inner.clone();
270        async move { inner.call(request).await }
271            .instrument(span)
272            .boxed()
273    }
274}
275
276struct MetadataInjector<'a>(&'a mut MetadataMap);
277
278impl Injector for MetadataInjector<'_> {
279    fn set(&mut self, key: &str, value: String) {
280        match MetadataKey::from_bytes(key.as_bytes()) {
281            Ok(key) => match MetadataValue::try_from(&value) {
282                Ok(value) => {
283                    self.0.insert(key, value);
284                }
285                Err(error) => {
286                    warn!(
287                        value,
288                        error = format!("{error:#}"),
289                        "failed to parse metadata value"
290                    );
291                }
292            },
293            Err(error) => {
294                warn!(
295                    key,
296                    error = format!("{error:#}"),
297                    "failed to parse metadata key"
298                );
299            }
300        }
301    }
302}
303
304struct MetadataExtractor<'a>(&'a MetadataMap);
305
306impl Extractor for MetadataExtractor<'_> {
307    fn get(&self, key: &str) -> Option<&str> {
308        self.0.get(key).and_then(|value| value.to_str().ok())
309    }
310
311    fn keys(&self) -> Vec<&str> {
312        self.0
313            .keys()
314            .filter_map(|key| match key {
315                tonic::metadata::KeyRef::Ascii(key) => Some(key.as_str()),
316                tonic::metadata::KeyRef::Binary(_) => None,
317            })
318            .collect()
319    }
320}
321
322struct HttpHeaderExtractor<'a>(&'a http::HeaderMap);
323
324impl Extractor for HttpHeaderExtractor<'_> {
325    fn get(&self, key: &str) -> Option<&str> {
326        self.0.get(key).and_then(|v| v.to_str().ok())
327    }
328
329    fn keys(&self) -> Vec<&str> {
330        self.0.keys().map(|k| k.as_str()).collect()
331    }
332}
333
334#[cfg(test)]
335mod tests {
336    use opentelemetry::{baggage::BaggageExt, Key, KeyValue};
337
338    use super::*;
339
340    #[test]
341    fn test_inject_and_extract_baggage() {
342        use opentelemetry::propagation::TextMapCompositePropagator;
343        use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
344
345        let propagator = TextMapCompositePropagator::new(vec![
346            Box::new(TraceContextPropagator::new()),
347            Box::new(BaggagePropagator::new()),
348        ]);
349        global::set_text_map_propagator(propagator);
350
351        let cx = Context::current().with_baggage(vec![KeyValue::new(
352            Key::new(TRAFFIC_TYPE_KEY),
353            TRAFFIC_TYPE_SYNTHETIC,
354        )]);
355
356        let mut metadata = MetadataMap::new();
357        inject_context(&cx, &mut metadata);
358
359        assert!(
360            metadata.get("baggage").is_some(),
361            "baggage header should be present"
362        );
363
364        let extracted_cx = extract_context(&metadata);
365        let traffic_type = get_traffic_type(&extracted_cx);
366        assert_eq!(traffic_type, TRAFFIC_TYPE_SYNTHETIC);
367    }
368
369    #[test]
370    fn test_default_traffic_type_is_organic() {
371        let cx = Context::current();
372        assert_eq!(get_traffic_type(&cx), TRAFFIC_TYPE_ORGANIC);
373    }
374}
linera_rpc/propagation.rs

linera_rpc/
propagation.rs
