rust_zero_core/
telemetry.rs

1//! OpenTelemetry tracing and OTLP export.
2//!
3//! ```no_run
4//! use rust_zero_core::{OtlpTransport, Telemetry, TelemetryConfig};
5//! let telemetry = Telemetry::init(TelemetryConfig::new(
6//!     "users-api", "http://127.0.0.1:4317", OtlpTransport::Grpc,
7//! ))?;
8//! telemetry.force_flush()?;
9//! # Ok::<(), Box<dyn std::error::Error>>(())
10//! ```
11
12use std::{fmt, time::Duration};
13
14use opentelemetry::{
15    global,
16    propagation::{Extractor, TextMapPropagator},
17    trace::{SpanKind, Status, TraceContextExt, Tracer},
18    Context, KeyValue,
19};
20use opentelemetry_otlp::WithExportConfig;
21use opentelemetry_sdk::{
22    propagation::TraceContextPropagator,
23    trace::{Sampler, SdkTracerProvider},
24    Resource,
25};
26
27use crate::TraceContext;
28
29/// OTLP wire transport used to export spans.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum OtlpTransport {
32    Grpc,
33    HttpBinary,
34}
35
36/// OpenTelemetry tracer-provider configuration.
37#[derive(Debug, Clone)]
38pub struct TelemetryConfig {
39    pub service_name: String,
40    pub endpoint: String,
41    pub transport: OtlpTransport,
42    pub sample_ratio: f64,
43    pub export_timeout: Duration,
44}
45
46impl TelemetryConfig {
47    pub fn new(
48        service_name: impl Into<String>,
49        endpoint: impl Into<String>,
50        transport: OtlpTransport,
51    ) -> Self {
52        Self {
53            service_name: service_name.into(),
54            endpoint: endpoint.into(),
55            transport,
56            sample_ratio: 1.0,
57            export_timeout: Duration::from_secs(10),
58        }
59    }
60
61    pub fn with_sample_ratio(mut self, ratio: f64) -> Self {
62        self.sample_ratio = ratio;
63        self
64    }
65
66    pub fn with_export_timeout(mut self, timeout: Duration) -> Self {
67        self.export_timeout = timeout;
68        self
69    }
70}
71
72/// Owns the global OpenTelemetry tracer provider and flushes it during shutdown.
73#[derive(Debug)]
74pub struct Telemetry {
75    provider: SdkTracerProvider,
76}
77
78impl Telemetry {
79    /// Configures a batched OTLP exporter and installs it as the global tracer provider.
80    pub fn init(config: TelemetryConfig) -> Result<Self, TelemetryError> {
81        validate_config(&config)?;
82
83        let exporter = match config.transport {
84            OtlpTransport::Grpc => opentelemetry_otlp::SpanExporter::builder()
85                .with_tonic()
86                .with_endpoint(config.endpoint)
87                .with_timeout(config.export_timeout)
88                .build(),
89            OtlpTransport::HttpBinary => opentelemetry_otlp::SpanExporter::builder()
90                .with_http()
91                .with_endpoint(config.endpoint)
92                .with_timeout(config.export_timeout)
93                .build(),
94        }
95        .map_err(|error| TelemetryError::Exporter(error.to_string()))?;
96
97        Ok(Self::install(
98            &config.service_name,
99            config.sample_ratio,
100            Some(exporter),
101        ))
102    }
103
104    /// Installs a recording provider without an exporter.
105    ///
106    /// This is useful for propagation-only deployments and deterministic middleware tests.
107    pub fn local(
108        service_name: impl Into<String>,
109        sample_ratio: f64,
110    ) -> Result<Self, TelemetryError> {
111        let service_name = service_name.into();
112        validate_service_and_ratio(&service_name, sample_ratio)?;
113        Ok(Self::install(&service_name, sample_ratio, None))
114    }
115
116    fn install(
117        service_name: &str,
118        sample_ratio: f64,
119        exporter: Option<opentelemetry_otlp::SpanExporter>,
120    ) -> Self {
121        let builder = SdkTracerProvider::builder()
122            .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(
123                sample_ratio,
124            ))))
125            .with_resource(
126                Resource::builder()
127                    .with_service_name(service_name.to_owned())
128                    .build(),
129            );
130        let provider = match exporter {
131            Some(exporter) => builder.with_batch_exporter(exporter).build(),
132            None => builder.build(),
133        };
134        global::set_text_map_propagator(TraceContextPropagator::new());
135        global::set_tracer_provider(provider.clone());
136        Self { provider }
137    }
138
139    pub fn force_flush(&self) -> Result<(), TelemetryError> {
140        self.provider
141            .force_flush()
142            .map_err(|error| TelemetryError::Flush(error.to_string()))
143    }
144
145    pub fn shutdown(&self) -> Result<(), TelemetryError> {
146        self.provider
147            .shutdown()
148            .map_err(|error| TelemetryError::Shutdown(error.to_string()))
149    }
150}
151
152impl Drop for Telemetry {
153    fn drop(&mut self) {
154        let _ = self.provider.shutdown();
155    }
156}
157
158/// Semantic kind of an exported span.
159#[derive(Debug, Clone, Copy, PartialEq, Eq)]
160pub enum TelemetrySpanKind {
161    Client,
162    Server,
163    Internal,
164    Producer,
165    Consumer,
166}
167
168/// An exportable span that also exposes rust-zero's W3C context representation.
169pub struct TelemetrySpan {
170    context: Context,
171    trace_context: Option<TraceContext>,
172    ended: bool,
173}
174
175impl fmt::Debug for TelemetrySpan {
176    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
177        formatter
178            .debug_struct("TelemetrySpan")
179            .field("trace_context", &self.trace_context)
180            .field("ended", &self.ended)
181            .finish_non_exhaustive()
182    }
183}
184
185impl TelemetrySpan {
186    pub fn start(
187        name: impl Into<String>,
188        kind: TelemetrySpanKind,
189        parent: Option<&TraceContext>,
190        attributes: impl IntoIterator<Item = (&'static str, String)>,
191    ) -> Self {
192        let parent_context = parent.map(parent_otel_context).unwrap_or_default();
193        let tracer = global::tracer("rust-zero");
194        let builder = tracer
195            .span_builder(name.into())
196            .with_kind(match kind {
197                TelemetrySpanKind::Client => SpanKind::Client,
198                TelemetrySpanKind::Server => SpanKind::Server,
199                TelemetrySpanKind::Internal => SpanKind::Internal,
200                TelemetrySpanKind::Producer => SpanKind::Producer,
201                TelemetrySpanKind::Consumer => SpanKind::Consumer,
202            })
203            .with_attributes(
204                attributes
205                    .into_iter()
206                    .map(|(key, value)| KeyValue::new(key, value)),
207            );
208        let span = tracer.build_with_context(builder, &parent_context);
209        let context = Context::current_with_span(span);
210        let trace_context = to_rust_zero_context(&context);
211        Self {
212            context,
213            trace_context,
214            ended: false,
215        }
216    }
217
218    pub fn trace_context(&self) -> Option<&TraceContext> {
219        self.trace_context.as_ref()
220    }
221
222    pub fn set_attribute(&self, key: &'static str, value: impl Into<String>) {
223        self.context
224            .span()
225            .set_attribute(KeyValue::new(key, value.into()));
226    }
227
228    pub fn set_error(&self, description: impl Into<String>) {
229        self.context
230            .span()
231            .set_status(Status::error(description.into()));
232    }
233
234    pub fn end(mut self) {
235        self.finish();
236    }
237
238    fn finish(&mut self) {
239        if !self.ended {
240            self.context.span().end();
241            self.ended = true;
242        }
243    }
244}
245
246impl Drop for TelemetrySpan {
247    fn drop(&mut self) {
248        self.finish();
249    }
250}
251
252struct TraceParentCarrier<'a>(&'a str);
253
254impl Extractor for TraceParentCarrier<'_> {
255    fn get(&self, key: &str) -> Option<&str> {
256        key.eq_ignore_ascii_case("traceparent").then_some(self.0)
257    }
258
259    fn keys(&self) -> Vec<&str> {
260        vec!["traceparent"]
261    }
262}
263
264fn parent_otel_context(parent: &TraceContext) -> Context {
265    TraceContextPropagator::new().extract(&TraceParentCarrier(&parent.traceparent()))
266}
267
268fn to_rust_zero_context(context: &Context) -> Option<TraceContext> {
269    let span = context.span();
270    let span_context = span.span_context();
271    if !span_context.is_valid() {
272        return None;
273    }
274    TraceContext::parse(&format!(
275        "00-{}-{}-{:02x}",
276        span_context.trace_id(),
277        span_context.span_id(),
278        span_context.trace_flags().to_u8()
279    ))
280    .ok()
281}
282
283fn validate_config(config: &TelemetryConfig) -> Result<(), TelemetryError> {
284    validate_service_and_ratio(&config.service_name, config.sample_ratio)?;
285    if config.endpoint.trim().is_empty() {
286        return Err(TelemetryError::EmptyEndpoint);
287    }
288    if config.export_timeout.is_zero() {
289        return Err(TelemetryError::InvalidTimeout);
290    }
291    Ok(())
292}
293
294fn validate_service_and_ratio(service_name: &str, sample_ratio: f64) -> Result<(), TelemetryError> {
295    if service_name.trim().is_empty() {
296        return Err(TelemetryError::EmptyServiceName);
297    }
298    if !(0.0..=1.0).contains(&sample_ratio) || !sample_ratio.is_finite() {
299        return Err(TelemetryError::InvalidSampleRatio);
300    }
301    Ok(())
302}
303
304#[derive(Debug, Clone, PartialEq, Eq)]
305pub enum TelemetryError {
306    EmptyServiceName,
307    EmptyEndpoint,
308    InvalidSampleRatio,
309    InvalidTimeout,
310    Exporter(String),
311    Flush(String),
312    Shutdown(String),
313}
314
315impl fmt::Display for TelemetryError {
316    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
317        match self {
318            Self::EmptyServiceName => formatter.write_str("telemetry service name cannot be empty"),
319            Self::EmptyEndpoint => formatter.write_str("telemetry endpoint cannot be empty"),
320            Self::InvalidSampleRatio => {
321                formatter.write_str("telemetry sample ratio must be between zero and one")
322            }
323            Self::InvalidTimeout => {
324                formatter.write_str("telemetry export timeout must be greater than zero")
325            }
326            Self::Exporter(error) => write!(formatter, "telemetry exporter error: {error}"),
327            Self::Flush(error) => write!(formatter, "telemetry flush error: {error}"),
328            Self::Shutdown(error) => write!(formatter, "telemetry shutdown error: {error}"),
329        }
330    }
331}
332
333impl std::error::Error for TelemetryError {}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338
339    #[test]
340    fn validates_configuration() {
341        assert_eq!(
342            Telemetry::local("", 1.0).unwrap_err(),
343            TelemetryError::EmptyServiceName
344        );
345        assert_eq!(
346            Telemetry::local("api", 1.1).unwrap_err(),
347            TelemetryError::InvalidSampleRatio
348        );
349    }
350
351    #[test]
352    fn creates_exportable_child_contexts() {
353        let telemetry = Telemetry::local("users-api", 1.0).unwrap();
354        let parent =
355            TraceContext::parse("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01").unwrap();
356        let span = TelemetrySpan::start(
357            "GET /users",
358            TelemetrySpanKind::Server,
359            Some(&parent),
360            [("http.request.method", "GET".to_owned())],
361        );
362        let context = span.trace_context().unwrap();
363
364        assert_eq!(context.trace_id(), parent.trace_id());
365        assert_ne!(context.span_id(), parent.span_id());
366        span.end();
367        telemetry.force_flush().unwrap();
368    }
369}
rust_zero_core/telemetry.rs

rust_zero_core/
telemetry.rs
