use async_trait::async_trait;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct TraceContext {
pub traceparent: Option<String>,
pub tracestate: Option<String>,
}
impl TraceContext {
pub fn new() -> Self {
Self::default()
}
pub fn from_traceparent(traceparent: impl Into<String>) -> Self {
Self::new().with_traceparent(traceparent)
}
pub fn with_traceparent(mut self, traceparent: impl Into<String>) -> Self {
self.traceparent = Some(traceparent.into());
self
}
pub fn with_tracestate(mut self, tracestate: impl Into<String>) -> Self {
self.tracestate = Some(tracestate.into());
self
}
pub fn is_empty(&self) -> bool {
self.traceparent.is_none() && self.tracestate.is_none()
}
}
#[async_trait]
pub trait TraceContextProvider: Send + Sync + 'static {
async fn get_trace_context(&self) -> TraceContext;
}
pub(crate) fn inject_trace_context(params: &mut serde_json::Value, ctx: &TraceContext) {
if let Some(tp) = &ctx.traceparent {
params["traceparent"] = serde_json::Value::String(tp.clone());
}
if let Some(ts) = &ctx.tracestate {
params["tracestate"] = serde_json::Value::String(ts.clone());
}
}
#[cfg(test)]
mod tests {
use super::TraceContext;
#[test]
fn new_yields_empty_context() {
let ctx = TraceContext::new();
assert!(ctx.is_empty());
assert!(ctx.traceparent.is_none());
assert!(ctx.tracestate.is_none());
}
#[test]
fn builder_composes_traceparent_and_tracestate() {
let ctx = TraceContext::new()
.with_traceparent("00-trace-span-01")
.with_tracestate("vendor=key");
assert_eq!(ctx.traceparent.as_deref(), Some("00-trace-span-01"));
assert_eq!(ctx.tracestate.as_deref(), Some("vendor=key"));
assert!(!ctx.is_empty());
}
#[test]
fn from_traceparent_matches_builder() {
let direct = TraceContext::from_traceparent("00-trace-span-01");
let chained = TraceContext::new().with_traceparent("00-trace-span-01");
assert_eq!(direct, chained);
}
}