1use std::fs::File;
2use std::io::{self, Read, Seek, SeekFrom, Write};
3use std::sync::Mutex;
4use std::sync::atomic::{AtomicBool, Ordering};
5
6#[derive(Clone, Debug, Default, PartialEq, Eq)]
8pub enum Display {
9 #[default]
11 Stream,
12 Prefixed(String),
14 Static {
18 id: String,
19 group: bool,
20 cached: bool,
21 },
22 Headed { id: String },
24 Failures { id: String },
27 Hidden,
29}
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum Shown {
34 Success,
35 Failure,
36 LocalCache,
37 RemoteCache,
38 Kept,
40}
41
42impl Display {
43 pub fn for_task(
47 style: OutputStyle,
48 id: &str,
49 project: &str,
50 continuous: bool,
51 root: bool,
52 ) -> Self {
53 match style {
54 OutputStyle::StreamWithoutPrefixes => Self::Stream,
55 OutputStyle::RunOne if root && !continuous => Self::Headed { id: id.to_owned() },
56 OutputStyle::Static | OutputStyle::RunOne if !continuous => Self::Static {
57 id: id.to_owned(),
58 group: std::env::var_os("GITHUB_ACTIONS").is_some()
59 && std::env::var_os("NX_SKIP_LOG_GROUPING").is_none_or(|value| value != "true"),
60 cached: style == OutputStyle::Static,
61 },
62 OutputStyle::Quiet => Self::Failures { id: id.to_owned() },
63 OutputStyle::Static | OutputStyle::RunOne | OutputStyle::Stream => {
64 Self::Prefixed(prefix(project))
65 }
66 }
67 }
68}
69
70#[derive(Clone, Copy, Debug, PartialEq, Eq)]
72pub enum OutputStyle {
73 Stream,
74 StreamWithoutPrefixes,
75 Static,
76 RunOne,
80 Quiet,
82}
83
84fn prefix(project: &str) -> String {
87 const COLOURS: [u8; 10] = [32, 92, 34, 94, 36, 96, 33, 93, 35, 95];
90 if !colour() {
91 return format!("{project}:");
92 }
93 let code: u32 = project.encode_utf16().map(u32::from).sum();
94 let colour = COLOURS[code as usize % COLOURS.len()];
95 format!("\x1b[1m\x1b[{colour}m{project}:\x1b[39m\x1b[22m")
96}
97
98#[derive(Clone, Debug, Default, PartialEq, Eq)]
100pub struct Decoration {
101 pub prefix: Option<String>,
102 pub prefix_color: Option<String>,
103 pub color: Option<String>,
104 pub bg_color: Option<String>,
105}
106
107const COLOURS: [&str; 8] = [
108 "black", "red", "green", "yellow", "blue", "magenta", "cyan", "white",
109];
110
111impl Decoration {
112 pub fn is_empty(&self) -> bool {
113 self == &Self::default()
114 }
115
116 pub fn valid(name: &str, background: bool) -> bool {
119 match name.strip_prefix("bg") {
120 Some(rest) if background => {
121 let mut chars = rest.chars();
122 chars.next().is_some_and(char::is_uppercase)
123 && COLOURS.contains(&rest.to_lowercase().as_str())
124 }
125 _ => !background && COLOURS.contains(&name),
126 }
127 }
128
129 fn apply(&self, bytes: &[u8]) -> Vec<u8> {
132 let colours = colour();
133 let paint = |text: &str, name: Option<&str>, background: bool| -> String {
134 let index = name.and_then(|name| {
135 let name = if background {
136 name.get(2..)?.to_lowercase()
137 } else {
138 name.to_owned()
139 };
140 COLOURS.iter().position(|colour| *colour == name)
141 });
142 match index {
143 Some(index) if colours => {
144 let (open, close) = if background { (40, 49) } else { (30, 39) };
145 format!("\x1b[{}m{text}\x1b[{close}m", open + index)
146 }
147 _ => text.to_owned(),
148 }
149 };
150 let mut text = String::from_utf8_lossy(bytes).into_owned();
151 if let Some(prefix) = &self.prefix {
152 let mut prefix = paint(prefix, self.prefix_color.as_deref(), false);
153 if colours {
154 prefix = format!("\x1b[1m{prefix}\x1b[22m");
155 }
156 text = text
157 .split('\n')
158 .map(|line| {
159 if line.trim().is_empty() {
160 line.to_owned()
161 } else {
162 format!("{prefix} {line}")
163 }
164 })
165 .collect::<Vec<_>>()
166 .join("\n");
167 }
168 let text = paint(&text, self.color.as_deref(), false);
169 paint(&text, self.bg_color.as_deref(), true).into_bytes()
170 }
171}
172
173const MEMORY_OUTPUT_BYTES: usize = 64 * 1024;
176
177#[derive(Default)]
178struct Spool {
179 memory: Vec<u8>,
180 file: Option<File>,
181 last: Option<u8>,
182}
183
184impl Spool {
185 fn append(&mut self, bytes: &[u8]) -> io::Result<()> {
186 if self.file.is_none() && self.memory.len() + bytes.len() > MEMORY_OUTPUT_BYTES {
187 let mut file = tempfile::tempfile()?;
188 file.write_all(&self.memory)?;
189 self.file = Some(file);
190 self.memory = Vec::new();
191 }
192 if let Some(file) = &mut self.file {
193 file.write_all(bytes)?;
194 } else {
195 self.memory.extend_from_slice(bytes);
196 }
197 if let Some(last) = bytes.last() {
198 self.last = Some(*last);
199 }
200 Ok(())
201 }
202
203 fn visit(&mut self, mut consume: impl FnMut(&[u8]) -> io::Result<()>) -> io::Result<()> {
204 if let Some(file) = &mut self.file {
205 file.seek(SeekFrom::Start(0))?;
206 let mut buffer = [0; 16 * 1024];
207 loop {
208 let count = file.read(&mut buffer)?;
209 if count == 0 {
210 break;
211 }
212 consume(&buffer[..count])?;
213 }
214 file.seek(SeekFrom::End(0))?;
216 } else {
217 for bytes in self.memory.chunks(16 * 1024) {
218 consume(bytes)?;
219 }
220 }
221 Ok(())
222 }
223 fn prefixed(&mut self, stderr: bool, prefix: &str) -> io::Result<()> {
224 if self.file.is_none() {
225 return write_to(stderr, &prefixed(prefix, &self.memory));
226 }
227 crate::report::output_batch(|output| {
228 let mut started = false;
229 self.visit(|bytes| {
230 for piece in bytes.split_inclusive(|byte| matches!(byte, b'\n' | b'\r')) {
231 let complete = piece
232 .last()
233 .is_some_and(|byte| matches!(byte, b'\n' | b'\r'));
234 let content = if complete {
235 &piece[..piece.len() - 1]
236 } else {
237 piece
238 };
239 if !content.is_empty() {
240 if !started {
241 output(stderr, prefix.as_bytes());
242 output(stderr, b" ");
243 started = true;
244 }
245 output(stderr, content);
246 }
247 if complete && started {
248 output(stderr, b"\n");
249 started = false;
250 }
251 }
252 Ok(())
253 })?;
254 if started {
255 output(stderr, b"\n");
256 }
257 Ok(())
258 })
259 }
260}
261
262fn colour() -> bool {
265 use std::io::IsTerminal;
266 let set = |name| std::env::var_os(name).is_some();
267 !set("NO_COLOR")
268 && (set("FORCE_COLOR")
269 || cfg!(windows)
270 || set("CI")
271 || (io::stdout().is_terminal()
272 && std::env::var("TERM").is_ok_and(|term| term != "dumb")))
273}
274
275struct Printer {
277 display: Display,
278 lines: Mutex<[Spool; 2]>,
280 held: Mutex<Spool>,
281}
282
283impl Printer {
284 fn new(display: Display) -> Self {
285 Self {
286 display,
287 lines: Mutex::default(),
288 held: Mutex::default(),
289 }
290 }
291
292 fn write(&self, stderr: bool, bytes: &[u8]) -> io::Result<()> {
293 match &self.display {
294 Display::Stream | Display::Headed { .. } => write_to(stderr, bytes),
295 Display::Hidden => Ok(()),
296 Display::Static { .. } | Display::Failures { .. } => {
297 self.held.lock().unwrap().append(bytes)
298 }
299 Display::Prefixed(prefix) => {
300 let mut lines = self.lines.lock().unwrap();
301 let buffer = &mut lines[usize::from(stderr)];
302 let end = bytes
303 .iter()
304 .rposition(|byte| matches!(byte, b'\n' | b'\r'))
305 .map(|end| end + 1);
306 if let Some(end) = end {
307 buffer.append(&bytes[..end])?;
308 let mut complete = std::mem::take(buffer);
309 complete.prefixed(stderr, prefix)?;
310 buffer.append(&bytes[end..])?;
311 } else {
312 buffer.append(bytes)?;
313 }
314 Ok(())
315 }
316 }
317 }
318
319 fn begin(&self, shown: Option<Shown>) -> io::Result<()> {
321 match &self.display {
322 Display::Headed { id } => {
323 let header = header(id, shown.unwrap_or(Shown::Success));
324 write_to(false, format!("\n{header}\n\n").as_bytes())
325 }
326 _ => Ok(()),
327 }
328 }
329
330 fn finish(&self, shown: Shown) -> io::Result<()> {
331 match &self.display {
332 Display::Stream | Display::Headed { .. } | Display::Hidden => Ok(()),
333 Display::Failures { id } => {
334 let mut held = std::mem::take(&mut *self.held.lock().unwrap());
335 if shown == Shown::Failure {
336 crate::report::output_batch(|output| -> io::Result<()> {
337 output(false, format!("\nā qk run {id} failed\n\n").as_bytes());
338 held.visit(|bytes| {
339 output(false, bytes);
340 Ok(())
341 })?;
342 if held.last != Some(b'\n') {
343 output(false, b"\n");
344 }
345 Ok(())
346 })?;
347 }
348 Ok(())
349 }
350 Display::Prefixed(prefix) => {
351 let mut lines = self.lines.lock().unwrap();
352 for stderr in [false, true] {
353 let mut rest = std::mem::take(&mut lines[usize::from(stderr)]);
354 rest.prefixed(stderr, prefix)?;
355 }
356 Ok(())
357 }
358 Display::Static { id, group, cached } => {
359 let mut held = std::mem::take(&mut *self.held.lock().unwrap());
360 if !cached && matches!(shown, Shown::LocalCache | Shown::RemoteCache | Shown::Kept)
361 {
362 held = Spool::default();
363 }
364 crate::report::output_batch(|output| -> io::Result<()> {
365 output(false, b"\n");
366 if *group {
367 let icon = match shown {
368 Shown::Success => "ā
",
369 Shown::Failure => "ā",
370 Shown::LocalCache | Shown::RemoteCache => "š",
371 Shown::Kept => "ā©",
372 };
373 output(false, format!("::group::{icon} ").as_bytes());
374 }
375 output(false, format!("{}\n\n", header(id, shown)).as_bytes());
376 held.visit(|bytes| {
377 output(false, bytes);
378 Ok(())
379 })?;
380 if *group {
381 if held.last.is_some() && held.last != Some(b'\n') {
382 output(false, b"\n");
383 }
384 output(false, b"::endgroup::\n");
385 }
386 Ok(())
387 })
388 }
389 }
390 }
391}
392
393fn header(id: &str, shown: Shown) -> String {
395 let status = match shown {
396 Shown::LocalCache => " [local cache]",
397 Shown::RemoteCache => " [remote cache]",
398 Shown::Kept => " [existing outputs match the cache, left as is]",
399 Shown::Success | Shown::Failure => "",
400 };
401 format!("> qk run {id}{status}")
402}
403
404fn prefixed(prefix: &str, bytes: &[u8]) -> Vec<u8> {
406 let mut result = Vec::new();
407 for line in bytes.split(|byte| matches!(byte, b'\n' | b'\r')) {
408 if !line.is_empty() {
409 result.extend_from_slice(prefix.as_bytes());
410 result.push(b' ');
411 result.extend_from_slice(line);
412 result.push(b'\n');
413 }
414 }
415 result
416}
417
418fn write_to(stderr: bool, bytes: &[u8]) -> io::Result<()> {
419 crate::report::output(stderr, bytes);
420 Ok(())
421}
422
423pub struct Capture {
426 file: Option<Mutex<File>>,
427 retained: Mutex<Option<RecordedLog>>,
428 printer: Printer,
429 healthy: AtomicBool,
430 readiness: Option<Readiness>,
431}
432
433pub const MAX_LOG_BYTES: usize = 4 * 1024 * 1024;
435
436#[derive(Default)]
438pub struct RecordedLog {
439 pub data: Vec<u8>,
440 pub truncated: bool,
442}
443
444impl RecordedLog {
445 fn append(&mut self, stderr: bool, bytes: &[u8]) {
447 if self.truncated {
448 return;
449 }
450 let remaining = MAX_LOG_BYTES.saturating_sub(self.data.len());
451 let count = bytes.len().min(remaining.saturating_sub(5));
452 if count > 0 {
453 self.data.push(u8::from(stderr));
454 self.data.extend_from_slice(&(count as u32).to_le_bytes());
455 self.data.extend_from_slice(&bytes[..count]);
456 }
457 self.truncated = count != bytes.len();
458 }
459}
460
461struct Readiness {
463 patterns: Vec<String>,
464 state: Mutex<(Vec<bool>, [Vec<u8>; 2])>,
467 ready: std::sync::Arc<AtomicBool>,
468}
469
470impl Readiness {
471 fn see(&self, stderr: bool, bytes: &[u8]) {
472 let Ok(mut state) = self.state.lock() else {
473 return;
474 };
475 let (found, tails) = &mut *state;
476 let tail = &mut tails[usize::from(stderr)];
477 tail.extend_from_slice(bytes);
478 let text = String::from_utf8_lossy(tail);
479 for (pattern, found) in self.patterns.iter().zip(found.iter_mut()) {
480 *found |= text.contains(pattern.as_str());
481 }
482 let keep = self.patterns.iter().map(String::len).max().unwrap_or(0);
483 if tail.len() > keep {
484 tail.drain(..tail.len() - keep);
485 }
486 if found.iter().all(|found| *found) {
487 self.ready.store(true, Ordering::SeqCst);
488 }
489 }
490}
491
492impl Capture {
493 pub fn new(file: Option<File>, display: Display) -> Self {
496 let printer = Printer::new(display);
497 let _ = printer.begin(None);
498 Self {
499 file: file.map(Mutex::new),
500 retained: Mutex::new(None),
501 printer,
502 healthy: AtomicBool::new(true),
503 readiness: None,
504 }
505 }
506
507 pub fn retain_log(self) -> Self {
509 *self.retained.lock().unwrap() = Some(RecordedLog::default());
510 self
511 }
512
513 pub fn take_log(&self) -> Option<RecordedLog> {
515 self.retained.lock().unwrap().take()
516 }
517
518 pub fn ready_when(mut self, patterns: &[String], ready: std::sync::Arc<AtomicBool>) -> Self {
520 if !patterns.is_empty() {
521 self.readiness = Some(Readiness {
522 patterns: patterns.to_vec(),
523 state: Mutex::new((vec![false; patterns.len()], [Vec::new(), Vec::new()])),
524 ready,
525 });
526 }
527 self
528 }
529
530 pub fn healthy(&self) -> bool {
531 self.healthy.load(Ordering::Relaxed)
532 }
533
534 pub(crate) fn copy(
536 &self,
537 mut reader: impl Read,
538 stderr: bool,
539 decoration: Option<&Decoration>,
540 ) -> io::Result<()> {
541 let decoration = decoration.filter(|decoration| !decoration.is_empty());
542 let mut buffer = [0; 16 * 1024];
543 let mut partial = Vec::new();
545 loop {
546 let count = reader.read(&mut buffer)?;
547 let chunk = &buffer[..count];
548 match decoration {
549 None if count == 0 => return Ok(()),
550 None => self.emit(stderr, chunk)?,
551 Some(decoration) if decoration.prefix.is_some() => {
552 partial.extend_from_slice(chunk);
553 let end = if count == 0 {
554 partial.len()
555 } else {
556 partial
557 .iter()
558 .rposition(|byte| *byte == b'\n')
559 .map_or(0, |end| end + 1)
560 };
561 if end > 0 {
562 let lines: Vec<u8> = partial.drain(..end).collect();
563 self.emit(stderr, &decoration.apply(&lines))?;
564 }
565 if count == 0 {
566 return Ok(());
567 }
568 }
569 Some(_) if count == 0 => return Ok(()),
570 Some(decoration) => self.emit(stderr, &decoration.apply(chunk))?,
571 }
572 }
573 }
574
575 fn emit(&self, stderr: bool, bytes: &[u8]) -> io::Result<()> {
577 for frame in bytes.chunks(16 * 1024) {
578 let mut retained = self
580 .retained
581 .lock()
582 .map_err(|_| io::Error::other("retained log lock poisoned"))?;
583 if let Some(log) = retained.as_mut() {
584 log.append(stderr, frame);
585 }
586 if let Some(file) = &self.file {
587 let mut file = file
588 .lock()
589 .map_err(|_| io::Error::other("output capture lock poisoned"))?;
590 if self.healthy() {
591 let result = (|| {
592 file.write_all(&[u8::from(stderr)])?;
593 file.write_all(&(frame.len() as u32).to_le_bytes())?;
594 file.write_all(frame)
595 })();
596 if result.is_err() {
597 self.healthy.store(false, Ordering::Relaxed);
598 }
599 }
600 }
601 self.printer.write(stderr, frame)?;
602 if let Some(readiness) = &self.readiness {
603 readiness.see(stderr, frame);
604 }
605 }
606 Ok(())
607 }
608
609 pub(crate) fn finish(&self, shown: Shown) -> io::Result<()> {
610 self.printer.finish(shown)
611 }
612}
613
614pub fn read_capture(
615 mut reader: impl Read,
616 mut consume: impl FnMut(bool, &[u8]) -> io::Result<()>,
617) -> io::Result<()> {
618 loop {
619 let mut stream = [0];
620 if reader.read(&mut stream)? == 0 {
621 return Ok(());
622 }
623 if stream[0] > 1 {
624 return Err(io::Error::other("invalid output capture stream"));
625 }
626 let mut length = [0; 4];
627 reader.read_exact(&mut length)?;
628 let length = u32::from_le_bytes(length) as usize;
629 if length > 16 * 1024 {
630 return Err(io::Error::other("invalid output capture frame"));
631 }
632 let mut bytes = vec![0; length];
633 reader.read_exact(&mut bytes)?;
634 consume(stream[0] == 1, &bytes)?;
635 }
636}
637
638pub fn replay(reader: impl Read, display: &Display, shown: Shown) -> io::Result<()> {
640 let printer = Printer::new(display.clone());
641 printer.begin(Some(shown))?;
642 read_capture(reader, |stderr, bytes| printer.write(stderr, bytes))?;
643 printer.finish(shown)
644}
645
646#[cfg(test)]
647mod tests {
648 use super::*;
649
650 #[test]
651 fn held_output_spills_without_losing_binary_stream_order() {
652 let printer = Printer::new(Display::Failures { id: "large".into() });
653 let chunk = [0, 255, b'x', b'\n'].repeat(4096);
654 for index in 0..32 {
655 printer.write(index % 2 == 1, &chunk).unwrap();
656 }
657 let mut held = printer.held.lock().unwrap();
658 assert!(held.file.is_some());
659 assert!(held.memory.is_empty());
660 let mut count = 0;
661 held.visit(|bytes| {
662 assert_eq!(bytes, chunk.as_slice());
663 count += 1;
664 Ok(())
665 })
666 .unwrap();
667 assert_eq!(count, 32);
668 drop(held);
670 printer.finish(Shown::Success).unwrap();
671 assert!(printer.held.lock().unwrap().file.is_none());
672 }
673
674 #[test]
675 fn unfinished_prefixed_lines_spill_and_keep_streams_separate() {
676 let printer = Printer::new(Display::Prefixed("app:".into()));
677 let bytes = vec![b'x'; MEMORY_OUTPUT_BYTES + 1];
678 printer.write(false, &bytes).unwrap();
679 printer.write(true, b"partial stderr").unwrap();
680 let mut lines = printer.lines.lock().unwrap();
681 assert!(lines[0].file.is_some());
682 assert!(lines[0].memory.is_empty());
683 let mut length = 0;
684 lines[0]
685 .visit(|chunk| {
686 assert!(chunk.len() <= 16 * 1024);
687 assert!(chunk.iter().all(|byte| *byte == b'x'));
688 length += chunk.len();
689 Ok(())
690 })
691 .unwrap();
692 assert_eq!(length, bytes.len());
693 assert_eq!(lines[1].memory, b"partial stderr");
694 }
695
696 #[test]
697 fn small_spools_stay_in_memory_and_can_replay_then_append() {
698 let mut spool = Spool::default();
699 spool.append(b"small").unwrap();
700 assert!(spool.file.is_none());
701 let mut output = Vec::new();
702 spool
703 .visit(|chunk| {
704 output.extend_from_slice(chunk);
705 Ok(())
706 })
707 .unwrap();
708 assert_eq!(output, b"small");
709 spool.append(&vec![b'x'; MEMORY_OUTPUT_BYTES]).unwrap();
710 assert!(spool.file.is_some());
711 spool.visit(|_| Ok(())).unwrap();
712 spool.append(b"last").unwrap();
713 output.clear();
714 spool
715 .visit(|chunk| {
716 output.extend_from_slice(chunk);
717 Ok(())
718 })
719 .unwrap();
720 assert!(output.starts_with(b"small"));
721 assert!(output.ends_with(b"last"));
722 }
723
724 #[test]
726 fn retained_logs_are_ordered_and_bounded() {
727 let capture = Capture::new(None, Display::Hidden).retain_log();
728 capture.emit(false, b"stdout\0\xff").unwrap();
729 capture.emit(true, b"stderr").unwrap();
730 let data = vec![b'x'; MAX_LOG_BYTES];
731 capture.emit(false, &data).unwrap();
732 capture.emit(true, b"omitted").unwrap();
733 let log = capture.take_log().unwrap();
734 assert!(log.truncated);
735 assert!(log.data.len() <= MAX_LOG_BYTES);
736 let mut frames = Vec::new();
737 read_capture(log.data.as_slice(), |stderr, bytes| {
738 frames.push((stderr, bytes.to_vec()));
739 Ok(())
740 })
741 .unwrap();
742 assert_eq!(frames[0], (false, b"stdout\0\xff".to_vec()));
743 assert_eq!(frames[1], (true, b"stderr".to_vec()));
744 assert!(
745 frames[2..]
746 .iter()
747 .all(|(stderr, bytes)| !stderr && bytes.iter().all(|b| *b == b'x'))
748 );
749 assert!(capture.take_log().is_none());
750 }
751
752 #[test]
753 fn prefixes_non_empty_lines_like_nx() {
754 assert_eq!(
755 prefixed("web:", b"one\n\ntwo\r\nthree"),
756 b"web: one\nweb: two\nweb: three\n"
757 );
758 }
759}