snix_tracing/propagate/
tonic.rs

1#[cfg(feature = "otlp")]
2use opentelemetry::{global, propagation::Injector};
3#[cfg(feature = "otlp")]
4use opentelemetry_http::HeaderExtractor;
5#[cfg(feature = "otlp")]
6use tracing_opentelemetry::OpenTelemetrySpanExt;
7
8/// Trace context propagation: associate the current span with the otlp trace of the given request,
9/// if any and valid. This only sets the parent trace if the otlp feature is also enabled.
10pub fn accept_trace<B>(request: http::Request<B>) -> http::Request<B> {
11    // we only extract and set a parent trace if otlp feature is enabled, otherwise this feature is
12    // an noop and we return the request as is
13    #[cfg(feature = "otlp")]
14    {
15        // Current context, if no or invalid data is received.
16        let parent_context = global::get_text_map_propagator(|propagator| {
17            propagator.extract(&HeaderExtractor(request.headers()))
18        });
19        tracing::Span::current().set_parent(parent_context);
20    }
21    request
22}
23
24#[cfg(feature = "otlp")]
25struct MetadataInjector<'a>(&'a mut tonic::metadata::MetadataMap);
26
27#[cfg(feature = "otlp")]
28impl Injector for MetadataInjector<'_> {
29    fn set(&mut self, key: &str, value: String) {
30        use tonic::metadata::{MetadataKey, MetadataValue};
31        use tracing::warn;
32
33        match MetadataKey::from_bytes(key.as_bytes()) {
34            Ok(key) => match MetadataValue::try_from(&value) {
35                Ok(value) => {
36                    self.0.insert(key, value);
37                }
38                Err(error) => warn!(value, error = format!("{error:#}"), "parse metadata value"),
39            },
40            Err(error) => warn!(key, error = format!("{error:#}"), "parse metadata key"),
41        }
42    }
43}
44
45/// Trace context propagation: send the trace context by injecting it into the metadata of the given
46/// request. This only injects the current span if the otlp feature is also enabled.
47#[allow(unused_mut)]
48pub fn send_trace<T>(mut request: tonic::Request<T>) -> Result<tonic::Request<T>, tonic::Status> {
49    #[cfg(feature = "otlp")]
50    {
51        global::get_text_map_propagator(|propagator| {
52            let context = tracing::Span::current().context();
53            propagator.inject_context(&context, &mut MetadataInjector(request.metadata_mut()))
54        });
55    }
56    Ok(request)
57}