llm_optimizer_processor/telemetry/
mod.rs

1//! OpenTelemetry instrumentation for the stream processor.
2//!
3//! This module provides comprehensive distributed tracing and metrics export
4//! using OpenTelemetry. It includes:
5//!
6//! - **Tracing**: Distributed tracing with OTLP exporter, span instrumentation,
7//!   and multiple sampling strategies.
8//! - **Metrics**: OTLP metrics export with support for counters, histograms,
9//!   and gauges.
10//! - **Context Propagation**: W3C Trace Context and Baggage propagation for
11//!   distributed tracing across service boundaries.
12//! - **Resource Attributes**: Automatic detection and configuration of service
13//!   metadata following OpenTelemetry semantic conventions.
14//!
15//! # Usage
16//!
17//! ## Initializing Tracing
18//!
19//! ```ignore
20//! use processor::telemetry::{TraceConfig, init_tracer_provider};
21//!
22//! let config = TraceConfig {
23//!     enabled: true,
24//!     service_name: "my-processor".to_string(),
25//!     otlp_endpoint: "http://localhost:4317".to_string(),
26//!     ..Default::default()
27//! };
28//!
29//! let provider = init_tracer_provider(config)?;
30//! ```
31//!
32//! ## Creating Spans
33//!
34//! ```ignore
35//! use processor::telemetry::SpanBuilder;
36//! use opentelemetry::trace::SpanKind;
37//!
38//! let span = SpanBuilder::new("process_event")
39//!     .with_kind(SpanKind::Internal)
40//!     .with_attribute("event.type", "metric")
41//!     .start();
42//! ```
43//!
44//! Or using the macro:
45//!
46//! ```ignore
47//! use processor::span;
48//!
49//! let span = span!("process_event", "event.type" => "metric");
50//! ```
51//!
52//! ## Context Propagation
53//!
54//! ```ignore
55//! use processor::telemetry::{inject_into_headers, extract_from_headers};
56//! use opentelemetry::Context;
57//!
58//! // Inject context into outgoing request
59//! let mut headers = HashMap::new();
60//! inject_into_headers(&Context::current(), &mut headers);
61//!
62//! // Extract context from incoming request
63//! let context = extract_from_headers(&headers);
64//! ```
65//!
66//! ## Recording Metrics
67//!
68//! ```ignore
69//! use processor::telemetry::{MetricsRecorder, metric_names};
70//! use opentelemetry::KeyValue;
71//!
72//! let recorder = MetricsRecorder::new()
73//!     .with_attribute("service", "processor");
74//!
75//! recorder.record_counter(
76//!     "processor",
77//!     metric_names::EVENTS_PROCESSED,
78//!     1,
79//!     vec![KeyValue::new("event.type", "metric")],
80//! );
81//! ```
82//!
83//! ## Integration with Existing Code
84//!
85//! The telemetry module integrates seamlessly with the existing `tracing` crate
86//! used throughout the processor. Use `tracing-opentelemetry` to bridge the two:
87//!
88//! ```ignore
89//! use tracing_subscriber::layer::SubscriberExt;
90//! use tracing_subscriber::Registry;
91//! use tracing_opentelemetry::OpenTelemetryLayer;
92//!
93//! let tracer = global::tracer("processor");
94//! let telemetry = OpenTelemetryLayer::new(tracer);
95//!
96//! let subscriber = Registry::default().with(telemetry);
97//! tracing::subscriber::set_global_default(subscriber)?;
98//! ```
99
100pub mod context;
101pub mod metrics;
102pub mod resource;
103pub mod tracing;
104
105// Re-export commonly used types
106pub use context::{
107    extract_context, extract_from_headers, inject_context, inject_into_headers,
108    context_with_span, get_baggage, propagate_context, with_baggage, with_context,
109    BaggageItem, HeaderCarrier,
110};
111
112pub use metrics::{
113    init_meter_provider, shutdown_meter_provider, metric_attributes, metric_names,
114    ExportConfig, MetricsConfig, MetricsRecorder, Temporality,
115};
116
117pub use resource::{create_processor_resource, ResourceBuilder};
118
119pub use tracing::{
120    init_tracer_provider, record_exception, shutdown_tracer_provider, BatchConfiguration,
121    OtlpProtocol, SamplingStrategy, SpanBuilder, TraceConfig,
122};
123
124// Re-export OpenTelemetry types for convenience
125pub use opentelemetry::{
126    global,
127    trace::{Span, SpanKind, Status, Tracer},
128    Context, KeyValue,
129};
130
131/// Combined configuration for both tracing and metrics.
132#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
133pub struct TelemetryConfig {
134    /// Tracing configuration.
135    pub tracing: TraceConfig,
136
137    /// Metrics configuration.
138    pub metrics: MetricsConfig,
139}
140
141impl Default for TelemetryConfig {
142    fn default() -> Self {
143        Self {
144            tracing: TraceConfig::default(),
145            metrics: MetricsConfig::default(),
146        }
147    }
148}
149
150impl TelemetryConfig {
151    /// Creates a new telemetry configuration.
152    pub fn new(service_name: impl Into<String>) -> Self {
153        let service_name = service_name.into();
154        Self {
155            tracing: TraceConfig {
156                service_name: service_name.clone(),
157                ..Default::default()
158            },
159            metrics: MetricsConfig {
160                service_name,
161                ..Default::default()
162            },
163        }
164    }
165
166    /// Sets the OTLP endpoint for both tracing and metrics.
167    pub fn with_otlp_endpoint(mut self, endpoint: impl Into<String>) -> Self {
168        let endpoint = endpoint.into();
169        self.tracing.otlp_endpoint = endpoint.clone();
170        self.metrics.otlp_endpoint = endpoint;
171        self
172    }
173
174    /// Sets the service version for both tracing and metrics.
175    pub fn with_service_version(mut self, version: impl Into<String>) -> Self {
176        let version = version.into();
177        self.tracing.service_version = version.clone();
178        self.metrics.service_version = version;
179        self
180    }
181
182    /// Sets the deployment environment for both tracing and metrics.
183    pub fn with_environment(mut self, env: impl Into<String>) -> Self {
184        let env = Some(env.into());
185        self.tracing.environment = env.clone();
186        self.metrics.environment = env;
187        self
188    }
189
190    /// Enables or disables tracing.
191    pub fn with_tracing_enabled(mut self, enabled: bool) -> Self {
192        self.tracing.enabled = enabled;
193        self
194    }
195
196    /// Enables or disables metrics.
197    pub fn with_metrics_enabled(mut self, enabled: bool) -> Self {
198        self.metrics.enabled = enabled;
199        self
200    }
201}
202
203/// Initializes both tracing and metrics providers.
204///
205/// # Errors
206///
207/// Returns an error if either provider fails to initialize.
208pub fn init_telemetry(
209    config: TelemetryConfig,
210) -> Result<TelemetryProviders, Box<dyn std::error::Error>> {
211    let tracer_provider = if config.tracing.enabled {
212        Some(init_tracer_provider(config.tracing)?)
213    } else {
214        None
215    };
216
217    let meter_provider = if config.metrics.enabled {
218        Some(init_meter_provider(config.metrics)?)
219    } else {
220        None
221    };
222
223    Ok(TelemetryProviders {
224        tracer_provider,
225        meter_provider,
226    })
227}
228
229/// Container for initialized telemetry providers.
230pub struct TelemetryProviders {
231    pub tracer_provider: Option<opentelemetry_sdk::trace::TracerProvider>,
232    pub meter_provider: Option<opentelemetry_sdk::metrics::SdkMeterProvider>,
233}
234
235impl TelemetryProviders {
236    /// Shuts down both providers and flushes pending data.
237    pub async fn shutdown(self) -> Result<(), Box<dyn std::error::Error>> {
238        if let Some(_provider) = self.tracer_provider {
239            shutdown_tracer_provider().await;
240        }
241
242        if let Some(_provider) = self.meter_provider {
243            shutdown_meter_provider()?;
244        }
245
246        Ok(())
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253
254    #[test]
255    fn test_telemetry_config_default() {
256        let config = TelemetryConfig::default();
257        assert!(config.tracing.enabled);
258        assert!(config.metrics.enabled);
259    }
260
261    #[test]
262    fn test_telemetry_config_new() {
263        let config = TelemetryConfig::new("test-service");
264        assert_eq!(config.tracing.service_name, "test-service");
265        assert_eq!(config.metrics.service_name, "test-service");
266    }
267
268    #[test]
269    fn test_telemetry_config_builder() {
270        let config = TelemetryConfig::new("test-service")
271            .with_otlp_endpoint("http://collector:4317")
272            .with_service_version("1.0.0")
273            .with_environment("production")
274            .with_tracing_enabled(true)
275            .with_metrics_enabled(false);
276
277        assert_eq!(config.tracing.otlp_endpoint, "http://collector:4317");
278        assert_eq!(config.metrics.otlp_endpoint, "http://collector:4317");
279        assert_eq!(config.tracing.service_version, "1.0.0");
280        assert_eq!(config.metrics.service_version, "1.0.0");
281        assert_eq!(
282            config.tracing.environment,
283            Some("production".to_string())
284        );
285        assert_eq!(
286            config.metrics.environment,
287            Some("production".to_string())
288        );
289        assert!(config.tracing.enabled);
290        assert!(!config.metrics.enabled);
291    }
292
293    #[test]
294    fn test_telemetry_config_serialization() {
295        let config = TelemetryConfig::new("test");
296        let json = serde_json::to_string(&config).unwrap();
297        let deserialized: TelemetryConfig = serde_json::from_str(&json).unwrap();
298
299        assert_eq!(config.tracing.service_name, deserialized.tracing.service_name);
300        assert_eq!(config.metrics.service_name, deserialized.metrics.service_name);
301    }
302}
llm_optimizer_processor/telemetry/mod.rs

llm_optimizer_processor/telemetry/
mod.rs
