Skip to main content

microclaw_engine/internal/observability/
traces.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5use opentelemetry::trace::{
6    Span as _, SpanContext, SpanId, SpanKind, Status as OtelStatus, TraceContextExt, TraceFlags,
7    TraceId, TraceState, Tracer, TracerProvider,
8};
9use opentelemetry::{Context, KeyValue};
10use opentelemetry_otlp::{Protocol, WithExportConfig, WithHttpConfig};
11use opentelemetry_proto::tonic::trace::v1::Status;
12use opentelemetry_sdk::runtime;
13use opentelemetry_sdk::trace::span_processor_with_async_runtime::BatchSpanProcessor;
14use opentelemetry_sdk::trace::{BatchConfigBuilder, SdkTracerProvider, Tracer as SdkTracer};
15use serde_yaml::Value as YamlValue;
16use tracing::{debug, info, warn};
17use uuid::Uuid;
18
19use crate::internal::observability::adapters;
20use crate::internal::observability::sdk::{
21    get_bool, get_trimmed, get_u64, parse_headers, OTelSdkContext,
22};
23
24#[derive(Debug, Clone)]
25pub struct SpanData {
26    pub trace_id: Vec<u8>,
27    pub span_id: Vec<u8>,
28    pub parent_span_id: Vec<u8>,
29    pub name: String,
30    pub start_time_unix_nano: u64,
31    pub end_time_unix_nano: u64,
32    pub attributes: Vec<KeyValue>,
33    pub status: Option<Status>,
34    pub kind: i32,
35}
36
37#[derive(Clone)]
38pub struct OtlpTraceExporter {
39    tracer: SdkTracer,
40    provider: Arc<SdkTracerProvider>,
41}
42
43pub fn new_trace_id() -> Vec<u8> {
44    Uuid::new_v4().as_bytes().to_vec()
45}
46
47pub fn new_span_id() -> Vec<u8> {
48    Uuid::new_v4().as_bytes()[0..8].to_vec()
49}
50
51pub fn now_unix_nano() -> u64 {
52    std::time::SystemTime::now()
53        .duration_since(std::time::UNIX_EPOCH)
54        .unwrap_or_default()
55        .as_nanos() as u64
56}
57
58pub fn kv(key: &str, value: &str) -> KeyValue {
59    KeyValue::new(key.to_string(), value.to_string())
60}
61
62pub fn kv_int(key: &str, value: i64) -> KeyValue {
63    KeyValue::new(key.to_string(), value)
64}
65
66impl OtlpTraceExporter {
67    pub fn from_observability(observability: Option<&YamlValue>) -> Option<Arc<Self>> {
68        let map = observability?.as_mapping()?;
69        let mut target = adapters::TraceTargetConfig {
70            endpoint: get_trimmed(map, "otlp_tracing_endpoint").map(ToOwned::to_owned),
71            headers: parse_headers(map),
72        };
73
74        adapters::langfuse::apply(map, &mut target);
75        adapters::agentops::apply(map, &mut target);
76
77        let enabled = get_bool(map, "otlp_tracing_enabled").unwrap_or(target.endpoint.is_some());
78        if !enabled {
79            debug!("otlp trace exporter disabled by config");
80            return None;
81        }
82
83        let endpoint = target.endpoint?;
84        let sdk = OTelSdkContext::from_observability(map);
85        let max_queue_size = get_u64(map, "otlp_tracing_max_queue_size")
86            .map(|v| v.clamp(512, 65536))
87            .unwrap_or(8192) as usize;
88        let max_export_batch_size = get_u64(map, "otlp_tracing_max_export_batch_size")
89            .map(|v| v.clamp(64, max_queue_size as u64))
90            .unwrap_or(512) as usize;
91        let scheduled_delay_ms = get_u64(map, "otlp_tracing_scheduled_delay_ms")
92            .map(|v| v.clamp(100, 10000))
93            .unwrap_or(5000);
94
95        let mut exporter_builder = opentelemetry_otlp::SpanExporter::builder()
96            .with_http()
97            .with_protocol(Protocol::HttpBinary)
98            .with_http_client(reqwest::Client::new())
99            .with_endpoint(endpoint.clone());
100        if !target.headers.is_empty() {
101            exporter_builder = exporter_builder
102                .with_headers(target.headers.into_iter().collect::<HashMap<_, _>>());
103        }
104        let exporter = match exporter_builder.build() {
105            Ok(exporter) => exporter,
106            Err(err) => {
107                warn!(
108                    endpoint = %endpoint,
109                    "failed to build otlp trace exporter: {}",
110                    err
111                );
112                return None;
113            }
114        };
115
116        let batch_config = BatchConfigBuilder::default()
117            .with_max_queue_size(max_queue_size)
118            .with_max_export_batch_size(max_export_batch_size)
119            .with_scheduled_delay(Duration::from_millis(scheduled_delay_ms))
120            .build();
121        let batch_processor = BatchSpanProcessor::builder(exporter, runtime::Tokio)
122            .with_batch_config(batch_config)
123            .build();
124        let provider = SdkTracerProvider::builder()
125            .with_resource(sdk.resource)
126            .with_span_processor(batch_processor)
127            .build();
128        let tracer = provider.tracer("microclaw.observability.traces");
129        info!(
130            endpoint = %endpoint,
131            service_name = %sdk.service_name,
132            max_queue_size,
133            max_export_batch_size,
134            scheduled_delay_ms,
135            "otlp trace exporter initialized"
136        );
137
138        Some(Arc::new(Self {
139            tracer,
140            provider: Arc::new(provider),
141        }))
142    }
143
144    pub fn send_span(&self, span: SpanData) {
145        let parent = build_parent_context(&span);
146        let SpanData {
147            trace_id,
148            span_id,
149            name,
150            kind,
151            start_time_unix_nano,
152            end_time_unix_nano,
153            attributes,
154            status,
155            ..
156        } = span;
157        let mut builder = self
158            .tracer
159            .span_builder(name)
160            .with_kind(to_span_kind(kind))
161            .with_start_time(unix_nano_to_time(start_time_unix_nano))
162            .with_attributes(attributes);
163        if trace_id.len() == 16 {
164            if let Ok(bytes) = trace_id.as_slice().try_into() {
165                builder = builder.with_trace_id(TraceId::from_bytes(bytes));
166            }
167        } else {
168            warn!(
169                trace_id_len = trace_id.len(),
170                "invalid trace_id length, expected 16 bytes"
171            );
172        }
173        if span_id.len() == 8 {
174            if let Ok(bytes) = span_id.as_slice().try_into() {
175                builder = builder.with_span_id(SpanId::from_bytes(bytes));
176            }
177        } else {
178            warn!(
179                span_id_len = span_id.len(),
180                "invalid span_id length, expected 8 bytes"
181            );
182        }
183
184        let mut sdk_span = if let Some(parent) = parent {
185            builder.start_with_context(&self.tracer, &parent)
186        } else {
187            builder.start(&self.tracer)
188        };
189
190        if let Some(status) = status {
191            if status.code == 2 {
192                sdk_span.set_status(OtelStatus::error(status.message));
193            } else if status.code == 1 {
194                sdk_span.set_status(OtelStatus::Ok);
195            }
196        }
197
198        sdk_span.end_with_timestamp(unix_nano_to_time(end_time_unix_nano));
199        debug!("trace span submitted to otel sdk");
200    }
201}
202
203impl Drop for OtlpTraceExporter {
204    fn drop(&mut self) {
205        if let Err(err) = self.provider.force_flush() {
206            warn!("otlp trace exporter force_flush failed: {}", err);
207        } else {
208            debug!("otlp trace exporter force_flush succeeded");
209        }
210    }
211}
212
213fn to_span_kind(kind: i32) -> SpanKind {
214    match kind {
215        2 => SpanKind::Server,
216        3 => SpanKind::Client,
217        4 => SpanKind::Producer,
218        5 => SpanKind::Consumer,
219        _ => SpanKind::Internal,
220    }
221}
222
223fn unix_nano_to_time(value: u64) -> SystemTime {
224    UNIX_EPOCH + Duration::from_nanos(value)
225}
226
227fn build_parent_context(span: &SpanData) -> Option<Context> {
228    if span.trace_id.len() != 16 || span.parent_span_id.len() != 8 {
229        return None;
230    }
231    let trace_id = TraceId::from_bytes(span.trace_id.as_slice().try_into().ok()?);
232    let span_id = SpanId::from_bytes(span.parent_span_id.as_slice().try_into().ok()?);
233    let span_context = SpanContext::new(
234        trace_id,
235        span_id,
236        TraceFlags::SAMPLED,
237        true,
238        TraceState::default(),
239    );
240    Some(Context::new().with_remote_span_context(span_context))
241}