1use std::sync::{Arc, Once};
4use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
5
6use crate::span_exporter::{AdkSpanExporter, AdkSpanLayer};
7
8static INIT: Once = Once::new();
9
10#[derive(Debug, thiserror::Error)]
12pub enum TelemetryError {
13 #[error("telemetry init failed: {0}")]
15 Init(String),
16}
17
18pub fn init_telemetry(service_name: &str) -> Result<(), TelemetryError> {
29 INIT.call_once(|| {
30 let filter = EnvFilter::try_from_default_env()
31 .or_else(|_| EnvFilter::try_new("info"))
32 .unwrap_or_else(|_| EnvFilter::new("info"));
33
34 tracing_subscriber::registry()
35 .with(filter)
36 .with(
37 tracing_subscriber::fmt::layer()
38 .with_target(true)
39 .with_thread_ids(true)
40 .with_line_number(true),
41 )
42 .init();
43
44 tracing::info!(service.name = service_name, "telemetry initialized");
45 });
46
47 Ok(())
48}
49
50pub fn init_with_otlp(service_name: &str, endpoint: &str) -> Result<(), TelemetryError> {
65 use opentelemetry::trace::TracerProvider;
66 use opentelemetry_otlp::WithExportConfig;
67 use tracing_opentelemetry::OpenTelemetryLayer;
68
69 let endpoint = endpoint.to_string();
70 let service_name = service_name.to_string();
71
72 let init_error: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
73
74 INIT.call_once(|| {
75 let resource = opentelemetry_sdk::Resource::builder_empty()
76 .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.clone())])
77 .build();
78
79 let span_exporter = match opentelemetry_otlp::SpanExporter::builder()
81 .with_tonic()
82 .with_endpoint(&endpoint)
83 .build()
84 {
85 Ok(e) => e,
86 Err(e) => {
87 *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
88 Some(format!("failed to build OTLP span exporter: {e}"));
89 return;
90 }
91 };
92
93 let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
95 .with_batch_exporter(span_exporter)
96 .with_resource(resource.clone())
97 .build();
98
99 let tracer = tracer_provider.tracer("adk-telemetry");
100 opentelemetry::global::set_tracer_provider(tracer_provider);
101
102 let metric_exporter = match opentelemetry_otlp::MetricExporter::builder()
104 .with_tonic()
105 .with_endpoint(&endpoint)
106 .build()
107 {
108 Ok(e) => e,
109 Err(e) => {
110 *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
111 Some(format!("failed to build OTLP metric exporter: {e}"));
112 return;
113 }
114 };
115
116 let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
117 .with_periodic_exporter(metric_exporter)
118 .with_resource(resource)
119 .build();
120
121 opentelemetry::global::set_meter_provider(meter_provider);
122
123 let telemetry_layer = OpenTelemetryLayer::new(tracer);
124
125 let filter = EnvFilter::try_from_default_env()
126 .or_else(|_| EnvFilter::try_new("info"))
127 .unwrap_or_else(|_| EnvFilter::new("info"));
128
129 tracing_subscriber::registry()
130 .with(filter)
131 .with(
132 tracing_subscriber::fmt::layer()
133 .with_target(true)
134 .with_thread_ids(true)
135 .with_line_number(true),
136 )
137 .with(telemetry_layer)
138 .init();
139
140 tracing::info!(
141 service.name = service_name,
142 otlp.endpoint = %endpoint,
143 "telemetry initialized with OpenTelemetry"
144 );
145 });
146
147 if let Some(err) = init_error.lock().unwrap_or_else(|p| p.into_inner()).take() {
148 return Err(TelemetryError::Init(err));
149 }
150
151 Ok(())
152}
153
154pub fn build_otlp_layer<S>(
189 service_name: &str,
190 endpoint: &str,
191) -> Result<Box<dyn tracing_subscriber::Layer<S> + Send + Sync>, TelemetryError>
192where
193 S: tracing::Subscriber
194 + for<'span> tracing_subscriber::registry::LookupSpan<'span>
195 + Send
196 + Sync,
197{
198 use opentelemetry::trace::TracerProvider;
199 use opentelemetry_otlp::WithExportConfig;
200 use tracing_opentelemetry::OpenTelemetryLayer;
201
202 let resource = opentelemetry_sdk::Resource::builder_empty()
203 .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.to_string())])
204 .build();
205
206 let span_exporter = opentelemetry_otlp::SpanExporter::builder()
208 .with_tonic()
209 .with_endpoint(endpoint)
210 .build()
211 .map_err(|e| TelemetryError::Init(format!("failed to build OTLP span exporter: {e}")))?;
212
213 let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
215 .with_batch_exporter(span_exporter)
216 .with_resource(resource.clone())
217 .build();
218
219 let tracer = tracer_provider.tracer("adk-telemetry");
220 opentelemetry::global::set_tracer_provider(tracer_provider);
221
222 let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
224 .with_tonic()
225 .with_endpoint(endpoint)
226 .build()
227 .map_err(|e| TelemetryError::Init(format!("failed to build OTLP metric exporter: {e}")))?;
228
229 let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
230 .with_periodic_exporter(metric_exporter)
231 .with_resource(resource)
232 .build();
233
234 opentelemetry::global::set_meter_provider(meter_provider);
235
236 Ok(Box::new(OpenTelemetryLayer::new(tracer)))
237}
238
239pub fn shutdown_telemetry() {
245 opentelemetry::global::set_tracer_provider(
249 opentelemetry::trace::noop::NoopTracerProvider::new(),
250 );
251}
252
253pub fn init_with_adk_exporter(service_name: &str) -> Result<Arc<AdkSpanExporter>, TelemetryError> {
258 let exporter = Arc::new(AdkSpanExporter::new());
259 let exporter_clone = exporter.clone();
260
261 INIT.call_once(|| {
262 let filter = EnvFilter::try_from_default_env()
263 .or_else(|_| EnvFilter::try_new("info"))
264 .unwrap_or_else(|_| EnvFilter::new("info"));
265
266 let adk_layer = AdkSpanLayer::new(exporter_clone);
267
268 tracing_subscriber::registry()
269 .with(filter)
270 .with(
271 tracing_subscriber::fmt::layer()
272 .with_target(true)
273 .with_thread_ids(true)
274 .with_line_number(true),
275 )
276 .with(adk_layer)
277 .init();
278
279 tracing::info!(service.name = service_name, "telemetry initialized with ADK span exporter");
280 });
281
282 Ok(exporter)
283}