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
50#[cfg(feature = "otlp")]
65pub fn init_with_otlp(service_name: &str, endpoint: &str) -> Result<(), TelemetryError> {
66 use opentelemetry::trace::TracerProvider;
67 use opentelemetry_otlp::WithExportConfig;
68 use tracing_opentelemetry::OpenTelemetryLayer;
69
70 let endpoint = endpoint.to_string();
71 let service_name = service_name.to_string();
72
73 let init_error: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
74
75 INIT.call_once(|| {
76 let resource = opentelemetry_sdk::Resource::builder_empty()
77 .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.clone())])
78 .build();
79
80 let span_exporter = match opentelemetry_otlp::SpanExporter::builder()
82 .with_tonic()
83 .with_endpoint(&endpoint)
84 .build()
85 {
86 Ok(e) => e,
87 Err(e) => {
88 *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
89 Some(format!("failed to build OTLP span exporter: {e}"));
90 return;
91 }
92 };
93
94 let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
96 .with_batch_exporter(span_exporter)
97 .with_resource(resource.clone())
98 .build();
99
100 let tracer = tracer_provider.tracer("adk-telemetry");
101 opentelemetry::global::set_tracer_provider(tracer_provider);
102
103 let metric_exporter = match opentelemetry_otlp::MetricExporter::builder()
105 .with_tonic()
106 .with_endpoint(&endpoint)
107 .build()
108 {
109 Ok(e) => e,
110 Err(e) => {
111 *init_error.lock().unwrap_or_else(|p| p.into_inner()) =
112 Some(format!("failed to build OTLP metric exporter: {e}"));
113 return;
114 }
115 };
116
117 let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
118 .with_periodic_exporter(metric_exporter)
119 .with_resource(resource)
120 .build();
121
122 opentelemetry::global::set_meter_provider(meter_provider);
123
124 let telemetry_layer = OpenTelemetryLayer::new(tracer);
125
126 let filter = EnvFilter::try_from_default_env()
127 .or_else(|_| EnvFilter::try_new("info"))
128 .unwrap_or_else(|_| EnvFilter::new("info"));
129
130 tracing_subscriber::registry()
131 .with(filter)
132 .with(
133 tracing_subscriber::fmt::layer()
134 .with_target(true)
135 .with_thread_ids(true)
136 .with_line_number(true),
137 )
138 .with(telemetry_layer)
139 .init();
140
141 tracing::info!(
142 service.name = service_name,
143 otlp.endpoint = %endpoint,
144 "telemetry initialized with OpenTelemetry"
145 );
146 });
147
148 if let Some(err) = init_error.lock().unwrap_or_else(|p| p.into_inner()).take() {
149 return Err(TelemetryError::Init(err));
150 }
151
152 Ok(())
153}
154
155#[cfg(feature = "otlp")]
190pub fn build_otlp_layer<S>(
191 service_name: &str,
192 endpoint: &str,
193) -> Result<Box<dyn tracing_subscriber::Layer<S> + Send + Sync>, TelemetryError>
194where
195 S: tracing::Subscriber
196 + for<'span> tracing_subscriber::registry::LookupSpan<'span>
197 + Send
198 + Sync,
199{
200 use opentelemetry::trace::TracerProvider;
201 use opentelemetry_otlp::WithExportConfig;
202 use tracing_opentelemetry::OpenTelemetryLayer;
203
204 let resource = opentelemetry_sdk::Resource::builder_empty()
205 .with_attributes([opentelemetry::KeyValue::new("service.name", service_name.to_string())])
206 .build();
207
208 let span_exporter = opentelemetry_otlp::SpanExporter::builder()
210 .with_tonic()
211 .with_endpoint(endpoint)
212 .build()
213 .map_err(|e| TelemetryError::Init(format!("failed to build OTLP span exporter: {e}")))?;
214
215 let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
217 .with_batch_exporter(span_exporter)
218 .with_resource(resource.clone())
219 .build();
220
221 let tracer = tracer_provider.tracer("adk-telemetry");
222 opentelemetry::global::set_tracer_provider(tracer_provider);
223
224 let metric_exporter = opentelemetry_otlp::MetricExporter::builder()
226 .with_tonic()
227 .with_endpoint(endpoint)
228 .build()
229 .map_err(|e| TelemetryError::Init(format!("failed to build OTLP metric exporter: {e}")))?;
230
231 let meter_provider = opentelemetry_sdk::metrics::SdkMeterProvider::builder()
232 .with_periodic_exporter(metric_exporter)
233 .with_resource(resource)
234 .build();
235
236 opentelemetry::global::set_meter_provider(meter_provider);
237
238 Ok(Box::new(OpenTelemetryLayer::new(tracer)))
239}
240
241pub fn shutdown_telemetry() {
247 #[cfg(feature = "otlp")]
248 {
249 opentelemetry::global::set_tracer_provider(
253 opentelemetry::trace::noop::NoopTracerProvider::new(),
254 );
255 }
256}
257
258pub fn init_with_adk_exporter(service_name: &str) -> Result<Arc<AdkSpanExporter>, TelemetryError> {
263 let exporter = Arc::new(AdkSpanExporter::new());
264 let exporter_clone = exporter.clone();
265
266 INIT.call_once(|| {
267 let filter = EnvFilter::try_from_default_env()
268 .or_else(|_| EnvFilter::try_new("info"))
269 .unwrap_or_else(|_| EnvFilter::new("info"));
270
271 let adk_layer = AdkSpanLayer::new(exporter_clone);
272
273 tracing_subscriber::registry()
274 .with(filter)
275 .with(
276 tracing_subscriber::fmt::layer()
277 .with_target(true)
278 .with_thread_ids(true)
279 .with_line_number(true),
280 )
281 .with(adk_layer)
282 .init();
283
284 tracing::info!(service.name = service_name, "telemetry initialized with ADK span exporter");
285 });
286
287 Ok(exporter)
288}