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
10pub fn init_telemetry(service_name: &str) -> Result<(), Box<dyn std::error::Error>> {
21 INIT.call_once(|| {
22 let filter = EnvFilter::try_from_default_env()
23 .or_else(|_| EnvFilter::try_new("info"))
24 .expect("Failed to create env filter");
25
26 tracing_subscriber::registry()
27 .with(filter)
28 .with(
29 tracing_subscriber::fmt::layer()
30 .with_target(true)
31 .with_thread_ids(true)
32 .with_line_number(true),
33 )
34 .init();
35
36 tracing::info!(service.name = service_name, "Telemetry initialized");
37 });
38
39 Ok(())
40}
41
42pub fn init_with_otlp(
57 service_name: &str,
58 endpoint: &str,
59) -> Result<(), Box<dyn std::error::Error>> {
60 use opentelemetry_otlp::WithExportConfig;
61 use tracing_opentelemetry::OpenTelemetryLayer;
62
63 INIT.call_once(|| {
64 let tracer = opentelemetry_otlp::new_pipeline()
66 .tracing()
67 .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
68 .with_trace_config(opentelemetry_sdk::trace::config().with_resource(
69 opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
70 "service.name",
71 service_name.to_string(),
72 )]),
73 ))
74 .install_batch(opentelemetry_sdk::runtime::Tokio)
75 .expect("Failed to install OTLP pipeline");
76
77 let meter_provider = opentelemetry_otlp::new_pipeline()
79 .metrics(opentelemetry_sdk::runtime::Tokio)
80 .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
81 .with_resource(opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
82 "service.name",
83 service_name.to_string(),
84 )]))
85 .build()
86 .expect("Failed to build meter provider");
87
88 opentelemetry::global::set_meter_provider(meter_provider);
89
90 let telemetry_layer = OpenTelemetryLayer::new(tracer);
91
92 let filter = EnvFilter::try_from_default_env()
93 .or_else(|_| EnvFilter::try_new("info"))
94 .expect("Failed to create env filter");
95
96 tracing_subscriber::registry()
97 .with(filter)
98 .with(
99 tracing_subscriber::fmt::layer()
100 .with_target(true)
101 .with_thread_ids(true)
102 .with_line_number(true),
103 )
104 .with(telemetry_layer)
105 .init();
106
107 tracing::info!(
108 service.name = service_name,
109 otlp.endpoint = endpoint,
110 "Telemetry initialized with OpenTelemetry"
111 );
112 });
113
114 Ok(())
115}
116
117pub fn shutdown_telemetry() {
121 opentelemetry::global::shutdown_tracer_provider();
122}
123
124pub fn init_with_adk_exporter(
128 service_name: &str,
129) -> Result<Arc<AdkSpanExporter>, Box<dyn std::error::Error>> {
130 let exporter = Arc::new(AdkSpanExporter::new());
131 let exporter_clone = exporter.clone();
132
133 INIT.call_once(|| {
134 let filter = EnvFilter::try_from_default_env()
135 .or_else(|_| EnvFilter::try_new("info"))
136 .expect("Failed to create env filter");
137
138 let adk_layer = AdkSpanLayer::new(exporter_clone);
139
140 tracing_subscriber::registry()
141 .with(filter)
142 .with(
143 tracing_subscriber::fmt::layer()
144 .with_target(true)
145 .with_thread_ids(true)
146 .with_line_number(true),
147 )
148 .with(adk_layer)
149 .init();
150
151 tracing::info!(service.name = service_name, "Telemetry initialized with ADK span exporter");
152 });
153
154 Ok(exporter)
155}