steamroom_cli/daemon/
tracing_layer.rs1use 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
17pub 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
70struct 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 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}