arminius_lib_nats/
opentelemetry.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Utilities for tracing over NATS
6
7mod tests;
8
9use async_nats::Message;
10use tracing_opentelemetry::OpenTelemetrySpanExt as _;
11
12/// OTel Injector to inject tracing contexts to NATS headers
13struct NatsHeaderInjector<'a>(pub &'a mut async_nats::HeaderMap);
14
15impl opentelemetry::propagation::Injector for NatsHeaderInjector<'_> {
16	fn set(&mut self, key: &str, value: String) {
17		self.0.insert(key, value);
18	}
19}
20
21/// OTel Extractor to extract tracing contexts from NATS headers
22struct NatsHeaderExtractor<'a>(pub &'a Option<async_nats::HeaderMap>);
23
24impl opentelemetry::propagation::Extractor for NatsHeaderExtractor<'_> {
25	/// Get a value for a key from the `HeaderMap` as a String
26	fn get(&self, key: &str) -> Option<&str> {
27		match self.0 {
28			Some(map) => map.get(key).map(AsRef::as_ref),
29			None => None,
30		}
31	}
32
33	/// Collect all the keys from the `HeaderMap`.
34	fn keys(&self) -> Vec<&str> {
35		match self.0 {
36			Some(map) => map.iter().map(|(name, _value)| name.as_ref()).collect::<Vec<_>>(),
37			None => Vec::new(),
38		}
39	}
40}
41
42/// Extracts the OpenTelemetry context carried in a NATS message's headers.
43///
44/// Uses the globally configured text-map propagator. Falls back to
45/// [`Context::current()`](opentelemetry::Context::current) when the message
46/// carries no headers.
47pub(crate) fn extract_otel_context(message: &Message) -> opentelemetry::Context {
48	if message.headers.is_some() {
49		opentelemetry::global::get_text_map_propagator(|propagator| {
50			propagator.extract(&NatsHeaderExtractor(&message.headers))
51		})
52	} else {
53		opentelemetry::Context::current()
54	}
55}
56
57/// Extract the current span from a NATS message.
58///
59/// The parent span is set when possible. (Remember to set it when
60/// sending the NATS message.)
61/// The span is not automatically entered, so you will either have to instrument
62/// a function manually or enter it yourself.
63#[must_use]
64pub fn span_from_nats_message(message: &Message) -> tracing::span::Span {
65	let span = tracing::span!(
66		tracing::Level::INFO,
67		"NATS message",
68		nats.subject = message.subject.as_str(),
69		nats.length = message.length,
70		nats.message = tracing::field::debug(&message.payload),
71		nats.headers = tracing::field::debug(&message.headers)
72	);
73
74	let _ = span.set_parent(extract_otel_context(message));
75
76	span
77}
78
79/// Create NATS message headers, that include the current tracing span
80#[must_use]
81pub fn span_to_nats_headers() -> async_nats::HeaderMap {
82	let mut headers = async_nats::HeaderMap::new();
83	let ctx = tracing::Span::current().context();
84	opentelemetry::global::get_text_map_propagator(|propagator| {
85		propagator.inject_context(&ctx, &mut NatsHeaderInjector(&mut headers));
86	});
87
88	headers
89}
arminius_lib_nats/opentelemetry.rs

arminius_lib_nats/
opentelemetry.rs
