revoke_trace/
span.rs

1use opentelemetry::trace::{SpanKind as OtelSpanKind, Status};
2use std::collections::HashMap;
3use std::time::SystemTime;
4use tracing::Level;
5
6/// Span 类型
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum SpanKind {
9    /// 服务端处理请求
10    Server,
11    /// 客户端发送请求
12    Client,
13    /// 生产者发送消息
14    Producer,
15    /// 消费者接收消息
16    Consumer,
17    /// 内部操作
18    Internal,
19}
20
21impl From<SpanKind> for OtelSpanKind {
22    fn from(kind: SpanKind) -> Self {
23        match kind {
24            SpanKind::Server => OtelSpanKind::Server,
25            SpanKind::Client => OtelSpanKind::Client,
26            SpanKind::Producer => OtelSpanKind::Producer,
27            SpanKind::Consumer => OtelSpanKind::Consumer,
28            SpanKind::Internal => OtelSpanKind::Internal,
29        }
30    }
31}
32
33/// Span 状态
34#[derive(Debug, Clone)]
35pub struct SpanStatus {
36    code: SpanStatusCode,
37    message: String,
38}
39
40/// Span 状态码
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SpanStatusCode {
43    /// 成功
44    Ok,
45    /// 错误
46    Error,
47    /// 未设置
48    Unset,
49}
50
51impl SpanStatus {
52    /// 创建成功状态
53    pub fn ok() -> Self {
54        Self {
55            code: SpanStatusCode::Ok,
56            message: String::new(),
57        }
58    }
59
60    /// 创建错误状态
61    pub fn error(message: impl Into<String>) -> Self {
62        Self {
63            code: SpanStatusCode::Error,
64            message: message.into(),
65        }
66    }
67
68    /// 创建未设置状态
69    pub fn unset() -> Self {
70        Self {
71            code: SpanStatusCode::Unset,
72            message: String::new(),
73        }
74    }
75}
76
77impl From<SpanStatus> for Status {
78    fn from(status: SpanStatus) -> Self {
79        match status.code {
80            SpanStatusCode::Ok => Status::Ok,
81            SpanStatusCode::Error => Status::error(status.message),
82            SpanStatusCode::Unset => Status::Unset,
83        }
84    }
85}
86
87/// Span 构建器
88pub struct SpanBuilder {
89    name: String,
90    kind: SpanKind,
91    attributes: HashMap<String, String>,
92    events: Vec<SpanEvent>,
93    links: Vec<SpanLink>,
94    start_time: Option<SystemTime>,
95    level: Level,
96}
97
98impl SpanBuilder {
99    /// 创建新的 Span 构建器
100    pub fn new(name: impl Into<String>) -> Self {
101        Self {
102            name: name.into(),
103            kind: SpanKind::Internal,
104            attributes: HashMap::new(),
105            events: Vec::new(),
106            links: Vec::new(),
107            start_time: None,
108            level: Level::INFO,
109        }
110    }
111
112    /// 设置 Span 类型
113    pub fn with_kind(mut self, kind: SpanKind) -> Self {
114        self.kind = kind;
115        self
116    }
117
118    /// 添加属性
119    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
120        self.attributes.insert(key.into(), value.into());
121        self
122    }
123
124    /// 批量添加属性
125    pub fn with_attributes<I, K, V>(mut self, attributes: I) -> Self
126    where
127        I: IntoIterator<Item = (K, V)>,
128        K: Into<String>,
129        V: Into<String>,
130    {
131        for (key, value) in attributes {
132            self.attributes.insert(key.into(), value.into());
133        }
134        self
135    }
136
137    /// 添加事件
138    pub fn with_event(mut self, event: SpanEvent) -> Self {
139        self.events.push(event);
140        self
141    }
142
143    /// 添加链接
144    pub fn with_link(mut self, link: SpanLink) -> Self {
145        self.links.push(link);
146        self
147    }
148
149    /// 设置开始时间
150    pub fn with_start_time(mut self, time: SystemTime) -> Self {
151        self.start_time = Some(time);
152        self
153    }
154
155    /// 设置日志级别
156    pub fn with_level(mut self, level: Level) -> Self {
157        self.level = level;
158        self
159    }
160
161    /// 构建并开始 Span
162    pub fn start(self) -> tracing::Span {
163        match self.level {
164            Level::TRACE => {
165                tracing::trace_span!(
166                    target: "revoke_trace",
167                    "{}",
168                    self.name,
169                    otel.kind = ?self.kind
170                )
171            }
172            Level::DEBUG => {
173                tracing::debug_span!(
174                    target: "revoke_trace",
175                    "{}",
176                    self.name,
177                    otel.kind = ?self.kind
178                )
179            }
180            Level::INFO => {
181                tracing::info_span!(
182                    target: "revoke_trace",
183                    "{}",
184                    self.name,
185                    otel.kind = ?self.kind
186                )
187            }
188            Level::WARN => {
189                tracing::warn_span!(
190                    target: "revoke_trace",
191                    "{}",
192                    self.name,
193                    otel.kind = ?self.kind
194                )
195            }
196            Level::ERROR => {
197                tracing::error_span!(
198                    target: "revoke_trace",
199                    "{}",
200                    self.name,
201                    otel.kind = ?self.kind
202                )
203            }
204        }
205    }
206}
207
208/// Span 事件
209#[derive(Debug, Clone)]
210pub struct SpanEvent {
211    name: String,
212    timestamp: SystemTime,
213    attributes: HashMap<String, String>,
214}
215
216impl SpanEvent {
217    /// 创建新的事件
218    pub fn new(name: impl Into<String>) -> Self {
219        Self {
220            name: name.into(),
221            timestamp: SystemTime::now(),
222            attributes: HashMap::new(),
223        }
224    }
225
226    /// 添加属性
227    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
228        self.attributes.insert(key.into(), value.into());
229        self
230    }
231
232    /// 设置时间戳
233    pub fn with_timestamp(mut self, timestamp: SystemTime) -> Self {
234        self.timestamp = timestamp;
235        self
236    }
237}
238
239/// Span 链接
240#[derive(Debug, Clone)]
241pub struct SpanLink {
242    #[allow(dead_code)]
243    context: crate::context::TraceContext,
244    attributes: HashMap<String, String>,
245}
246
247impl SpanLink {
248    /// 创建新的链接
249    pub fn new(context: crate::context::TraceContext) -> Self {
250        Self {
251            context,
252            attributes: HashMap::new(),
253        }
254    }
255
256    /// 添加属性
257    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
258        self.attributes.insert(key.into(), value.into());
259        self
260    }
261}
262
263/// Span 扩展功能
264pub trait SpanExt {
265    /// 记录异常
266    fn record_exception(&self, error: &dyn std::error::Error);
267
268    /// 设置状态
269    fn set_status(&self, status: SpanStatus);
270
271    /// 添加事件
272    fn add_event(&self, event: SpanEvent);
273}
274
275impl SpanExt for tracing::Span {
276    fn record_exception(&self, error: &dyn std::error::Error) {
277        self.record("exception.type", &error.to_string());
278        self.record("exception.message", &error.to_string());
279
280        if let Some(source) = error.source() {
281            self.record("exception.stacktrace", &source.to_string());
282        }
283    }
284
285    fn set_status(&self, status: SpanStatus) {
286        self.record("otel.status_code", format!("{:?}", status.code).as_str());
287        if !status.message.is_empty() {
288            self.record("otel.status_message", &status.message);
289        }
290    }
291
292    fn add_event(&self, event: SpanEvent) {
293        tracing::event!(
294            target: "revoke_trace",
295            parent: self,
296            Level::INFO,
297            name = %event.name,
298            timestamp = ?event.timestamp,
299            ?event.attributes,
300            "span event"
301        );
302    }
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308
309    #[test]
310    fn test_span_builder() {
311        let _span = SpanBuilder::new("test_operation")
312            .with_kind(SpanKind::Server)
313            .with_attribute("http.method", "GET")
314            .with_attribute("http.url", "/api/users")
315            .with_level(Level::DEBUG)
316            .start();
317
318        // 如果没有 panic,说明成功创建
319    }
320
321    #[test]
322    fn test_span_status() {
323        let ok_status = SpanStatus::ok();
324        let error_status = SpanStatus::error("Something went wrong");
325
326        // 只验证转换不会 panic
327        let _otel_status: Status = ok_status.into();
328        let _otel_error: Status = error_status.into();
329    }
330}