use std::collections::HashMap;
use std::fmt;
use std::time::Duration;
use opentelemetry::baggage::BaggageExt;
use opentelemetry::metrics::{Counter, Histogram, Meter};
use opentelemetry::trace::Span as _;
use opentelemetry::{Context, ContextGuard, KeyValue};
use opentelemetry_sdk::error::OTelSdkResult;
use opentelemetry_sdk::trace::{Span, SpanData, SpanProcessor};
use turnframe_core::observe::{Observer, Signal, SignalKind, SignalLabels};
use crate::metrics::{description, metric_labels};
use crate::tracing::{ProviderCall, TraceGrouping};
const UNIT_MILLIS: &str = "ms";
const UNIT_MICROS: &str = "us";
const UNIT_COUNT: &str = "1";
pub struct OtelObserver {
counters: HashMap<Signal, Counter<u64>>,
histograms: HashMap<Signal, Histogram<f64>>,
}
impl OtelObserver {
#[must_use]
pub fn new(meter: &Meter) -> Self {
let mut counters = HashMap::new();
let mut histograms = HashMap::new();
for signal in Signal::ALL {
match signal.kind() {
SignalKind::DurationMillis | SignalKind::DurationMicros => {
let unit = if matches!(signal.kind(), SignalKind::DurationMicros) {
UNIT_MICROS
} else {
UNIT_MILLIS
};
let histogram = meter
.f64_histogram(signal.name())
.with_description(description(signal))
.with_unit(unit)
.build();
histograms.insert(signal, histogram);
}
_ => {
let counter = meter
.u64_counter(signal.name())
.with_description(description(signal))
.with_unit(UNIT_COUNT)
.build();
counters.insert(signal, counter);
}
}
}
Self {
counters,
histograms,
}
}
#[must_use]
pub fn instrument_count(&self) -> usize {
self.counters.len() + self.histograms.len()
}
#[must_use]
pub fn attributes(signal: Signal, labels: &SignalLabels) -> Vec<KeyValue> {
metric_labels(signal, labels)
.into_iter()
.map(|label| {
let (key, value) = label.into_parts();
KeyValue::new(key.into_owned(), value.into_owned())
})
.collect()
}
fn record(&self, signal: Signal, labels: &SignalLabels, duration: Option<Duration>) {
let attributes = Self::attributes(signal, labels);
match signal.kind() {
SignalKind::DurationMillis => {
if let (Some(histogram), Some(duration)) = (self.histograms.get(&signal), duration)
{
histogram.record(duration.as_secs_f64() * 1_000.0, &attributes);
}
}
SignalKind::DurationMicros => {
if let (Some(histogram), Some(duration)) = (self.histograms.get(&signal), duration)
{
histogram.record(duration.as_secs_f64() * 1_000_000.0, &attributes);
}
}
_ => {
if let Some(counter) = self.counters.get(&signal) {
counter.add(1, &attributes);
}
}
}
}
}
impl fmt::Debug for OtelObserver {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OtelObserver")
.field("counters", &self.counters.len())
.field("histograms", &self.histograms.len())
.finish()
}
}
impl Observer for OtelObserver {
fn observe(&self, signal: &Signal) {
self.record(*signal, &SignalLabels::none(), None);
}
fn observe_labeled(&self, signal: &Signal, labels: &SignalLabels) {
self.record(*signal, labels, None);
}
fn observe_duration(&self, signal: &Signal, duration: Duration, labels: &SignalLabels) {
self.record(*signal, labels, Some(duration));
}
}
#[must_use]
pub fn provider_call_attributes(call: &ProviderCall) -> Vec<KeyValue> {
call.attributes()
.into_iter()
.map(|(key, value)| KeyValue::new(key, value))
.collect()
}
#[must_use]
pub fn grouping_attributes(grouping: &TraceGrouping) -> Vec<KeyValue> {
grouping
.fields()
.into_iter()
.map(|(key, value)| KeyValue::new(key, value))
.collect()
}
#[must_use]
pub fn attach_grouping(grouping: &TraceGrouping) -> ContextGuard {
Context::map_current(|current| current.with_baggage(grouping_attributes(grouping))).attach()
}
#[must_use]
pub fn grouping_from_baggage(context: &Context) -> Vec<KeyValue> {
let baggage = context.baggage();
crate::attrs::GROUPING_KEYS
.iter()
.filter_map(|key| {
baggage
.get(*key)
.map(|value| KeyValue::new(*key, value.as_str().to_owned()))
})
.collect()
}
#[derive(Debug, Clone, Copy, Default)]
pub struct GroupingSpanProcessor;
impl GroupingSpanProcessor {
#[must_use]
pub const fn new() -> Self {
Self
}
}
impl SpanProcessor for GroupingSpanProcessor {
fn on_start(&self, span: &mut Span, cx: &Context) {
for attribute in grouping_from_baggage(cx) {
span.set_attribute(attribute);
}
}
fn on_end(&self, _span: SpanData) {}
fn force_flush(&self) -> OTelSdkResult {
Ok(())
}
fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex, PoisonError};
use opentelemetry::metrics::{MeterProvider, NoopMeterProvider};
use opentelemetry::trace::{TraceContextExt, Tracer, TracerProvider};
use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
use turnframe_core::command::RiskClass;
use turnframe_core::ids::{AccountId, ConversationId, WorkflowKey};
use super::*;
#[test]
fn every_signal_gets_exactly_one_instrument() {
let observer = OtelObserver::new(&NoopMeterProvider::new().meter("turnframe"));
assert_eq!(observer.instrument_count(), Signal::ALL.len());
for signal in Signal::ALL {
let counted = observer.counters.contains_key(&signal);
let measured = observer.histograms.contains_key(&signal);
assert!(counted ^ measured, "{signal:?}");
}
}
#[test]
fn attributes_match_the_metrics_label_set() {
let labels = SignalLabels::workflow(WorkflowKey::from("trip"))
.with_risk(RiskClass::Destructive)
.with_error_code("trip.traveler_missing");
for signal in Signal::ALL {
let expected = metric_labels(signal, &labels);
let attributes = OtelObserver::attributes(signal, &labels);
assert_eq!(attributes.len(), expected.len(), "{signal:?}");
for (attribute, label) in attributes.iter().zip(expected.iter()) {
assert_eq!(attribute.key.as_str(), label.key(), "{signal:?}");
assert_eq!(attribute.value.to_string(), label.value(), "{signal:?}");
}
}
}
#[test]
fn recording_through_a_noop_meter_is_harmless() {
let observer = OtelObserver::new(&NoopMeterProvider::new().meter("turnframe"));
for signal in Signal::ALL {
observer.observe(&signal);
observer.observe_labeled(&signal, &SignalLabels::none());
observer.observe_duration(&signal, Duration::from_millis(2), &SignalLabels::none());
}
assert!(format!("{observer:?}").starts_with("OtelObserver"));
}
#[test]
fn provider_call_attributes_mirror_the_tracing_ones() {
let call = ProviderCall::new("openai", "extract", "gpt-x")
.with_temperature(0.0)
.with_response("resp-1", "gpt-x")
.with_finish_reason("stop")
.with_cached_usage(1_000, 800, 120);
let expected = call.attributes();
let attributes = provider_call_attributes(&call);
assert_eq!(attributes.len(), expected.len());
for (attribute, (key, value)) in attributes.iter().zip(expected.iter()) {
assert_eq!(attribute.key.as_str(), *key);
assert_eq!(attribute.value.to_string(), *value);
}
let input = attributes
.iter()
.find(|kv| kv.key.as_str() == crate::attrs::GEN_AI_USAGE_INPUT_TOKENS)
.expect("input tokens");
assert_eq!(
input.value.to_string(),
"200",
"input tokens must be net of cache"
);
}
#[derive(Debug, Clone, Default)]
struct CapturingProcessor {
spans: Arc<Mutex<Vec<SpanData>>>,
}
impl SpanProcessor for CapturingProcessor {
fn on_start(&self, _span: &mut Span, _cx: &Context) {}
fn on_end(&self, span: SpanData) {
self.spans
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(span);
}
fn force_flush(&self) -> OTelSdkResult {
Ok(())
}
fn shutdown_with_timeout(&self, _timeout: Duration) -> OTelSdkResult {
Ok(())
}
}
fn provider_with_grouping() -> (SdkTracerProvider, Arc<Mutex<Vec<SpanData>>>) {
let capture = CapturingProcessor::default();
let spans = Arc::clone(&capture.spans);
let provider = SdkTracerProvider::builder()
.with_sampler(Sampler::AlwaysOn)
.with_span_processor(GroupingSpanProcessor::new())
.with_span_processor(capture)
.build();
(provider, spans)
}
fn span_named(spans: &Arc<Mutex<Vec<SpanData>>>, name: &str) -> SpanData {
spans
.lock()
.unwrap_or_else(PoisonError::into_inner)
.iter()
.find(|span| span.name == name)
.unwrap_or_else(|| panic!("span {name} was never finished"))
.clone()
}
fn attributes_of(spans: &Arc<Mutex<Vec<SpanData>>>, name: &str) -> Vec<KeyValue> {
span_named(spans, name).attributes
}
fn attribute_value(attributes: &[KeyValue], key: &str) -> Option<String> {
attributes
.iter()
.find(|kv| kv.key.as_str() == key)
.map(|kv| kv.value.to_string())
}
#[test]
fn the_grouping_lands_on_a_child_span_not_only_on_the_root() {
let (provider, spans) = provider_with_grouping();
let tracer = provider.tracer("turnframe");
let grouping = TraceGrouping::for_conversation(ConversationId::nil())
.with_account(&AccountId::from("acct-1"))
.with_environment("production")
.with_release("v0.1.0")
.with_tag("trip");
{
let _grouping_guard = attach_grouping(&grouping);
let parent = tracer.start("turnframe.turn");
let parent_guard = Context::current_with_span(parent).attach();
let mut child = tracer.start("turnframe.stage");
child.end();
drop(parent_guard);
}
let parent = span_named(&spans, "turnframe.turn");
let child = span_named(&spans, "turnframe.stage");
assert_eq!(
child.parent_span_id,
parent.span_context.span_id(),
"the second span must really be a child of the first"
);
for (key, value) in grouping.fields() {
assert_eq!(
attribute_value(&child.attributes, key).as_deref(),
Some(value.as_str()),
"{key} did not reach the child span"
);
}
assert!(
child
.attributes
.iter()
.all(|kv| kv.value.to_string() != "acct-1")
);
assert_eq!(
attribute_value(&parent.attributes, crate::attrs::SESSION_ID).as_deref(),
Some(ConversationId::nil().to_string().as_str())
);
}
#[test]
fn unrelated_baggage_never_becomes_a_span_attribute() {
let (provider, spans) = provider_with_grouping();
let tracer = provider.tracer("turnframe");
{
let context = Context::map_current(|current| {
current.with_baggage([
KeyValue::new(crate::attrs::SESSION_ID, "sess-1"),
KeyValue::new("internal.traveler_email", "someone@example.com"),
])
});
let _guard = context.attach();
let mut span = tracer.start("turnframe.stage");
span.end();
}
let attributes = attributes_of(&spans, "turnframe.stage");
assert_eq!(
attribute_value(&attributes, crate::attrs::SESSION_ID).as_deref(),
Some("sess-1"),
"the grouping key should still be copied"
);
assert_eq!(
attribute_value(&attributes, "internal.traveler_email"),
None,
"baggage outside the grouping keys must not leak onto spans"
);
assert!(
attributes
.iter()
.all(|kv| crate::attrs::GROUPING_KEYS.contains(&kv.key.as_str()))
);
}
#[test]
fn a_span_opened_before_the_grouping_is_attached_carries_nothing() {
let (provider, spans) = provider_with_grouping();
let tracer = provider.tracer("turnframe");
let grouping = TraceGrouping::for_conversation(ConversationId::nil());
let mut early = tracer.start("turnframe.early");
let guard = attach_grouping(&grouping);
early.end();
drop(guard);
assert!(attributes_of(&spans, "turnframe.early").is_empty());
}
#[test]
fn grouping_travels_through_baggage_and_comes_back_whole() {
let grouping = TraceGrouping::for_conversation(ConversationId::nil())
.with_account(&AccountId::from("acct-1"))
.with_environment("production")
.with_release("v0.1.0")
.with_tag("trip");
let guard = attach_grouping(&grouping);
let recovered = Context::map_current(grouping_from_baggage);
drop(guard);
let expected = grouping_attributes(&grouping);
assert_eq!(recovered.len(), expected.len());
for (key, value) in expected
.iter()
.map(|kv| (kv.key.as_str(), kv.value.to_string()))
{
let found = recovered
.iter()
.find(|kv| kv.key.as_str() == key)
.unwrap_or_else(|| panic!("{key} missing from baggage"));
assert_eq!(found.value.to_string(), value);
}
assert!(
recovered
.iter()
.all(|kv| crate::attrs::GROUPING_KEYS.contains(&kv.key.as_str()))
);
}
}