opentelemetry_tonic/
lib.rs

1use std::str::FromStr;
2
3use opentelemetry::{global, Context, ContextGuard};
4use opentelemetry::propagation::{Extractor, Injector};
5
6use tonic::metadata::{MetadataKey, KeyRef, MetadataMap};
7use tonic::Request;
8
9// extend tracing::Span with context()
10use tracing_opentelemetry::OpenTelemetrySpanExt;
11
12
13pub struct MetadataInjector<'a>(&'a mut MetadataMap);
14
15impl<'a> Injector for MetadataInjector<'a> {
16    /// Set a key and value in the MetadataMap.  Does nothing if the key or value are not valid inputs
17    fn set(&mut self, key: &str, value: String) {
18        if let Ok(key) = MetadataKey::from_str(key) {
19            if let Ok(val) = value.parse() {
20                self.0.insert(key, val);
21            }
22        }
23    }
24}
25
26
27pub struct MetadataExtractor<'a>(&'a MetadataMap);
28
29impl<'a> Extractor for MetadataExtractor<'a> {
30    /// Get a value for a key from the MetadataMap.  If the value can't be converted to &str, returns None
31    fn get(&self, key: &str) -> Option<&str> {
32        self.0.get(key).and_then(|metadata| metadata.to_str().ok())
33    }
34
35    /// Collect all the keys from the MetadataMap.
36    fn keys(&self) -> Vec<&str> {
37        self.0
38            .keys()
39            .map(|key| match key {
40                KeyRef::Ascii(v) => v.as_str(),
41                KeyRef::Binary(v) => v.as_str(),
42            })
43            .collect::<Vec<_>>()
44    }
45}
46
47
48// pre-requisite:
49// global::set_text_map_propagator(TraceContextPropagator::new());
50pub fn tracing_parent_span_from_req<T>(request: &Request<T>){
51		let cx = global::get_text_map_propagator(|propagator| {
52				propagator.extract(&MetadataExtractor(request.metadata()))
53		});
54
55		tracing::Span::current().set_parent(cx);
56}
57
58// pre-requisite:
59// global::set_text_map_propagator(TraceContextPropagator::new());
60pub fn tracing_current_span_to_req<T>(request: &mut Request<T>){
61		let cx = tracing::Span::current().context();
62		global::get_text_map_propagator(|propagator| {
63				propagator.inject_context(&cx, &mut MetadataInjector(request.metadata_mut()))
64		});
65}
66
67// pre-requisite:
68// global::set_text_map_propagator(TraceContextPropagator::new());
69// context is bind to thread, not like tracing::Span
70// need to hold returned guard(e.g. let _xx = ) for this new context to take effect
71// context will restore when guard dropped
72pub fn otel_thread_cx_from_req<T>(request: &Request<T>)  -> ContextGuard {
73		let cx = global::get_text_map_propagator(|propagator| {
74				propagator.extract(&MetadataExtractor(request.metadata()))
75		});
76		cx.attach()
77}
78
79// pre-requisite:
80// global::set_text_map_propagator(TraceContextPropagator::new());
81pub fn otel_thread_cx_to_req<T>(request: &mut Request<T>){
82		let cx = Context::current();
83		global::get_text_map_propagator(|propagator| {
84				propagator.inject_context(&cx, &mut MetadataInjector(request.metadata_mut()))
85		});
86}
87
88		
89#[cfg(test)]
90mod tests {
91		use opentelemetry::{global, Context};
92		use opentelemetry::sdk::{
93				propagation::TraceContextPropagator,
94				export::trace::stdout
95		};
96		use opentelemetry::trace::{Tracer, TraceContextExt};
97		
98		use super::MetadataExtractor;
99
100		use super::MetadataInjector;
101
102    #[test]
103    fn inject() {
104				global::set_text_map_propagator(TraceContextPropagator::new());
105				let tracer = stdout::new_pipeline()
106						.install_simple();
107
108				let span = tracer.start("client-span");
109
110				let cx = Context::current_with_span(span);
111
112				let mut request = tonic::Request::new(1);
113
114				global::get_text_map_propagator(|propagator| {
115						propagator.inject_context(&cx, &mut MetadataInjector(request.metadata_mut()))
116				});
117    }
118
119		#[test]
120    fn extract() {
121				global::set_text_map_propagator(TraceContextPropagator::new());
122				let tracer = stdout::new_pipeline()
123						.install_simple();
124
125				let request = tonic::Request::new(1);
126
127				let cx = global::get_text_map_propagator(|propagator| {
128						propagator.extract(&MetadataExtractor(request.metadata()))
129				});
130
131				let span = tracer.start_with_context("server-span", &cx);
132				
133    }
134}
opentelemetry_tonic/lib.rs

opentelemetry_tonic/
lib.rs
