rskit_observability/
tracer.rs

1use std::time::Duration;
2
3use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
4use opentelemetry_sdk::propagation::TraceContextPropagator;
5use serde::{Deserialize, Serialize};
6use tracing::Span;
7use tracing_opentelemetry::OpenTelemetrySpanExt;
8
9use rskit_errors::{AppError, AppResult, ErrorCode};
10
11use crate::attribute::set_span_attribute;
12
13/// Stable OTel semantic-convention key for service name.
14pub const SERVICE_NAME: &str = "service.name";
15
16/// OTLP exporter protocol.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
18#[non_exhaustive]
19pub enum OtlpProtocol {
20    /// OTLP over gRPC/Tonic, usually port 4317.
21    #[default]
22    Grpc,
23    /// OTLP over HTTP/protobuf, usually port 4318.
24    HttpBinary,
25}
26
27/// Configuration for an OTLP trace exporter.
28#[derive(Debug, Clone, Deserialize)]
29pub struct TracingConfig {
30    /// Logical service name reported in every span.
31    pub service_name: String,
32    /// OTLP collector endpoint for the selected exporter protocol.
33    pub endpoint: String,
34    /// Sampling probability in `0.0..=1.0`.
35    pub sample_rate: f64,
36    /// Maximum time to wait when flushing spans on shutdown.
37    pub export_timeout: Duration,
38}
39
40/// Injectable tracer provider handle.
41pub struct TracerGuard {
42    provider: opentelemetry_sdk::trace::SdkTracerProvider,
43    tracer: opentelemetry_sdk::trace::SdkTracer,
44}
45
46impl Drop for TracerGuard {
47    fn drop(&mut self) {
48        if let Err(e) = self.provider.shutdown() {
49            tracing::warn!(error = %e, "failed to shutdown tracer provider");
50        }
51    }
52}
53
54impl TracerGuard {
55    /// Return the SDK tracer provider for constructor injection.
56    #[must_use]
57    pub fn provider(&self) -> opentelemetry_sdk::trace::SdkTracerProvider {
58        self.provider.clone()
59    }
60
61    /// Return a tracer that can be injected into `tracing_opentelemetry::layer()`.
62    #[must_use]
63    pub fn tracer(&self) -> opentelemetry_sdk::trace::SdkTracer {
64        self.tracer.clone()
65    }
66}
67
68/// Build an injectable OpenTelemetry tracer provider without touching global state.
69pub fn tracer_provider(cfg: &TracingConfig) -> AppResult<TracerGuard> {
70    tracer_provider_with_protocol(cfg, OtlpProtocol::Grpc)
71}
72
73/// Build an injectable OpenTelemetry tracer provider with an explicit OTLP protocol.
74pub fn tracer_provider_with_protocol(
75    cfg: &TracingConfig,
76    protocol: OtlpProtocol,
77) -> AppResult<TracerGuard> {
78    #[cfg(not(feature = "otlp"))]
79    {
80        let _ = (cfg, protocol);
81        Err(AppError::new(
82            ErrorCode::InvalidInput,
83            "OTLP tracing exporter requires the `otlp` feature",
84        ))
85    }
86
87    #[cfg(feature = "otlp")]
88    {
89        use opentelemetry::trace::TracerProvider as _;
90        use opentelemetry::{InstrumentationScope, KeyValue};
91        use opentelemetry_otlp::{SpanExporter, WithExportConfig};
92        use opentelemetry_sdk::Resource;
93        use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
94        let exporter = match protocol {
95            OtlpProtocol::Grpc => SpanExporter::builder()
96                .with_tonic()
97                .with_endpoint(&cfg.endpoint)
98                .with_timeout(cfg.export_timeout)
99                .build(),
100            OtlpProtocol::HttpBinary => SpanExporter::builder()
101                .with_http()
102                .with_protocol(opentelemetry_otlp::Protocol::HttpBinary)
103                .with_endpoint(&cfg.endpoint)
104                .with_timeout(cfg.export_timeout)
105                .build(),
106        }
107        .map_err(|e| AppError::new(ErrorCode::Internal, format!("OTLP span exporter: {e}")))?;
108
109        let sampler = if (cfg.sample_rate - 1.0).abs() < f64::EPSILON {
110            Sampler::AlwaysOn
111        } else if cfg.sample_rate <= 0.0 {
112            Sampler::AlwaysOff
113        } else {
114            Sampler::TraceIdRatioBased(cfg.sample_rate)
115        };
116
117        let resource = Resource::builder_empty()
118            .with_attributes([KeyValue::new(SERVICE_NAME, cfg.service_name.clone())])
119            .build();
120
121        let provider = SdkTracerProvider::builder()
122            .with_batch_exporter(exporter)
123            .with_sampler(sampler)
124            .with_resource(resource)
125            .build();
126
127        let scope = InstrumentationScope::builder(cfg.service_name.clone()).build();
128        let tracer = provider.tracer_with_scope(scope);
129        Ok(TracerGuard { provider, tracer })
130    }
131}
132
133/// Attach bounded operation attributes to an active tracing span.
134pub fn set_operation_attributes(
135    span: &Span,
136    service_name: &str,
137    operation_name: &str,
138    request_id: &str,
139) {
140    set_span_attribute(span, SERVICE_NAME, service_name);
141    set_span_attribute(span, "operation.name", operation_name);
142    set_span_attribute(span, "request.id", request_id);
143}
144
145/// Inject the current trace context into HTTP headers (W3C Trace-Context).
146pub fn inject_trace_context(headers: &mut http::HeaderMap) {
147    let propagator = TraceContextPropagator::new();
148    let cx = Span::current().context();
149    propagator.inject_context(&cx, &mut HeaderMapCarrier(headers));
150}
151
152/// Extract a trace context from incoming HTTP headers.
153pub fn extract_trace_context(headers: &http::HeaderMap) -> opentelemetry::Context {
154    let propagator = TraceContextPropagator::new();
155    propagator.extract(&HeaderMapExtractor(headers))
156}
157
158// ---- carriers ---------------------------------------------------------------
159
160struct HeaderMapCarrier<'a>(&'a mut http::HeaderMap);
161
162impl Injector for HeaderMapCarrier<'_> {
163    fn set(&mut self, key: &str, value: String) {
164        if let Ok(name) = http::header::HeaderName::from_bytes(key.as_bytes())
165            && let Ok(val) = http::header::HeaderValue::from_str(&value)
166        {
167            self.0.insert(name, val);
168        }
169    }
170}
171
172struct HeaderMapExtractor<'a>(&'a http::HeaderMap);
173
174impl Extractor for HeaderMapExtractor<'_> {
175    fn get(&self, key: &str) -> Option<&str> {
176        self.0.get(key).and_then(|v| v.to_str().ok())
177    }
178
179    fn keys(&self) -> Vec<&str> {
180        self.0.keys().map(|k| k.as_str()).collect()
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use std::time::Duration;
187
188    use super::*;
189
190    fn config() -> TracingConfig {
191        TracingConfig {
192            service_name: "test-service".to_string(),
193            endpoint: "http://127.0.0.1:4317".to_string(),
194            sample_rate: 1.0,
195            export_timeout: Duration::from_millis(10),
196        }
197    }
198
199    #[test]
200    fn otlp_protocol_round_trips_through_serde() {
201        let grpc = serde_json::to_string(&OtlpProtocol::Grpc).expect("serialize grpc");
202        let http = serde_json::to_string(&OtlpProtocol::HttpBinary).expect("serialize http");
203
204        assert_eq!(
205            serde_json::from_str::<OtlpProtocol>(&grpc).expect("deserialize grpc"),
206            OtlpProtocol::Grpc
207        );
208        assert_eq!(
209            serde_json::from_str::<OtlpProtocol>(&http).expect("deserialize http"),
210            OtlpProtocol::HttpBinary
211        );
212    }
213
214    #[cfg(not(feature = "otlp"))]
215    #[test]
216    fn tracer_provider_requires_otlp_feature_for_exporter() {
217        let error = match tracer_provider_with_protocol(&config(), OtlpProtocol::Grpc) {
218            Ok(_) => panic!("otlp disabled should reject exporter configuration"),
219            Err(error) => error,
220        };
221
222        assert_eq!(error.code(), ErrorCode::InvalidInput);
223        assert!(error.message().contains("otlp"));
224    }
225
226    #[cfg(feature = "otlp")]
227    #[tokio::test]
228    async fn tracer_provider_builds_exporters_for_supported_protocols() {
229        for (protocol, sample_rate) in [
230            (OtlpProtocol::Grpc, 1.0),
231            (OtlpProtocol::HttpBinary, 0.0),
232            (OtlpProtocol::HttpBinary, 0.5),
233        ] {
234            let mut cfg = config();
235            cfg.sample_rate = sample_rate;
236            let guard =
237                tracer_provider_with_protocol(&cfg, protocol).expect("build tracer provider");
238            let _provider = guard.provider();
239            let _tracer = guard.tracer();
240        }
241    }
242
243    #[test]
244    fn operation_attributes_and_trace_context_helpers_are_tolerant() {
245        let span = tracing::info_span!("operation");
246        set_operation_attributes(&span, "svc", "op", "req-1");
247
248        let mut headers = http::HeaderMap::new();
249        headers.insert(
250            "traceparent",
251            http::HeaderValue::from_static(
252                "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
253            ),
254        );
255        headers.insert("tracestate", http::HeaderValue::from_static("vendor=value"));
256        let _context = extract_trace_context(&headers);
257
258        headers.insert(
259            "traceparent",
260            http::HeaderValue::from_bytes(b"\xff").expect("invalid utf8 header value is allowed"),
261        );
262        let _context = extract_trace_context(&headers);
263
264        let _entered = span.enter();
265        let mut injected = http::HeaderMap::new();
266        inject_trace_context(&mut injected);
267    }
268
269    #[test]
270    fn header_carriers_ignore_invalid_injected_values_and_list_valid_keys() {
271        let mut headers = http::HeaderMap::new();
272        let mut injector = HeaderMapCarrier(&mut headers);
273        injector.set("UPPERCASE", "ok".to_string());
274        injector.set("valid-header", "\ninvalid".to_string());
275        assert!(headers.contains_key("uppercase"));
276        assert!(!headers.contains_key("valid-header"));
277
278        let extractor = HeaderMapExtractor(&headers);
279        assert_eq!(extractor.get("uppercase"), Some("ok"));
280        assert!(extractor.keys().contains(&"uppercase"));
281    }
282}
rskit_observability/tracer.rs

rskit_observability/
tracer.rs
