krabka_trace_context/
propagation.rs

1//! W3C Trace Context propagation over Kafka record headers.
2//!
3//! The telemetry initialisation of the service, `krabka_telemetry::init`,
4//! installs a global `TraceContextPropagator`. A producer can then serialise
5//! the current span's `traceparent` and `tracestate` into record headers with
6//! [`current_trace_headers`]. A consumer can rebuild that context with
7//! [`extract_context`] and [`set_remote_parent`], and make its own span a child
8//! of the producer's span. One distributed trace then covers all the services
9//! that the Kafka WAL and topics connect.
10//!
11//! The helpers work with `(key, value)` string and byte pairs, so this crate
12//! stays independent of the concrete `krabka_client_{producer,consumer}` header
13//! types. Callers convert to and from their own `Header` at the edge.
14
15use std::{collections::HashMap, str::from_utf8};
16
17use opentelemetry::{
18    Context, global,
19    propagation::{Extractor, Injector},
20};
21use tracing_opentelemetry::OpenTelemetrySpanExt as _;
22
23/// The canonical W3C trace-context header keys.
24///
25/// A producer should attach these keys, in lower case, to Kafka record headers.
26pub const TRACEPARENT: &str = "traceparent";
27pub const TRACESTATE: &str = "tracestate";
28
29/// Collects injected key and value pairs into a `Vec`.
30///
31/// The caller converts each pair into its own header type, such as
32/// `krabka_client_producer::Header`.
33struct VecInjector<'a>(&'a mut Vec<(String, String)>);
34
35impl Injector for VecInjector<'_> {
36    fn set(&mut self, key: &str, value: String) {
37        self.0.push((key.to_owned(), value));
38    }
39}
40
41/// Reads W3C headers back out of a borrowed key→value map.
42struct MapExtractor<'a>(&'a HashMap<String, String>);
43
44impl Extractor for MapExtractor<'_> {
45    fn get(&self, key: &str) -> Option<&str> {
46        self.0.get(key).map(String::as_str)
47    }
48
49    fn keys(&self) -> Vec<&str> {
50        self.0.keys().map(String::as_str).collect()
51    }
52}
53
54/// The W3C trace-context headers for the **current** `tracing` span.
55///
56/// Returns the `traceparent`, and the `tracestate` when the span has one, as
57/// `(key, value)` pairs for Kafka record headers. Returns an empty `Vec` when
58/// there is no active span, when the span is not sampled or recorded, or when
59/// OTLP is disabled. It is always safe to call.
60///
61/// ```no_run
62/// let span = tracing::info_span!("produce_order");
63/// let _g = span.enter();
64/// let carriers = krabka_trace_context::current_trace_headers();
65/// // convert each (key, value) into your producer's Header and attach it
66/// ```
67#[must_use]
68pub fn current_trace_headers() -> Vec<(String, String)> {
69    let cx = tracing::Span::current().context();
70    let mut out = Vec::new();
71    global::get_text_map_propagator(|prop| {
72        prop.inject_context(&cx, &mut VecInjector(&mut out));
73    });
74    out
75}
76
77/// Rebuild an OpenTelemetry [`Context`] from Kafka record headers.
78///
79/// Each header is a key plus a raw byte value. This function skips a non-UTF-8
80/// value. Use the result as the remote parent of a consumer-side span: see
81/// [`set_remote_parent`].
82#[must_use]
83pub fn extract_context<'a, I, V>(headers: I) -> Context
84where
85    I: IntoIterator<Item = (&'a str, V)>,
86    V: AsRef<[u8]>,
87{
88    let map: HashMap<String, String> = headers
89        .into_iter()
90        .filter_map(|(k, v)| {
91            from_utf8(v.as_ref())
92                .ok()
93                .map(|s| (k.to_owned(), s.to_owned()))
94        })
95        .collect();
96    global::get_text_map_propagator(|prop| prop.extract(&MapExtractor(&map)))
97}
98
99/// Make `span` a child of the trace that `headers` carries.
100///
101/// The headers hold the producer's context, so consumer-side work then appears
102/// in the same distributed trace. This function does nothing when the headers
103/// carry no valid trace context.
104pub fn set_remote_parent<'a, I, V>(span: &tracing::Span, headers: I)
105where
106    I: IntoIterator<Item = (&'a str, V)>,
107    V: AsRef<[u8]>,
108{
109    // `set_parent` returns the previous `Context`; we don't need it.
110    let _ = span.set_parent(extract_context(headers));
111}
112
113#[cfg(test)]
114mod tests {
115
116    use std::iter;
117
118    use opentelemetry::trace::{TraceContextExt as _, TraceId};
119    use opentelemetry_sdk::propagation::TraceContextPropagator;
120    use tracing::subscriber;
121
122    use super::*;
123
124    #[test]
125    fn extract_context_roundtrips_w3c_traceparent() {
126        // The global propagator is a no-op until one is installed; install the
127        // W3C propagator (idempotent for our purposes) so extraction works even
128        // when `init()` has not run in this test process.
129        global::set_text_map_propagator(TraceContextPropagator::new());
130
131        let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
132        let cx = extract_context([(TRACEPARENT, traceparent.as_bytes())]);
133        let sc = cx.span().span_context().clone();
134
135        assert2::assert!(sc.is_valid());
136        assert2::assert!(
137            sc.trace_id() == TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap()
138        );
139        assert2::assert!(sc.is_sampled());
140    }
141
142    #[test]
143    fn extract_context_skips_non_utf8_and_missing() {
144        global::set_text_map_propagator(TraceContextPropagator::new());
145        // Invalid UTF-8 header value is dropped, leaving no valid context.
146        let cx = extract_context([(TRACEPARENT, [0xff, 0xfe].as_slice())]);
147        assert2::assert!(!cx.span().span_context().is_valid());
148    }
149
150    /// Run `f` under a subscriber wired to a real `OTel` tracer with
151    /// `AlwaysOn` sampling and no exporter.
152    ///
153    /// Spans created inside then carry a valid, *sampled* `OTel` context. The
154    /// inject and re-parent helpers do nothing observable without it.
155    fn with_otel_subscriber(f: impl FnOnce()) {
156        use opentelemetry::trace::TracerProvider as _;
157        use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
158        use tracing_subscriber::prelude::*;
159
160        global::set_text_map_propagator(TraceContextPropagator::new());
161        let provider = SdkTracerProvider::builder()
162            .with_sampler(Sampler::AlwaysOn)
163            .build();
164        let tracer = provider.tracer("propagation-test");
165        let subscriber =
166            tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
167        subscriber::with_default(subscriber, f);
168    }
169
170    #[test]
171    fn vec_injector_collects_key_value_pairs() {
172        // The Injector must actually record each (key, value); a no-op set()
173        // would silently drop the trace headers.
174        let mut out = Vec::new();
175        VecInjector(&mut out).set(TRACEPARENT, "abc".to_owned());
176        assert2::assert!(out == vec![(TRACEPARENT.to_owned(), "abc".to_owned())]);
177    }
178
179    #[test]
180    fn map_extractor_reads_values_and_lists_keys() {
181        // get() returns the stored value (None when absent); keys() must list the
182        // real header keys, not an empty/placeholder set.
183        let map: HashMap<String, String> =
184            iter::once((TRACEPARENT.to_owned(), "v".to_owned())).collect();
185        let ex = MapExtractor(&map);
186        assert2::assert!(ex.get(TRACEPARENT) == Some("v"));
187        assert2::assert!(ex.get("absent").is_none());
188        assert2::assert!(ex.keys() == vec![TRACEPARENT]);
189    }
190
191    #[test]
192    fn current_trace_headers_emits_traceparent_for_active_span() {
193        with_otel_subscriber(|| {
194            let span = tracing::info_span!("producer");
195            let _g = span.enter();
196
197            let headers = current_trace_headers();
198            let tp = headers.iter().find(|(k, _)| k == TRACEPARENT);
199            assert2::assert!(tp.is_some());
200
201            // The injected value carries *this* span's trace id and the sampled
202            // flag — pinning both the key and the content so a wrong-key or
203            // placeholder-value mutant is caught.
204            let (_, value) = tp.unwrap();
205            let trace_id = tracing::Span::current()
206                .context()
207                .span()
208                .span_context()
209                .trace_id();
210            assert2::assert!(value.contains(&trace_id.to_string()));
211            assert2::assert!(value.ends_with("-01"));
212        });
213    }
214
215    #[test]
216    fn current_trace_headers_empty_without_active_span() {
217        // No subscriber ⇒ no OTel context ⇒ nothing to inject. Safe to call.
218        assert2::assert!(current_trace_headers().is_empty());
219    }
220
221    #[test]
222    fn set_remote_parent_reparents_span_into_header_trace() {
223        with_otel_subscriber(|| {
224            let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
225            let span = tracing::info_span!("consumer");
226            set_remote_parent(&span, [(TRACEPARENT, traceparent.as_bytes())]);
227
228            // The span now belongs to the producer's trace (shares its trace id).
229            let sc = span.context().span().span_context().clone();
230            assert2::assert!(
231                sc.trace_id() == TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap()
232            );
233        });
234    }
235}
krabka_trace_context/propagation.rs

krabka_trace_context/
propagation.rs
