Skip to main content

steamroom_cli/daemon/
tracing_layer.rs

1//! `tracing_subscriber::Layer` that intercepts events emitted inside a
2//! `job_id`-tagged span and republishes them as `Event::Log` for that
3//! job. Off-span events have `job_id: None` and only land in the daemon
4//! log file (handled by the wrapping fmt layer).
5
6use tokio::sync::broadcast::Sender;
7use tracing::Event as TracingEvent;
8use tracing::Subscriber;
9use tracing_subscriber::Layer;
10use tracing_subscriber::layer::Context;
11use tracing_subscriber::registry::LookupSpan;
12
13use crate::daemon::proto::Event;
14use crate::daemon::proto::JobId;
15use crate::daemon::proto::LogLevel;
16
17/// Span field name carrying the job id. The worker sets this when entering
18/// each job's span. Pass the numeric value directly for the most reliable
19/// path: `tracing::info_span!("job", job_id = job.0)`. The `%`-formatter
20/// (Display) is also handled but requires a successful parse from the
21/// formatted string.
22pub const JOB_ID_FIELD: &str = "job_id";
23
24pub struct JobScopedLogLayer {
25    pub events: Sender<Event>,
26}
27
28impl JobScopedLogLayer {
29    pub fn new(events: Sender<Event>) -> Self {
30        Self { events }
31    }
32}
33
34impl<S> Layer<S> for JobScopedLogLayer
35where
36    S: Subscriber + for<'a> LookupSpan<'a>,
37{
38    fn on_event(&self, event: &TracingEvent<'_>, ctx: Context<'_, S>) {
39        let job_id = find_job_id_in_scope(event, &ctx);
40        let mut visitor = MessageVisitor::default();
41        event.record(&mut visitor);
42        let message = visitor.message.unwrap_or_default();
43
44        let level: LogLevel = (*event.metadata().level()).into();
45        let target = event.metadata().target().to_string();
46
47        let _ = self.events.send(Event::Log {
48            job_id: job_id.map(JobId),
49            level,
50            target,
51            message,
52        });
53    }
54}
55
56fn find_job_id_in_scope<S>(event: &TracingEvent<'_>, ctx: &Context<'_, S>) -> Option<u64>
57where
58    S: Subscriber + for<'a> LookupSpan<'a>,
59{
60    let scope = ctx.event_scope(event)?;
61    for span in scope.from_root() {
62        let ext = span.extensions();
63        if let Some(id) = ext.get::<JobIdAttachment>() {
64            return Some(id.0);
65        }
66    }
67    None
68}
69
70/// Attached to a span by `on_new_span` when the span's recorded fields
71/// include `job_id`. Pure data; no Arc.
72struct JobIdAttachment(u64);
73
74pub struct JobIdAttachmentInstaller;
75
76impl<S> Layer<S> for JobIdAttachmentInstaller
77where
78    S: Subscriber + for<'a> LookupSpan<'a>,
79{
80    fn on_new_span(
81        &self,
82        attrs: &tracing::span::Attributes<'_>,
83        id: &tracing::span::Id,
84        ctx: Context<'_, S>,
85    ) {
86        let mut v = JobIdFieldVisitor::default();
87        attrs.record(&mut v);
88        if let Some(jid) = v.job_id
89            && let Some(span) = ctx.span(id)
90        {
91            span.extensions_mut().insert(JobIdAttachment(jid));
92        }
93    }
94}
95
96#[derive(Default)]
97struct JobIdFieldVisitor {
98    job_id: Option<u64>,
99}
100
101impl tracing::field::Visit for JobIdFieldVisitor {
102    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
103        if field.name() == JOB_ID_FIELD {
104            // Tracing's %/?formatters route through record_debug. Try to
105            // parse a numeric form (works for `%id` where id is u64-Display).
106            if let Ok(n) = format!("{value:?}").parse::<u64>() {
107                self.job_id = Some(n);
108            }
109        }
110    }
111    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
112        if field.name() == JOB_ID_FIELD {
113            self.job_id = Some(value);
114        }
115    }
116    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
117        if field.name() == JOB_ID_FIELD && value >= 0 {
118            self.job_id = Some(value as u64);
119        }
120    }
121}
122
123#[derive(Default)]
124struct MessageVisitor {
125    message: Option<String>,
126}
127
128impl tracing::field::Visit for MessageVisitor {
129    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
130        if field.name() == "message" {
131            self.message = Some(format!("{value:?}").trim_matches('"').to_string());
132        }
133    }
134    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
135        if field.name() == "message" {
136            self.message = Some(value.to_string());
137        }
138    }
139}