Skip to main content

studio_worker/
job_log.rs

1//! Per-job log capture.
2//!
3//! A job runs inside a tracing span named `job` carrying a `job_id` field
4//! ([`job_span`]).  [`JobLogLayer`] copies every event emitted inside such a
5//! span, or carrying a `job_id` field itself, into that job's log in a
6//! bounded [`JobLogStore`].  Engines need not know about jobs: their events
7//! land in the right log because they run inside the job's span.
8//!
9//! The daemon installs the layer next to the stderr formatter (`main.rs`),
10//! so the job log sees the same, `RUST_LOG`-filtered, events.
11
12use std::collections::VecDeque;
13use std::fmt::Write as _;
14use std::sync::{Arc, OnceLock};
15
16use chrono::{DateTime, Utc};
17use parking_lot::Mutex;
18use serde::{Deserialize, Serialize};
19use tracing::field::{Field, Visit};
20use tracing::span::{Attributes, Id};
21use tracing::{Event, Subscriber};
22use tracing_subscriber::layer::Context;
23use tracing_subscriber::registry::LookupSpan;
24use tracing_subscriber::Layer;
25
26/// Tracing target for job bookkeeping (started / finished / thumbnail).
27pub const TRACE_TARGET: &str = "studio_worker::job";
28
29/// Logs of this many most recent jobs are kept.  Covers both 50-job rings
30/// plus the running jobs, with slack.
31pub const JOB_LOG_JOBS_CAP: usize = 128;
32
33/// Lines kept per job; older lines are dropped (and counted).
34pub const JOB_LOG_LINES_CAP: usize = 400;
35
36/// A longer message is clipped to this many characters.
37pub const JOB_LOG_LINE_CHARS: usize = 2_000;
38
39/// One captured log line.
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
41#[serde(rename_all = "camelCase")]
42pub struct JobLogLine {
43    pub ts: DateTime<Utc>,
44    pub level: String,
45    pub target: String,
46    pub message: String,
47}
48
49/// A job's captured log.
50#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(rename_all = "camelCase")]
52pub struct JobLog {
53    pub lines: Vec<JobLogLine>,
54    /// Lines dropped from the front because the job logged more than
55    /// [`JOB_LOG_LINES_CAP`].
56    pub dropped: u64,
57}
58
59#[derive(Default)]
60struct Entry {
61    job_id: String,
62    lines: VecDeque<JobLogLine>,
63    dropped: u64,
64}
65
66/// Bounded store of per-job logs, oldest job evicted first.  Cheap to clone.
67#[derive(Clone, Default)]
68pub struct JobLogStore {
69    inner: Arc<Mutex<VecDeque<Entry>>>,
70}
71
72impl JobLogStore {
73    /// Append `line` to `job_id`'s log.
74    pub fn push(&self, job_id: &str, mut line: JobLogLine) {
75        if line.message.chars().count() > JOB_LOG_LINE_CHARS {
76            line.message = line.message.chars().take(JOB_LOG_LINE_CHARS).collect();
77            line.message.push('…');
78        }
79        let mut jobs = self.inner.lock();
80        let index = match jobs.iter().rposition(|e| e.job_id == job_id) {
81            Some(index) => index,
82            None => {
83                jobs.push_back(Entry {
84                    job_id: job_id.to_string(),
85                    ..Default::default()
86                });
87                if jobs.len() > JOB_LOG_JOBS_CAP {
88                    jobs.pop_front();
89                }
90                jobs.len() - 1
91            }
92        };
93        let entry = &mut jobs[index];
94        entry.lines.push_back(line);
95        if entry.lines.len() > JOB_LOG_LINES_CAP {
96            entry.lines.pop_front();
97            entry.dropped += 1;
98        }
99    }
100
101    /// The log of `job_id`, if any line was captured for it.
102    pub fn get(&self, job_id: &str) -> Option<JobLog> {
103        let jobs = self.inner.lock();
104        jobs.iter().rfind(|e| e.job_id == job_id).map(|e| JobLog {
105            lines: e.lines.iter().cloned().collect(),
106            dropped: e.dropped,
107        })
108    }
109
110    /// Number of jobs with a log.
111    pub fn len(&self) -> usize {
112        self.inner.lock().len()
113    }
114
115    /// True when no job has a log.
116    pub fn is_empty(&self) -> bool {
117        self.len() == 0
118    }
119}
120
121/// The process-wide store the installed [`JobLogLayer`] writes to.
122pub fn global() -> &'static JobLogStore {
123    static STORE: OnceLock<JobLogStore> = OnceLock::new();
124    STORE.get_or_init(JobLogStore::default)
125}
126
127/// The span a job runs in.  Every event inside it lands in the job's log.
128pub fn job_span(job_id: &str) -> tracing::Span {
129    tracing::info_span!(target: TRACE_TARGET, "job", job_id = %job_id)
130}
131
132/// The job id a span carries, kept in the span's extensions.
133struct SpanJobId(String);
134
135/// Keep a new span's `job_id` in its extensions, once (both layers call
136/// this; an extension type may only be inserted once per span).
137fn remember_job_id<S>(attrs: &Attributes<'_>, id: &Id, ctx: &Context<'_, S>)
138where
139    S: Subscriber + for<'a> LookupSpan<'a>,
140{
141    let mut fields = Fields::default();
142    attrs.record(&mut fields);
143    if let (Some(job_id), Some(span)) = (fields.job_id, ctx.span(id)) {
144        let mut extensions = span.extensions_mut();
145        if extensions.get_mut::<SpanJobId>().is_none() {
146            extensions.insert(SpanJobId(job_id));
147        }
148    }
149}
150
151/// Collects the `job_id` field and renders the rest as `message k=v …`.
152#[derive(Default)]
153struct Fields {
154    job_id: Option<String>,
155    message: String,
156    rest: String,
157}
158
159impl Visit for Fields {
160    fn record_str(&mut self, field: &Field, value: &str) {
161        match field.name() {
162            "job_id" => self.job_id = Some(value.to_string()),
163            "message" => self.message.push_str(value),
164            name => {
165                let _ = write!(self.rest, " {name}={value:?}");
166            }
167        }
168    }
169
170    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
171        match field.name() {
172            "job_id" => self.job_id = Some(format!("{value:?}")),
173            "message" => {
174                let _ = write!(self.message, "{value:?}");
175            }
176            name => {
177                let _ = write!(self.rest, " {name}={value:?}");
178            }
179        }
180    }
181}
182
183impl Fields {
184    fn rendered(self) -> String {
185        format!("{}{}", self.message, self.rest)
186    }
187}
188
189/// The worker log ring the Logs tab shows: the entries and the sequence
190/// number of the newest (see `runtime::recent_logs_after`).
191#[derive(Clone, Default)]
192pub struct WorkerLogRing {
193    pub entries: Arc<Mutex<VecDeque<crate::types::LogEntry>>>,
194    pub seq: Arc<std::sync::atomic::AtomicU64>,
195}
196
197impl WorkerLogRing {
198    /// Append `entry`, keeping the newest [`crate::runtime::RECENT_LOGS_CAP`].
199    pub fn push(&self, entry: crate::types::LogEntry) {
200        let mut ring = self.entries.lock();
201        self.seq.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
202        ring.push_back(entry);
203        while ring.len() > crate::runtime::RECENT_LOGS_CAP {
204            ring.pop_front();
205        }
206    }
207}
208
209/// The process-wide worker log ring the installed [`WorkerLogLayer`]
210/// writes to; the daemon's observers share it.
211pub fn global_worker_log() -> &'static WorkerLogRing {
212    static RING: OnceLock<WorkerLogRing> = OnceLock::new();
213    RING.get_or_init(WorkerLogRing::default)
214}
215
216/// Copies the worker's own info / warn / error events into a
217/// [`WorkerLogRing`], so the Logs tab shows everything the daemon does,
218/// not only the studio session's breadcrumbs.  Events on the bare
219/// `studio_worker` target are skipped: `runtime::push_log` writes those
220/// into the ring itself.
221pub struct WorkerLogLayer {
222    ring: WorkerLogRing,
223}
224
225impl WorkerLogLayer {
226    pub fn new(ring: WorkerLogRing) -> Self {
227        Self { ring }
228    }
229
230    /// A layer writing to [`global_worker_log`].
231    pub fn global() -> Self {
232        Self::new(global_worker_log().clone())
233    }
234}
235
236impl<S> Layer<S> for WorkerLogLayer
237where
238    S: Subscriber + for<'a> LookupSpan<'a>,
239{
240    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
241        remember_job_id(attrs, id, &ctx);
242    }
243
244    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
245        let meta = event.metadata();
246        let target = meta.target();
247        if *meta.level() > tracing::Level::INFO
248            || target == "studio_worker"
249            || !target.starts_with("studio_worker")
250        {
251            return;
252        }
253        let mut fields = Fields::default();
254        event.record(&mut fields);
255        let job_id = fields.job_id.take().or_else(|| {
256            ctx.event_scope(event)?
257                .find_map(|span| span.extensions().get::<SpanJobId>().map(|j| j.0.clone()))
258        });
259        let mut message = fields.rendered();
260        if message.chars().count() > JOB_LOG_LINE_CHARS {
261            message = message.chars().take(JOB_LOG_LINE_CHARS).collect();
262            message.push('…');
263        }
264        self.ring.push(crate::types::LogEntry {
265            ts: Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
266            level: meta.level().as_str().to_ascii_lowercase(),
267            category: target
268                .strip_prefix("studio_worker::")
269                .unwrap_or(target)
270                .to_string(),
271            message,
272            job_id,
273        });
274    }
275}
276
277/// Copies job-scoped events into a [`JobLogStore`].
278pub struct JobLogLayer {
279    store: JobLogStore,
280}
281
282impl JobLogLayer {
283    /// A layer writing to `store`.
284    pub fn new(store: JobLogStore) -> Self {
285        Self { store }
286    }
287
288    /// A layer writing to the process-wide [`global`] store.
289    pub fn global() -> Self {
290        Self::new(global().clone())
291    }
292}
293
294impl<S> Layer<S> for JobLogLayer
295where
296    S: Subscriber + for<'a> LookupSpan<'a>,
297{
298    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
299        remember_job_id(attrs, id, &ctx);
300    }
301
302    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
303        let mut fields = Fields::default();
304        event.record(&mut fields);
305        let job_id = fields.job_id.take().or_else(|| {
306            ctx.event_scope(event)?
307                .find_map(|span| span.extensions().get::<SpanJobId>().map(|j| j.0.clone()))
308        });
309        let Some(job_id) = job_id else { return };
310        let meta = event.metadata();
311        self.store.push(
312            &job_id,
313            JobLogLine {
314                ts: Utc::now(),
315                level: meta.level().as_str().to_ascii_lowercase(),
316                target: meta.target().to_string(),
317                message: fields.rendered(),
318            },
319        );
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use tracing_subscriber::layer::SubscriberExt as _;
327
328    fn line(message: &str) -> JobLogLine {
329        JobLogLine {
330            ts: Utc::now(),
331            level: "info".into(),
332            target: "t".into(),
333            message: message.into(),
334        }
335    }
336
337    /// Run `f` with a subscriber whose only layer captures into `store`.
338    fn with_layer(store: &JobLogStore, f: impl FnOnce()) {
339        let subscriber = tracing_subscriber::registry().with(JobLogLayer::new(store.clone()));
340        tracing::subscriber::with_default(subscriber, || {
341            tracing::callsite::rebuild_interest_cache();
342            f();
343        });
344    }
345
346    #[test]
347    fn events_inside_a_job_span_land_in_that_jobs_log() {
348        let store = JobLogStore::default();
349        with_layer(&store, || {
350            let span = job_span("job-a");
351            let _entered = span.enter();
352            tracing::info!(target: "studio_worker::engine", op = "download", bytes = 42, "fetching weights");
353        });
354        let log = store.get("job-a").expect("job-a has a log");
355        assert_eq!(log.lines.len(), 1);
356        assert_eq!(log.lines[0].level, "info");
357        assert_eq!(log.lines[0].target, "studio_worker::engine");
358        assert_eq!(
359            log.lines[0].message,
360            "fetching weights op=\"download\" bytes=42"
361        );
362    }
363
364    #[test]
365    fn nested_spans_resolve_to_the_enclosing_job() {
366        let store = JobLogStore::default();
367        with_layer(&store, || {
368            let job = job_span("job-b");
369            let _job = job.enter();
370            let inner = tracing::info_span!("inner-work");
371            let _inner = inner.enter();
372            tracing::warn!("deep inside");
373        });
374        let log = store.get("job-b").expect("job-b has a log");
375        assert_eq!(log.lines[0].level, "warn");
376        assert_eq!(log.lines[0].message, "deep inside");
377    }
378
379    #[test]
380    fn an_event_carrying_a_job_id_field_lands_in_that_jobs_log() {
381        let store = JobLogStore::default();
382        with_layer(&store, || {
383            tracing::info!(job_id = "job-c", "[ws] accepted");
384        });
385        let log = store.get("job-c").expect("job-c has a log");
386        assert_eq!(log.lines[0].message, "[ws] accepted");
387    }
388
389    #[test]
390    fn an_explicit_job_id_field_wins_over_the_enclosing_span() {
391        let store = JobLogStore::default();
392        with_layer(&store, || {
393            let span = job_span("outer");
394            let _entered = span.enter();
395            tracing::info!(job_id = "other", "for the other job");
396        });
397        assert!(store.get("outer").is_none());
398        assert_eq!(store.get("other").unwrap().lines.len(), 1);
399    }
400
401    #[test]
402    fn events_outside_any_job_are_ignored() {
403        let store = JobLogStore::default();
404        with_layer(&store, || {
405            tracing::info!("background chatter");
406        });
407        assert!(store.is_empty());
408    }
409
410    #[test]
411    fn the_worker_log_takes_the_workers_own_info_and_up() {
412        let ring = WorkerLogRing::default();
413        let subscriber = tracing_subscriber::registry().with(WorkerLogLayer::new(ring.clone()));
414        tracing::subscriber::with_default(subscriber, || {
415            tracing::callsite::rebuild_interest_cache();
416            let span = job_span("job-w");
417            let _entered = span.enter();
418            tracing::info!(target: "studio_worker::host", op = "load", "model loaded");
419            tracing::warn!(target: "studio_worker::local_api", "denied");
420            tracing::debug!(target: "studio_worker::host", "too chatty");
421            tracing::info!(target: "studio_worker", "[ws] pushed by push_log");
422            tracing::info!(target: "hyper::client", "not ours");
423        });
424        let entries: Vec<_> = ring.entries.lock().iter().cloned().collect();
425        assert_eq!(entries.len(), 2, "{entries:?}");
426        assert_eq!(entries[0].category, "host");
427        assert_eq!(entries[0].message, "model loaded op=\"load\"");
428        assert_eq!(entries[0].job_id.as_deref(), Some("job-w"));
429        assert_eq!(entries[1].level, "warn");
430        assert_eq!(ring.seq.load(std::sync::atomic::Ordering::SeqCst), 2);
431    }
432
433    #[test]
434    fn both_layers_together_share_the_span_job_id() {
435        let store = JobLogStore::default();
436        let ring = WorkerLogRing::default();
437        let subscriber = tracing_subscriber::registry()
438            .with(JobLogLayer::new(store.clone()))
439            .with(WorkerLogLayer::new(ring.clone()));
440        tracing::subscriber::with_default(subscriber, || {
441            tracing::callsite::rebuild_interest_cache();
442            let span = job_span("job-both");
443            let _entered = span.enter();
444            tracing::info!(target: "studio_worker::host", "loaded");
445        });
446        assert_eq!(store.get("job-both").unwrap().lines.len(), 1);
447        assert_eq!(ring.entries.lock()[0].job_id.as_deref(), Some("job-both"));
448    }
449
450    #[test]
451    fn the_worker_log_keeps_only_the_newest_entries() {
452        let ring = WorkerLogRing::default();
453        for i in 0..(crate::runtime::RECENT_LOGS_CAP + 2) {
454            ring.push(crate::types::LogEntry {
455                ts: String::new(),
456                level: "info".into(),
457                category: "c".into(),
458                message: format!("m{i}"),
459                job_id: None,
460            });
461        }
462        assert_eq!(ring.entries.lock().len(), crate::runtime::RECENT_LOGS_CAP);
463        assert_eq!(ring.entries.lock()[0].message, "m2");
464    }
465
466    #[test]
467    fn a_job_keeps_only_its_newest_lines_and_counts_the_drop() {
468        let store = JobLogStore::default();
469        for i in 0..(JOB_LOG_LINES_CAP + 3) {
470            store.push("busy", line(&format!("line {i}")));
471        }
472        let log = store.get("busy").unwrap();
473        assert_eq!(log.lines.len(), JOB_LOG_LINES_CAP);
474        assert_eq!(log.dropped, 3);
475        assert_eq!(log.lines[0].message, "line 3");
476    }
477
478    #[test]
479    fn only_the_most_recent_jobs_keep_a_log() {
480        let store = JobLogStore::default();
481        for i in 0..(JOB_LOG_JOBS_CAP + 2) {
482            store.push(&format!("job-{i}"), line("x"));
483        }
484        assert_eq!(store.len(), JOB_LOG_JOBS_CAP);
485        assert!(store.get("job-0").is_none());
486        assert!(store.get("job-1").is_none());
487        assert!(store
488            .get(&format!("job-{}", JOB_LOG_JOBS_CAP + 1))
489            .is_some());
490    }
491
492    #[test]
493    fn a_long_message_is_clipped() {
494        let store = JobLogStore::default();
495        store.push("long", line(&"x".repeat(JOB_LOG_LINE_CHARS + 50)));
496        let message = &store.get("long").unwrap().lines[0].message;
497        assert_eq!(message.chars().count(), JOB_LOG_LINE_CHARS + 1);
498        assert!(message.ends_with('…'));
499    }
500
501    #[test]
502    fn a_job_log_round_trips_through_json() {
503        let log = JobLog {
504            lines: vec![line("hello")],
505            dropped: 2,
506        };
507        let json = serde_json::to_value(&log).unwrap();
508        assert_eq!(json["dropped"], 2);
509        assert_eq!(json["lines"][0]["message"], "hello");
510        let back: JobLog = serde_json::from_value(json).unwrap();
511        assert_eq!(back, log);
512    }
513
514    #[test]
515    fn the_global_store_is_one_instance() {
516        global().push("global-probe", line("g"));
517        assert!(JobLogLayer::global().store.get("global-probe").is_some());
518    }
519}