1#![warn(missing_docs)]
30
31pub mod logging;
32
33use rustc_hash::FxHashMap;
34use std::fmt;
35use std::time::{Duration, Instant};
36use parking_lot::Mutex;
37use serde::ser::{Serialize, Serializer, SerializeMap};
38use tracing::{
39 event::Event,
40 field::{Visit, Field},
41 Level,
42 span::{Attributes, Id, Record},
43 subscriber::Subscriber,
44};
45use tracing_subscriber::{
46 CurrentSpan,
47 layer::{Layer, Context},
48};
49use tc_telemetry::{telemetry, TETCORE_INFO};
50use tetcore_tracing::{WASM_NAME_KEY, WASM_TARGET_KEY, WASM_TRACE_IDENTIFIER};
51
52#[doc(hidden)]
53pub use tracing;
54
55const ZERO_DURATION: Duration = Duration::from_nanos(0);
56
57pub struct ProfilingLayer {
59 targets: Vec<(String, Level)>,
60 trace_handler: Box<dyn TraceHandler>,
61 span_data: Mutex<FxHashMap<Id, SpanDatum>>,
62 current_span: CurrentSpan,
63}
64
65#[derive(Debug, Clone)]
67pub enum TracingReceiver {
68 Log,
70 Telemetry,
72}
73
74impl Default for TracingReceiver {
75 fn default() -> Self {
76 Self::Log
77 }
78}
79
80pub trait TraceHandler: Send + Sync {
82 fn handle_span(&self, span: SpanDatum);
84 fn handle_event(&self, event: TraceEvent);
86}
87
88#[derive(Debug)]
90pub struct TraceEvent {
91 pub name: &'static str,
93 pub target: String,
95 pub level: Level,
97 pub values: Values,
99 pub parent_id: Option<Id>,
101}
102
103#[derive(Debug)]
105pub struct SpanDatum {
106 pub id: Id,
108 pub parent_id: Option<Id>,
110 pub name: String,
112 pub target: String,
114 pub level: Level,
116 pub line: u32,
118 pub start_time: Instant,
120 pub overall_time: Duration,
122 pub values: Values,
124}
125
126#[derive(Default, Clone, Debug)]
128pub struct Values {
129 pub bool_values: FxHashMap<String, bool>,
131 pub i64_values: FxHashMap<String, i64>,
133 pub u64_values: FxHashMap<String, u64>,
135 pub string_values: FxHashMap<String, String>,
137}
138
139impl Values {
140 pub fn new() -> Self {
142 Default::default()
143 }
144
145 pub fn is_empty(&self) -> bool {
147 self.bool_values.is_empty() &&
148 self.i64_values.is_empty() &&
149 self.u64_values.is_empty() &&
150 self.string_values.is_empty()
151 }
152}
153
154impl Visit for Values {
155 fn record_i64(&mut self, field: &Field, value: i64) {
156 self.i64_values.insert(field.name().to_string(), value);
157 }
158
159 fn record_u64(&mut self, field: &Field, value: u64) {
160 self.u64_values.insert(field.name().to_string(), value);
161 }
162
163 fn record_bool(&mut self, field: &Field, value: bool) {
164 self.bool_values.insert(field.name().to_string(), value);
165 }
166
167 fn record_str(&mut self, field: &Field, value: &str) {
168 self.string_values.insert(field.name().to_string(), value.to_owned());
169 }
170
171 fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
172 self.string_values.insert(field.name().to_string(), format!("{:?}", value).to_owned());
173 }
174}
175
176impl Serialize for Values {
177 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
178 where S: Serializer,
179 {
180 let len = self.bool_values.len() + self.i64_values.len() + self.u64_values.len() + self.string_values.len();
181 let mut map = serializer.serialize_map(Some(len))?;
182 for (k, v) in &self.bool_values {
183 map.serialize_entry(k, v)?;
184 }
185 for (k, v) in &self.i64_values {
186 map.serialize_entry(k, v)?;
187 }
188 for (k, v) in &self.u64_values {
189 map.serialize_entry(k, v)?;
190 }
191 for (k, v) in &self.string_values {
192 map.serialize_entry(k, v)?;
193 }
194 map.end()
195 }
196}
197
198impl fmt::Display for Values {
199 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
200 let bool_iter = self.bool_values.iter().map(|(k, v)| format!("{}={}", k, v));
201 let i64_iter = self.i64_values.iter().map(|(k, v)| format!("{}={}", k, v));
202 let u64_iter = self.u64_values.iter().map(|(k, v)| format!("{}={}", k, v));
203 let string_iter = self.string_values.iter().map(|(k, v)| format!("{}=\"{}\"", k, v));
204 let values = bool_iter.chain(i64_iter).chain(u64_iter).chain(string_iter).collect::<Vec<String>>().join(", ");
205 write!(f, "{}", values)
206 }
207}
208
209impl ProfilingLayer {
210 pub fn new(receiver: TracingReceiver, targets: &str) -> Self {
215 match receiver {
216 TracingReceiver::Log => Self::new_with_handler(Box::new(LogTraceHandler), targets),
217 TracingReceiver::Telemetry => Self::new_with_handler(
218 Box::new(TelemetryTraceHandler),
219 targets,
220 ),
221 }
222 }
223
224 pub fn new_with_handler(trace_handler: Box<dyn TraceHandler>, targets: &str) -> Self {
230 let targets: Vec<_> = targets.split(',').map(|s| parse_target(s)).collect();
231 Self {
232 targets,
233 trace_handler,
234 span_data: Mutex::new(FxHashMap::default()),
235 current_span: Default::default(),
236 }
237 }
238
239 fn check_target(&self, target: &str, level: &Level) -> bool {
240 for t in &self.targets {
241 if target.starts_with(t.0.as_str()) && level <= &t.1 {
242 return true;
243 }
244 }
245 false
246 }
247}
248
249fn parse_target(s: &str) -> (String, Level) {
252 match s.find('=') {
253 Some(i) => {
254 let target = s[0..i].to_string();
255 if s.len() > i {
256 let level = s[i + 1..s.len()].parse::<Level>().unwrap_or(Level::TRACE);
257 (target, level)
258 } else {
259 (target, Level::TRACE)
260 }
261 }
262 None => (s.to_string(), Level::TRACE)
263 }
264}
265
266impl<S: Subscriber> Layer<S> for ProfilingLayer {
267 fn new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<S>) {
268 let mut values = Values::default();
269 attrs.record(&mut values);
270 let span_datum = SpanDatum {
271 id: id.clone(),
272 parent_id: attrs.parent().cloned().or_else(|| self.current_span.id()),
273 name: attrs.metadata().name().to_owned(),
274 target: attrs.metadata().target().to_owned(),
275 level: attrs.metadata().level().clone(),
276 line: attrs.metadata().line().unwrap_or(0),
277 start_time: Instant::now(),
278 overall_time: ZERO_DURATION,
279 values,
280 };
281 self.span_data.lock().insert(id.clone(), span_datum);
282 }
283
284 fn on_record(&self, span: &Id, values: &Record<'_>, _ctx: Context<S>) {
285 let mut span_data = self.span_data.lock();
286 if let Some(s) = span_data.get_mut(span) {
287 values.record(&mut s.values);
288 }
289 }
290
291 fn on_event(&self, event: &Event<'_>, _ctx: Context<S>) {
292 let mut values = Values::default();
293 event.record(&mut values);
294 let trace_event = TraceEvent {
295 name: event.metadata().name(),
296 target: event.metadata().target().to_owned(),
297 level: event.metadata().level().clone(),
298 values,
299 parent_id: event.parent().cloned().or_else(|| self.current_span.id()),
300 };
301 self.trace_handler.handle_event(trace_event);
302 }
303
304 fn on_enter(&self, span: &Id, _ctx: Context<S>) {
305 self.current_span.enter(span.clone());
306 let mut span_data = self.span_data.lock();
307 let start_time = Instant::now();
308 if let Some(mut s) = span_data.get_mut(&span) {
309 s.start_time = start_time;
310 }
311 }
312
313 fn on_exit(&self, span: &Id, _ctx: Context<S>) {
314 self.current_span.exit();
315 let end_time = Instant::now();
316 let span_datum = {
317 let mut span_data = self.span_data.lock();
318 span_data.remove(&span)
319 };
320
321 if let Some(mut span_datum) = span_datum {
322 span_datum.overall_time += end_time - span_datum.start_time;
323 if span_datum.name == WASM_TRACE_IDENTIFIER {
324 span_datum.values.bool_values.insert("wasm".to_owned(), true);
325 if let Some(n) = span_datum.values.string_values.remove(WASM_NAME_KEY) {
326 span_datum.name = n;
327 }
328 if let Some(t) = span_datum.values.string_values.remove(WASM_TARGET_KEY) {
329 span_datum.target = t;
330 }
331 if self.check_target(&span_datum.target, &span_datum.level) {
332 self.trace_handler.handle_span(span_datum);
333 }
334 } else {
335 self.trace_handler.handle_span(span_datum);
336 }
337 };
338 }
339
340 fn on_close(&self, span: Id, ctx: Context<S>) {
341 self.on_exit(&span, ctx)
342 }
343}
344
345pub struct LogTraceHandler;
347
348fn log_level(level: Level) -> log::Level {
349 match level {
350 Level::TRACE => log::Level::Trace,
351 Level::DEBUG => log::Level::Debug,
352 Level::INFO => log::Level::Info,
353 Level::WARN => log::Level::Warn,
354 Level::ERROR => log::Level::Error,
355 }
356}
357
358impl TraceHandler for LogTraceHandler {
359 fn handle_span(&self, span_datum: SpanDatum) {
360 if span_datum.values.is_empty() {
361 log::log!(
362 log_level(span_datum.level),
363 "{}: {}, time: {}, id: {}, parent_id: {:?}",
364 span_datum.target,
365 span_datum.name,
366 span_datum.overall_time.as_nanos(),
367 span_datum.id.into_u64(),
368 span_datum.parent_id.map(|s| s.into_u64()),
369 );
370 } else {
371 log::log!(
372 log_level(span_datum.level),
373 "{}: {}, time: {}, id: {}, parent_id: {:?}, values: {}",
374 span_datum.target,
375 span_datum.name,
376 span_datum.overall_time.as_nanos(),
377 span_datum.id.into_u64(),
378 span_datum.parent_id.map(|s| s.into_u64()),
379 span_datum.values,
380 );
381 }
382 }
383
384 fn handle_event(&self, event: TraceEvent) {
385 log::log!(
386 log_level(event.level),
387 "{}, parent_id: {:?}, {}",
388 event.target,
389 event.parent_id.map(|s| s.into_u64()),
390 event.values,
391 );
392 }
393}
394
395pub struct TelemetryTraceHandler;
399
400impl TraceHandler for TelemetryTraceHandler {
401 fn handle_span(&self, span_datum: SpanDatum) {
402 telemetry!(TETCORE_INFO; "tracing.profiling";
403 "name" => span_datum.name,
404 "target" => span_datum.target,
405 "time" => span_datum.overall_time.as_nanos(),
406 "id" => span_datum.id.into_u64(),
407 "parent_id" => span_datum.parent_id.as_ref().map(|i| i.into_u64()),
408 "values" => span_datum.values
409 );
410 }
411
412 fn handle_event(&self, event: TraceEvent) {
413 telemetry!(TETCORE_INFO; "tracing.event";
414 "name" => event.name,
415 "target" => event.target,
416 "parent_id" => event.parent_id.as_ref().map(|i| i.into_u64()),
417 "values" => event.values
418 );
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use std::sync::Arc;
426 use tracing_subscriber::layer::SubscriberExt;
427
428 struct TestTraceHandler {
429 spans: Arc<Mutex<Vec<SpanDatum>>>,
430 events: Arc<Mutex<Vec<TraceEvent>>>,
431 }
432
433 impl TraceHandler for TestTraceHandler {
434 fn handle_span(&self, sd: SpanDatum) {
435 self.spans.lock().push(sd);
436 }
437
438 fn handle_event(&self, event: TraceEvent) {
439 self.events.lock().push(event);
440 }
441 }
442
443 type TestSubscriber = tracing_subscriber::layer::Layered<
444 ProfilingLayer,
445 tracing_subscriber::fmt::Subscriber
446 >;
447
448 fn setup_subscriber() -> (TestSubscriber, Arc<Mutex<Vec<SpanDatum>>>, Arc<Mutex<Vec<TraceEvent>>>) {
449 let spans = Arc::new(Mutex::new(Vec::new()));
450 let events = Arc::new(Mutex::new(Vec::new()));
451 let handler = TestTraceHandler {
452 spans: spans.clone(),
453 events: events.clone(),
454 };
455 let layer = ProfilingLayer::new_with_handler(
456 Box::new(handler),
457 "test_target",
458 );
459 let subscriber = tracing_subscriber::fmt().finish().with(layer);
460 (subscriber, spans, events)
461 }
462
463 #[test]
464 fn test_span() {
465 let (sub, spans, events) = setup_subscriber();
466 let _sub_guard = tracing::subscriber::set_default(sub);
467 let span = tracing::info_span!(target: "test_target", "test_span1");
468 assert_eq!(spans.lock().len(), 0);
469 assert_eq!(events.lock().len(), 0);
470 let _guard = span.enter();
471 assert_eq!(spans.lock().len(), 0);
472 assert_eq!(events.lock().len(), 0);
473 drop(_guard);
474 drop(span);
475 assert_eq!(spans.lock().len(), 1);
476 assert_eq!(events.lock().len(), 0);
477 let sd = spans.lock().remove(0);
478 assert_eq!(sd.name, "test_span1");
479 assert_eq!(sd.target, "test_target");
480 let time: u128 = sd.overall_time.as_nanos();
481 assert!(time > 0);
482 }
483
484 #[test]
485 fn test_span_parent_id() {
486 let (sub, spans, _events) = setup_subscriber();
487 let _sub_guard = tracing::subscriber::set_default(sub);
488 let span1 = tracing::info_span!(target: "test_target", "test_span1");
489 let _guard1 = span1.enter();
490 let span2 = tracing::info_span!(target: "test_target", "test_span2");
491 let _guard2 = span2.enter();
492 drop(_guard2);
493 drop(span2);
494 let sd2 = spans.lock().remove(0);
495 drop(_guard1);
496 drop(span1);
497 let sd1 = spans.lock().remove(0);
498 assert_eq!(sd1.id, sd2.parent_id.unwrap())
499 }
500
501 #[test]
502 fn test_span_values() {
503 let (sub, spans, _events) = setup_subscriber();
504 let _sub_guard = tracing::subscriber::set_default(sub);
505 let test_bool = true;
506 let test_u64 = 1u64;
507 let test_i64 = 2i64;
508 let test_str = "test_str";
509 let span = tracing::info_span!(
510 target: "test_target",
511 "test_span1",
512 test_bool,
513 test_u64,
514 test_i64,
515 test_str
516 );
517 let _guard = span.enter();
518 drop(_guard);
519 drop(span);
520 let sd = spans.lock().remove(0);
521 assert_eq!(sd.name, "test_span1");
522 assert_eq!(sd.target, "test_target");
523 let values = sd.values;
524 assert_eq!(values.bool_values.get("test_bool").unwrap(), &test_bool);
525 assert_eq!(values.u64_values.get("test_u64").unwrap(), &test_u64);
526 assert_eq!(values.i64_values.get("test_i64").unwrap(), &test_i64);
527 assert_eq!(values.string_values.get("test_str").unwrap(), &test_str.to_owned());
528 }
529
530 #[test]
531 fn test_event() {
532 let (sub, _spans, events) = setup_subscriber();
533 let _sub_guard = tracing::subscriber::set_default(sub);
534 tracing::event!(target: "test_target", tracing::Level::INFO, "test_event");
535 let mut te1 = events.lock().remove(0);
536 assert_eq!(te1.values.string_values.remove(&"message".to_owned()).unwrap(), "test_event".to_owned());
537 }
538
539 #[test]
540 fn test_event_parent_id() {
541 let (sub, spans, events) = setup_subscriber();
542 let _sub_guard = tracing::subscriber::set_default(sub);
543
544 let span1 = tracing::info_span!(target: "test_target", "test_span1");
546 let _guard1 = span1.enter();
547
548 tracing::event!(target: "test_target", tracing::Level::INFO, "test_event");
550
551 drop(_guard1);
553 drop(span1);
554
555 let sd1 = spans.lock().remove(0);
556 let te1 = events.lock().remove(0);
557
558 assert_eq!(sd1.id, te1.parent_id.unwrap());
559 }
560
561 #[test]
562 fn test_parent_id_with_threads() {
563 use std::sync::mpsc;
564 use std::thread;
565
566 let (sub, spans, events) = setup_subscriber();
567 let _sub_guard = tracing::subscriber::set_global_default(sub);
568 let span1 = tracing::info_span!(target: "test_target", "test_span1");
569 let _guard1 = span1.enter();
570
571 let (tx, rx) = mpsc::channel();
572 let handle = thread::spawn(move || {
573 let span2 = tracing::info_span!(target: "test_target", "test_span2");
574 let _guard2 = span2.enter();
575 tracing::event!(target: "test_target", tracing::Level::INFO, "test_event1");
577 for msg in rx.recv() {
578 if msg == false {
579 break;
580 }
581 }
582 });
584
585 while events.lock().is_empty() {
587 thread::sleep(Duration::from_millis(1));
588 }
589
590 tracing::event!(target: "test_target", tracing::Level::INFO, "test_event2");
592
593 let _ = tx.send(false);
595 let _ = handle.join();
596
597 while spans.lock().is_empty() {
599 thread::sleep(Duration::from_millis(1));
600 }
601 let span2 = spans.lock().remove(0);
602 let event1 = events.lock().remove(0);
603 drop(_guard1);
604 drop(span1);
605
606 tracing::event!(target: "test_target", tracing::Level::INFO, "test_event3");
608
609 let span1 = spans.lock().remove(0);
610 let event2 = events.lock().remove(0);
611
612 assert_eq!(event1.values.string_values.get("message").unwrap(), "test_event1");
613 assert_eq!(event2.values.string_values.get("message").unwrap(), "test_event2");
614 assert!(span1.parent_id.is_none());
615 assert!(span2.parent_id.is_none());
616 assert_eq!(span2.id, event1.parent_id.unwrap());
617 assert_eq!(span1.id, event2.parent_id.unwrap());
618 assert_ne!(span2.id, span1.id);
619
620 let event3 = events.lock().remove(0);
621 assert!(event3.parent_id.is_none());
622 }
623}