byre/
telemetry.rs

1//! # Telemetry for Rust Applications
2//!
3//! This module provides a complete telemetry solution integrating:
4//!
5//! - **Tracing**: Distributed tracing for tracking request flows across services
6//! - **Metrics**: Quantitative measurements of your application's performance
7//! - **Logging**: Structured event logging for visibility into application behavior
8//!
9//! The implementation uses [OpenTelemetry](https://opentelemetry.io/) standards for compatibility
10//! with popular observability platforms like Jaeger, Prometheus, Grafana, etc.
11//!
12//! ## Why use telemetry?
13//!
14//! - **Troubleshooting**: Quickly identify and debug issues in production
15//! - **Performance optimizations**: Measure and improve application performance
16//! - **Service health monitoring**: Get alerts when services degrade
17//! - **Cross-service visibility**: Track requests across microservices
18//!
19//! ## Quick Start
20//!
21//! ```rust,no_run
22//! use byre::telemetry::{TelemetrySettings, TelemetryProviders};
23//! use byre::ServiceInfo;
24//!
25//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
26//! // 1. Configure telemetry in your settings
27//! let settings = TelemetrySettings::default();
28//! let service = ServiceInfo {
29//!     name: "my-service",
30//!     name_in_metrics: "my_service".to_string(),
31//!     version: "1.0.0",
32//!     author: "Author",
33//!     description: "My service description",
34//! };
35//!
36//! // 2. Initialize telemetry (keep the returned handle alive for the app lifetime!)
37//! let _telemetry: TelemetryProviders = byre::telemetry::init(&service, &settings)?;
38//!
39//! // 3. Use tracing macros as normal - they automatically go to OpenTelemetry
40//! tracing::info!("Application started");
41//! # Ok(())
42//! # }
43//! ```
44//!
45//! ## Common Patterns
46//!
47//! ### gRPC Metadata (linking incoming trace context)
48//!
49//! ```
50//! use byre::telemetry::{TraceContextCarrier, TraceContextExt};
51//!
52//! // In a gRPC handler, extract trace context from incoming metadata
53//! let mut metadata = tonic::metadata::MetadataMap::new();
54//! metadata.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
55//!
56//! // Extract the trace context
57//! let ctx = metadata.extract_trace_context();
58//!
59//! // Or link it directly to the current span
60//! let _ = metadata.link_distributed_trace();
61//! ```
62//!
63//! ### gRPC Metadata (propagating trace context)
64//!
65//! ```
66//! use byre::telemetry::{TraceContextCarrier, TraceContextExt};
67//!
68//! // Before making outgoing gRPC calls, inject trace context
69//! let mut metadata = tonic::metadata::MetadataMap::new();
70//! metadata.inject_trace_context();
71//! // metadata now contains traceparent header (if there's an active span)
72//! ```
73//!
74//! ### HTTP Headers (linking incoming trace context)
75//!
76//! ```
77//! use byre::telemetry::{TraceContextCarrier, TraceContextExt};
78//!
79//! // In an HTTP handler, extract trace context from incoming headers
80//! let mut headers = http::HeaderMap::new();
81//! headers.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
82//!
83//! // Extract the trace context
84//! let ctx = headers.extract_trace_context();
85//!
86//! // Or link it directly to the current span
87//! let _ = headers.link_distributed_trace();
88//! ```
89//!
90//! ### HTTP Headers (propagating trace context)
91//!
92//! ```
93//! use byre::telemetry::{TraceContextCarrier, TraceContextExt};
94//!
95//! // Before making outgoing HTTP calls, inject trace context
96//! let mut headers = http::HeaderMap::new();
97//! headers.inject_trace_context();
98//! // headers now contains traceparent header (if there's an active span)
99//! ```
100//!
101//! ### HashMap (for message queues)
102//!
103//! ```
104//! use std::collections::HashMap;
105//! use byre::telemetry::{TraceContextCarrier, TraceContextExt};
106//!
107//! // Producer: inject trace context into message headers
108//! let mut headers: HashMap<String, String> = HashMap::new();
109//! headers.inject_trace_context();
110//!
111//! // Consumer: extract trace context from message headers
112//! let ctx = headers.extract_trace_context();
113//!
114//! // Or link it directly to the current span
115//! let _ = headers.link_distributed_trace();
116//! ```
117
118use doku::Document;
119use opentelemetry::propagation::{Extractor, Injector};
120use opentelemetry::trace::TracerProvider as _;
121use opentelemetry::{global, KeyValue};
122use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
123use opentelemetry_otlp::{
124    ExporterBuildError, LogExporter, MetricExporter, SpanExporter, WithExportConfig,
125};
126use opentelemetry_sdk::logs::SdkLoggerProvider;
127use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
128use opentelemetry_sdk::propagation::TraceContextPropagator;
129use opentelemetry_sdk::{trace as sdktrace, Resource};
130use serde::{Deserialize, Serialize};
131use snafu::{ResultExt as _, Snafu};
132use tracing::Subscriber;
133use tracing_opentelemetry::OpenTelemetryLayer;
134use tracing_subscriber::prelude::*;
135use tracing_subscriber::EnvFilter;
136
137use crate::ServiceInfo;
138
139// ============================================================================
140// Trace Context Carrier Traits
141// ============================================================================
142
143/// Trait for types that can carry trace context (e.g., HTTP headers, gRPC metadata).
144///
145/// This trait provides a unified interface for extracting and injecting
146/// W3C Trace Context headers across different transport types.
147///
148/// Implementations are provided for:
149/// - `tonic::metadata::MetadataMap` (gRPC)
150/// - `http::HeaderMap` (HTTP)
151/// - `HashMap<String, String>` (message queues, generic use)
152pub trait TraceContextCarrier {
153    /// Extract trace context from this carrier.
154    ///
155    /// Returns an OpenTelemetry context that can be used to link spans
156    /// to an incoming distributed trace.
157    fn extract_trace_context(&self) -> opentelemetry::Context;
158
159    /// Inject the current span's trace context into this carrier.
160    ///
161    /// Call this before making outgoing requests to propagate the trace.
162    fn inject_trace_context(&mut self);
163}
164
165/// Extension trait providing convenient methods for trace context propagation.
166///
167/// This trait is automatically implemented for all types that implement
168/// [`TraceContextCarrier`]. Import this trait to use the extension methods:
169///
170/// ```
171/// use byre::telemetry::{TraceContextCarrier, TraceContextExt};
172///
173/// // Now you can call these methods on any carrier type:
174/// let mut headers = http::HeaderMap::new();
175/// let ctx = headers.extract_trace_context();
176/// headers.inject_trace_context();
177/// let _ = headers.link_distributed_trace();
178/// ```
179pub trait TraceContextExt: TraceContextCarrier {
180    /// Link the current tracing span to an incoming distributed trace.
181    ///
182    /// This is a convenience method that extracts the trace context and
183    /// sets it as the parent of the current span. Call this at the start
184    /// of your handler after the `#[tracing::instrument]` span is created.
185    ///
186    /// Returns `Ok(())` if successful, or an error if the span context
187    /// couldn't be set. Most callers will want to ignore the error:
188    ///
189    /// ```
190    /// use byre::telemetry::{TraceContextCarrier, TraceContextExt};
191    ///
192    /// let headers = http::HeaderMap::new();
193    /// let _ = headers.link_distributed_trace();
194    /// ```
195    fn link_distributed_trace(&self) -> Result<(), Error>;
196}
197
198impl<T: TraceContextCarrier> TraceContextExt for T {
199    fn link_distributed_trace(&self) -> Result<(), Error> {
200        use tracing_opentelemetry::OpenTelemetrySpanExt;
201        let parent_cx = self.extract_trace_context();
202        tracing::Span::current()
203            .set_parent(parent_cx)
204            .map_err(|e| Error::LinkDistributedTrace {
205                source: Box::new(e),
206            })
207    }
208}
209
210/// Errors initializing telemetry
211#[derive(Debug, Snafu)]
212pub enum Error {
213    /// Could not initialize the logger
214    #[snafu(display("Could not initialize logging: {source}"))]
215    InitLog {
216        /// The error from initializing the gRPC connection
217        source: ExporterBuildError,
218    },
219
220    /// Could not initialize metrics
221    #[snafu(display("Could not initialize metrics: {source}"))]
222    InitMetric {
223        /// The error from initializing the gRPC connection
224        source: ExporterBuildError,
225    },
226
227    /// Could not initialize tracing
228    #[snafu(display("Could not initialize tracing: {source}"))]
229    InitTrace {
230        /// The error from initializing the gRPC connection
231        source: ExporterBuildError,
232    },
233
234    /// Could not link distributed trace context to current span
235    #[snafu(display("Could not link distributed trace: {source}"))]
236    LinkDistributedTrace {
237        /// The underlying error from tracing-opentelemetry
238        source: Box<dyn std::error::Error + Send + Sync>,
239    },
240}
241
242/// Settings for metrics collection and export.
243///
244/// Metrics provide quantitative measurements about your application's performance and behavior.
245/// Examples include request counts, error rates, response times, and resource usage.
246#[derive(Debug, Default, Serialize, Deserialize, Document)]
247pub struct MetricSettings {
248    /// gRPC endpoint to send metrics to. Omit to disable opentelemetry metrics.
249    #[doku(example = "http://localhost:4318/v1/metrics")]
250    pub endpoint: Option<String>,
251}
252
253/// Settings for logging configuration.
254///
255/// Logs provide contextual information about application events and are essential
256/// for debugging and monitoring application behavior.
257///
258/// Note: `otel_level` will filter the logs before they are sent to the console, so if `otel_level` is `warn`, then `console_level` can only be `warn`, `error`, or `off`.
259#[derive(Debug, Default, Serialize, Deserialize, Document)]
260pub struct LogSettings {
261    /// log level used when filtering console logs. Uses env-logger style syntax. Set to "off" to disable console logging.
262    /// `console_level` is limited by `otel_level`, so if `otel_level` is `warn`, then `console_level` can only be `warn`, `error`, or `off`.
263    #[doku(example = "debug,yourcrate=info")]
264    pub console_level: String,
265
266    /// log level used when filtering opentelemetry logs. Uses env-logger style syntax.
267    #[doku(example = "warn,yourcrate=debug")]
268    pub otel_level: String,
269
270    /// gRPC endpoint to send the opentelemetry logs. Omit to disable opentelemetry logs, will not disable console logs.
271    #[doku(example = "http://localhost:4317")]
272    pub endpoint: Option<String>,
273}
274
275/// Settings for distributed tracing.
276///
277/// Traces track the flow of requests as they propagate through your system, helping you
278/// understand the execution path and identify performance bottlenecks.
279#[derive(Debug, Default, Serialize, Deserialize, Document)]
280pub struct TraceSettings {
281    /// gRPC endpoint to send opentelemetry traces to, omit to disable.
282    #[doku(example = "http://localhost:4317")]
283    pub endpoint: Option<String>,
284}
285
286/**
287Settings for tracing, logging, and metrics.
288
289Use `TelemetrySettings` as a member in your own `Settings` object.
290
291# Example
292
293```rust
294use doku::Document;
295use serde::Deserialize;
296
297#[derive(Deserialize, Document)]
298/// Data Archive Settings
299pub struct Settings {
300    /// Server Settings
301    pub application: Application,
302    /// Telemetry settings.
303    pub telemetry: byre::telemetry::TelemetrySettings,
304}
305
306#[derive(Deserialize, Document)]
307pub struct Application {
308    #[doku(example = "localhost")]
309    pub listen_host: String,
310    #[doku(example = "8080")]
311    pub listen_port: u16,
312}
313``` */
314#[derive(Debug, Default, Serialize, Deserialize, Document)]
315pub struct TelemetrySettings {
316    /// Settings for tracing
317    pub trace: TraceSettings,
318    /// Settings for logging
319    pub log: LogSettings,
320    /// Settings for metrics
321    pub metric: MetricSettings,
322}
323
324/// Container for the initialized telemetry providers.
325///
326/// This struct owns the telemetry providers and ensures they are properly
327/// shut down when dropped. You must keep this value alive for the duration
328/// of your application; dropping it will shut down all telemetry.
329#[derive(Debug, Default)]
330#[must_use = "dropping TelemetryProviders will shut down all telemetry"]
331pub struct TelemetryProviders {
332    meter: Option<SdkMeterProvider>,
333    tracer: Option<sdktrace::SdkTracerProvider>,
334    logger: Option<SdkLoggerProvider>,
335}
336
337impl Drop for TelemetryProviders {
338    fn drop(&mut self) {
339        if let Some(tracer_provider) = self.tracer.take() {
340            if let Err(err) = tracer_provider.shutdown() {
341                eprintln!("Error shutting down Telemetry tracer provider: {err}");
342            }
343        }
344        if let Some(logger_provider) = self.logger.take() {
345            if let Err(err) = logger_provider.shutdown() {
346                eprintln!("Error shutting down Telemetry logger provider: {err}");
347            }
348        }
349        if let Some(meter_provider) = self.meter.take() {
350            if let Err(err) = meter_provider.shutdown() {
351                eprintln!("Error shutting down Telemetry meter provider: {err}");
352            }
353        }
354    }
355}
356
357fn init_traces(
358    service_info: &ServiceInfo,
359    settings: &TraceSettings,
360) -> Result<Option<sdktrace::SdkTracerProvider>, ExporterBuildError> {
361    match &settings.endpoint {
362        Some(endpoint) => {
363            let exporter = SpanExporter::builder()
364                .with_tonic()
365                .with_endpoint(endpoint)
366                .build()?;
367
368            let resource = Resource::builder()
369                .with_attribute(KeyValue::new(
370                    opentelemetry_semantic_conventions::resource::SERVICE_NAME,
371                    service_info.name_in_metrics.clone(),
372                ))
373                .build();
374
375            Ok(Some(
376                sdktrace::SdkTracerProvider::builder()
377                    .with_resource(resource)
378                    .with_batch_exporter(exporter)
379                    .build(),
380            ))
381        }
382        None => Ok(None),
383    }
384}
385
386fn init_metrics(
387    service_info: &ServiceInfo,
388    setting: &MetricSettings,
389) -> Result<Option<opentelemetry_sdk::metrics::SdkMeterProvider>, ExporterBuildError> {
390    match &setting.endpoint {
391        Some(endpoint) => {
392            let exporter = MetricExporter::builder()
393                .with_tonic()
394                .with_endpoint(endpoint)
395                .build()?;
396            let reader = PeriodicReader::builder(exporter).build();
397
398            let resource = Resource::builder()
399                .with_attribute(KeyValue::new(
400                    opentelemetry_semantic_conventions::resource::SERVICE_NAME,
401                    service_info.name_in_metrics.clone(),
402                ))
403                .build();
404
405            Ok(Some(
406                SdkMeterProvider::builder()
407                    .with_reader(reader)
408                    .with_resource(resource)
409                    .build(),
410            ))
411        }
412
413        None => Ok(None),
414    }
415}
416
417fn init_otel_logs<S>(
418    service_info: &ServiceInfo,
419    settings: &LogSettings,
420) -> Result<
421    (
422        Option<opentelemetry_sdk::logs::SdkLoggerProvider>,
423        Option<impl tracing_subscriber::layer::Layer<S> + use<S>>,
424    ),
425    Error,
426>
427where
428    S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
429{
430    match &settings.endpoint {
431        None => Ok((None, None)),
432
433        Some(endpoint) => {
434            let builder = init_otel_logs_builder(service_info, endpoint)?;
435
436            let logger_provider = builder.build();
437
438            // Create a new OpenTelemetryTracingBridge using the above LoggerProvider.
439            let otel_layer = OpenTelemetryTracingBridge::new(&logger_provider);
440
441            // For the OpenTelemetry layer, add a tracing filter to filter events from
442            // OpenTelemetry and its dependent crates (opentelemetry-otlp uses crates
443            // like reqwest/tonic etc.) from being sent back to OTel itself, thus
444            // preventing infinite telemetry generation. The filter levels are set as
445            // follows:
446            // - Allow `info` level and above by default.
447            // - Restrict `opentelemetry`, `hyper`, `tonic`, and `reqwest` completely.
448            // Note: This will also drop events from crates like `tonic` etc. even when
449            // they are used outside the OTLP Exporter. For more details, see:
450            // https://github.com/open-telemetry/opentelemetry-rust/issues/761
451            // FIXME: the directives below should be noted in the documentation!
452            let filter_otel = EnvFilter::new(&settings.otel_level)
453                .add_directive("hyper=off".parse().unwrap())
454                .add_directive("opentelemetry=off".parse().unwrap())
455                .add_directive("tonic=off".parse().unwrap())
456                .add_directive("h2=off".parse().unwrap())
457                .add_directive("reqwest=off".parse().unwrap());
458            let otel_layer = otel_layer.with_filter(filter_otel);
459
460            Ok((Some(logger_provider), Some(otel_layer)))
461        }
462    }
463}
464
465fn init_otel_logs_builder(
466    service_info: &ServiceInfo,
467    endpoint: &String,
468) -> Result<opentelemetry_sdk::logs::LoggerProviderBuilder, Error> {
469    let builder = SdkLoggerProvider::builder();
470    let exporter = LogExporter::builder()
471        .with_tonic()
472        .with_endpoint(endpoint)
473        .build()
474        .with_context(|_| InitLogSnafu {})?;
475    let resource = Resource::builder()
476        .with_attribute(KeyValue::new(
477            opentelemetry_semantic_conventions::resource::SERVICE_NAME,
478            service_info.name_in_metrics.clone(),
479        ))
480        .build();
481    let builder = builder
482        .with_resource(resource)
483        .with_batch_exporter(exporter);
484    Ok(builder)
485}
486
487/// Builder for configuring and initializing the logging/tracing subscriber.
488///
489/// This builder separates configuration from initialization, making it easier
490/// to test the subscriber configuration without installing it globally.
491struct LogSubscriberBuilder<'a> {
492    service_info: &'a ServiceInfo,
493    settings: &'a LogSettings,
494    tracer_provider: Option<&'a sdktrace::SdkTracerProvider>,
495}
496
497/// The built subscriber components, ready to be installed or used for testing.
498struct BuiltSubscriber<S> {
499    /// The logger provider (if OTel logging endpoint was configured)
500    logger_provider: Option<opentelemetry_sdk::logs::SdkLoggerProvider>,
501    /// The fully configured subscriber
502    subscriber: S,
503}
504
505impl<'a> LogSubscriberBuilder<'a> {
506    /// Create a new builder with the required configuration.
507    fn new(service_info: &'a ServiceInfo, settings: &'a LogSettings) -> Self {
508        Self {
509            service_info,
510            settings,
511            tracer_provider: None,
512        }
513    }
514
515    /// Set the tracer provider for OpenTelemetry trace integration.
516    fn with_tracer_provider(mut self, provider: &'a sdktrace::SdkTracerProvider) -> Self {
517        self.tracer_provider = Some(provider);
518        self
519    }
520
521    /// Build the subscriber without installing it globally.
522    /// Use this for testing with `tracing::subscriber::with_default`.
523    fn build(
524        self,
525    ) -> Result<
526        BuiltSubscriber<
527            impl Subscriber // + for<'span> tracing_subscriber::registry::LookupSpan<'span>
528                // + Send
529                // + Sync
530                + use<'a>,
531        >,
532        Error,
533    > {
534        let (logger_provider, otel_log_layer) = init_otel_logs(self.service_info, self.settings)?;
535
536        // Create the OpenTelemetry tracing layer if a tracer provider is configured.
537        // This bridges tracing spans to OpenTelemetry traces.
538        let otel_trace_layer = self.tracer_provider.map(|provider| {
539            let tracer = provider.tracer(self.service_info.name_in_metrics.clone());
540            let filter = EnvFilter::new(&self.settings.otel_level)
541                .add_directive("hyper=off".parse().unwrap())
542                .add_directive("opentelemetry=off".parse().unwrap())
543                .add_directive("opentelemetry_sdk=off".parse().unwrap())
544                .add_directive("tonic=off".parse().unwrap())
545                .add_directive("h2=off".parse().unwrap())
546                .add_directive("reqwest=off".parse().unwrap());
547            OpenTelemetryLayer::new(tracer).with_filter(filter)
548        });
549
550        // Create a new tracing::Fmt layer to print the logs to stdout.
551        let filter_fmt = EnvFilter::new(&self.settings.console_level);
552        let fmt_layer = tracing_subscriber::fmt::layer()
553            .with_thread_names(true)
554            .with_filter(filter_fmt);
555
556        // Build the subscriber with all layers (but don't install it)
557        let subscriber = tracing_subscriber::registry()
558            .with(otel_log_layer)
559            .with(otel_trace_layer)
560            .with(fmt_layer);
561
562        Ok(BuiltSubscriber {
563            logger_provider,
564            subscriber,
565        })
566    }
567
568    /// Build and install the subscriber globally.
569    /// Returns the logger provider if OTel logging was configured.
570    fn init(self) -> Result<Option<opentelemetry_sdk::logs::SdkLoggerProvider>, Error> {
571        let built = self.build()?;
572        built.subscriber.init();
573        Ok(built.logger_provider)
574    }
575}
576
577fn init_logs(
578    service_info: &ServiceInfo,
579    settings: &LogSettings,
580    tracer_provider: Option<&sdktrace::SdkTracerProvider>,
581) -> Result<Option<opentelemetry_sdk::logs::SdkLoggerProvider>, Error> {
582    let mut builder = LogSubscriberBuilder::new(service_info, settings);
583    if let Some(provider) = tracer_provider {
584        builder = builder.with_tracer_provider(provider);
585    }
586    builder.init()
587}
588
589/// Initializes the telemetry backend for your application.
590///
591/// This function sets up tracing, metrics, and logging according to the provided settings.
592/// It integrates with OpenTelemetry to provide a complete observability solution.
593///
594/// # Errors
595///
596/// - `InitLog` if the logger provider cannot be initialized.
597/// - `InitTrace` if the tracer provider cannot be initialized.
598/// - `InitMetric` if the metric provider cannot be initialized.
599/// Initializes the telemetry backend for your application.
600///
601/// This function sets up tracing, metrics, logging, and the W3C Trace Context
602/// propagator for distributed tracing according to the provided settings.
603/// It integrates with OpenTelemetry to provide a complete observability solution.
604///
605/// # Errors
606///
607/// - `InitLog` if the logger provider cannot be initialized.
608/// - `InitTrace` if the tracer provider cannot be initialized.
609/// - `InitMetric` if the metric provider cannot be initialized.
610#[must_use]
611pub fn init(
612    service_info: &ServiceInfo,
613    settings: &TelemetrySettings,
614) -> Result<TelemetryProviders, Error> {
615    // Initialize the W3C Trace Context propagator for distributed tracing
616    init_propagator();
617    // Initialize traces first so we can pass the provider to init_logs for the tracing layer
618    let tracer_provider =
619        init_traces(service_info, &settings.trace).with_context(|_| InitTraceSnafu {})?;
620    if let Some(tracer_provider) = &tracer_provider {
621        global::set_tracer_provider(tracer_provider.clone());
622    }
623
624    // Initialize logs with the tracer provider to enable span export via tracing-opentelemetry
625    let logger_provider = init_logs(service_info, &settings.log, tracer_provider.as_ref())?;
626
627    let meter_provider =
628        init_metrics(service_info, &settings.metric).with_context(|_| InitMetricSnafu {})?;
629    if let Some(meter_provider) = &meter_provider {
630        global::set_meter_provider(meter_provider.clone());
631    }
632
633    Ok(TelemetryProviders {
634        meter: meter_provider,
635        tracer: tracer_provider,
636        logger: logger_provider,
637    })
638}
639
640// ============================================================================
641// Distributed Tracing Propagation
642// ============================================================================
643
644/// Wrapper for tonic::metadata::MetadataMap to implement Extractor trait.
645/// Used for extracting trace context from incoming gRPC requests.
646pub struct MetadataExtractor<'a>(pub &'a tonic::metadata::MetadataMap);
647
648impl Extractor for MetadataExtractor<'_> {
649    fn get(&self, key: &str) -> Option<&str> {
650        self.0.get(key).and_then(|v| v.to_str().ok())
651    }
652
653    fn keys(&self) -> Vec<&str> {
654        // W3C Trace Context only uses "traceparent" and optionally "tracestate".
655        // Only return the keys that actually exist in the metadata.
656        ["traceparent", "tracestate"]
657            .into_iter()
658            .filter(|k| self.0.get(*k).is_some())
659            .collect()
660    }
661}
662
663/// Wrapper for tonic::metadata::MetadataMap to implement Injector trait.
664/// Used for injecting trace context into outgoing gRPC requests.
665pub struct MetadataInjector<'a>(pub &'a mut tonic::metadata::MetadataMap);
666
667impl Injector for MetadataInjector<'_> {
668    fn set(&mut self, key: &str, value: String) {
669        if let Ok(key) = tonic::metadata::MetadataKey::from_bytes(key.as_bytes()) {
670            if let Ok(value) = tonic::metadata::MetadataValue::try_from(&value) {
671                self.0.insert(key, value);
672            }
673        }
674    }
675}
676
677impl TraceContextCarrier for tonic::metadata::MetadataMap {
678    fn extract_trace_context(&self) -> opentelemetry::Context {
679        global::get_text_map_propagator(|propagator| propagator.extract(&MetadataExtractor(self)))
680    }
681
682    fn inject_trace_context(&mut self) {
683        use tracing_opentelemetry::OpenTelemetrySpanExt;
684        let cx = tracing::Span::current().context();
685        global::get_text_map_propagator(|propagator| {
686            propagator.inject_context(&cx, &mut MetadataInjector(self));
687        });
688    }
689}
690
691/// Extract trace context from incoming gRPC request metadata.
692///
693/// Returns the extracted OpenTelemetry context. Use [`link_distributed_trace`] for a more
694/// convenient way to extract and link the trace context in one call.
695///
696/// # Example
697///
698/// ```
699/// let mut metadata = tonic::metadata::MetadataMap::new();
700/// metadata.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
701///
702/// let parent_cx = byre::telemetry::extract_trace_context(&metadata);
703/// let _guard = parent_cx.attach();
704/// // Spans created here will be children of the incoming trace
705/// ```
706pub fn extract_trace_context(metadata: &tonic::metadata::MetadataMap) -> opentelemetry::Context {
707    global::get_text_map_propagator(|propagator| propagator.extract(&MetadataExtractor(metadata)))
708}
709
710/// Link the current span to an incoming distributed trace from gRPC metadata.
711///
712/// This is a convenience function that extracts the trace context from the
713/// incoming request metadata and sets it as the parent of the current span.
714/// Call this at the start of your gRPC handler after the `#[tracing::instrument]` span is created.
715///
716/// Returns `Ok(())` if successful, or an error if the span context couldn't be set.
717/// Most callers will want to ignore the error with `let _ = link_distributed_trace(...)`.
718///
719/// # Example
720///
721/// ```
722/// let mut metadata = tonic::metadata::MetadataMap::new();
723/// metadata.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
724///
725/// let _ = byre::telemetry::link_distributed_trace(&metadata);
726/// // Current span is now part of the distributed trace
727/// ```
728pub fn link_distributed_trace(metadata: &tonic::metadata::MetadataMap) -> Result<(), Error> {
729    use tracing_opentelemetry::OpenTelemetrySpanExt;
730    let parent_cx = extract_trace_context(metadata);
731    tracing::Span::current()
732        .set_parent(parent_cx)
733        .map_err(|e| Error::LinkDistributedTrace {
734            source: Box::new(e),
735        })
736}
737
738/// Inject trace context into outgoing gRPC request metadata.
739///
740/// Call this before making outgoing gRPC calls to propagate the trace context.
741///
742/// # Example
743///
744/// ```
745/// let mut metadata = tonic::metadata::MetadataMap::new();
746/// byre::telemetry::inject_trace_context(&mut metadata);
747/// // metadata now contains traceparent header (if there's an active span)
748/// ```
749pub fn inject_trace_context(metadata: &mut tonic::metadata::MetadataMap) {
750    use tracing_opentelemetry::OpenTelemetrySpanExt;
751    // Get the OpenTelemetry context from the current tracing span
752    let cx = tracing::Span::current().context();
753    global::get_text_map_propagator(|propagator| {
754        propagator.inject_context(&cx, &mut MetadataInjector(metadata));
755    });
756}
757
758/// Initialize the global text map propagator for W3C Trace Context.
759///
760/// This is called automatically by `init()`, but can be called manually if needed.
761pub fn init_propagator() {
762    global::set_text_map_propagator(TraceContextPropagator::new());
763}
764
765// ============================================================================
766// HTTP Header Propagation (for HTTP proxies and clients)
767// ============================================================================
768
769/// Wrapper for http::HeaderMap to implement Extractor trait.
770/// Used for extracting trace context from incoming HTTP requests.
771pub struct HttpHeaderExtractor<'a>(pub &'a http::HeaderMap);
772
773impl Extractor for HttpHeaderExtractor<'_> {
774    fn get(&self, key: &str) -> Option<&str> {
775        self.0.get(key).and_then(|v| v.to_str().ok())
776    }
777
778    fn keys(&self) -> Vec<&str> {
779        // W3C Trace Context only uses "traceparent" and optionally "tracestate".
780        // Only return the keys that actually exist in the headers.
781        ["traceparent", "tracestate"]
782            .into_iter()
783            .filter(|k| self.0.get(*k).is_some())
784            .collect()
785    }
786}
787
788/// Wrapper for http::HeaderMap to implement Injector trait.
789/// Used for injecting trace context into outgoing HTTP requests.
790pub struct HttpHeaderInjector<'a>(pub &'a mut http::HeaderMap);
791
792impl Injector for HttpHeaderInjector<'_> {
793    fn set(&mut self, key: &str, value: String) {
794        if let Ok(key) = http::header::HeaderName::from_bytes(key.as_bytes()) {
795            if let Ok(value) = http::header::HeaderValue::from_str(&value) {
796                self.0.insert(key, value);
797            }
798        }
799    }
800}
801
802impl TraceContextCarrier for http::HeaderMap {
803    fn extract_trace_context(&self) -> opentelemetry::Context {
804        global::get_text_map_propagator(|propagator| propagator.extract(&HttpHeaderExtractor(self)))
805    }
806
807    fn inject_trace_context(&mut self) {
808        use tracing_opentelemetry::OpenTelemetrySpanExt;
809        let cx = tracing::Span::current().context();
810        global::get_text_map_propagator(|propagator| {
811            propagator.inject_context(&cx, &mut HttpHeaderInjector(self));
812        });
813    }
814}
815
816/// Extract trace context from incoming HTTP request headers.
817///
818/// Returns the extracted OpenTelemetry context. Use [`link_distributed_trace_http`] for a more
819/// convenient way to extract and link the trace context in one call.
820///
821/// # Example
822///
823/// ```
824/// let mut headers = http::HeaderMap::new();
825/// headers.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
826///
827/// let parent_cx = byre::telemetry::extract_trace_context_http(&headers);
828/// let _guard = parent_cx.attach();
829/// // Spans created here will be children of the incoming trace
830/// ```
831pub fn extract_trace_context_http(headers: &http::HeaderMap) -> opentelemetry::Context {
832    global::get_text_map_propagator(|propagator| propagator.extract(&HttpHeaderExtractor(headers)))
833}
834
835/// Link the current span to an incoming distributed trace from HTTP headers.
836///
837/// This is a convenience function that extracts the trace context from the
838/// incoming request headers and sets it as the parent of the current span.
839/// Call this at the start of your HTTP handler after the `#[tracing::instrument]` span is created.
840///
841/// Returns `Ok(())` if successful, or an error if the span context couldn't be set.
842/// Most callers will want to ignore the error with `let _ = link_distributed_trace_http(...)`.
843///
844/// # Example
845///
846/// ```
847/// let mut headers = http::HeaderMap::new();
848/// headers.insert("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".parse().unwrap());
849///
850/// let _ = byre::telemetry::link_distributed_trace_http(&headers);
851/// // Current span is now part of the distributed trace
852/// ```
853pub fn link_distributed_trace_http(headers: &http::HeaderMap) -> Result<(), Error> {
854    use tracing_opentelemetry::OpenTelemetrySpanExt;
855    let parent_cx = extract_trace_context_http(headers);
856    tracing::Span::current()
857        .set_parent(parent_cx)
858        .map_err(|e| Error::LinkDistributedTrace {
859            source: Box::new(e),
860        })
861}
862
863/// Inject trace context into outgoing HTTP request headers.
864///
865/// Call this before making outgoing HTTP calls to propagate the trace context.
866///
867/// # Example
868///
869/// ```
870/// let mut headers = http::HeaderMap::new();
871/// byre::telemetry::inject_trace_context_http(&mut headers);
872/// // headers now contains traceparent header (if there's an active span)
873/// ```
874pub fn inject_trace_context_http(headers: &mut http::HeaderMap) {
875    use tracing_opentelemetry::OpenTelemetrySpanExt;
876    // Get the OpenTelemetry context from the current tracing span
877    let cx = tracing::Span::current().context();
878    global::get_text_map_propagator(|propagator| {
879        propagator.inject_context(&cx, &mut HttpHeaderInjector(headers));
880    });
881}
882
883// ============================================================================
884// Tower Layer for Distributed Trace Context (gRPC/tonic)
885// ============================================================================
886
887/// A Tower layer that extracts distributed trace context from incoming gRPC requests
888/// and creates a parent span for all downstream handlers.
889///
890/// This layer should be added to tonic services to enable distributed tracing.
891/// It extracts the W3C Trace Context headers from incoming requests and creates
892/// a span that becomes the parent of all spans created within the handler.
893///
894/// # Example
895///
896/// ```
897/// use byre::telemetry::GrpcTraceContextLayer;
898///
899/// // Create the layer
900/// let layer = GrpcTraceContextLayer::new("my-service");
901///
902/// // Use with tonic Server::builder().layer(layer)
903/// ```
904#[derive(Clone)]
905pub struct GrpcTraceContextLayer {
906    service_name: &'static str,
907}
908
909impl GrpcTraceContextLayer {
910    /// Create a new layer with the given service name.
911    /// The service name is used to identify spans in the trace.
912    pub fn new(service_name: &'static str) -> Self {
913        Self { service_name }
914    }
915}
916
917impl<S> tower::Layer<S> for GrpcTraceContextLayer {
918    type Service = GrpcTraceContextService<S>;
919
920    fn layer(&self, inner: S) -> Self::Service {
921        GrpcTraceContextService {
922            inner,
923            service_name: self.service_name,
924        }
925    }
926}
927
928/// The service that wraps inner services with trace context extraction.
929#[derive(Clone)]
930pub struct GrpcTraceContextService<S> {
931    inner: S,
932    service_name: &'static str,
933}
934
935impl<S, B> tower::Service<http::Request<B>> for GrpcTraceContextService<S>
936where
937    S: tower::Service<http::Request<B>> + Clone + Send + 'static,
938    S::Future: Send,
939    B: Send + 'static,
940{
941    type Response = S::Response;
942    type Error = S::Error;
943    type Future = std::pin::Pin<
944        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
945    >;
946
947    fn poll_ready(
948        &mut self,
949        cx: &mut std::task::Context<'_>,
950    ) -> std::task::Poll<Result<(), Self::Error>> {
951        self.inner.poll_ready(cx)
952    }
953
954    fn call(&mut self, request: http::Request<B>) -> Self::Future {
955        use tracing::Instrument;
956        use tracing_opentelemetry::OpenTelemetrySpanExt;
957
958        // Extract trace context from incoming HTTP/2 headers (gRPC uses HTTP/2)
959        let parent_cx = extract_trace_context_http(request.headers());
960
961        // Create a tracing span and link it to the incoming OpenTelemetry context.
962        // This makes all child spans (from #[tracing::instrument]) part of the distributed trace.
963        let span = tracing::info_span!("grpc_request", service = self.service_name);
964        let _ = span.set_parent(parent_cx);
965
966        // Clone inner service for use in async block
967        let mut inner = self.inner.clone();
968
969        // Instrument the future with our span so it stays active for the entire request
970        Box::pin(async move { inner.call(request).await }.instrument(span))
971    }
972}
973
974// ============================================================================
975// Message Queue Trace Context Propagation (for Iggy and similar systems)
976// ============================================================================
977
978impl TraceContextCarrier for std::collections::HashMap<String, String> {
979    fn extract_trace_context(&self) -> opentelemetry::Context {
980        global::get_text_map_propagator(|propagator| propagator.extract(self))
981    }
982
983    fn inject_trace_context(&mut self) {
984        use tracing_opentelemetry::OpenTelemetrySpanExt;
985        let cx = tracing::Span::current().context();
986        global::get_text_map_propagator(|propagator| {
987            propagator.inject_context(&cx, self);
988        });
989    }
990}
991
992/// Inject the current trace context into a HashMap suitable for message queue headers.
993///
994/// This is useful for propagating trace context through message queues like Iggy
995/// where headers are stored as a `HashMap<HeaderKey, HeaderValue>`.
996///
997/// # Example
998///
999/// ```
1000/// use std::collections::HashMap;
1001///
1002/// let mut headers: HashMap<String, String> = HashMap::new();
1003/// byre::telemetry::inject_trace_context_map(&mut headers);
1004/// // headers now contains traceparent key (if there's an active span)
1005/// ```
1006pub fn inject_trace_context_map(headers: &mut std::collections::HashMap<String, String>) {
1007    use tracing_opentelemetry::OpenTelemetrySpanExt;
1008    // Get the OpenTelemetry context from the current tracing span
1009    let cx = tracing::Span::current().context();
1010    global::get_text_map_propagator(|propagator| {
1011        propagator.inject_context(&cx, headers);
1012    });
1013}
1014
1015/// Extract trace context from a HashMap of message queue headers.
1016///
1017/// This is useful for extracting trace context from message queues like Iggy
1018/// where headers are stored as a `HashMap<HeaderKey, HeaderValue>`.
1019///
1020/// # Example
1021///
1022/// ```
1023/// use std::collections::HashMap;
1024///
1025/// let mut headers: HashMap<String, String> = HashMap::new();
1026/// headers.insert("traceparent".to_string(), "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string());
1027///
1028/// let parent_cx = byre::telemetry::extract_trace_context_map(&headers);
1029/// ```
1030pub fn extract_trace_context_map(
1031    headers: &std::collections::HashMap<String, String>,
1032) -> opentelemetry::Context {
1033    global::get_text_map_propagator(|propagator| propagator.extract(headers))
1034}
1035
1036/// Link the current span to an incoming distributed trace from message queue headers.
1037///
1038/// This is a convenience function that extracts the trace context from the
1039/// message headers and sets it as the parent of the current span.
1040///
1041/// Returns `Ok(())` if successful, or an error if the span context couldn't be set.
1042/// Most callers will want to ignore the error with `let _ = link_distributed_trace_map(...)`.
1043///
1044/// # Example
1045///
1046/// ```
1047/// use std::collections::HashMap;
1048///
1049/// let mut headers: HashMap<String, String> = HashMap::new();
1050/// headers.insert("traceparent".to_string(), "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string());
1051///
1052/// let _ = byre::telemetry::link_distributed_trace_map(&headers);
1053/// // Current span is now part of the distributed trace
1054/// ```
1055pub fn link_distributed_trace_map(
1056    headers: &std::collections::HashMap<String, String>,
1057) -> Result<(), Error> {
1058    use tracing_opentelemetry::OpenTelemetrySpanExt;
1059    let parent_cx = extract_trace_context_map(headers);
1060    tracing::Span::current()
1061        .set_parent(parent_cx)
1062        .map_err(|e| Error::LinkDistributedTrace {
1063            source: Box::new(e),
1064        })
1065}
1066
1067/// Set a span's parent from an OpenTelemetry context.
1068///
1069/// This links the given tracing span to a distributed trace context,
1070/// making it a child of the remote span.
1071///
1072/// # Example
1073///
1074/// ```
1075/// use std::collections::HashMap;
1076///
1077/// let mut headers: HashMap<String, String> = HashMap::new();
1078/// headers.insert("traceparent".to_string(), "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string());
1079///
1080/// let parent_cx = byre::telemetry::extract_trace_context_map(&headers);
1081/// let span = tracing::info_span!("process_message", message_id = 42);
1082/// byre::telemetry::set_span_parent(&span, parent_cx);
1083/// let _enter = span.enter();
1084/// ```
1085pub fn set_span_parent(span: &tracing::Span, parent_cx: opentelemetry::Context) {
1086    use tracing_opentelemetry::OpenTelemetrySpanExt;
1087    let _ = span.set_parent(parent_cx);
1088}
1089
1090// ============================================================================
1091// Prelude - convenient re-exports for common use
1092// ============================================================================
1093
1094/// Convenient re-exports for common telemetry usage.
1095///
1096/// Import with:
1097/// ```rust
1098/// use byre::telemetry::prelude::*;
1099/// ```
1100///
1101/// This gives you access to:
1102/// - [`init`] - Initialize telemetry
1103/// - [`TelemetrySettings`] - Configuration for telemetry
1104/// - [`TelemetryProviders`] - Handle to keep telemetry alive
1105/// - [`TraceContextCarrier`] - Trait for types that carry trace context
1106/// - [`TraceContextExt`] - Extension methods for trace context propagation
1107/// - [`GrpcTraceContextLayer`] - Tower layer for gRPC distributed tracing
1108pub mod prelude {
1109    pub use super::{
1110        init, GrpcTraceContextLayer, TelemetryProviders, TelemetrySettings, TraceContextCarrier,
1111        TraceContextExt,
1112    };
1113}
1114
1115#[cfg(test)]
1116mod tests {
1117    use super::*;
1118    use opentelemetry::trace::{
1119        SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState, Tracer,
1120    };
1121    use std::collections::HashMap;
1122
1123    /// Initialize the W3C TraceContext propagator for tests
1124    fn init_test_propagator() {
1125        use opentelemetry::propagation::TextMapCompositePropagator;
1126        use opentelemetry_sdk::propagation::TraceContextPropagator;
1127
1128        let propagator =
1129            TextMapCompositePropagator::new(vec![Box::new(TraceContextPropagator::new())]);
1130        global::set_text_map_propagator(propagator);
1131    }
1132
1133    #[test]
1134    fn test_inject_and_extract_trace_context_roundtrip() {
1135        use tracing_opentelemetry::OpenTelemetrySpanExt;
1136
1137        let _provider = init_tracing_with_otel();
1138
1139        // Create a span context with known values
1140        let trace_id = TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap();
1141        let span_id = SpanId::from_hex("b7ad6b7169203331").unwrap();
1142        let span_context = SpanContext::new(
1143            trace_id,
1144            span_id,
1145            TraceFlags::SAMPLED,
1146            false,
1147            TraceState::default(),
1148        );
1149
1150        // Create an OpenTelemetry context with our span context
1151        let parent_cx = opentelemetry::Context::new().with_remote_span_context(span_context);
1152
1153        // Create a tracing span and set the parent context
1154        let span = tracing::info_span!("test_roundtrip_span");
1155        let _ = span.set_parent(parent_cx);
1156        let _enter = span.enter();
1157
1158        // Inject the trace context into headers
1159        let mut headers: HashMap<String, String> = HashMap::new();
1160        inject_trace_context_map(&mut headers);
1161
1162        // Verify traceparent header was injected
1163        assert!(
1164            headers.contains_key("traceparent"),
1165            "traceparent header should be present"
1166        );
1167        let traceparent = headers.get("traceparent").unwrap();
1168        assert!(
1169            traceparent.starts_with("00-0af7651916cd43dd8448eb211c80319c-"),
1170            "traceparent should contain the correct trace ID"
1171        );
1172
1173        // Extract the trace context from headers
1174        let extracted_cx = extract_trace_context_map(&headers);
1175        let extracted_span = extracted_cx.span();
1176        let extracted_span_context = extracted_span.span_context();
1177
1178        // Verify the extracted context matches the original
1179        assert_eq!(
1180            extracted_span_context.trace_id(),
1181            trace_id,
1182            "trace ID should match"
1183        );
1184        assert!(
1185            extracted_span_context.trace_flags().is_sampled(),
1186            "sampled flag should be preserved"
1187        );
1188    }
1189
1190    #[test]
1191    fn test_extract_empty_headers_returns_empty_context() {
1192        init_test_propagator();
1193
1194        let headers: HashMap<String, String> = HashMap::new();
1195        let extracted_cx = extract_trace_context_map(&headers);
1196
1197        // An empty context should have an invalid span context
1198        let extracted_span = extracted_cx.span();
1199        let span_context = extracted_span.span_context();
1200        assert!(
1201            !span_context.is_valid(),
1202            "extracted context from empty headers should be invalid"
1203        );
1204    }
1205
1206    #[test]
1207    fn test_extract_invalid_traceparent_returns_empty_context() {
1208        init_test_propagator();
1209
1210        let mut headers: HashMap<String, String> = HashMap::new();
1211        headers.insert(
1212            "traceparent".to_string(),
1213            "invalid-traceparent-value".to_string(),
1214        );
1215
1216        let extracted_cx = extract_trace_context_map(&headers);
1217        let extracted_span = extracted_cx.span();
1218        let span_context = extracted_span.span_context();
1219
1220        // Invalid traceparent should result in invalid span context
1221        assert!(
1222            !span_context.is_valid(),
1223            "extracted context from invalid traceparent should be invalid"
1224        );
1225    }
1226
1227    #[test]
1228    fn test_inject_without_active_span_produces_no_traceparent() {
1229        init_test_propagator();
1230
1231        // Without an active span, inject should not add traceparent
1232        let mut headers: HashMap<String, String> = HashMap::new();
1233        inject_trace_context_map(&mut headers);
1234
1235        // When there's no active span with a valid context, traceparent may be empty or missing
1236        // The behavior depends on whether there's a valid span context
1237        if let Some(traceparent) = headers.get("traceparent") {
1238            // If present, it should start with "00-00000000000000000000000000000000" (invalid trace)
1239            assert!(
1240                traceparent.contains("00000000000000000000000000000000") || traceparent.is_empty(),
1241                "traceparent without active span should be empty or have zero trace ID"
1242            );
1243        }
1244    }
1245
1246    #[test]
1247    fn test_trace_context_preserves_tracestate() {
1248        init_test_propagator();
1249
1250        // Create headers with both traceparent and tracestate
1251        let mut headers: HashMap<String, String> = HashMap::new();
1252        headers.insert(
1253            "traceparent".to_string(),
1254            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
1255        );
1256        headers.insert("tracestate".to_string(), "congo=t61rcWkgMzE".to_string());
1257
1258        let extracted_cx = extract_trace_context_map(&headers);
1259        let extracted_span = extracted_cx.span();
1260        let span_context = extracted_span.span_context();
1261
1262        assert!(span_context.is_valid(), "span context should be valid");
1263        // TraceState doesn't have is_empty, check the header value instead
1264        assert!(
1265            !span_context.trace_state().header().is_empty(),
1266            "tracestate should be preserved"
1267        );
1268    }
1269
1270    #[test]
1271    fn test_inject_extract_with_real_span() {
1272        init_test_propagator();
1273
1274        // Create a real tracer and span
1275        let tracer = global::tracer("test-tracer");
1276        let span = tracer.start("test-span");
1277        let cx = opentelemetry::Context::current_with_span(span);
1278
1279        // Get the span context for comparison before attaching
1280        let original_span = cx.span();
1281        let original_trace_id = original_span.span_context().trace_id();
1282
1283        let _guard = cx.attach();
1284
1285        // Inject
1286        let mut headers: HashMap<String, String> = HashMap::new();
1287        inject_trace_context_map(&mut headers);
1288
1289        // Extract
1290        let extracted_cx = extract_trace_context_map(&headers);
1291        let extracted_span = extracted_cx.span();
1292        let extracted_span_context = extracted_span.span_context();
1293
1294        // The trace ID should match (span ID may differ as it's a child reference)
1295        assert_eq!(
1296            extracted_span_context.trace_id(),
1297            original_trace_id,
1298            "trace ID should be preserved through inject/extract cycle"
1299        );
1300    }
1301
1302    // ========================================================================
1303    // Tests for inject_trace_context from tracing spans
1304    // ========================================================================
1305
1306    /// Initialize a tracing subscriber with OpenTelemetry layer for tests.
1307    /// This uses try_init() which only works once per process - use
1308    /// `with_otel_subscriber` for tests that need isolated subscribers.
1309    fn init_tracing_with_otel() -> opentelemetry_sdk::trace::SdkTracerProvider {
1310        use opentelemetry::trace::TracerProvider as _;
1311        use opentelemetry_sdk::trace::SdkTracerProvider;
1312        use tracing_subscriber::layer::SubscriberExt;
1313        use tracing_subscriber::util::SubscriberInitExt;
1314
1315        init_test_propagator();
1316
1317        // Create a simple tracer provider
1318        let provider = SdkTracerProvider::builder().build();
1319        let tracer = provider.tracer("test-tracer");
1320
1321        // Create the OpenTelemetry layer
1322        let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
1323
1324        // Build and set the subscriber
1325        let subscriber = tracing_subscriber::registry().with(otel_layer);
1326
1327        // Use try_init to avoid panics if already initialized
1328        let _ = subscriber.try_init();
1329
1330        // Return provider to keep it alive
1331        provider
1332    }
1333
1334    /// Run a test closure with an isolated OpenTelemetry-enabled tracing subscriber.
1335    /// This uses `with_default` to avoid global subscriber conflicts between tests.
1336    fn with_otel_subscriber<F, R>(f: F) -> R
1337    where
1338        F: FnOnce() -> R,
1339    {
1340        use opentelemetry::trace::TracerProvider as _;
1341        use opentelemetry_sdk::trace::SdkTracerProvider;
1342        use tracing_subscriber::layer::SubscriberExt;
1343
1344        init_test_propagator();
1345
1346        // Create a simple tracer provider
1347        let provider = SdkTracerProvider::builder().build();
1348        let tracer = provider.tracer("test-tracer");
1349
1350        // Create the OpenTelemetry layer
1351        let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
1352
1353        // Build the subscriber (don't set globally)
1354        let subscriber = tracing_subscriber::registry().with(otel_layer);
1355
1356        // Run the test with this subscriber as the default for this scope only
1357        tracing::subscriber::with_default(subscriber, f)
1358    }
1359
1360    /// Helper to assert a traceparent value is valid (non-empty trace ID)
1361    fn assert_valid_traceparent(traceparent: &str) {
1362        assert!(
1363            traceparent.starts_with("00-"),
1364            "traceparent should start with version 00"
1365        );
1366        assert!(
1367            !traceparent.contains("00000000000000000000000000000000"),
1368            "traceparent should have a valid (non-zero) trace ID"
1369        );
1370    }
1371
1372    #[test]
1373    fn test_inject_trace_context_from_tracing_span() {
1374        let _provider = init_tracing_with_otel();
1375        let span = tracing::info_span!("test_span_for_injection");
1376        let _enter = span.enter();
1377
1378        // Test HashMap injection
1379        let mut map_headers: HashMap<String, String> = HashMap::new();
1380        inject_trace_context_map(&mut map_headers);
1381        assert!(map_headers.contains_key("traceparent"));
1382        assert_valid_traceparent(map_headers.get("traceparent").unwrap());
1383
1384        // Test gRPC metadata injection
1385        let mut grpc_metadata = tonic::metadata::MetadataMap::new();
1386        inject_trace_context(&mut grpc_metadata);
1387        let grpc_traceparent = grpc_metadata.get("traceparent").unwrap().to_str().unwrap();
1388        assert_valid_traceparent(grpc_traceparent);
1389
1390        // Test HTTP header injection
1391        let mut http_headers = http::HeaderMap::new();
1392        inject_trace_context_http(&mut http_headers);
1393        let http_traceparent = http_headers.get("traceparent").unwrap().to_str().unwrap();
1394        assert_valid_traceparent(http_traceparent);
1395    }
1396
1397    #[test]
1398    fn test_nested_tracing_spans_propagate_trace_id() {
1399        let _provider = init_tracing_with_otel();
1400
1401        // Create parent span
1402        let parent_span = tracing::info_span!("parent_span");
1403        let _parent_enter = parent_span.enter();
1404
1405        // Get trace context from parent
1406        let mut parent_headers: HashMap<String, String> = HashMap::new();
1407        inject_trace_context_map(&mut parent_headers);
1408        let parent_traceparent = parent_headers.get("traceparent").unwrap().clone();
1409
1410        // Extract trace ID from parent (format: 00-{trace_id}-{span_id}-{flags})
1411        let parent_trace_id: String = parent_traceparent.split('-').nth(1).unwrap().to_string();
1412
1413        // Create child span
1414        let child_span = tracing::info_span!("child_span");
1415        let _child_enter = child_span.enter();
1416
1417        // Get trace context from child
1418        let mut child_headers: HashMap<String, String> = HashMap::new();
1419        inject_trace_context_map(&mut child_headers);
1420        let child_traceparent = child_headers.get("traceparent").unwrap();
1421
1422        // Extract trace ID from child
1423        let child_trace_id: String = child_traceparent.split('-').nth(1).unwrap().to_string();
1424
1425        // Both spans should have the same trace ID
1426        assert_eq!(
1427            parent_trace_id, child_trace_id,
1428            "nested spans should share the same trace ID"
1429        );
1430
1431        // But span IDs should be different
1432        let parent_span_id: String = parent_traceparent.split('-').nth(2).unwrap().to_string();
1433        let child_span_id: String = child_traceparent.split('-').nth(2).unwrap().to_string();
1434
1435        assert_ne!(
1436            parent_span_id, child_span_id,
1437            "nested spans should have different span IDs"
1438        );
1439    }
1440
1441    #[test]
1442    fn test_inject_trace_context_roundtrip_with_tracing_span() {
1443        let _provider = init_tracing_with_otel();
1444
1445        // Create a tracing span
1446        let span = tracing::info_span!("roundtrip_test_span");
1447        let _enter = span.enter();
1448
1449        // Inject into HashMap
1450        let mut headers: HashMap<String, String> = HashMap::new();
1451        inject_trace_context_map(&mut headers);
1452
1453        // Extract the context
1454        let extracted_cx = extract_trace_context_map(&headers);
1455        let extracted_span = extracted_cx.span();
1456        let extracted_span_context = extracted_span.span_context();
1457
1458        // Verify we got a valid span context
1459        assert!(
1460            extracted_span_context.is_valid(),
1461            "extracted span context should be valid"
1462        );
1463
1464        // Get the original trace ID from the injected headers
1465        let traceparent = headers.get("traceparent").unwrap();
1466        let original_trace_id: String = traceparent.split('-').nth(1).unwrap().to_string();
1467
1468        // Compare with extracted trace ID
1469        let extracted_trace_id = format!("{:032x}", extracted_span_context.trace_id());
1470
1471        assert_eq!(
1472            original_trace_id, extracted_trace_id,
1473            "trace ID should survive inject/extract roundtrip"
1474        );
1475    }
1476
1477    // ========================================================================
1478    // Tests for MetadataExtractor (gRPC metadata)
1479    // ========================================================================
1480
1481    #[test]
1482    fn test_metadata_extractor_get_returns_value() {
1483        let mut metadata = tonic::metadata::MetadataMap::new();
1484        metadata.insert("traceparent", "00-abc123-def456-01".parse().unwrap());
1485
1486        let extractor = MetadataExtractor(&metadata);
1487        let value = extractor.get("traceparent");
1488
1489        assert!(value.is_some(), "get should return Some for existing key");
1490        assert_eq!(value.unwrap(), "00-abc123-def456-01");
1491    }
1492
1493    #[test]
1494    fn test_metadata_extractor_get_returns_none_for_missing() {
1495        let metadata = tonic::metadata::MetadataMap::new();
1496        let extractor = MetadataExtractor(&metadata);
1497
1498        let value = extractor.get("nonexistent");
1499        assert!(value.is_none(), "get should return None for missing key");
1500    }
1501
1502    #[test]
1503    fn test_metadata_extractor_keys_returns_trace_context_keys() {
1504        let mut metadata = tonic::metadata::MetadataMap::new();
1505        metadata.insert("traceparent", "value1".parse().unwrap());
1506        metadata.insert("tracestate", "value2".parse().unwrap());
1507        metadata.insert("custom-header", "value3".parse().unwrap());
1508
1509        let extractor = MetadataExtractor(&metadata);
1510        let keys = extractor.keys();
1511
1512        // keys() only returns W3C Trace Context keys, not all headers
1513        assert_eq!(keys.len(), 2, "keys should only contain trace context keys");
1514        assert!(
1515            keys.contains(&"traceparent"),
1516            "keys should contain traceparent"
1517        );
1518        assert!(
1519            keys.contains(&"tracestate"),
1520            "keys should contain tracestate"
1521        );
1522    }
1523
1524    #[test]
1525    fn test_metadata_extractor_keys_returns_only_present_keys() {
1526        let mut metadata = tonic::metadata::MetadataMap::new();
1527        metadata.insert("traceparent", "value1".parse().unwrap());
1528        // tracestate not present
1529
1530        let extractor = MetadataExtractor(&metadata);
1531        let keys = extractor.keys();
1532
1533        assert_eq!(
1534            keys.len(),
1535            1,
1536            "keys should only contain present trace context keys"
1537        );
1538        assert!(
1539            keys.contains(&"traceparent"),
1540            "keys should contain traceparent"
1541        );
1542    }
1543
1544    // ========================================================================
1545    // Tests for HttpHeaderExtractor (HTTP headers)
1546    // ========================================================================
1547
1548    #[test]
1549    fn test_http_header_extractor_get_returns_value() {
1550        let mut headers = http::HeaderMap::new();
1551        headers.insert("traceparent", "00-abc123-def456-01".parse().unwrap());
1552
1553        let extractor = HttpHeaderExtractor(&headers);
1554        let value = extractor.get("traceparent");
1555
1556        assert!(value.is_some(), "get should return Some for existing key");
1557        assert_eq!(value.unwrap(), "00-abc123-def456-01");
1558    }
1559
1560    #[test]
1561    fn test_http_header_extractor_get_returns_none_for_missing() {
1562        let headers = http::HeaderMap::new();
1563        let extractor = HttpHeaderExtractor(&headers);
1564
1565        let value = extractor.get("nonexistent");
1566        assert!(value.is_none(), "get should return None for missing key");
1567    }
1568
1569    #[test]
1570    fn test_http_header_extractor_keys_returns_trace_context_keys() {
1571        let mut headers = http::HeaderMap::new();
1572        headers.insert("traceparent", "value1".parse().unwrap());
1573        headers.insert("tracestate", "value2".parse().unwrap());
1574        headers.insert("x-custom-header", "value3".parse().unwrap());
1575
1576        let extractor = HttpHeaderExtractor(&headers);
1577        let keys = extractor.keys();
1578
1579        // keys() only returns W3C Trace Context keys, not all headers
1580        assert_eq!(keys.len(), 2, "keys should only contain trace context keys");
1581        assert!(
1582            keys.contains(&"traceparent"),
1583            "keys should contain traceparent"
1584        );
1585        assert!(
1586            keys.contains(&"tracestate"),
1587            "keys should contain tracestate"
1588        );
1589    }
1590
1591    #[test]
1592    fn test_http_header_extractor_keys_returns_only_present_keys() {
1593        let mut headers = http::HeaderMap::new();
1594        headers.insert("traceparent", "value1".parse().unwrap());
1595        // tracestate not present
1596
1597        let extractor = HttpHeaderExtractor(&headers);
1598        let keys = extractor.keys();
1599
1600        assert_eq!(
1601            keys.len(),
1602            1,
1603            "keys should only contain present trace context keys"
1604        );
1605        assert!(
1606            keys.contains(&"traceparent"),
1607            "keys should contain traceparent"
1608        );
1609    }
1610
1611    // ========================================================================
1612    // Tests for extract_trace_context functions
1613    // ========================================================================
1614
1615    /// Helper to verify extracted context has expected trace ID
1616    fn assert_extracted_trace_id(context: &opentelemetry::Context, expected_trace_id: &str) {
1617        let span = context.span();
1618        let span_context = span.span_context();
1619        assert!(span_context.is_valid(), "span context should be valid");
1620        assert_eq!(
1621            format!("{:032x}", span_context.trace_id()),
1622            expected_trace_id,
1623            "trace ID should match"
1624        );
1625    }
1626
1627    #[test]
1628    fn test_extract_trace_context_grpc_with_valid_headers() {
1629        init_test_propagator();
1630
1631        let mut metadata = tonic::metadata::MetadataMap::new();
1632        metadata.insert(
1633            "traceparent",
1634            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
1635                .parse()
1636                .unwrap(),
1637        );
1638
1639        let context = extract_trace_context(&metadata);
1640        assert_extracted_trace_id(&context, "0af7651916cd43dd8448eb211c80319c");
1641    }
1642
1643    #[test]
1644    fn test_extract_trace_context_http_with_valid_headers() {
1645        init_test_propagator();
1646
1647        let mut headers = http::HeaderMap::new();
1648        headers.insert(
1649            "traceparent",
1650            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
1651                .parse()
1652                .unwrap(),
1653        );
1654
1655        let context = extract_trace_context_http(&headers);
1656        assert_extracted_trace_id(&context, "0af7651916cd43dd8448eb211c80319c");
1657    }
1658
1659    #[test]
1660    fn test_extract_trace_context_map_with_valid_headers() {
1661        init_test_propagator();
1662
1663        let mut headers: HashMap<String, String> = HashMap::new();
1664        headers.insert(
1665            "traceparent".to_string(),
1666            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01".to_string(),
1667        );
1668
1669        let context = extract_trace_context_map(&headers);
1670        assert_extracted_trace_id(&context, "0af7651916cd43dd8448eb211c80319c");
1671    }
1672
1673    // ========================================================================
1674    // Tests for set_span_parent
1675    // ========================================================================
1676
1677    #[test]
1678    fn test_set_span_parent_actually_links() {
1679        // Use with_otel_subscriber for isolated per-test subscriber
1680        with_otel_subscriber(|| {
1681            let trace_id = TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap();
1682            let span_id = SpanId::from_hex("b7ad6b7169203331").unwrap();
1683            let remote_span_context = SpanContext::new(
1684                trace_id,
1685                span_id,
1686                TraceFlags::SAMPLED,
1687                true,
1688                TraceState::default(),
1689            );
1690            let parent_cx =
1691                opentelemetry::Context::new().with_remote_span_context(remote_span_context);
1692
1693            let span = tracing::info_span!("test_set_parent");
1694            set_span_parent(&span, parent_cx);
1695            let _enter = span.enter();
1696
1697            let mut headers: HashMap<String, String> = HashMap::new();
1698            inject_trace_context_map(&mut headers);
1699
1700            let traceparent = headers.get("traceparent").unwrap();
1701            assert!(
1702                traceparent.contains("0af7651916cd43dd8448eb211c80319c"),
1703                "trace ID should be preserved after set_span_parent"
1704            );
1705        });
1706    }
1707
1708    #[test]
1709    fn test_link_distributed_trace_grpc_actually_links() {
1710        // Test that link_distributed_trace extracts context and calls set_parent.
1711        // Uses TraceContextExt trait which works on a specific span reference.
1712        with_otel_subscriber(|| {
1713            let mut metadata = tonic::metadata::MetadataMap::new();
1714            metadata.insert(
1715                "traceparent",
1716                "00-11111111111111111111111111111111-aaaaaaaaaaaaaaaa-01"
1717                    .parse()
1718                    .unwrap(),
1719            );
1720
1721            // Use the trait method which extracts and links
1722            let span = tracing::info_span!("test_link_grpc");
1723
1724            // First extract the context
1725            let parent_cx = metadata.extract_trace_context();
1726
1727            // Then set it as parent (same as what link_distributed_trace does internally)
1728            set_span_parent(&span, parent_cx);
1729            let _enter = span.enter();
1730
1731            // Verify the span is now linked by injecting and checking trace ID
1732            let mut verify: HashMap<String, String> = HashMap::new();
1733            inject_trace_context_map(&mut verify);
1734
1735            let traceparent = verify.get("traceparent").unwrap();
1736            assert!(
1737                traceparent.contains("11111111111111111111111111111111"),
1738                "link_distributed_trace should link the span to trace 11111111111111111111111111111111, got {traceparent}"
1739            );
1740        });
1741    }
1742
1743    #[test]
1744    fn test_link_distributed_trace_http_actually_links() {
1745        with_otel_subscriber(|| {
1746            let mut headers = http::HeaderMap::new();
1747            headers.insert(
1748                "traceparent",
1749                "00-22222222222222222222222222222222-bbbbbbbbbbbbbbbb-01"
1750                    .parse()
1751                    .unwrap(),
1752            );
1753
1754            let span = tracing::info_span!("test_link_http");
1755            let parent_cx = headers.extract_trace_context();
1756            set_span_parent(&span, parent_cx);
1757            let _enter = span.enter();
1758
1759            let mut verify: HashMap<String, String> = HashMap::new();
1760            inject_trace_context_map(&mut verify);
1761
1762            let traceparent = verify.get("traceparent").unwrap();
1763            assert!(
1764                traceparent.contains("22222222222222222222222222222222"),
1765                "link_distributed_trace_http should link the span"
1766            );
1767        });
1768    }
1769
1770    #[test]
1771    fn test_link_distributed_trace_map_actually_links() {
1772        with_otel_subscriber(|| {
1773            let mut headers: HashMap<String, String> = HashMap::new();
1774            headers.insert(
1775                "traceparent".to_string(),
1776                "00-33333333333333333333333333333333-cccccccccccccccc-01".to_string(),
1777            );
1778
1779            let span = tracing::info_span!("test_link_map");
1780            let parent_cx = headers.extract_trace_context();
1781            set_span_parent(&span, parent_cx);
1782            let _enter = span.enter();
1783
1784            let mut verify: HashMap<String, String> = HashMap::new();
1785            inject_trace_context_map(&mut verify);
1786
1787            let traceparent = verify.get("traceparent").unwrap();
1788            assert!(
1789                traceparent.contains("33333333333333333333333333333333"),
1790                "link_distributed_trace_map should link the span"
1791            );
1792        });
1793    }
1794
1795    // ========================================================================
1796    // Tests for init_propagator
1797    // ========================================================================
1798
1799    #[test]
1800    fn test_init_propagator_enables_trace_context_propagation() {
1801        // Call init_propagator to set the W3C TraceContext propagator
1802        init_propagator();
1803
1804        // Create a context with a known trace ID
1805        let trace_id = TraceId::from_hex("1234567890abcdef1234567890abcdef").unwrap();
1806        let span_id = SpanId::from_hex("fedcba0987654321").unwrap();
1807        let span_context = SpanContext::new(
1808            trace_id,
1809            span_id,
1810            TraceFlags::SAMPLED,
1811            false,
1812            TraceState::default(),
1813        );
1814        let cx = opentelemetry::Context::new().with_remote_span_context(span_context);
1815
1816        // Inject using the propagator
1817        let mut headers: HashMap<String, String> = HashMap::new();
1818        global::get_text_map_propagator(|propagator| {
1819            propagator.inject_context(&cx, &mut headers);
1820        });
1821
1822        // Verify traceparent was injected (this would fail if init_propagator did nothing)
1823        assert!(
1824            headers.contains_key("traceparent"),
1825            "init_propagator should enable traceparent injection"
1826        );
1827        let traceparent = headers.get("traceparent").unwrap();
1828        assert!(
1829            traceparent.contains("1234567890abcdef1234567890abcdef"),
1830            "traceparent should contain the trace ID"
1831        );
1832    }
1833
1834    // ========================================================================
1835    // Tests for TraceContextCarrier::extract_trace_context implementations
1836    // ========================================================================
1837
1838    #[test]
1839    fn test_metadata_map_extract_trace_context_returns_valid_context() {
1840        init_test_propagator();
1841
1842        let mut metadata = tonic::metadata::MetadataMap::new();
1843        metadata.insert(
1844            "traceparent",
1845            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
1846                .parse()
1847                .unwrap(),
1848        );
1849
1850        // Use the TraceContextCarrier trait method
1851        let context = TraceContextCarrier::extract_trace_context(&metadata);
1852        let span = context.span();
1853        let span_context = span.span_context();
1854
1855        // Verify this is NOT a default context - it has the trace ID from headers
1856        assert!(span_context.is_valid(), "span context should be valid");
1857        assert_eq!(
1858            format!("{:032x}", span_context.trace_id()),
1859            "0af7651916cd43dd8448eb211c80319c",
1860            "trace ID should be extracted from headers, not default"
1861        );
1862    }
1863
1864    #[test]
1865    fn test_http_header_map_extract_trace_context_returns_valid_context() {
1866        init_test_propagator();
1867
1868        let mut headers = http::HeaderMap::new();
1869        headers.insert(
1870            "traceparent",
1871            "00-1234567890abcdef1234567890abcdef-b7ad6b7169203331-01"
1872                .parse()
1873                .unwrap(),
1874        );
1875
1876        // Use the TraceContextCarrier trait method
1877        let context = TraceContextCarrier::extract_trace_context(&headers);
1878        let span = context.span();
1879        let span_context = span.span_context();
1880
1881        // Verify this is NOT a default context - it has the trace ID from headers
1882        assert!(span_context.is_valid(), "span context should be valid");
1883        assert_eq!(
1884            format!("{:032x}", span_context.trace_id()),
1885            "1234567890abcdef1234567890abcdef",
1886            "trace ID should be extracted from headers, not default"
1887        );
1888    }
1889
1890    #[test]
1891    fn test_hashmap_extract_trace_context_returns_valid_context() {
1892        init_test_propagator();
1893
1894        let mut headers: HashMap<String, String> = HashMap::new();
1895        headers.insert(
1896            "traceparent".to_string(),
1897            "00-abcdef1234567890abcdef1234567890-b7ad6b7169203331-01".to_string(),
1898        );
1899
1900        // Use the TraceContextCarrier trait method
1901        let context = TraceContextCarrier::extract_trace_context(&headers);
1902        let span = context.span();
1903        let span_context = span.span_context();
1904
1905        // Verify this is NOT a default context - it has the trace ID from headers
1906        assert!(span_context.is_valid(), "span context should be valid");
1907        assert_eq!(
1908            format!("{:032x}", span_context.trace_id()),
1909            "abcdef1234567890abcdef1234567890",
1910            "trace ID should be extracted from headers, not default"
1911        );
1912    }
1913
1914    // ========================================================================
1915    // Tests for TraceContextCarrier::inject_trace_context implementations
1916    // ========================================================================
1917
1918    #[test]
1919    fn test_metadata_map_inject_trace_context_modifies_carrier() {
1920        let _provider = init_tracing_with_otel();
1921
1922        // Create a span with a known trace ID
1923        let trace_id = TraceId::from_hex("fedcba9876543210fedcba9876543210").unwrap();
1924        let span_id = SpanId::from_hex("1234567890abcdef").unwrap();
1925        let span_context = SpanContext::new(
1926            trace_id,
1927            span_id,
1928            TraceFlags::SAMPLED,
1929            false,
1930            TraceState::default(),
1931        );
1932        let parent_cx = opentelemetry::Context::new().with_remote_span_context(span_context);
1933
1934        let span = tracing::info_span!("test_grpc_inject");
1935        {
1936            use tracing_opentelemetry::OpenTelemetrySpanExt;
1937            let _ = span.set_parent(parent_cx);
1938        }
1939        let _enter = span.enter();
1940
1941        // Use the TraceContextCarrier trait method
1942        let mut metadata = tonic::metadata::MetadataMap::new();
1943        assert!(
1944            metadata.get("traceparent").is_none(),
1945            "metadata should start empty"
1946        );
1947
1948        TraceContextCarrier::inject_trace_context(&mut metadata);
1949
1950        // Verify injection actually modified the carrier
1951        assert!(
1952            metadata.get("traceparent").is_some(),
1953            "inject_trace_context should add traceparent"
1954        );
1955        let traceparent = metadata.get("traceparent").unwrap().to_str().unwrap();
1956        assert!(
1957            traceparent.contains("fedcba9876543210fedcba9876543210"),
1958            "injected traceparent should contain the trace ID"
1959        );
1960    }
1961
1962    #[test]
1963    fn test_http_header_map_inject_trace_context_modifies_carrier() {
1964        let _provider = init_tracing_with_otel();
1965
1966        // Create a span with a known trace ID
1967        let trace_id = TraceId::from_hex("11111111111111111111111111111111").unwrap();
1968        let span_id = SpanId::from_hex("2222222222222222").unwrap();
1969        let span_context = SpanContext::new(
1970            trace_id,
1971            span_id,
1972            TraceFlags::SAMPLED,
1973            false,
1974            TraceState::default(),
1975        );
1976        let parent_cx = opentelemetry::Context::new().with_remote_span_context(span_context);
1977
1978        let span = tracing::info_span!("test_http_inject");
1979        {
1980            use tracing_opentelemetry::OpenTelemetrySpanExt;
1981            let _ = span.set_parent(parent_cx);
1982        }
1983        let _enter = span.enter();
1984
1985        // Use the TraceContextCarrier trait method
1986        let mut headers = http::HeaderMap::new();
1987        assert!(
1988            headers.get("traceparent").is_none(),
1989            "headers should start empty"
1990        );
1991
1992        TraceContextCarrier::inject_trace_context(&mut headers);
1993
1994        // Verify injection actually modified the carrier
1995        assert!(
1996            headers.get("traceparent").is_some(),
1997            "inject_trace_context should add traceparent"
1998        );
1999        let traceparent = headers.get("traceparent").unwrap().to_str().unwrap();
2000        assert!(
2001            traceparent.contains("11111111111111111111111111111111"),
2002            "injected traceparent should contain the trace ID"
2003        );
2004    }
2005
2006    #[test]
2007    fn test_hashmap_inject_trace_context_modifies_carrier() {
2008        let _provider = init_tracing_with_otel();
2009
2010        // Create a span with a known trace ID
2011        let trace_id = TraceId::from_hex("33333333333333333333333333333333").unwrap();
2012        let span_id = SpanId::from_hex("4444444444444444").unwrap();
2013        let span_context = SpanContext::new(
2014            trace_id,
2015            span_id,
2016            TraceFlags::SAMPLED,
2017            false,
2018            TraceState::default(),
2019        );
2020        let parent_cx = opentelemetry::Context::new().with_remote_span_context(span_context);
2021
2022        let span = tracing::info_span!("test_map_inject");
2023        {
2024            use tracing_opentelemetry::OpenTelemetrySpanExt;
2025            let _ = span.set_parent(parent_cx);
2026        }
2027        let _enter = span.enter();
2028
2029        // Use the TraceContextCarrier trait method
2030        let mut headers: HashMap<String, String> = HashMap::new();
2031        assert!(
2032            headers.get("traceparent").is_none(),
2033            "headers should start empty"
2034        );
2035
2036        TraceContextCarrier::inject_trace_context(&mut headers);
2037
2038        // Verify injection actually modified the carrier
2039        assert!(
2040            headers.get("traceparent").is_some(),
2041            "inject_trace_context should add traceparent"
2042        );
2043        let traceparent = headers.get("traceparent").unwrap();
2044        assert!(
2045            traceparent.contains("33333333333333333333333333333333"),
2046            "injected traceparent should contain the trace ID"
2047        );
2048    }
2049
2050    // ========================================================================
2051    // Tests for link_distributed_trace functions
2052    // ========================================================================
2053
2054    #[test]
2055    fn test_link_distributed_trace_grpc_calls_set_parent() {
2056        // This test verifies that link_distributed_trace actually calls set_parent
2057        // by using a mock-like approach: we verify the extraction happens and
2058        // the function attempts to link (even if it errors due to no OTel layer)
2059        init_test_propagator();
2060
2061        let mut metadata = tonic::metadata::MetadataMap::new();
2062        metadata.insert(
2063            "traceparent",
2064            "00-aaaabbbbccccddddaaaabbbbccccdddd-1111222233334444-01"
2065                .parse()
2066                .unwrap(),
2067        );
2068
2069        // Call link_distributed_trace - it extracts context and calls set_parent
2070        // The result depends on whether an OTel layer is registered
2071        let _ = link_distributed_trace(&metadata);
2072
2073        // To verify the function does something (not just returns Ok(())),
2074        // we verify that extract_trace_context (which it calls internally)
2075        // returns the correct context. If the mutation replaced the body with
2076        // Ok(()), the function wouldn't extract or link anything.
2077        //
2078        // The actual linking verification is done by test_set_span_parent_links_trace
2079        // which tests the same code path with a properly initialized OTel layer.
2080        let ctx = extract_trace_context(&metadata);
2081        let span_ref = ctx.span();
2082        let span_context = span_ref.span_context();
2083        assert!(span_context.is_valid());
2084        assert_eq!(
2085            format!("{:032x}", span_context.trace_id()),
2086            "aaaabbbbccccddddaaaabbbbccccdddd"
2087        );
2088    }
2089
2090    #[test]
2091    fn test_link_distributed_trace_http_extracts_and_links() {
2092        init_test_propagator();
2093
2094        let mut headers = http::HeaderMap::new();
2095        headers.insert(
2096            "traceparent",
2097            "00-55556666777788885555666677778888-9999aaaabbbbcccc-01"
2098                .parse()
2099                .unwrap(),
2100        );
2101
2102        // The function should execute the extraction and linking logic
2103        let _ = link_distributed_trace_http(&headers);
2104
2105        // Verify that extraction happened
2106        let ctx = extract_trace_context_http(&headers);
2107        let span = ctx.span();
2108        let span_context = span.span_context();
2109        assert!(span_context.is_valid());
2110        assert_eq!(
2111            format!("{:032x}", span_context.trace_id()),
2112            "55556666777788885555666677778888"
2113        );
2114    }
2115
2116    #[test]
2117    fn test_link_distributed_trace_map_extracts_and_links() {
2118        init_test_propagator();
2119
2120        let mut headers: HashMap<String, String> = HashMap::new();
2121        headers.insert(
2122            "traceparent".to_string(),
2123            "00-ddddeeeeffff0000ddddeeeeffff0000-1234567890abcdef-01".to_string(),
2124        );
2125
2126        // The function should execute the extraction and linking logic
2127        let _ = link_distributed_trace_map(&headers);
2128
2129        // Verify that extraction happened
2130        let ctx = extract_trace_context_map(&headers);
2131        let span = ctx.span();
2132        let span_context = span.span_context();
2133        assert!(span_context.is_valid());
2134        assert_eq!(
2135            format!("{:032x}", span_context.trace_id()),
2136            "ddddeeeeffff0000ddddeeeeffff0000"
2137        );
2138    }
2139
2140    #[test]
2141    fn test_trace_context_ext_link_distributed_trace_extracts_context() {
2142        init_test_propagator();
2143
2144        let mut headers = http::HeaderMap::new();
2145        headers.insert(
2146            "traceparent",
2147            "00-11112222333344441111222233334444-5555666677778888-01"
2148                .parse()
2149                .unwrap(),
2150        );
2151
2152        // Use TraceContextExt trait method - it may error without OTel layer
2153        let _ = super::TraceContextExt::link_distributed_trace(&headers);
2154
2155        // Verify extraction works via the trait method
2156        let ctx = TraceContextCarrier::extract_trace_context(&headers);
2157        let span = ctx.span();
2158        let span_context = span.span_context();
2159        assert!(span_context.is_valid());
2160        assert_eq!(
2161            format!("{:032x}", span_context.trace_id()),
2162            "11112222333344441111222233334444"
2163        );
2164    }
2165
2166    // ========================================================================
2167    // Tests for MetadataInjector
2168    // ========================================================================
2169
2170    #[test]
2171    fn test_metadata_injector_set_adds_header() {
2172        let mut metadata = tonic::metadata::MetadataMap::new();
2173        {
2174            let mut injector = MetadataInjector(&mut metadata);
2175            injector.set("traceparent", "00-test-value-01".to_string());
2176        }
2177
2178        assert!(
2179            metadata.get("traceparent").is_some(),
2180            "set should add the header"
2181        );
2182        assert_eq!(
2183            metadata.get("traceparent").unwrap().to_str().unwrap(),
2184            "00-test-value-01"
2185        );
2186    }
2187
2188    #[test]
2189    fn test_metadata_injector_set_handles_invalid_key_gracefully() {
2190        let mut metadata = tonic::metadata::MetadataMap::new();
2191        {
2192            let mut injector = MetadataInjector(&mut metadata);
2193            // Invalid header name with spaces - should not panic
2194            injector.set("invalid key with spaces", "value".to_string());
2195        }
2196
2197        // The invalid key should not be added
2198        assert!(
2199            metadata.is_empty(),
2200            "invalid header keys should be handled gracefully"
2201        );
2202    }
2203
2204    // ========================================================================
2205    // Tests for HttpHeaderInjector
2206    // ========================================================================
2207
2208    #[test]
2209    fn test_http_header_injector_set_adds_header() {
2210        let mut headers = http::HeaderMap::new();
2211        {
2212            let mut injector = HttpHeaderInjector(&mut headers);
2213            injector.set("traceparent", "00-http-test-01".to_string());
2214        }
2215
2216        assert!(
2217            headers.get("traceparent").is_some(),
2218            "set should add the header"
2219        );
2220        assert_eq!(
2221            headers.get("traceparent").unwrap().to_str().unwrap(),
2222            "00-http-test-01"
2223        );
2224    }
2225
2226    #[test]
2227    fn test_http_header_injector_set_handles_invalid_key_gracefully() {
2228        let mut headers = http::HeaderMap::new();
2229        {
2230            let mut injector = HttpHeaderInjector(&mut headers);
2231            // Invalid header name with spaces - should not panic
2232            injector.set("invalid key", "value".to_string());
2233        }
2234
2235        // The invalid key should not be added
2236        assert!(
2237            headers.is_empty(),
2238            "invalid header keys should be handled gracefully"
2239        );
2240    }
2241
2242    // ========================================================================
2243    // Tests for init_otel_logs_builder
2244    // ========================================================================
2245
2246    #[test]
2247    fn test_init_otel_logs_builder_returns_configured_builder() {
2248        // Create a tokio runtime for the async exporter
2249        let rt = tokio::runtime::Runtime::new().unwrap();
2250        rt.block_on(async {
2251            let service_info = crate::ServiceInfo {
2252                name: "test-service",
2253                name_in_metrics: "test_service".to_string(),
2254                version: "1.0.0",
2255                author: "Test",
2256                description: "Test service",
2257            };
2258
2259            // Use a dummy endpoint - the builder doesn't connect until export
2260            let endpoint = "http://localhost:4317".to_string();
2261
2262            let result = super::init_otel_logs_builder(&service_info, &endpoint);
2263
2264            // The function should succeed and return a configured builder
2265            assert!(
2266                result.is_ok(),
2267                "init_otel_logs_builder should return Ok with valid endpoint"
2268            );
2269
2270            // Build the provider to verify configuration was applied
2271            let builder = result.unwrap();
2272            let provider = builder.build();
2273
2274            // If the builder was Default::default(), the provider wouldn't have
2275            // the exporter or resource configured. We can verify by checking
2276            // that shutdown succeeds (it would fail differently if misconfigured)
2277            let shutdown_result = provider.shutdown();
2278            assert!(
2279                shutdown_result.is_ok(),
2280                "provider built from configured builder should shutdown cleanly"
2281            );
2282        });
2283    }
2284
2285    // ========================================================================
2286    // Tests for init_traces and init_metrics
2287    // ========================================================================
2288
2289    #[test]
2290    fn test_init_traces_with_endpoint_returns_provider() {
2291        let rt = tokio::runtime::Runtime::new().unwrap();
2292        rt.block_on(async {
2293            let service_info = crate::ServiceInfo {
2294                name: "test-service",
2295                name_in_metrics: "test_service".to_string(),
2296                version: "1.0.0",
2297                author: "Test",
2298                description: "Test service",
2299            };
2300
2301            let settings = TraceSettings {
2302                endpoint: Some("http://localhost:4317".to_string()),
2303            };
2304
2305            let result = super::init_traces(&service_info, &settings);
2306
2307            assert!(result.is_ok(), "init_traces should succeed");
2308            let provider = result.unwrap();
2309            assert!(
2310                provider.is_some(),
2311                "init_traces should return Some(provider) when endpoint is configured"
2312            );
2313
2314            // Clean up
2315            if let Some(p) = provider {
2316                let _ = p.shutdown();
2317            }
2318        });
2319    }
2320
2321    #[test]
2322    fn test_init_traces_without_endpoint_returns_none() {
2323        let service_info = crate::ServiceInfo {
2324            name: "test-service",
2325            name_in_metrics: "test_service".to_string(),
2326            version: "1.0.0",
2327            author: "Test",
2328            description: "Test service",
2329        };
2330
2331        let settings = TraceSettings { endpoint: None };
2332
2333        let result = super::init_traces(&service_info, &settings);
2334
2335        assert!(result.is_ok(), "init_traces should succeed");
2336        let provider = result.unwrap();
2337        assert!(
2338            provider.is_none(),
2339            "init_traces should return None when endpoint is not configured"
2340        );
2341    }
2342
2343    #[test]
2344    fn test_init_metrics_with_endpoint_returns_provider() {
2345        let rt = tokio::runtime::Runtime::new().unwrap();
2346        rt.block_on(async {
2347            let service_info = crate::ServiceInfo {
2348                name: "test-service",
2349                name_in_metrics: "test_service".to_string(),
2350                version: "1.0.0",
2351                author: "Test",
2352                description: "Test service",
2353            };
2354
2355            let settings = MetricSettings {
2356                endpoint: Some("http://localhost:4317".to_string()),
2357            };
2358
2359            let result = super::init_metrics(&service_info, &settings);
2360
2361            assert!(result.is_ok(), "init_metrics should succeed");
2362            let provider = result.unwrap();
2363            assert!(
2364                provider.is_some(),
2365                "init_metrics should return Some(provider) when endpoint is configured"
2366            );
2367
2368            // Clean up
2369            if let Some(p) = provider {
2370                let _ = p.shutdown();
2371            }
2372        });
2373    }
2374
2375    #[test]
2376    fn test_init_metrics_without_endpoint_returns_none() {
2377        let service_info = crate::ServiceInfo {
2378            name: "test-service",
2379            name_in_metrics: "test_service".to_string(),
2380            version: "1.0.0",
2381            author: "Test",
2382            description: "Test service",
2383        };
2384
2385        let settings = MetricSettings { endpoint: None };
2386
2387        let result = super::init_metrics(&service_info, &settings);
2388
2389        assert!(result.is_ok(), "init_metrics should succeed");
2390        let provider = result.unwrap();
2391        assert!(
2392            provider.is_none(),
2393            "init_metrics should return None when endpoint is not configured"
2394        );
2395    }
2396
2397    // ========================================================================
2398    // Tests for LogSubscriberBuilder
2399    // ========================================================================
2400
2401    #[test]
2402    fn test_log_subscriber_builder_new_sets_fields() {
2403        let service_info = crate::ServiceInfo {
2404            name: "test-service",
2405            name_in_metrics: "test_service".to_string(),
2406            version: "1.0.0",
2407            author: "Test",
2408            description: "Test service",
2409        };
2410
2411        let settings = LogSettings {
2412            console_level: "info".to_string(),
2413            otel_level: "info".to_string(),
2414            endpoint: None,
2415        };
2416
2417        let builder = super::LogSubscriberBuilder::new(&service_info, &settings);
2418
2419        // Verify the builder captured the references
2420        assert_eq!(builder.service_info.name, "test-service");
2421        assert_eq!(builder.settings.console_level, "info");
2422        assert!(builder.tracer_provider.is_none());
2423    }
2424
2425    #[test]
2426    fn test_log_subscriber_builder_with_tracer_provider_sets_provider() {
2427        use opentelemetry_sdk::trace::SdkTracerProvider;
2428
2429        let service_info = crate::ServiceInfo {
2430            name: "test-service",
2431            name_in_metrics: "test_service".to_string(),
2432            version: "1.0.0",
2433            author: "Test",
2434            description: "Test service",
2435        };
2436
2437        let settings = LogSettings {
2438            console_level: "info".to_string(),
2439            otel_level: "info".to_string(),
2440            endpoint: None,
2441        };
2442
2443        let tracer_provider = SdkTracerProvider::builder().build();
2444
2445        let builder = super::LogSubscriberBuilder::new(&service_info, &settings)
2446            .with_tracer_provider(&tracer_provider);
2447
2448        // Verify the tracer provider was set
2449        assert!(builder.tracer_provider.is_some());
2450
2451        let _ = tracer_provider.shutdown();
2452    }
2453
2454    #[test]
2455    fn test_log_subscriber_builder_build_returns_working_subscriber() {
2456        // Test that build() returns a subscriber that can be used with with_default.
2457        // This catches mutations like "replace init_logs body with Ok(None)" because
2458        // if build() returned a broken/default subscriber, this test would fail.
2459        let service_info = crate::ServiceInfo {
2460            name: "test-service",
2461            name_in_metrics: "test_service".to_string(),
2462            version: "1.0.0",
2463            author: "Test",
2464            description: "Test service",
2465        };
2466
2467        let settings = LogSettings {
2468            console_level: "info".to_string(),
2469            otel_level: "info".to_string(),
2470            endpoint: None, // No OTel endpoint - just console logging
2471        };
2472
2473        let result = super::LogSubscriberBuilder::new(&service_info, &settings).build();
2474        assert!(result.is_ok(), "build() should succeed");
2475
2476        let built = result.unwrap();
2477
2478        // logger_provider should be None when no endpoint is configured
2479        assert!(
2480            built.logger_provider.is_none(),
2481            "logger_provider should be None without OTel endpoint"
2482        );
2483
2484        // The subscriber should be usable with with_default
2485        // This verifies that build() actually built something, not just returned default
2486        use std::sync::atomic::{AtomicBool, Ordering};
2487        static LOG_RECEIVED: AtomicBool = AtomicBool::new(false);
2488
2489        // Create a simple layer that sets a flag when it receives events
2490        struct TestLayer;
2491        impl<S: Subscriber> tracing_subscriber::Layer<S> for TestLayer {
2492            fn on_event(
2493                &self,
2494                _event: &tracing::Event<'_>,
2495                _ctx: tracing_subscriber::layer::Context<'_, S>,
2496            ) {
2497                LOG_RECEIVED.store(true, Ordering::SeqCst);
2498            }
2499        }
2500
2501        // Use the built subscriber with an additional test layer
2502        use tracing_subscriber::layer::SubscriberExt;
2503        let subscriber_with_test = built.subscriber.with(TestLayer);
2504
2505        tracing::subscriber::with_default(subscriber_with_test, || {
2506            tracing::info!("test log message");
2507        });
2508
2509        // The test layer should have received the event, proving the subscriber works
2510        assert!(
2511            LOG_RECEIVED.load(Ordering::SeqCst),
2512            "subscriber from build() should process log events"
2513        );
2514    }
2515
2516    #[test]
2517    fn test_log_subscriber_builder_build_with_tracer_provider() {
2518        // Test that when a tracer_provider is passed, the subscriber includes the OTel trace layer.
2519        use opentelemetry_sdk::trace::SdkTracerProvider;
2520
2521        let service_info = crate::ServiceInfo {
2522            name: "test-service",
2523            name_in_metrics: "test_service".to_string(),
2524            version: "1.0.0",
2525            author: "Test",
2526            description: "Test service",
2527        };
2528
2529        let settings = LogSettings {
2530            console_level: "info".to_string(),
2531            otel_level: "info".to_string(),
2532            endpoint: None,
2533        };
2534
2535        let tracer_provider = SdkTracerProvider::builder().build();
2536
2537        let result = super::LogSubscriberBuilder::new(&service_info, &settings)
2538            .with_tracer_provider(&tracer_provider)
2539            .build();
2540        assert!(
2541            result.is_ok(),
2542            "build() should succeed with tracer_provider"
2543        );
2544
2545        let built = result.unwrap();
2546
2547        // Use the subscriber and verify spans work (they're processed by the OTel layer)
2548        tracing::subscriber::with_default(built.subscriber, || {
2549            let span = tracing::info_span!("test_span_with_otel");
2550            let _enter = span.enter();
2551            tracing::info!("inside span");
2552        });
2553
2554        // Clean up
2555        let _ = tracer_provider.shutdown();
2556    }
2557}