1use chrono::prelude::*;
2
3static SPACE: usize = 4;
4
5#[derive(Debug, Serialize)]
6pub struct Context {
7 id: String,
8 key: String,
9 #[serde(serialize_with = "datetime_to_millis")]
10 pub created_on: DateTime<Utc>,
11 pub span_stack: std::cell::RefCell<Vec<crate::Span>>,
12}
13
14fn datetime_to_millis<S>(d: &DateTime<Utc>, ser: S) -> Result<S::Ok, S::Error>
15where
16 S: serde::Serializer,
17{
18 ser.serialize_i64(d.timestamp_millis())
19}
20
21thread_local! {
22 static CONTEXT: std::cell::RefCell<Option<Context>> = std::cell::RefCell::new(None);
23}
24
25impl Context {
26 pub fn new(id: String) -> Context {
27 Context {
28 id,
29 key: uuid::Uuid::new_v4().to_string(),
30 created_on: Utc::now(),
31 span_stack: std::cell::RefCell::new(vec![crate::Span::new("main")]),
32 }
33 }
34
35 pub fn id(&self) -> String {
36 self.id.to_string()
37 }
38
39 pub fn start_span(&self, id: &str) {
40 self.span_stack.borrow_mut().push(crate::Span::new(id));
41 }
42
43 #[allow(dead_code)]
44 pub(crate) fn observe_span_id(&self, id: &str) {
45 let frame = self.span_stack.borrow_mut().pop();
46 if let Some(mut frame) = frame {
47 frame.set_id(id);
48 self.span_stack.borrow_mut().push(frame);
49 }
50 }
51
52 pub(crate) fn observe_span_field(&self, key: &'static str, value: serde_json::Value) {
53 let frame = self.span_stack.borrow_mut().pop();
54 if let Some(mut frame) = frame {
55 frame.add_breadcrumbs(key, value);
56 self.span_stack.borrow_mut().push(frame);
57 }
58 }
59
60 #[allow(dead_code)]
61 pub(crate) fn observe_query(
62 &self,
63 query: String,
64 bind: Option<String>,
65 result: Result<usize, String>,
66 ) {
67 let frame = self.span_stack.borrow_mut().pop();
68 if let Some(mut frame) = frame {
69 frame.add_query(query, bind, result);
70 self.span_stack.borrow_mut().push(frame);
71 }
72 }
73
74 pub(crate) fn observe_span_transient_field(&self, key: &'static str, value: serde_json::Value) {
75 let frame = self.span_stack.borrow_mut().pop();
76 if let Some(mut frame) = frame {
77 frame.add_transient_field(key, value);
78 self.span_stack.borrow_mut().push(frame);
79 }
80 }
81
82 pub(crate) fn observe_span_result(&self, value: impl serde::Serialize) {
83 let frame = self.span_stack.borrow_mut().pop();
84 if let Some(mut frame) = frame {
85 frame.set_result(value);
86 self.span_stack.borrow_mut().push(frame);
87 }
88 }
89
90 pub(crate) fn span_log(&self, value: &'static str) {
91 let frame = self.span_stack.borrow_mut().pop();
92 if let Some(mut frame) = frame {
93 frame.add_log(value);
94 self.span_stack.borrow_mut().push(frame);
95 }
96 }
97
98 pub fn end_span(&self, _is_critical: bool, err: Option<String>) {
99 let child = self.span_stack.borrow_mut().pop();
100 let parent = self.span_stack.borrow_mut().pop();
101 if let Some(mut child_frame) = child {
102 child_frame.set_success(err.is_none()).set_err(err).end();
103 if let Some(mut parent_frame) = parent {
104 parent_frame.add_sub_frame(child_frame.created_on, child_frame);
105 self.span_stack.borrow_mut().push(parent_frame);
106 } else {
107 self.span_stack.borrow_mut().push(child_frame);
108 }
109 }
110 }
111
112 pub(crate) fn end_ctx_frame(&self) {
113 let frame = self.span_stack.borrow_mut().pop();
114 if let Some(mut frame) = frame {
115 frame.end();
116 self.span_stack.borrow_mut().push(frame);
117 }
118 }
119
120 pub fn finalise(&self) {
121 self.end_ctx_frame();
122 }
123
124 pub fn get_key(&self) -> String {
125 self.key.clone()
126 }
127
128 pub fn trace_without_data(&self, print_values: bool) -> String {
129 let mut buffer = "".to_string();
130 for span in self.span_stack.borrow().iter() {
131 print_span_without_data(&mut buffer, span, 0, print_values);
132 }
133 buffer
134 }
135}
136
137pub fn print_span_without_data(
138 buffer: &mut String,
139 span: &crate::Span,
140 space: usize,
141 print_values: bool,
142) {
143 if space != 0 {
144 buffer.push_str(&format!("{}\n", span.id.as_str()));
145 }
146
147 for (_, item) in span.items.iter() {
148 match item {
149 crate::SpanItem::Log { message } => {
150 buffer.push_str(&format!(
151 "{:_>space$}- {message}\n",
152 "",
153 message = message,
154 space = space,
155 ));
156 }
157 crate::SpanItem::Field { name, value } => {
158 if print_values {
159 buffer.push_str(&format!(
160 "{:_>space$}- {name}={value}\n",
161 "",
162 name = name,
163 value = if let serde_json::Value::String(s) = value {
164 s.to_string()
165 } else {
166 value.to_string()
167 },
168 space = space
169 ));
170 } else {
171 buffer.push_str(&format!(
172 "{:_>space$}- {name}: observed\n",
173 "",
174 name = name,
175 space = space
176 ));
177 }
178 }
179 crate::SpanItem::TransientField { name, value: _ } => {
180 buffer.push_str(&format!(
181 "{:_>space$}- {name}: observed\n",
182 "",
183 name = name,
184 space = space
185 ));
186 }
187 crate::SpanItem::Query {
188 query,
189 bind,
190 result,
191 } => {
192 buffer.push_str(&format!(
193 "{:_>space$}- query: {}\n",
194 "",
195 query,
196 space = space
197 ));
198 if bind.is_some() {
199 buffer.push_str(&format!(
200 "{:_>space$} bind: observed\n",
201 "",
202 space = space
203 ));
204 }
205 match result {
206 Ok(rows) => {
207 if print_values {
208 buffer.push_str(&format!(
209 "{:_>space$} rows: {}\n",
210 "",
211 rows,
212 space = space
213 ))
214 } else {
215 buffer.push_str(&format!(
216 "{:_>space$} rows: observed\n",
217 "",
218 space = space
219 ))
220 }
221 }
222 Err(e) => {
223 buffer.push_str(&format!("{:_>space$} error: {}\n", "", e, space = space))
224 }
225 };
226 }
227 crate::SpanItem::Frame(inner) => {
228 buffer.push_str(&format!("{:_>space$}- ", "", space = space));
229 print_span_without_data(buffer, inner, space + SPACE, print_values);
230 }
231 }
232 }
233}