1use opentelemetry::trace::{SpanKind as OtelSpanKind, Status};
2use std::collections::HashMap;
3use std::time::SystemTime;
4use tracing::Level;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum SpanKind {
9 Server,
11 Client,
13 Producer,
15 Consumer,
17 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#[derive(Debug, Clone)]
35pub struct SpanStatus {
36 code: SpanStatusCode,
37 message: String,
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SpanStatusCode {
43 Ok,
45 Error,
47 Unset,
49}
50
51impl SpanStatus {
52 pub fn ok() -> Self {
54 Self {
55 code: SpanStatusCode::Ok,
56 message: String::new(),
57 }
58 }
59
60 pub fn error(message: impl Into<String>) -> Self {
62 Self {
63 code: SpanStatusCode::Error,
64 message: message.into(),
65 }
66 }
67
68 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
87pub 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 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 pub fn with_kind(mut self, kind: SpanKind) -> Self {
114 self.kind = kind;
115 self
116 }
117
118 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 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 pub fn with_event(mut self, event: SpanEvent) -> Self {
139 self.events.push(event);
140 self
141 }
142
143 pub fn with_link(mut self, link: SpanLink) -> Self {
145 self.links.push(link);
146 self
147 }
148
149 pub fn with_start_time(mut self, time: SystemTime) -> Self {
151 self.start_time = Some(time);
152 self
153 }
154
155 pub fn with_level(mut self, level: Level) -> Self {
157 self.level = level;
158 self
159 }
160
161 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#[derive(Debug, Clone)]
210pub struct SpanEvent {
211 name: String,
212 timestamp: SystemTime,
213 attributes: HashMap<String, String>,
214}
215
216impl SpanEvent {
217 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 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 pub fn with_timestamp(mut self, timestamp: SystemTime) -> Self {
234 self.timestamp = timestamp;
235 self
236 }
237}
238
239#[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 pub fn new(context: crate::context::TraceContext) -> Self {
250 Self {
251 context,
252 attributes: HashMap::new(),
253 }
254 }
255
256 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
263pub trait SpanExt {
265 fn record_exception(&self, error: &dyn std::error::Error);
267
268 fn set_status(&self, status: SpanStatus);
270
271 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 }
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 let _otel_status: Status = ok_status.into();
328 let _otel_error: Status = error_status.into();
329 }
330}