wasmcloud_tracing/
context.rs

1//! Contains helpers and code for enabling [OpenTelemetry](https://opentelemetry.io/) tracing for
2//! wasmcloud. Please note that right now this is only supported for providers. This module is only
3//! available with the `otel` feature enabled
4
5use std::collections::HashMap;
6use std::ops::Deref;
7
8use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
9use opentelemetry::trace::TraceContextExt;
10use opentelemetry_sdk::propagation::TraceContextPropagator;
11use tracing::span::Span;
12use tracing_opentelemetry::OpenTelemetrySpanExt;
13use wasmcloud_core::TraceContext;
14
15/// A convenience type that wraps an invocation [`TraceContext`] and implements the [`Extractor`] trait
16#[derive(Debug)]
17pub struct TraceContextExtractor<'a> {
18    inner: &'a TraceContext,
19}
20
21impl<'a> TraceContextExtractor<'a> {
22    /// Creates a new extractor using the given [`TraceContext`]
23    #[must_use]
24    pub fn new(context: &'a TraceContext) -> Self {
25        TraceContextExtractor { inner: context }
26    }
27}
28
29impl Extractor for TraceContextExtractor<'_> {
30    fn get(&self, key: &str) -> Option<&str> {
31        // NOTE(thomastaylor312): I don't like that we have to iterate to find this, but I didn't
32        // want to allocate hashmap for now. If this starts to cause performance issues, we can see
33        // what the tradeoff is for increasing space usage for a faster lookup
34        self.inner
35            .iter()
36            .find_map(|(k, v)| (k == key).then_some(v.as_str()))
37    }
38
39    fn keys(&self) -> Vec<&str> {
40        self.inner.iter().map(|(k, _)| k.as_str()).collect()
41    }
42}
43
44/// A convenience type that wraps an invocation [`TraceContext`] and implements the [`Injector`] trait
45#[derive(Clone, Debug, Default)]
46pub struct TraceContextInjector {
47    inner: HashMap<String, String>,
48}
49
50impl TraceContextInjector {
51    /// Creates a new injector using the given [`TraceContext`]
52    #[must_use]
53    pub fn new(headers: TraceContext) -> Self {
54        // NOTE(thomastaylor312): Same point here with performance, technically we aren't allocating anything here except the hashmap, but we could do more optimization here if needed
55        // Manually constructing the map here so we are sure we're only allocating once
56        let mut inner = HashMap::with_capacity(headers.len());
57        inner.extend(headers);
58        TraceContextInjector { inner }
59    }
60
61    /// Convenience constructor that returns a new injector with the current span context already
62    /// injected into the given header map
63    #[must_use]
64    pub fn new_with_span(headers: TraceContext) -> Self {
65        let mut header_map = Self::new(headers);
66        header_map.inject_context();
67        header_map
68    }
69
70    // Creates a new injector with the context extracted from the given extractor. If the context is empty, it will use the current span's context
71    pub fn new_with_extractor(extractor: &dyn Extractor) -> Self {
72        let mut header_map = Self::default();
73        let ctx_propagator = TraceContextPropagator::new();
74        let context = ctx_propagator.extract(extractor);
75
76        // Check if the extracted context is empty and use the current span's context if necessary
77        if !context.span().span_context().is_valid() {
78            ctx_propagator.inject_context(&Span::current().context(), &mut header_map);
79        } else {
80            ctx_propagator.inject_context(&context, &mut header_map);
81        }
82
83        header_map
84    }
85
86    /// Convenience constructor that returns a new injector with the current span context already
87    /// injected into a default [`TraceContext`]
88    #[must_use]
89    pub fn default_with_span() -> Self {
90        let mut header_map = Self::default();
91        header_map.inject_context();
92        header_map
93    }
94
95    /// Injects the context from the current span into the headers
96    pub fn inject_context(&mut self) {
97        let ctx_propagator = TraceContextPropagator::new();
98        ctx_propagator.inject_context(&Span::current().context(), self);
99    }
100
101    /// Injects the context from the given span into the headers
102    pub fn inject_context_from_span(&mut self, span: &Span) {
103        let ctx_propagator = TraceContextPropagator::new();
104        ctx_propagator.inject_context(&span.context(), self);
105    }
106}
107
108impl Injector for TraceContextInjector {
109    fn set(&mut self, key: &str, value: String) {
110        self.inner.insert(key.to_owned(), value);
111    }
112}
113
114impl AsRef<HashMap<String, String>> for TraceContextInjector {
115    fn as_ref(&self) -> &HashMap<String, String> {
116        &self.inner
117    }
118}
119
120impl Deref for TraceContextInjector {
121    type Target = HashMap<String, String>;
122
123    fn deref(&self) -> &Self::Target {
124        &self.inner
125    }
126}
127
128impl From<TraceContext> for TraceContextInjector {
129    fn from(context: TraceContext) -> Self {
130        TraceContextInjector::new(context)
131    }
132}
133
134impl From<TraceContextInjector> for TraceContext {
135    fn from(inj: TraceContextInjector) -> Self {
136        inj.inner.into_iter().collect()
137    }
138}
139
140/// A convenience function that will extract the [`opentelemetry::Context`] from the given
141/// [`TraceContext`]. If you want to do something more advanced, use the [`TraceContextExtractor`]
142pub fn get_span_context(trace_context: &TraceContext) -> opentelemetry::Context {
143    let ctx_propagator = TraceContextPropagator::new();
144    let extractor = TraceContextExtractor::new(trace_context);
145    ctx_propagator.extract(&extractor)
146}
147
148/// A convenience function that will extract from an incoming context and set the parent span for
149/// the current tracing Span. If you want to do something more advanced, use the
150/// [`TraceContextExtractor`] type directly
151///
152/// **WARNING**: To avoid performance issues, this function does not check if you have empty tracing
153/// headers. **If you pass an empty Extractor to this function, you will orphan the current span
154/// hierarchy.**
155#[allow(clippy::module_name_repetitions)]
156pub fn attach_span_context(trace_context: &TraceContext) {
157    let parent_ctx = get_span_context(trace_context);
158    Span::current().set_parent(parent_ctx);
159}
wasmcloud_tracing/context.rs

wasmcloud_tracing/
context.rs
