common/
remote_tracing.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Clone, Debug, Deserialize, Serialize)]
4pub enum TracingMessage {
5    Log {
6        timestamp: std::time::SystemTime,
7        level: String,
8        target: String,
9        message: String,
10    },
11    Progress(crate::progress::SerializableProgress),
12}
13
14#[derive(Debug)]
15pub struct RemoteTracingLayer {
16    pub sender: tokio::sync::mpsc::UnboundedSender<TracingMessage>,
17}
18
19impl RemoteTracingLayer {
20    #[must_use]
21    pub fn new() -> (
22        Self,
23        tokio::sync::mpsc::UnboundedSender<TracingMessage>,
24        tokio::sync::mpsc::UnboundedReceiver<TracingMessage>,
25    ) {
26        let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
27        (
28            Self {
29                sender: sender.clone(),
30            },
31            sender,
32            receiver,
33        )
34    }
35}
36
37pub fn send_progress_update(
38    sender: &tokio::sync::mpsc::UnboundedSender<TracingMessage>,
39    progress: &crate::progress::Progress,
40) -> anyhow::Result<()> {
41    let serializable_progress = crate::progress::SerializableProgress::from(progress);
42    sender.send(TracingMessage::Progress(serializable_progress))?;
43    Ok(())
44}
45
46struct FieldVisitor {
47    fields: std::collections::HashMap<String, String>,
48    message: Option<String>,
49}
50
51impl FieldVisitor {
52    fn new() -> Self {
53        Self {
54            fields: std::collections::HashMap::new(),
55            message: None,
56        }
57    }
58}
59
60impl tracing_subscriber::field::Visit for FieldVisitor {
61    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
62        let value_str = format!("{value:?}");
63        if field.name() == "message" {
64            self.message = Some(value_str);
65        } else {
66            self.fields.insert(field.name().to_string(), value_str);
67        }
68    }
69
70    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
71        if field.name() == "message" {
72            self.message = Some(value.to_string());
73        } else {
74            self.fields
75                .insert(field.name().to_string(), value.to_string());
76        }
77    }
78
79    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
80        self.fields
81            .insert(field.name().to_string(), value.to_string());
82    }
83
84    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
85        self.fields
86            .insert(field.name().to_string(), value.to_string());
87    }
88
89    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
90        self.fields
91            .insert(field.name().to_string(), value.to_string());
92    }
93
94    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
95        self.fields
96            .insert(field.name().to_string(), value.to_string());
97    }
98}
99
100impl<S> tracing_subscriber::Layer<S> for RemoteTracingLayer
101where
102    S: tracing::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
103{
104    fn on_event(
105        &self,
106        event: &tracing::Event<'_>,
107        _ctx: tracing_subscriber::layer::Context<'_, S>,
108    ) {
109        let mut visitor = FieldVisitor::new();
110        event.record(&mut visitor);
111        let message = visitor.message.unwrap_or_else(|| {
112            if visitor.fields.is_empty() {
113                String::new()
114            } else {
115                format!("{:?}", visitor.fields)
116            }
117        });
118        let tracing_message = TracingMessage::Log {
119            timestamp: std::time::SystemTime::now(),
120            level: event.metadata().level().to_string(),
121            target: event.metadata().target().to_string(),
122            message,
123        };
124        if self.sender.send(tracing_message).is_err() {
125            // If we can't send the tracing message, there's not much we can do
126            // The receiver has probably been dropped
127        }
128    }
129}