use std::collections::HashMap;
use std::path::PathBuf;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum TraceLevel {
Off,
Error,
Warn,
#[default]
Info,
Debug,
Trace,
}
impl TraceLevel {
pub fn parse(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"off" | "none" => Some(Self::Off),
"error" => Some(Self::Error),
"warn" | "warning" => Some(Self::Warn),
"info" => Some(Self::Info),
"debug" => Some(Self::Debug),
"trace" => Some(Self::Trace),
_ => None,
}
}
pub fn name(&self) -> &'static str {
match self {
Self::Off => "OFF",
Self::Error => "ERROR",
Self::Warn => "WARN",
Self::Info => "INFO",
Self::Debug => "DEBUG",
Self::Trace => "TRACE",
}
}
pub fn should_log(&self, level: TraceLevel) -> bool {
level <= *self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SpanStatus {
#[default]
Unset,
Ok,
Error,
}
impl SpanStatus {
pub fn name(&self) -> &'static str {
match self {
Self::Unset => "UNSET",
Self::Ok => "OK",
Self::Error => "ERROR",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SpanKind {
#[default]
Internal,
Server,
Client,
Producer,
Consumer,
}
impl SpanKind {
pub fn name(&self) -> &'static str {
match self {
Self::Internal => "INTERNAL",
Self::Server => "SERVER",
Self::Client => "CLIENT",
Self::Producer => "PRODUCER",
Self::Consumer => "CONSUMER",
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum AttributeValue {
String(String),
Int(i64),
Float(f64),
Bool(bool),
StringArray(Vec<String>),
IntArray(Vec<i64>),
}
impl AttributeValue {
pub fn string(s: impl Into<String>) -> Self {
Self::String(s.into())
}
pub fn int(i: i64) -> Self {
Self::Int(i)
}
pub fn float(f: f64) -> Self {
Self::Float(f)
}
pub fn bool(b: bool) -> Self {
Self::Bool(b)
}
pub fn to_json(&self) -> String {
match self {
Self::String(s) => format!("\"{}\"", escape_json(s)),
Self::Int(i) => i.to_string(),
Self::Float(f) => f.to_string(),
Self::Bool(b) => b.to_string(),
Self::StringArray(arr) => {
let items: Vec<String> = arr
.iter()
.map(|s| format!("\"{}\"", escape_json(s)))
.collect();
format!("[{}]", items.join(", "))
}
Self::IntArray(arr) => {
let items: Vec<String> = arr.iter().map(|i| i.to_string()).collect();
format!("[{}]", items.join(", "))
}
}
}
}
#[derive(Debug, Clone)]
pub struct Span {
pub span_id: String,
pub parent_id: Option<String>,
pub trace_id: String,
pub name: String,
pub kind: SpanKind,
pub start_time: SystemTime,
pub end_time: Option<SystemTime>,
pub status: SpanStatus,
pub status_message: Option<String>,
pub attributes: HashMap<String, AttributeValue>,
pub events: Vec<SpanEvent>,
}
impl Span {
pub fn new(name: &str, trace_id: &str) -> Self {
Self {
span_id: generate_id(),
parent_id: None,
trace_id: trace_id.to_string(),
name: name.to_string(),
kind: SpanKind::Internal,
start_time: SystemTime::now(),
end_time: None,
status: SpanStatus::Unset,
status_message: None,
attributes: HashMap::new(),
events: Vec::new(),
}
}
pub fn child(&self, name: &str) -> Self {
let mut span = Self::new(name, &self.trace_id);
span.parent_id = Some(self.span_id.clone());
span
}
pub fn with_parent(mut self, parent_id: &str) -> Self {
self.parent_id = Some(parent_id.to_string());
self
}
pub fn with_kind(mut self, kind: SpanKind) -> Self {
self.kind = kind;
self
}
pub fn set_attribute(&mut self, key: &str, value: AttributeValue) {
self.attributes.insert(key.to_string(), value);
}
pub fn set_string(&mut self, key: &str, value: impl Into<String>) {
self.set_attribute(key, AttributeValue::string(value));
}
pub fn set_int(&mut self, key: &str, value: i64) {
self.set_attribute(key, AttributeValue::int(value));
}
pub fn set_bool(&mut self, key: &str, value: bool) {
self.set_attribute(key, AttributeValue::bool(value));
}
pub fn add_event(&mut self, name: &str) {
self.events.push(SpanEvent::new(name));
}
pub fn add_event_with_attrs(&mut self, name: &str, attrs: HashMap<String, AttributeValue>) {
self.events.push(SpanEvent::with_attributes(name, attrs));
}
pub fn end_ok(&mut self) {
self.end_time = Some(SystemTime::now());
self.status = SpanStatus::Ok;
}
pub fn end_error(&mut self, message: &str) {
self.end_time = Some(SystemTime::now());
self.status = SpanStatus::Error;
self.status_message = Some(message.to_string());
}
pub fn duration(&self) -> Option<Duration> {
self.end_time.map(|end| {
end.duration_since(self.start_time)
.unwrap_or(Duration::ZERO)
})
}
pub fn is_completed(&self) -> bool {
self.end_time.is_some()
}
pub fn to_json(&self) -> String {
let mut json = String::from("{\n");
json.push_str(&format!(" \"traceId\": \"{}\",\n", self.trace_id));
json.push_str(&format!(" \"spanId\": \"{}\",\n", self.span_id));
if let Some(ref parent) = self.parent_id {
json.push_str(&format!(" \"parentSpanId\": \"{}\",\n", parent));
}
json.push_str(&format!(" \"name\": \"{}\",\n", escape_json(&self.name)));
json.push_str(&format!(" \"kind\": \"{}\",\n", self.kind.name()));
json.push_str(&format!(
" \"startTimeUnixNano\": {},\n",
time_to_nanos(self.start_time)
));
if let Some(end) = self.end_time {
json.push_str(&format!(" \"endTimeUnixNano\": {},\n", time_to_nanos(end)));
}
json.push_str(&format!(
" \"status\": {{\n \"code\": \"{}\"\n }}",
self.status.name()
));
if !self.attributes.is_empty() {
json.push_str(",\n \"attributes\": {\n");
let attrs: Vec<String> = self
.attributes
.iter()
.map(|(k, v)| format!(" \"{}\": {}", escape_json(k), v.to_json()))
.collect();
json.push_str(&attrs.join(",\n"));
json.push_str("\n }");
}
if !self.events.is_empty() {
json.push_str(",\n \"events\": [\n");
let events: Vec<String> = self
.events
.iter()
.map(|e| format!(" {}", e.to_json()))
.collect();
json.push_str(&events.join(",\n"));
json.push_str("\n ]");
}
json.push_str("\n}");
json
}
}
include!("tracing_spanevent.rs");