adk_telemetry/
init.rs

1//! Telemetry initialization and configuration
2
3use std::sync::Once;
4use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
5
6static INIT: Once = Once::new();
7
8/// Initialize basic telemetry with console logging
9///
10/// # Arguments
11/// * `service_name` - Name of the service for trace identification
12///
13/// # Example
14/// ```
15/// use adk_telemetry::init_telemetry;
16/// init_telemetry("my-agent-service").expect("Failed to initialize telemetry");
17/// ```
18pub fn init_telemetry(service_name: &str) -> Result<(), Box<dyn std::error::Error>> {
19    INIT.call_once(|| {
20        let filter = EnvFilter::try_from_default_env()
21            .or_else(|_| EnvFilter::try_new("info"))
22            .expect("Failed to create env filter");
23
24        tracing_subscriber::registry()
25            .with(filter)
26            .with(
27                tracing_subscriber::fmt::layer()
28                    .with_target(true)
29                    .with_thread_ids(true)
30                    .with_line_number(true),
31            )
32            .init();
33
34        tracing::info!(service.name = service_name, "Telemetry initialized");
35    });
36
37    Ok(())
38}
39
40/// Initialize telemetry with OpenTelemetry OTLP export
41///
42/// Enables distributed tracing by exporting spans to an OTLP collector.
43///
44/// # Arguments
45/// * `service_name` - Name of the service for trace identification
46/// * `endpoint` - OTLP collector endpoint (e.g., "http://localhost:4317")
47///
48/// # Example
49/// ```no_run
50/// use adk_telemetry::init_with_otlp;
51/// init_with_otlp("my-agent", "http://localhost:4317")
52///     .expect("Failed to initialize telemetry");
53/// ```
54pub fn init_with_otlp(
55    service_name: &str,
56    endpoint: &str,
57) -> Result<(), Box<dyn std::error::Error>> {
58    use opentelemetry_otlp::WithExportConfig;
59    use tracing_opentelemetry::OpenTelemetryLayer;
60
61    INIT.call_once(|| {
62        // install_batch returns a Tracer directly
63        let tracer = opentelemetry_otlp::new_pipeline()
64            .tracing()
65            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
66            .with_trace_config(opentelemetry_sdk::trace::config().with_resource(
67                opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
68                    "service.name",
69                    service_name.to_string(),
70                )]),
71            ))
72            .install_batch(opentelemetry_sdk::runtime::Tokio)
73            .expect("Failed to install OTLP pipeline");
74
75        // Initialize metrics
76        let meter_provider = opentelemetry_otlp::new_pipeline()
77            .metrics(opentelemetry_sdk::runtime::Tokio)
78            .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint(endpoint))
79            .with_resource(opentelemetry_sdk::Resource::new(vec![opentelemetry::KeyValue::new(
80                "service.name",
81                service_name.to_string(),
82            )]))
83            .build()
84            .expect("Failed to build meter provider");
85
86        opentelemetry::global::set_meter_provider(meter_provider);
87
88        let telemetry_layer = OpenTelemetryLayer::new(tracer);
89
90        let filter = EnvFilter::try_from_default_env()
91            .or_else(|_| EnvFilter::try_new("info"))
92            .expect("Failed to create env filter");
93
94        tracing_subscriber::registry()
95            .with(filter)
96            .with(
97                tracing_subscriber::fmt::layer()
98                    .with_target(true)
99                    .with_thread_ids(true)
100                    .with_line_number(true),
101            )
102            .with(telemetry_layer)
103            .init();
104
105        tracing::info!(
106            service.name = service_name,
107            otlp.endpoint = endpoint,
108            "Telemetry initialized with OpenTelemetry"
109        );
110    });
111
112    Ok(())
113}
114
115/// Shutdown telemetry and flush any pending spans
116///
117/// Should be called before application exit to ensure all telemetry data is sent.
118pub fn shutdown_telemetry() {
119    opentelemetry::global::shutdown_tracer_provider();
120}