lux_lib/progress/
layer.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    sync::{
5        atomic::{AtomicI32, Ordering},
6        Arc, Mutex,
7    },
8};
9use tracing::{
10    field::{Field, Visit},
11    span::{Attributes, Id},
12    Subscriber,
13};
14use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
15
16use crate::progress::client::{LspClient, ProgressMessage, CLIENT};
17
18pub struct LspProgressLayer {
19    span_ids: Mutex<HashMap<Id, i32>>,
20    next: AtomicI32,
21}
22
23impl LspProgressLayer {
24    pub fn new() -> Self {
25        Self {
26            span_ids: Mutex::new(HashMap::new()),
27            next: AtomicI32::new(1),
28        }
29    }
30
31    pub fn next_id(&self) -> i32 {
32        self.next.fetch_add(1, Ordering::Relaxed)
33    }
34}
35
36impl Default for LspProgressLayer {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl<S> Layer<S> for LspProgressLayer
43where
44    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
45{
46    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, _ctx: Context<'_, S>) {
47        if *attrs.metadata().level() > tracing::Level::INFO {
48            return;
49        }
50
51        with_client(|client| {
52            let pid = self.next_id();
53            if let Ok(mut ids) = self.span_ids.lock() {
54                ids.insert(id.clone(), pid);
55            }
56            client.send(&ProgressMessage::Begin {
57                id: pid,
58                title: attrs.metadata().name().to_string(),
59            });
60        });
61    }
62
63    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
64        if *event.metadata().level() > tracing::Level::INFO {
65            return;
66        }
67
68        let pid = ctx
69            .event_span(event)
70            .and_then(|span_ref| self.span_ids.lock().ok()?.get(&span_ref.id()).copied());
71
72        if let Some(pid) = pid {
73            let mut visitor = MessageVisitor { message: None };
74            event.record(&mut visitor);
75
76            if let Some(message) = visitor.message {
77                with_client(|client| {
78                    client.send(&ProgressMessage::Report { id: pid, message });
79                });
80            }
81        }
82    }
83
84    fn on_close(&self, id: Id, _ctx: Context<'_, S>) {
85        if let Some(pid) = self
86            .span_ids
87            .lock()
88            .ok()
89            .and_then(|mut ids| ids.remove(&id))
90        {
91            with_client(|client| {
92                client.send(&ProgressMessage::End { id: pid });
93            });
94        }
95    }
96}
97
98struct MessageVisitor {
99    message: Option<String>,
100}
101
102impl Visit for MessageVisitor {
103    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
104        if field.name() == "message" {
105            self.message = Some(format!("{value:?}"));
106        }
107    }
108
109    fn record_str(&mut self, field: &Field, value: &str) {
110        if field.name() == "message" {
111            self.message = Some(value.to_string());
112        }
113    }
114}
115
116fn with_client<F>(f: F)
117where
118    F: FnOnce(&LspClient),
119{
120    let client = CLIENT
121        .read()
122        .ok()
123        .and_then(|guard| guard.as_ref().map(Arc::clone));
124    if let Some(ref c) = client {
125        f(c);
126    }
127}
lux_lib/progress/layer.rs

lux_lib/progress/
layer.rs
