lsp_bin/
timings.rs

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        // Spawn a single background task
66        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                        // tracing::info!("🔍 Top 10 Longest Spans:");
96                        // for (_, timing) in longest_spans.iter().take(10) {
97                        //     tracing::info!(
98                        //         "⏳ Span: {} | Count: {} | Mean: {:.3}ms | Time: {:.3}ms",
99                        //         timing.name,
100                        //         timing.count,
101                        //         timing.elapsed.as_millis() / timing.count,
102                        //         timing.elapsed.as_millis()
103                        //     );
104                        // }
105                        // let names = longest_spans.iter().map(|(x, _)| x.as_str()).fold(String::new(), |mut acc, item| {
106                        //     acc += ", ";
107                        //     acc += item;
108                        //     acc
109                        // });
110                        //
111                        //     tracing::info!(
112                        //         "span names {}",
113                        //         names
114                        //     );
115                    }
116                }
117            }
118        });
119
120        Self { sender }
121    }
122}