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