• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 use crate::layer::WithContext;
2 use opentelemetry::{trace::SpanContext, Context, KeyValue};
3 
4 /// Utility functions to allow tracing [`Span`]s to accept and return
5 /// [OpenTelemetry] [`Context`]s.
6 ///
7 /// [`Span`]: tracing::Span
8 /// [OpenTelemetry]: https://opentelemetry.io
9 /// [`Context`]: opentelemetry::Context
10 pub trait OpenTelemetrySpanExt {
11     /// Associates `self` with a given OpenTelemetry trace, using the provided
12     /// parent [`Context`].
13     ///
14     /// [`Context`]: opentelemetry::Context
15     ///
16     /// # Examples
17     ///
18     /// ```rust
19     /// use opentelemetry::{propagation::TextMapPropagator, trace::TraceContextExt};
20     /// use opentelemetry::sdk::propagation::TraceContextPropagator;
21     /// use tracing_opentelemetry::OpenTelemetrySpanExt;
22     /// use std::collections::HashMap;
23     /// use tracing::Span;
24     ///
25     /// // Example carrier, could be a framework header map that impls otel's `Extractor`.
26     /// let mut carrier = HashMap::new();
27     ///
28     /// // Propagator can be swapped with b3 propagator, jaeger propagator, etc.
29     /// let propagator = TraceContextPropagator::new();
30     ///
31     /// // Extract otel parent context via the chosen propagator
32     /// let parent_context = propagator.extract(&carrier);
33     ///
34     /// // Generate a tracing span as usual
35     /// let app_root = tracing::span!(tracing::Level::INFO, "app_start");
36     ///
37     /// // Assign parent trace from external context
38     /// app_root.set_parent(parent_context.clone());
39     ///
40     /// // Or if the current span has been created elsewhere:
41     /// Span::current().set_parent(parent_context);
42     /// ```
set_parent(&self, cx: Context)43     fn set_parent(&self, cx: Context);
44 
45     /// Associates `self` with a given OpenTelemetry trace, using the provided
46     /// followed span [`SpanContext`].
47     ///
48     /// [`SpanContext`]: opentelemetry::trace::SpanContext
49     ///
50     /// # Examples
51     ///
52     /// ```rust
53     /// use opentelemetry::{propagation::TextMapPropagator, trace::TraceContextExt};
54     /// use opentelemetry::sdk::propagation::TraceContextPropagator;
55     /// use tracing_opentelemetry::OpenTelemetrySpanExt;
56     /// use std::collections::HashMap;
57     /// use tracing::Span;
58     ///
59     /// // Example carrier, could be a framework header map that impls otel's `Extractor`.
60     /// let mut carrier = HashMap::new();
61     ///
62     /// // Propagator can be swapped with b3 propagator, jaeger propagator, etc.
63     /// let propagator = TraceContextPropagator::new();
64     ///
65     /// // Extract otel context of linked span via the chosen propagator
66     /// let linked_span_otel_context = propagator.extract(&carrier);
67     ///
68     /// // Extract the linked span context from the otel context
69     /// let linked_span_context = linked_span_otel_context.span().span_context().clone();
70     ///
71     /// // Generate a tracing span as usual
72     /// let app_root = tracing::span!(tracing::Level::INFO, "app_start");
73     ///
74     /// // Assign linked trace from external context
75     /// app_root.add_link(linked_span_context);
76     ///
77     /// // Or if the current span has been created elsewhere:
78     /// let linked_span_context = linked_span_otel_context.span().span_context().clone();
79     /// Span::current().add_link(linked_span_context);
80     /// ```
add_link(&self, cx: SpanContext)81     fn add_link(&self, cx: SpanContext);
82 
83     /// Associates `self` with a given OpenTelemetry trace, using the provided
84     /// followed span [`SpanContext`] and attributes.
85     ///
86     /// [`SpanContext`]: opentelemetry::trace::SpanContext
add_link_with_attributes(&self, cx: SpanContext, attributes: Vec<KeyValue>)87     fn add_link_with_attributes(&self, cx: SpanContext, attributes: Vec<KeyValue>);
88 
89     /// Extracts an OpenTelemetry [`Context`] from `self`.
90     ///
91     /// [`Context`]: opentelemetry::Context
92     ///
93     /// # Examples
94     ///
95     /// ```rust
96     /// use opentelemetry::Context;
97     /// use tracing_opentelemetry::OpenTelemetrySpanExt;
98     /// use tracing::Span;
99     ///
100     /// fn make_request(cx: Context) {
101     ///     // perform external request after injecting context
102     ///     // e.g. if the request's headers impl `opentelemetry::propagation::Injector`
103     ///     // then `propagator.inject_context(cx, request.headers_mut())`
104     /// }
105     ///
106     /// // Generate a tracing span as usual
107     /// let app_root = tracing::span!(tracing::Level::INFO, "app_start");
108     ///
109     /// // To include tracing context in client requests from _this_ app,
110     /// // extract the current OpenTelemetry context.
111     /// make_request(app_root.context());
112     ///
113     /// // Or if the current span has been created elsewhere:
114     /// make_request(Span::current().context())
115     /// ```
context(&self) -> Context116     fn context(&self) -> Context;
117 }
118 
119 impl OpenTelemetrySpanExt for tracing::Span {
set_parent(&self, cx: Context)120     fn set_parent(&self, cx: Context) {
121         let mut cx = Some(cx);
122         self.with_subscriber(move |(id, subscriber)| {
123             if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
124                 get_context.with_context(subscriber, id, move |data, _tracer| {
125                     if let Some(cx) = cx.take() {
126                         data.parent_cx = cx;
127                     }
128                 });
129             }
130         });
131     }
132 
add_link(&self, cx: SpanContext)133     fn add_link(&self, cx: SpanContext) {
134         self.add_link_with_attributes(cx, Vec::new())
135     }
136 
add_link_with_attributes(&self, cx: SpanContext, attributes: Vec<KeyValue>)137     fn add_link_with_attributes(&self, cx: SpanContext, attributes: Vec<KeyValue>) {
138         if cx.is_valid() {
139             let mut cx = Some(cx);
140             let mut att = Some(attributes);
141             self.with_subscriber(move |(id, subscriber)| {
142                 if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
143                     get_context.with_context(subscriber, id, move |data, _tracer| {
144                         if let Some(cx) = cx.take() {
145                             let attr = att.take().unwrap_or_default();
146                             let follows_link = opentelemetry::trace::Link::new(cx, attr);
147                             data.builder
148                                 .links
149                                 .get_or_insert_with(|| Vec::with_capacity(1))
150                                 .push(follows_link);
151                         }
152                     });
153                 }
154             });
155         }
156     }
157 
context(&self) -> Context158     fn context(&self) -> Context {
159         let mut cx = None;
160         self.with_subscriber(|(id, subscriber)| {
161             if let Some(get_context) = subscriber.downcast_ref::<WithContext>() {
162                 get_context.with_context(subscriber, id, |builder, tracer| {
163                     cx = Some(tracer.sampled_context(builder));
164                 })
165             }
166         });
167 
168         cx.unwrap_or_default()
169     }
170 }
171