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        with_client(|client| {
48            let pid = self.next_id();
49            if let Ok(mut ids) = self.span_ids.lock() {
50                ids.insert(id.clone(), pid);
51            }
52            client.send(&ProgressMessage::Begin {
53                id: pid,
54                title: attrs.metadata().name().to_string(),
55            });
56        });
57    }
58
59    fn on_event(&self, event: &tracing::Event<'_>, ctx: Context<'_, S>) {
60        if *event.metadata().level() > tracing::Level::INFO {
61            return;
62        }
63
64        let pid = ctx
65            .event_span(event)
66            .and_then(|span_ref| self.span_ids.lock().ok()?.get(&span_ref.id()).copied());
67
68        if let Some(pid) = pid {
69            let mut visitor = MessageVisitor { message: None };
70            event.record(&mut visitor);
71
72            if let Some(message) = visitor.message {
73                with_client(|client| {
74                    client.send(&ProgressMessage::Report { id: pid, message });
75                });
76            }
77        }
78    }
79
80    fn on_close(&self, id: Id, _ctx: Context<'_, S>) {
81        if let Some(pid) = self
82            .span_ids
83            .lock()
84            .ok()
85            .and_then(|mut ids| ids.remove(&id))
86        {
87            with_client(|client| {
88                client.send(&ProgressMessage::End { id: pid });
89            });
90        }
91    }
92}
93
94struct MessageVisitor {
95    message: Option<String>,
96}
97
98impl Visit for MessageVisitor {
99    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
100        if field.name() == "message" {
101            self.message = Some(format!("{value:?}"));
102        }
103    }
104
105    fn record_str(&mut self, field: &Field, value: &str) {
106        if field.name() == "message" {
107            self.message = Some(value.to_string());
108        }
109    }
110}
111
112fn with_client<F>(f: F)
113where
114    F: FnOnce(&LspClient),
115{
116    let client = CLIENT
117        .read()
118        .ok()
119        .and_then(|guard| guard.as_ref().map(Arc::clone));
120    if let Some(ref c) = client {
121        f(c);
122    }
123}
lux_lib/progress/layer.rs

lux_lib/progress/
layer.rs
