1use super::{failure, stop_after_failure, terminate_runner, LeaderTermination, RunnerPhase};
6use super::{StdinWriter, DRAIN_WAIT};
7use crate::{NativeOutputCapture, OwnedCommand};
8use std::{
9 io,
10 process::ExitStatus,
11 sync::{
12 atomic::{AtomicBool, Ordering},
13 mpsc::{self, RecvTimeoutError, SyncSender, TrySendError},
14 },
15 time::{Duration, Instant},
16};
17
18const KILL_REAP_WAIT: Duration = Duration::from_secs(5);
21
22const MAX_STEP: Duration = Duration::from_millis(25);
25
26const QUEUE_CHUNKS: usize = 64;
30
31const LINE_CAP: usize = 256 * 1024;
35
36#[derive(Clone, Debug, Eq, PartialEq)]
40#[non_exhaustive]
41pub struct StreamOptions {
42 pub drain_grace: Duration,
44 pub timeout: Option<Duration>,
47 pub poll: Duration,
51 pub stdout_tail_cap: usize,
54 pub stderr_tail_cap: usize,
57 pub no_window: bool,
59 pub input: Option<Vec<u8>>,
62}
63
64impl Default for StreamOptions {
65 fn default() -> Self {
66 Self {
67 drain_grace: DRAIN_WAIT,
68 timeout: None,
69 poll: Duration::from_millis(100),
70 stdout_tail_cap: 64 * 1024,
71 stderr_tail_cap: 64 * 1024,
72 no_window: true,
73 input: None,
74 }
75 }
76}
77
78impl StreamOptions {
79 pub fn drain_grace(mut self, grace: Duration) -> Self {
80 self.drain_grace = grace;
81 self
82 }
83 pub fn new() -> Self {
85 Self::default()
86 }
87
88 pub fn timeout(mut self, timeout: Duration) -> Self {
89 self.timeout = Some(timeout);
90 self
91 }
92
93 pub fn poll(mut self, poll: Duration) -> Self {
94 self.poll = poll;
95 self
96 }
97
98 pub fn stdout_tail_cap(mut self, bytes: usize) -> Self {
99 self.stdout_tail_cap = bytes;
100 self
101 }
102
103 pub fn stderr_tail_cap(mut self, bytes: usize) -> Self {
104 self.stderr_tail_cap = bytes;
105 self
106 }
107
108 pub fn no_window(mut self, hidden: bool) -> Self {
109 self.no_window = hidden;
110 self
111 }
112
113 pub fn input(mut self, bytes: impl Into<Vec<u8>>) -> Self {
114 self.input = Some(bytes.into());
115 self
116 }
117}
118
119#[derive(Clone, Debug, Eq, PartialEq)]
121#[non_exhaustive]
122pub struct StreamOutput {
123 pub status: ExitStatus,
126 pub cancelled: bool,
128 pub timed_out: bool,
131 pub stdout_tail: Vec<u8>,
133 pub stderr_tail: String,
137}
138
139pub fn run_streaming(
202 mut cmd: OwnedCommand,
203 options: &StreamOptions,
204 cancel: &dyn Fn() -> bool,
205 on_stdout: &mut dyn FnMut(&[u8]),
206 on_stderr_line: &mut dyn FnMut(&str),
207) -> io::Result<StreamOutput> {
208 cmd.pipe_stdio(options.input.is_some());
211 cmd.bound_runner_cleanup(super::RUNNER_REAP_WAIT);
212 if options.no_window {
213 cmd.windows_hide();
214 } else {
215 cmd.windows_visible();
216 }
217 let mut child = cmd
218 .spawn()
219 .map_err(|error| failure(RunnerPhase::Spawn, LeaderTermination::NotStarted, error))?;
220 let writer = match (options.input.as_deref(), child.take_stdin()) {
223 (Some(bytes), Some(stdin)) => Some(
224 StdinWriter::spawn(stdin, bytes.to_vec())
225 .map_err(|error| stop_after_failure(&mut child, RunnerPhase::Setup, error))?,
226 ),
227 _ => None,
228 };
229 let (sender, receiver) = mpsc::sync_channel::<Message>(QUEUE_CHUNKS);
230 let mut stdout = child
231 .take_stdout()
232 .map(|pipe| {
233 NativeOutputCapture::stdout_streaming(pipe, forward(sender.clone(), Message::Stdout))
234 })
235 .transpose()
236 .map_err(|error| stop_after_failure(&mut child, RunnerPhase::Setup, error))?;
237 let mut stderr = child
238 .take_stderr()
239 .map(|pipe| {
240 NativeOutputCapture::stderr_streaming(pipe, forward(sender.clone(), Message::Stderr))
241 })
242 .transpose()
243 .map_err(|error| stop_after_failure(&mut child, RunnerPhase::Setup, error))?;
244 drop(sender);
246
247 let mut sinks = Sinks {
248 on_stdout,
249 on_stderr_line,
250 stdout_tail: Tail::new(options.stdout_tail_cap),
251 stderr_tail: Tail::new(options.stderr_tail_cap),
252 splitter: LineSplitter::default(),
253 };
254 let poll = options.poll.max(Duration::from_millis(1));
255 let step = poll.min(MAX_STEP);
256 let started = Instant::now();
257 let deadline = options.timeout.and_then(|limit| started.checked_add(limit));
258 let mut last_cancel: Option<Instant> = None;
259 let mut last_exit_look = started;
260 let mut streams_done = false;
261 let ended = loop {
262 let now = Instant::now();
263 if last_cancel.is_none_or(|at| now.duration_since(at) >= poll) {
264 last_cancel = Some(now);
265 if cancel() {
266 break Ended::Cancelled;
267 }
268 }
269 if deadline.is_some_and(|at| now >= at) {
270 break Ended::TimedOut;
271 }
272 if streams_done {
273 if let Some(status) = child
275 .wait_exit_then_sweep(step)
276 .map_err(|error| stop_after_failure(&mut child, RunnerPhase::Wait, error))?
277 {
278 break Ended::Exited(status);
279 }
280 continue;
281 }
282 match receiver.recv_timeout(step) {
283 Ok(message) => sinks.handle(message),
284 Err(RecvTimeoutError::Timeout) => {}
285 Err(RecvTimeoutError::Disconnected) => {
286 streams_done = true;
287 continue;
288 }
289 }
290 if last_exit_look.elapsed() >= step {
291 last_exit_look = Instant::now();
292 if let Some(status) = child
293 .wait_exit_then_sweep(Duration::ZERO)
294 .map_err(|error| stop_after_failure(&mut child, RunnerPhase::Wait, error))?
295 {
296 break Ended::Exited(status);
297 }
298 }
299 };
300
301 let (status, cancelled, timed_out) = match ended {
302 Ended::Exited(status) => (Some(status), false, false),
303 Ended::Cancelled | Ended::TimedOut => {
304 let reaped = Some(terminate_runner(&mut child, KILL_REAP_WAIT)?);
305 (
306 reaped,
307 matches!(ended, Ended::Cancelled),
308 matches!(ended, Ended::TimedOut),
309 )
310 }
311 };
312
313 if !streams_done {
316 drain_output(&receiver, &mut sinks, options.drain_grace);
317 }
318 if let Some(capture) = stdout.as_mut() {
321 let _ = capture.cancel();
322 }
323 if let Some(capture) = stderr.as_mut() {
324 let _ = capture.cancel();
325 }
326 for message in receiver.try_iter() {
327 sinks.handle(message);
328 }
329 sinks.finish_stream();
330 let writer_result = writer.map(StdinWriter::finish).transpose();
331 let stdout_failure = super::capture_failure("stdout", stdout.as_ref());
334 let stderr_failure = super::capture_failure("stderr", stderr.as_ref());
335 let Some(status) = status else {
336 return Err(failure(
337 RunnerPhase::Cleanup,
338 LeaderTermination::Unconfirmed,
339 io::Error::new(
340 io::ErrorKind::TimedOut,
341 "the terminated process was not reaped within its bound",
342 ),
343 ));
344 };
345 stdout_failure?;
346 stderr_failure?;
347 writer_result
348 .map_err(|error| failure(RunnerPhase::Stdin, LeaderTermination::Confirmed, error))?;
349 Ok(StreamOutput {
350 status,
351 cancelled,
352 timed_out,
353 stdout_tail: sinks.stdout_tail.into_bytes(),
354 stderr_tail: sinks.stderr_tail.into_text(),
355 })
356}
357
358fn drain_output(receiver: &mpsc::Receiver<Message>, sinks: &mut Sinks<'_>, grace: Duration) {
359 let mut drain_end = Instant::now() + grace;
360 loop {
361 let left = drain_end.saturating_duration_since(Instant::now());
362 if left.is_zero() {
363 return;
364 }
365 match receiver.recv_timeout(left) {
366 Ok(message) => {
367 let callback_started = Instant::now();
368 sinks.handle(message);
369 drain_end += callback_started.elapsed();
370 }
371 Err(RecvTimeoutError::Timeout | RecvTimeoutError::Disconnected) => return,
372 }
373 }
374}
375
376#[cfg(test)]
377mod drain_tests {
378 use super::*;
379
380 #[test]
381 fn callback_longer_than_grace_does_not_cancel_a_backpressured_reader() {
382 let (sender, receiver) = mpsc::sync_channel(QUEUE_CHUNKS);
383 for _ in 0..QUEUE_CHUNKS {
386 sender.send(Message::Stdout(vec![b'x'])).unwrap();
387 }
388 let reader = std::thread::spawn(move || {
389 for _ in 0..QUEUE_CHUNKS {
390 if sender.send(Message::Stdout(vec![b'x'])).is_err() {
391 return;
392 }
393 }
394 });
395 let mut count = 0;
396 let mut on_stdout = |chunk: &[u8]| {
397 if count == 0 {
398 std::thread::sleep(Duration::from_millis(100));
399 }
400 count += chunk.len();
401 };
402 let mut on_stderr = |_: &str| {};
403 {
404 let mut sinks = Sinks {
405 on_stdout: &mut on_stdout,
406 on_stderr_line: &mut on_stderr,
407 stdout_tail: Tail::new(8),
408 stderr_tail: Tail::new(8),
409 splitter: LineSplitter::default(),
410 };
411 drain_output(&receiver, &mut sinks, Duration::from_millis(50));
412 }
413 drop(receiver);
415 reader.join().unwrap();
416 assert_eq!(count, 2 * QUEUE_CHUNKS);
417 }
418}
419
420#[derive(Clone, Copy)]
421enum Ended {
422 Exited(ExitStatus),
423 Cancelled,
424 TimedOut,
425}
426
427enum Message {
428 Stdout(Vec<u8>),
429 Stderr(Vec<u8>),
430}
431
432fn forward(
435 sender: SyncSender<Message>,
436 wrap: fn(Vec<u8>) -> Message,
437) -> crate::native_capture::ChunkSink {
438 Box::new(move |chunk: &[u8], stop: &AtomicBool| {
439 let mut message = wrap(chunk.to_vec());
440 loop {
441 match sender.try_send(message) {
442 Ok(()) | Err(TrySendError::Disconnected(_)) => return,
443 Err(TrySendError::Full(back)) => {
444 if stop.load(Ordering::Acquire) {
445 return;
446 }
447 message = back;
448 std::thread::sleep(Duration::from_millis(2));
449 }
450 }
451 }
452 })
453}
454
455struct Sinks<'a> {
457 on_stdout: &'a mut dyn FnMut(&[u8]),
458 on_stderr_line: &'a mut dyn FnMut(&str),
459 stdout_tail: Tail,
460 stderr_tail: Tail,
461 splitter: LineSplitter,
462}
463
464impl Sinks<'_> {
465 fn handle(&mut self, message: Message) {
466 match message {
467 Message::Stdout(chunk) => {
468 self.stdout_tail.push(&chunk);
469 (self.on_stdout)(&chunk);
470 }
471 Message::Stderr(chunk) => {
472 let Self {
473 on_stderr_line,
474 stderr_tail,
475 splitter,
476 ..
477 } = self;
478 splitter.feed(&chunk, &mut |line| {
479 stderr_tail.push_line(line);
480 (**on_stderr_line)(line);
481 });
482 }
483 }
484 }
485
486 fn finish_stream(&mut self) {
488 let Self {
489 on_stderr_line,
490 stderr_tail,
491 splitter,
492 ..
493 } = self;
494 splitter.finish(&mut |line| {
495 stderr_tail.push_line(line);
496 (**on_stderr_line)(line);
497 });
498 }
499}
500
501struct Tail {
503 cap: usize,
504 bytes: Vec<u8>,
505}
506
507impl Tail {
508 fn new(cap: usize) -> Self {
509 Self {
510 cap,
511 bytes: Vec::new(),
512 }
513 }
514
515 fn push(&mut self, chunk: &[u8]) {
516 if self.cap == 0 {
517 return;
518 }
519 if chunk.len() >= self.cap {
520 self.bytes.clear();
521 self.bytes
522 .extend_from_slice(&chunk[chunk.len() - self.cap..]);
523 return;
524 }
525 self.bytes.extend_from_slice(chunk);
526 if self.bytes.len() > self.cap.saturating_mul(2) {
527 self.trim();
528 }
529 }
530
531 fn push_line(&mut self, line: &str) {
533 self.push(line.as_bytes());
534 self.push(b"\n");
535 }
536
537 fn trim(&mut self) {
538 if self.bytes.len() > self.cap {
539 let cut = self.bytes.len() - self.cap;
540 self.bytes.drain(..cut);
541 }
542 }
543
544 fn into_bytes(mut self) -> Vec<u8> {
545 self.trim();
546 self.bytes
547 }
548
549 fn into_text(self) -> String {
552 let bytes = self.into_bytes();
553 let start = bytes
554 .iter()
555 .position(|&byte| byte & 0xC0 != 0x80)
556 .unwrap_or(bytes.len());
557 let mut text = String::from_utf8_lossy(&bytes[start..]).into_owned();
558 if text.ends_with('\n') {
559 text.pop();
560 }
561 text
562 }
563}
564
565#[derive(Default)]
567struct LineSplitter {
568 line: Vec<u8>,
569 skip_lf: bool,
571 dropping: bool,
573}
574
575impl LineSplitter {
576 fn feed(&mut self, bytes: &[u8], emit: &mut dyn FnMut(&str)) {
577 for &byte in bytes {
578 if std::mem::take(&mut self.skip_lf) && byte == b'\n' {
579 continue;
580 }
581 match byte {
582 b'\n' | b'\r' => {
583 self.skip_lf = byte == b'\r';
584 self.end_line(emit);
585 }
586 _ if self.dropping => {}
587 _ => {
588 self.line.push(byte);
589 if self.line.len() > LINE_CAP {
590 self.line.clear();
591 self.dropping = true;
592 }
593 }
594 }
595 }
596 }
597
598 fn finish(&mut self, emit: &mut dyn FnMut(&str)) {
600 if !self.line.is_empty() || self.dropping {
601 self.end_line(emit);
602 }
603 }
604
605 fn end_line(&mut self, emit: &mut dyn FnMut(&str)) {
606 if self.dropping {
607 self.dropping = false;
608 emit(&format!("[stderr line over {LINE_CAP} bytes dropped]"));
609 } else {
610 let text = String::from_utf8_lossy(&self.line).into_owned();
611 emit(&text);
612 }
613 self.line.clear();
614 }
615}
616
617#[cfg(test)]
618mod tests {
619 use super::*;
620
621 fn lines(chunks: &[&[u8]]) -> Vec<String> {
622 let mut splitter = LineSplitter::default();
623 let mut out = Vec::new();
624 for chunk in chunks {
625 splitter.feed(chunk, &mut |line| out.push(line.to_owned()));
626 }
627 splitter.finish(&mut |line| out.push(line.to_owned()));
628 out
629 }
630
631 #[test]
632 fn lines_end_at_lf_cr_and_crlf_across_chunks() {
633 assert_eq!(lines(&[b"a\rb\nc\r\nd"]), ["a", "b", "c", "d"]);
634 assert_eq!(lines(&[b"x\r", b"\ny\n\nz"]), ["x", "y", "", "z"]);
636 assert_eq!(lines(&[b"par", b"tial\n"]), ["partial"]);
637 assert!(lines(&[b""]).is_empty());
638 }
639
640 #[test]
641 fn an_overlong_line_becomes_one_marker() {
642 let long = vec![b'x'; LINE_CAP + 10];
643 let got = lines(&[&long, b"\nnext\n"]);
644 let marker = format!("[stderr line over {LINE_CAP} bytes dropped]");
645 assert_eq!(got, [marker.as_str(), "next"]);
646 }
647
648 #[test]
649 fn a_tail_keeps_the_last_bytes() {
650 let mut tail = Tail::new(5);
651 tail.push(b"ab");
652 tail.push(b"cdefgh");
653 tail.push(b"i");
654 assert_eq!(tail.into_bytes(), b"efghi");
655 let mut none = Tail::new(0);
656 none.push(b"abc");
657 assert!(none.into_bytes().is_empty());
658 let mut many = Tail::new(4);
660 for byte in 0..100u8 {
661 many.push(&[byte]);
662 assert!(many.bytes.len() <= 9);
663 }
664 assert_eq!(many.into_bytes(), [96, 97, 98, 99]);
665 }
666
667 #[test]
668 fn the_stderr_tail_is_joined_lines_without_a_split_character() {
669 let mut tail = Tail::new(6);
670 for line in ["l1", "l2", "l3", "l4"] {
671 tail.push_line(line);
672 }
673 assert_eq!(tail.into_text(), "l3\nl4");
674 let mut tail = Tail::new(4);
676 tail.push_line("a\u{e9}\u{e9}");
677 assert_eq!(tail.into_text(), "\u{e9}");
678 }
679
680 #[test]
681 fn options_default_and_build() {
682 let options = StreamOptions::new();
683 assert_eq!(options.timeout, None);
684 assert_eq!(options.poll, Duration::from_millis(100));
685 assert_eq!(options.stdout_tail_cap, 64 * 1024);
686 assert_eq!(options.stderr_tail_cap, 64 * 1024);
687 assert!(options.no_window);
688 assert!(options.input.is_none());
689 let options = options
690 .timeout(Duration::from_secs(3))
691 .poll(Duration::from_millis(10))
692 .stdout_tail_cap(1)
693 .stderr_tail_cap(2)
694 .no_window(false)
695 .input(b"in".to_vec());
696 assert_eq!(options.timeout, Some(Duration::from_secs(3)));
697 assert_eq!(options.poll, Duration::from_millis(10));
698 assert_eq!((options.stdout_tail_cap, options.stderr_tail_cap), (1, 2));
699 assert!(!options.no_window);
700 assert_eq!(options.input.as_deref(), Some(&b"in"[..]));
701 }
702}