1use 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
26pub const TRACE_TARGET: &str = "studio_worker::job";
28
29pub const JOB_LOG_JOBS_CAP: usize = 128;
32
33pub const JOB_LOG_LINES_CAP: usize = 400;
35
36pub const JOB_LOG_LINE_CHARS: usize = 2_000;
38
39#[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#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(rename_all = "camelCase")]
52pub struct JobLog {
53 pub lines: Vec<JobLogLine>,
54 pub dropped: u64,
57}
58
59#[derive(Default)]
60struct Entry {
61 job_id: String,
62 lines: VecDeque<JobLogLine>,
63 dropped: u64,
64}
65
66#[derive(Clone, Default)]
68pub struct JobLogStore {
69 inner: Arc<Mutex<VecDeque<Entry>>>,
70}
71
72impl JobLogStore {
73 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 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 pub fn len(&self) -> usize {
112 self.inner.lock().len()
113 }
114
115 pub fn is_empty(&self) -> bool {
117 self.len() == 0
118 }
119}
120
121pub fn global() -> &'static JobLogStore {
123 static STORE: OnceLock<JobLogStore> = OnceLock::new();
124 STORE.get_or_init(JobLogStore::default)
125}
126
127pub fn job_span(job_id: &str) -> tracing::Span {
129 tracing::info_span!(target: TRACE_TARGET, "job", job_id = %job_id)
130}
131
132struct SpanJobId(String);
134
135fn 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#[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#[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 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
209pub fn global_worker_log() -> &'static WorkerLogRing {
212 static RING: OnceLock<WorkerLogRing> = OnceLock::new();
213 RING.get_or_init(WorkerLogRing::default)
214}
215
216pub struct WorkerLogLayer {
222 ring: WorkerLogRing,
223}
224
225impl WorkerLogLayer {
226 pub fn new(ring: WorkerLogRing) -> Self {
227 Self { ring }
228 }
229
230 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
277pub struct JobLogLayer {
279 store: JobLogStore,
280}
281
282impl JobLogLayer {
283 pub fn new(store: JobLogStore) -> Self {
285 Self { store }
286 }
287
288 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 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}