use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use starweaver_core::TraceContext;
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum TraceLevel {
#[default]
Info,
Debug,
}
#[allow(clippy::trivially_copy_pass_by_ref)]
const fn is_default_trace_level(level: &TraceLevel) -> bool {
matches!(level, TraceLevel::Info)
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SpanKind {
#[default]
Internal,
Client,
Server,
}
#[allow(clippy::trivially_copy_pass_by_ref)]
const fn is_default_span_kind(kind: &SpanKind) -> bool {
matches!(kind, SpanKind::Internal)
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SpanStatus {
Open,
Ok,
Error {
error_type: String,
},
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct SpanEvent {
pub name: String,
#[serde(default, skip_serializing_if = "is_default_trace_level")]
pub level: TraceLevel,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub attributes: BTreeMap<String, Value>,
}
impl SpanEvent {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
level: TraceLevel::Info,
attributes: BTreeMap::new(),
}
}
#[must_use]
pub const fn debug(mut self) -> Self {
self.level = TraceLevel::Debug;
self
}
#[must_use]
pub fn with_attribute(mut self, key: impl Into<String>, value: Value) -> Self {
self.attributes.insert(key.into(), value);
self
}
}
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub struct SpanSpec {
pub name: String,
#[serde(default, skip_serializing_if = "is_default_span_kind")]
pub kind: SpanKind,
#[serde(default, skip_serializing_if = "is_default_trace_level")]
pub level: TraceLevel,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub attributes: BTreeMap<String, Value>,
}
impl SpanSpec {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
kind: SpanKind::Internal,
level: TraceLevel::Info,
attributes: BTreeMap::new(),
}
}
#[must_use]
pub const fn with_kind(mut self, kind: SpanKind) -> Self {
self.kind = kind;
self
}
#[must_use]
pub const fn debug(mut self) -> Self {
self.level = TraceLevel::Debug;
self
}
#[must_use]
pub fn with_attribute(mut self, key: impl Into<String>, value: Value) -> Self {
self.attributes.insert(key.into(), value);
self
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct RecordedSpan {
pub span_id: String,
pub trace_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_span_id: Option<String>,
pub name: String,
#[serde(default, skip_serializing_if = "is_default_span_kind")]
pub kind: SpanKind,
#[serde(default, skip_serializing_if = "is_default_trace_level")]
pub level: TraceLevel,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub attributes: BTreeMap<String, Value>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub events: Vec<SpanEvent>,
pub status: SpanStatus,
}
#[derive(Clone, Debug)]
pub struct SpanHandle {
context: TraceContext,
span_id: String,
}
impl SpanHandle {
pub(super) fn new(context: TraceContext, span_id: impl Into<String>) -> Self {
Self {
context,
span_id: span_id.into(),
}
}
#[must_use]
pub const fn context(&self) -> &TraceContext {
&self.context
}
#[must_use]
pub fn into_context(self) -> TraceContext {
self.context
}
#[must_use]
pub fn span_id(&self) -> &str {
&self.span_id
}
}