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