1use std::{
2 collections::HashMap,
3 sync::Arc,
4 time::{Duration, Instant},
5};
6
7use tokio::{
8 sync::{mpsc, Mutex},
9 task,
10};
11use tracing::{span, Subscriber};
12use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
13
14#[derive(Debug)]
15enum SpanEvent {
16 Enter { name: String, metadata: String },
17 Exit { name: String },
18}
19
20#[derive(Debug)]
21#[allow(unused)]
22struct SpanTiming {
23 name: String,
24 metadata: String,
25 start: Instant,
26 elapsed: Duration,
27 count: u128,
28}
29
30#[derive(Debug)]
31pub struct TracingLayer {
32 sender: mpsc::Sender<SpanEvent>,
33}
34
35impl<S> Layer<S> for TracingLayer
36where
37 S: Subscriber + for<'a> LookupSpan<'a>,
38{
39 fn on_enter(&self, id: &span::Id, ctx: Context<S>) {
40 if let Some(span) = ctx.span(id) {
41 let name = span.name().to_string();
42 let metadata = span
43 .extensions()
44 .get::<tracing::field::FieldSet>()
45 .map(|fields| format!("{:?}", fields))
46 .unwrap_or_else(|| "".to_string());
47
48 let _ = self.sender.try_send(SpanEvent::Enter { name, metadata });
49 }
50 }
51
52 fn on_exit(&self, id: &span::Id, ctx: Context<S>) {
53 if let Some(span) = ctx.span(id) {
54 let name = span.name().to_string();
55 let _ = self.sender.try_send(SpanEvent::Exit { name });
56 }
57 }
58}
59
60impl TracingLayer {
61 pub fn new() -> Self {
62 let (sender, mut receiver) = mpsc::channel(100);
63 let spans = Arc::new(Mutex::new(HashMap::new()));
64
65 let spans_clone = spans.clone();
67 task::spawn(async move {
68 loop {
69 tokio::select! {
70 Some(event) = receiver.recv() => {
71 let mut spans = spans_clone.lock().await;
72 match event {
73 SpanEvent::Enter { name, metadata, .. } => {
74 let e = spans.entry(name).or_insert_with_key(|name| SpanTiming {
75 name: name.clone(), metadata, start: Instant::now(), elapsed: Duration::ZERO, count: 0
76 });
77 e.start = Instant::now();
78 e.count += 1;
79 }
80 SpanEvent::Exit { name, .. } => {
81 if let Some(span) = spans.get_mut(&name) {
82 span.elapsed = span.start.elapsed();
83 span.count += 1;
84 }
85 }
86 }
87 }
88 _ = tokio::time::sleep(Duration::from_secs(5)) => {
89 let spans = spans_clone.lock().await;
90
91 let mut longest_spans: Vec<_> = spans.iter().collect();
92 longest_spans.sort_by_key(|(_, timing)| timing.elapsed);
93 longest_spans.reverse();
94
95 }
116 }
117 }
118 });
119
120 Self { sender }
121 }
122}