1use std::collections::HashMap;
17use std::fmt;
18use std::time::Duration;
19
20use opentelemetry::baggage::BaggageExt;
21use opentelemetry::metrics::{Counter, Histogram, Meter};
22use opentelemetry::trace::Span as _;
23use opentelemetry::{Context, ContextGuard, KeyValue};
24use opentelemetry_sdk::error::OTelSdkResult;
25use opentelemetry_sdk::trace::{Span, SpanData, SpanProcessor};
26use turnframe_core::observe::{Observer, Signal, SignalKind, SignalLabels};
27
28use crate::metrics::{description, metric_labels};
29use crate::tracing::{ProviderCall, TraceGrouping};
30
31const UNIT_MILLIS: &str = "ms";
33const UNIT_MICROS: &str = "us";
35const UNIT_COUNT: &str = "1";
37
38pub struct OtelObserver {
48 counters: HashMap<Signal, Counter<u64>>,
49 histograms: HashMap<Signal, Histogram<f64>>,
50}
51
52impl OtelObserver {
53 #[must_use]
59 pub fn new(meter: &Meter) -> Self {
60 let mut counters = HashMap::new();
61 let mut histograms = HashMap::new();
62
63 for signal in Signal::ALL {
64 match signal.kind() {
65 SignalKind::DurationMillis | SignalKind::DurationMicros => {
66 let unit = if matches!(signal.kind(), SignalKind::DurationMicros) {
67 UNIT_MICROS
68 } else {
69 UNIT_MILLIS
70 };
71 let histogram = meter
72 .f64_histogram(signal.name())
73 .with_description(description(signal))
74 .with_unit(unit)
75 .build();
76 histograms.insert(signal, histogram);
77 }
78 _ => {
79 let counter = meter
80 .u64_counter(signal.name())
81 .with_description(description(signal))
82 .with_unit(UNIT_COUNT)
83 .build();
84 counters.insert(signal, counter);
85 }
86 }
87 }
88
89 Self {
90 counters,
91 histograms,
92 }
93 }
94
95 #[must_use]
97 pub fn instrument_count(&self) -> usize {
98 self.counters.len() + self.histograms.len()
99 }
100
101 #[must_use]
104 pub fn attributes(signal: Signal, labels: &SignalLabels) -> Vec<KeyValue> {
105 metric_labels(signal, labels)
106 .into_iter()
107 .map(|label| {
108 let (key, value) = label.into_parts();
109 KeyValue::new(key.into_owned(), value.into_owned())
110 })
111 .collect()
112 }
113
114 fn record(&self, signal: Signal, labels: &SignalLabels, duration: Option<Duration>) {
115 let attributes = Self::attributes(signal, labels);
116 match signal.kind() {
117 SignalKind::DurationMillis => {
118 if let (Some(histogram), Some(duration)) = (self.histograms.get(&signal), duration)
119 {
120 histogram.record(duration.as_secs_f64() * 1_000.0, &attributes);
121 }
122 }
123 SignalKind::DurationMicros => {
124 if let (Some(histogram), Some(duration)) = (self.histograms.get(&signal), duration)
125 {
126 histogram.record(duration.as_secs_f64() * 1_000_000.0, &attributes);
127 }
128 }
129 _ => {
130 if let Some(counter) = self.counters.get(&signal) {
131 counter.add(1, &attributes);
132 }
133 }
134 }
135 }
136}
137
138impl fmt::Debug for OtelObserver {
139 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140 f.debug_struct("OtelObserver")
141 .field("counters", &self.counters.len())
142 .field("histograms", &self.histograms.len())
143 .finish()
144 }
145}
146
147impl Observer for OtelObserver {
148 fn observe(&self, signal: &Signal) {
149 self.record(*signal, &SignalLabels::none(), None);
150 }
151
152 fn observe_labeled(&self, signal: &Signal, labels: &SignalLabels) {
153 self.record(*signal, labels, None);
154 }
155
156 fn observe_duration(&self, signal: &Signal, duration: Duration, labels: &SignalLabels) {
157 self.record(*signal, labels, Some(duration));
158 }
159}
160
161#[must_use]
182pub fn provider_call_attributes(call: &ProviderCall) -> Vec<KeyValue> {
183 call.attributes()
184 .into_iter()
185 .map(|(key, value)| KeyValue::new(key, value))
186 .collect()
187}
188
189#[must_use]
191pub fn grouping_attributes(grouping: &TraceGrouping) -> Vec<KeyValue> {
192 grouping
193 .fields()
194 .into_iter()
195 .map(|(key, value)| KeyValue::new(key, value))
196 .collect()
197}
198
199#[must_use]
226pub fn attach_grouping(grouping: &TraceGrouping) -> ContextGuard {
227 Context::map_current(|current| current.with_baggage(grouping_attributes(grouping))).attach()
228}
229
230#[must_use]
237pub fn grouping_from_baggage(context: &Context) -> Vec<KeyValue> {
238 let baggage = context.baggage();
239 crate::attrs::GROUPING_KEYS
240 .iter()
241 .filter_map(|key| {
242 baggage
243 .get(*key)
244 .map(|value| KeyValue::new(*key, value.as_str().to_owned()))
245 })
246 .collect()
247}
248
249#[derive(Debug, Clone, Copy, Default)]
284pub struct GroupingSpanProcessor;
285
286impl GroupingSpanProcessor {
287 #[must_use]
289 pub const fn new() -> Self {
290 Self
291 }
292}
293
294impl SpanProcessor for GroupingSpanProcessor {
295 fn on_start(&self, span: &mut Span, cx: &Context) {
296 for attribute in grouping_from_baggage(cx) {
297 span.set_attribute(attribute);
298 }
299 }
300
301 fn on_end(&self, _span: SpanData) {}
305
306 fn force_flush(&self) -> OTelSdkResult {
307 Ok(())
308 }
309
310 fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
311 Ok(())
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use std::sync::{Arc, Mutex, PoisonError};
318
319 use opentelemetry::metrics::{MeterProvider, NoopMeterProvider};
320 use opentelemetry::trace::{TraceContextExt, Tracer, TracerProvider};
321 use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
322 use turnframe_core::command::RiskClass;
323 use turnframe_core::ids::{AccountId, ConversationId, WorkflowKey};
324
325 use super::*;
326
327 #[test]
328 fn every_signal_gets_exactly_one_instrument() {
329 let observer = OtelObserver::new(&NoopMeterProvider::new().meter("turnframe"));
330 assert_eq!(observer.instrument_count(), Signal::ALL.len());
331 for signal in Signal::ALL {
332 let counted = observer.counters.contains_key(&signal);
333 let measured = observer.histograms.contains_key(&signal);
334 assert!(counted ^ measured, "{signal:?}");
335 }
336 }
337
338 #[test]
339 fn attributes_match_the_metrics_label_set() {
340 let labels = SignalLabels::workflow(WorkflowKey::from("trip"))
341 .with_risk(RiskClass::Destructive)
342 .with_error_code("trip.traveler_missing");
343 for signal in Signal::ALL {
344 let expected = metric_labels(signal, &labels);
345 let attributes = OtelObserver::attributes(signal, &labels);
346 assert_eq!(attributes.len(), expected.len(), "{signal:?}");
347 for (attribute, label) in attributes.iter().zip(expected.iter()) {
348 assert_eq!(attribute.key.as_str(), label.key(), "{signal:?}");
349 assert_eq!(attribute.value.to_string(), label.value(), "{signal:?}");
350 }
351 }
352 }
353
354 #[test]
355 fn recording_through_a_noop_meter_is_harmless() {
356 let observer = OtelObserver::new(&NoopMeterProvider::new().meter("turnframe"));
357 for signal in Signal::ALL {
358 observer.observe(&signal);
359 observer.observe_labeled(&signal, &SignalLabels::none());
360 observer.observe_duration(&signal, Duration::from_millis(2), &SignalLabels::none());
361 }
362 assert!(format!("{observer:?}").starts_with("OtelObserver"));
363 }
364
365 #[test]
366 fn provider_call_attributes_mirror_the_tracing_ones() {
367 let call = ProviderCall::new("openai", "extract", "gpt-x")
368 .with_temperature(0.0)
369 .with_response("resp-1", "gpt-x")
370 .with_finish_reason("stop")
371 .with_cached_usage(1_000, 800, 120);
372
373 let expected = call.attributes();
374 let attributes = provider_call_attributes(&call);
375 assert_eq!(attributes.len(), expected.len());
376 for (attribute, (key, value)) in attributes.iter().zip(expected.iter()) {
377 assert_eq!(attribute.key.as_str(), *key);
378 assert_eq!(attribute.value.to_string(), *value);
379 }
380
381 let input = attributes
382 .iter()
383 .find(|kv| kv.key.as_str() == crate::attrs::GEN_AI_USAGE_INPUT_TOKENS)
384 .expect("input tokens");
385 assert_eq!(
386 input.value.to_string(),
387 "200",
388 "input tokens must be net of cache"
389 );
390 }
391
392 #[derive(Debug, Clone, Default)]
395 struct CapturingProcessor {
396 spans: Arc<Mutex<Vec<SpanData>>>,
397 }
398
399 impl SpanProcessor for CapturingProcessor {
400 fn on_start(&self, _span: &mut Span, _cx: &Context) {}
401
402 fn on_end(&self, span: SpanData) {
403 self.spans
404 .lock()
405 .unwrap_or_else(PoisonError::into_inner)
406 .push(span);
407 }
408
409 fn force_flush(&self) -> OTelSdkResult {
410 Ok(())
411 }
412
413 fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
414 Ok(())
415 }
416 }
417
418 fn provider_with_grouping() -> (SdkTracerProvider, Arc<Mutex<Vec<SpanData>>>) {
421 let capture = CapturingProcessor::default();
422 let spans = Arc::clone(&capture.spans);
423 let provider = SdkTracerProvider::builder()
424 .with_sampler(Sampler::AlwaysOn)
425 .with_span_processor(GroupingSpanProcessor::new())
426 .with_span_processor(capture)
427 .build();
428 (provider, spans)
429 }
430
431 fn span_named(spans: &Arc<Mutex<Vec<SpanData>>>, name: &str) -> SpanData {
432 spans
433 .lock()
434 .unwrap_or_else(PoisonError::into_inner)
435 .iter()
436 .find(|span| span.name == name)
437 .unwrap_or_else(|| panic!("span {name} was never finished"))
438 .clone()
439 }
440
441 fn attributes_of(spans: &Arc<Mutex<Vec<SpanData>>>, name: &str) -> Vec<KeyValue> {
442 span_named(spans, name).attributes
443 }
444
445 fn attribute_value(attributes: &[KeyValue], key: &str) -> Option<String> {
446 attributes
447 .iter()
448 .find(|kv| kv.key.as_str() == key)
449 .map(|kv| kv.value.to_string())
450 }
451
452 #[test]
453 fn the_grouping_lands_on_a_child_span_not_only_on_the_root() {
454 let (provider, spans) = provider_with_grouping();
455 let tracer = provider.tracer("turnframe");
456
457 let grouping = TraceGrouping::for_conversation(ConversationId::nil())
458 .with_account(&AccountId::from("acct-1"))
459 .with_environment("production")
460 .with_release("v0.1.0")
461 .with_tag("trip");
462
463 {
464 let _grouping_guard = attach_grouping(&grouping);
466 let parent = tracer.start("turnframe.turn");
467 let parent_guard = Context::current_with_span(parent).attach();
469 let mut child = tracer.start("turnframe.stage");
470 child.end();
471 drop(parent_guard);
472 }
473
474 let parent = span_named(&spans, "turnframe.turn");
475 let child = span_named(&spans, "turnframe.stage");
476 assert_eq!(
477 child.parent_span_id,
478 parent.span_context.span_id(),
479 "the second span must really be a child of the first"
480 );
481
482 for (key, value) in grouping.fields() {
483 assert_eq!(
484 attribute_value(&child.attributes, key).as_deref(),
485 Some(value.as_str()),
486 "{key} did not reach the child span"
487 );
488 }
489 assert!(
491 child
492 .attributes
493 .iter()
494 .all(|kv| kv.value.to_string() != "acct-1")
495 );
496
497 assert_eq!(
499 attribute_value(&parent.attributes, crate::attrs::SESSION_ID).as_deref(),
500 Some(ConversationId::nil().to_string().as_str())
501 );
502 }
503
504 #[test]
505 fn unrelated_baggage_never_becomes_a_span_attribute() {
506 let (provider, spans) = provider_with_grouping();
507 let tracer = provider.tracer("turnframe");
508
509 {
510 let context = Context::map_current(|current| {
511 current.with_baggage([
512 KeyValue::new(crate::attrs::SESSION_ID, "sess-1"),
513 KeyValue::new("internal.traveler_email", "someone@example.com"),
514 ])
515 });
516 let _guard = context.attach();
517 let mut span = tracer.start("turnframe.stage");
518 span.end();
519 }
520
521 let attributes = attributes_of(&spans, "turnframe.stage");
522 assert_eq!(
523 attribute_value(&attributes, crate::attrs::SESSION_ID).as_deref(),
524 Some("sess-1"),
525 "the grouping key should still be copied"
526 );
527 assert_eq!(
528 attribute_value(&attributes, "internal.traveler_email"),
529 None,
530 "baggage outside the grouping keys must not leak onto spans"
531 );
532 assert!(
533 attributes
534 .iter()
535 .all(|kv| crate::attrs::GROUPING_KEYS.contains(&kv.key.as_str()))
536 );
537 }
538
539 #[test]
540 fn a_span_opened_before_the_grouping_is_attached_carries_nothing() {
541 let (provider, spans) = provider_with_grouping();
542 let tracer = provider.tracer("turnframe");
543 let grouping = TraceGrouping::for_conversation(ConversationId::nil());
544
545 let mut early = tracer.start("turnframe.early");
548 let guard = attach_grouping(&grouping);
549 early.end();
550 drop(guard);
551
552 assert!(attributes_of(&spans, "turnframe.early").is_empty());
553 }
554
555 #[test]
556 fn grouping_travels_through_baggage_and_comes_back_whole() {
557 let grouping = TraceGrouping::for_conversation(ConversationId::nil())
558 .with_account(&AccountId::from("acct-1"))
559 .with_environment("production")
560 .with_release("v0.1.0")
561 .with_tag("trip");
562
563 let guard = attach_grouping(&grouping);
564 let recovered = Context::map_current(grouping_from_baggage);
565 drop(guard);
566
567 let expected = grouping_attributes(&grouping);
568 assert_eq!(recovered.len(), expected.len());
569 for (key, value) in expected
570 .iter()
571 .map(|kv| (kv.key.as_str(), kv.value.to_string()))
572 {
573 let found = recovered
574 .iter()
575 .find(|kv| kv.key.as_str() == key)
576 .unwrap_or_else(|| panic!("{key} missing from baggage"));
577 assert_eq!(found.value.to_string(), value);
578 }
579 assert!(
581 recovered
582 .iter()
583 .all(|kv| crate::attrs::GROUPING_KEYS.contains(&kv.key.as_str()))
584 );
585 }
586}