1use std::sync::atomic::{AtomicBool, Ordering};
18use std::sync::{Arc, Condvar, Mutex};
19use std::time::Duration;
20
21use serde::Serialize;
22use tokio::sync::Mutex as AsyncMutex;
23use tokio::sync::Notify;
24use tonic::{Code, Status, Streaming};
25
26use crate::error::SailError;
27use crate::pb::workerproxy::v1 as pb;
28use crate::worker::{
29 retry_deadline, rpc_attempt_timeout, should_invalidate_channel,
30 should_retry_transient_exec_rpc, sleep_before_retry, WorkerProxy,
31 EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS,
32};
33
34#[doc(hidden)]
37pub const EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS: f64 = 30.0;
38const STREAM_BUFFER_CAP_BYTES: usize = 1024 * 1024;
44const STDIN_WRITE_CHUNK_BYTES: usize = 256 * 1024;
47
48fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
52 mutex
53 .lock()
54 .unwrap_or_else(std::sync::PoisonError::into_inner)
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum OutputStream {
60 Stdout,
62 Stderr,
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
68pub enum ReadStep {
69 Chunk(Vec<u8>),
74 Eof,
76 Pending,
79}
80
81#[derive(Debug, Clone, Serialize)]
83#[non_exhaustive]
84#[allow(clippy::struct_excessive_bools)]
85pub struct ExecResult {
86 pub stdout: String,
91 pub stderr: String,
93 pub exit_code: i32,
95 pub timed_out: bool,
97 pub stdout_truncated: bool,
100 pub stderr_truncated: bool,
103 pub stdout_complete: bool,
109 pub stderr_complete: bool,
112 pub stdout_total_bytes: i64,
117 pub stderr_total_bytes: i64,
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub enum CancelSignal {
125 Interrupt,
127 Kill,
129}
130
131impl CancelSignal {
132 fn is_force(self) -> bool {
134 matches!(self, CancelSignal::Kill)
135 }
136}
137
138#[derive(Debug, Clone, Copy, PartialEq)]
141pub enum RetryBudget {
142 None,
144 Within(Duration),
146 Forever,
148}
149
150impl RetryBudget {
151 #[doc(hidden)]
154 pub fn as_secs_f64(self) -> f64 {
155 match self {
156 RetryBudget::None => 0.0,
157 RetryBudget::Within(d) => d.as_secs_f64(),
158 RetryBudget::Forever => f64::INFINITY,
159 }
160 }
161
162 #[doc(hidden)]
165 pub fn from_secs_f64(secs: f64) -> RetryBudget {
166 if secs <= 0.0 {
167 RetryBudget::None
168 } else if secs.is_finite() {
169 RetryBudget::Within(Duration::from_secs_f64(secs))
170 } else {
171 RetryBudget::Forever
172 }
173 }
174}
175
176#[derive(Debug, Clone)]
182pub struct ExecOptions {
183 pub timeout: Option<Duration>,
187 pub open_stdin: bool,
189 pub pty: bool,
191 pub term: String,
193 pub cols: u32,
195 pub rows: u32,
197 pub env: Vec<(String, String)>,
204 pub idempotency_key: String,
207 pub retry_timeout: RetryBudget,
212 pub cwd: Option<String>,
215 pub background: bool,
220}
221
222impl Default for ExecOptions {
223 fn default() -> ExecOptions {
224 ExecOptions {
225 timeout: None,
226 open_stdin: false,
227 pty: false,
228 term: String::new(),
229 cols: 0,
230 rows: 0,
231 env: Vec::new(),
232 idempotency_key: String::new(),
233 retry_timeout: RetryBudget::Within(Duration::from_secs_f64(
234 EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
235 )),
236 cwd: None,
237 background: false,
238 }
239 }
240}
241
242pub(crate) fn sh_quote(value: &str) -> String {
244 format!("'{}'", value.replace('\'', "'\\''"))
245}
246
247const PTY_ENV_WHITELIST: [&str; 3] = ["COLORTERM", "LANG", "TERM_PROGRAM"];
251
252fn pty_env_whitelisted(key: &str) -> bool {
253 PTY_ENV_WHITELIST.contains(&key) || key.starts_with("LC_")
254}
255
256pub(crate) fn pty_forward_env() -> Vec<(String, String)> {
258 pty_forward_env_from(
262 std::env::vars_os()
263 .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))),
264 )
265}
266
267fn pty_forward_env_from(vars: impl Iterator<Item = (String, String)>) -> Vec<(String, String)> {
268 vars.filter(|(key, _)| pty_env_whitelisted(key)).collect()
269}
270
271#[doc(hidden)]
277pub fn encode_env(
278 pairs: &[(String, String)],
279) -> Result<std::collections::HashMap<String, String>, SailError> {
280 let mut env = std::collections::HashMap::with_capacity(pairs.len());
281 for (key, value) in pairs {
282 if !is_portable_env_name(key) || value.contains('\0') {
286 return Err(SailError::InvalidArgument {
287 message: format!("invalid env entry {key:?}"),
288 });
289 }
290 env.insert(key.clone(), value.clone());
291 }
292 Ok(env)
293}
294
295fn is_portable_env_name(name: &str) -> bool {
300 let mut chars = name.chars();
301 match chars.next() {
302 Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
303 _ => return false,
304 }
305 chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
306}
307
308#[doc(hidden)]
313pub fn shell_argv(command: &str, options: &ExecOptions) -> Result<Vec<String>, SailError> {
314 let invalid = |message: &str| {
315 Err(SailError::InvalidArgument {
316 message: message.to_string(),
317 })
318 };
319 if command.is_empty() {
320 return invalid("command must be non-empty");
321 }
322 if options.background && (options.open_stdin || options.pty) {
323 return invalid("background is not supported with open_stdin or pty");
324 }
325 let mut command = command.to_string();
326 if let Some(cwd) = &options.cwd {
327 let cwd = cwd.trim();
328 if cwd.is_empty() {
329 return invalid("cwd must be non-empty");
330 }
331 command = format!(
332 "cd {} && exec /bin/sh -lc {}",
333 sh_quote(cwd),
334 sh_quote(&command)
335 );
336 }
337 if options.background {
338 command = format!(
339 "nohup /bin/sh -lc {} </dev/null >/dev/null 2>&1 &",
340 sh_quote(&command)
341 );
342 }
343 Ok(vec!["/bin/sh".to_string(), "-lc".to_string(), command])
344}
345
346#[doc(hidden)]
348#[derive(Debug, Clone)]
349pub struct ExecParams {
350 pub sailbox_id: String,
352 pub exec_endpoint: String,
354 pub argv: Vec<String>,
356 pub timeout_seconds: u32,
359 pub idempotency_key: String,
362 pub open_stdin: bool,
364 pub pty: bool,
366 pub term: String,
368 pub cols: u32,
370 pub rows: u32,
372 pub env: std::collections::HashMap<String, String>,
376 pub retry_timeout: f64,
379 pub extra_metadata: Vec<(String, String)>,
381}
382
383#[derive(Default)]
388struct Ring {
389 pieces: Vec<Vec<u8>>,
390 first_idx: usize,
391 size: usize,
392 dropped: bool,
393 front_clips: u64,
398 resets: u64,
403}
404
405impl Ring {
406 fn append(&mut self, data: Vec<u8>) {
407 self.size += data.len();
408 self.pieces.push(data);
409 while self.size > STREAM_BUFFER_CAP_BYTES {
410 let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
411 if self.pieces[0].len() <= overflow {
412 self.size -= self.pieces[0].len();
413 self.pieces.remove(0);
414 self.first_idx += 1;
415 } else {
416 self.pieces[0].drain(..overflow);
417 self.size -= overflow;
418 self.front_clips += 1;
419 }
420 self.dropped = true;
421 }
422 }
423
424 fn reset_to(&mut self, repaint: Vec<u8>) {
432 self.first_idx += self.pieces.len();
433 self.pieces.clear();
434 self.size = 0;
435 self.dropped = false;
436 self.resets += 1;
437 if !repaint.is_empty() {
438 self.append(repaint);
439 }
440 }
441
442 fn tail(&self) -> Vec<u8> {
443 self.pieces.concat()
444 }
445}
446
447fn lossy_tail(ring: &Ring) -> String {
453 String::from_utf8_lossy(&ring.tail()).replace('\0', "\u{FFFD}")
454}
455
456#[derive(Default)]
457struct State {
458 stdout: Ring,
459 stderr: Ring,
460 ended: bool,
461}
462
463impl State {
464 fn ring(&self, which: OutputStream) -> &Ring {
465 match which {
466 OutputStream::Stdout => &self.stdout,
467 OutputStream::Stderr => &self.stderr,
468 }
469 }
470}
471
472#[derive(Clone)]
474struct ExitInfo {
475 status: i32,
476 exit_code: i32,
477 timed_out: bool,
478 stdout_truncated: bool,
479 stderr_truncated: bool,
480 error_message: String,
481 stdout_seq: i64,
482 stderr_seq: i64,
483 stdout_total_bytes: i64,
484 stderr_total_bytes: i64,
485}
486
487#[derive(Default)]
488struct StdinState {
489 offset: i64,
490 eof_sent: bool,
491 broken: bool,
492 write_in_flight: bool,
500}
501
502struct ExecShared {
503 worker: Arc<WorkerProxy>,
504 params: ExecParams,
505 state: Mutex<State>,
506 cond: Condvar,
509 data_notify: Notify,
512 exit: Mutex<Option<ExitInfo>>,
513 high_seq: Mutex<(i64, i64)>,
515 stdin: AsyncMutex<StdinState>,
516 ended: AtomicBool,
517 ended_notify: Notify,
518 closing: AtomicBool,
519 close_notify: Notify,
520}
521
522pub struct ExecProcess {
525 shared: Arc<ExecShared>,
526 exec_request_id: String,
527}
528
529impl Drop for ExecProcess {
530 fn drop(&mut self) {
531 self.close();
536 }
537}
538
539impl ExecProcess {
540 #[doc(hidden)]
548 pub async fn start(
549 worker: Arc<WorkerProxy>,
550 mut params: ExecParams,
551 ) -> Result<ExecProcess, SailError> {
552 let key = params.idempotency_key.trim();
556 params.idempotency_key = if key.is_empty() {
557 format!("exec_{}", uuid::Uuid::new_v4())
558 } else {
559 key.to_string()
560 };
561 if params.pty {
564 for (key, value) in pty_forward_env() {
565 params.env.entry(key).or_insert(value);
566 }
567 }
568 let (exec_request_id, stream) = submit(
569 &worker, ¶ms, 0, 0,
570 )
571 .await?;
572 let shared = Arc::new(ExecShared {
573 worker,
574 params,
575 state: Mutex::new(State::default()),
576 cond: Condvar::new(),
577 data_notify: Notify::new(),
578 exit: Mutex::new(None),
579 high_seq: Mutex::new((0, 0)),
580 stdin: AsyncMutex::new(StdinState::default()),
581 ended: AtomicBool::new(false),
582 ended_notify: Notify::new(),
583 closing: AtomicBool::new(false),
584 close_notify: Notify::new(),
585 });
586 let pump_shared = shared.clone();
587 tokio::spawn(async move { pump(pump_shared, stream).await });
588 Ok(ExecProcess {
589 shared,
590 exec_request_id,
591 })
592 }
593
594 pub fn exec_request_id(&self) -> &str {
597 &self.exec_request_id
598 }
599
600 pub fn idempotency_key(&self) -> &str {
603 &self.shared.params.idempotency_key
604 }
605
606 pub fn reader(&self, which: OutputStream) -> StreamReader {
609 StreamReader {
610 shared: self.shared.clone(),
611 which,
612 cursor: 0,
613 dropped: false,
614 reset: false,
615 seen_front_clips: 0,
616 seen_resets: 0,
617 }
618 }
619
620 pub fn reader_async(&self, which: OutputStream) -> AsyncStreamReader {
623 AsyncStreamReader {
624 shared: self.shared.clone(),
625 which,
626 cursor: 0,
627 }
628 }
629
630 #[doc(hidden)]
636 pub fn buffered_output(&self, which: OutputStream) -> Vec<u8> {
637 lock(&self.shared.state).ring(which).tail()
638 }
639
640 pub fn poll(&self) -> Option<Result<i32, SailError>> {
644 let exit = lock(&self.shared.exit);
645 exit.as_ref()
646 .map(|exit| match terminal_status_error(exit.status) {
647 Some(err) => Err(err),
648 None => Ok(exit.exit_code),
649 })
650 }
651
652 pub fn close(&self) {
654 self.shared.closing.store(true, Ordering::SeqCst);
655 self.shared.close_notify.notify_one();
660 }
661
662 #[doc(hidden)]
667 pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
668 let _ = tokio::time::timeout(timeout, self.await_ended()).await;
669 self.shared.ended.load(Ordering::SeqCst)
670 }
671
672 async fn await_ended(&self) {
674 loop {
675 let notified = self.shared.ended_notify.notified();
676 if self.shared.ended.load(Ordering::SeqCst) {
677 return;
678 }
679 notified.await;
680 }
681 }
682
683 pub async fn wait(&self) -> Result<ExecResult, SailError> {
690 self.await_ended().await;
691 let exit = lock(&self.shared.exit).clone();
692 let (high_out, high_err) = *lock(&self.shared.high_seq);
693 let (out_dropped, err_dropped) = {
694 let state = lock(&self.shared.state);
695 (state.stdout.dropped, state.stderr.dropped)
696 };
697
698 let incomplete = exit
706 .as_ref()
707 .is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
708 let stdout_truncated_flag = exit
713 .as_ref()
714 .is_some_and(|exit| exit.stdout_truncated || out_dropped);
715 let stderr_truncated_flag = exit
716 .as_ref()
717 .is_some_and(|exit| exit.stderr_truncated || err_dropped);
718
719 let stdout_complete = exit
724 .as_ref()
725 .is_some_and(|exit| exit.stdout_seq <= high_out);
726 let stderr_complete = exit
727 .as_ref()
728 .is_some_and(|exit| exit.stderr_seq <= high_err);
729
730 let stdout_total_bytes = exit.as_ref().map_or(0, |exit| exit.stdout_total_bytes);
734 let stderr_total_bytes = exit.as_ref().map_or(0, |exit| exit.stderr_total_bytes);
735
736 if exit.is_none() || incomplete {
737 let outcome = self
738 .shared
739 .worker
740 .wait_exec(
741 &self.shared.params.exec_endpoint,
742 &self.shared.params.sailbox_id,
743 &self.exec_request_id,
744 self.shared.params.retry_timeout,
745 )
746 .await?;
747 {
750 let mut exit_slot = lock(&self.shared.exit);
751 if exit_slot.is_none() {
752 *exit_slot = Some(ExitInfo {
753 status: outcome.status,
754 exit_code: outcome.exit_code,
755 timed_out: outcome.timed_out,
756 stdout_truncated: outcome.stdout_truncated,
757 stderr_truncated: outcome.stderr_truncated,
758 error_message: String::new(),
759 stdout_seq: 0,
760 stderr_seq: 0,
761 stdout_total_bytes: 0,
763 stderr_total_bytes: 0,
764 });
765 }
766 }
767 if let Some(witnessed) = exit.as_ref() {
773 if outcome.status == pb::SailboxExecStatus::WorkerLost as i32 {
774 let state = lock(&self.shared.state);
775 let mut stderr = lossy_tail(&state.stderr);
776 if witnessed.status == pb::SailboxExecStatus::Failed as i32
777 && !witnessed.error_message.is_empty()
778 {
779 stderr = witnessed.error_message.clone();
780 }
781 return Ok(ExecResult {
782 stdout: lossy_tail(&state.stdout),
783 stderr,
784 exit_code: witnessed.exit_code,
785 timed_out: witnessed.timed_out,
786 stdout_truncated: witnessed.stdout_truncated
787 || out_dropped
788 || witnessed.stdout_seq > high_out,
789 stderr_truncated: witnessed.stderr_truncated
790 || err_dropped
791 || witnessed.stderr_seq > high_err,
792 stdout_complete,
793 stderr_complete,
794 stdout_total_bytes,
795 stderr_total_bytes,
796 });
797 }
798 }
799 if let Some(err) = terminal_status_error(outcome.status) {
800 return Err(err);
801 }
802 return Ok(ExecResult {
803 stdout: outcome.stdout,
804 stderr: outcome.stderr,
805 exit_code: outcome.exit_code,
806 timed_out: outcome.timed_out,
807 stdout_truncated: outcome.stdout_truncated,
808 stderr_truncated: outcome.stderr_truncated,
809 stdout_complete,
810 stderr_complete,
811 stdout_total_bytes,
812 stderr_total_bytes,
813 });
814 }
815
816 let exit = exit.expect("exit present on the clean path");
817 if let Some(err) = terminal_status_error(exit.status) {
818 return Err(err);
819 }
820 let state = lock(&self.shared.state);
821 let mut stderr = lossy_tail(&state.stderr);
822 if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
823 stderr = exit.error_message.clone();
825 }
826 Ok(ExecResult {
827 stdout: lossy_tail(&state.stdout),
828 stderr,
829 exit_code: exit.exit_code,
830 timed_out: exit.timed_out,
831 stdout_truncated: stdout_truncated_flag,
832 stderr_truncated: stderr_truncated_flag,
833 stdout_complete,
834 stderr_complete,
835 stdout_total_bytes,
836 stderr_total_bytes,
837 })
838 }
839
840 pub async fn cancel(&self, signal: CancelSignal, retry: RetryBudget) -> Result<(), SailError> {
843 self.shared
844 .worker
845 .cancel_exec(
846 &self.shared.params.exec_endpoint,
847 &self.shared.params.sailbox_id,
848 &self.exec_request_id,
849 signal.is_force(),
850 retry.as_secs_f64(),
851 )
852 .await
853 }
854
855 pub async fn resize(&self, cols: u32, rows: u32) {
859 let message = pb::ResizeSailboxExecRequest {
860 sailbox_id: self.shared.params.sailbox_id.clone(),
861 exec_request_id: self.exec_request_id.clone(),
862 cols,
863 rows,
864 };
865 let Ok(request) =
866 self.shared
867 .worker
868 .request_for(message, &[], Some(Duration::from_secs(5)))
869 else {
870 return;
871 };
872 if let Ok(mut client) = self
873 .shared
874 .worker
875 .client_for(&self.shared.params.exec_endpoint)
876 {
877 let _ = client.resize_sailbox_exec(request).await;
878 }
879 }
880
881 pub async fn resync(&self) {
889 let message = pb::ResyncSailboxExecRequest {
890 sailbox_id: self.shared.params.sailbox_id.clone(),
891 exec_request_id: self.exec_request_id.clone(),
892 };
893 let Ok(request) =
894 self.shared
895 .worker
896 .request_for(message, &[], Some(Duration::from_secs(5)))
897 else {
898 return;
899 };
900 if let Ok(mut client) = self
901 .shared
902 .worker
903 .client_for(&self.shared.params.exec_endpoint)
904 {
905 let _ = client.resync_sailbox_exec(request).await;
906 }
907 }
908
909 pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
913 let mut stdin = self.shared.stdin.lock().await;
914 if stdin.eof_sent {
915 return Err(SailError::BrokenPipe {
916 message: "stdin is closed".to_string(),
917 });
918 }
919 if stdin.broken {
920 return Err(SailError::BrokenPipe {
921 message: "an earlier stdin write failed".to_string(),
922 });
923 }
924 if stdin.write_in_flight {
925 stdin.broken = true;
929 return Err(SailError::BrokenPipe {
930 message: "an earlier stdin write was interrupted".to_string(),
931 });
932 }
933 if data.is_empty() {
934 return Ok(());
935 }
936 stdin.write_in_flight = true;
937 let result = self.send_stdin(&mut stdin, data, false).await;
938 stdin.write_in_flight = false;
941 result
942 }
943
944 pub async fn close_stdin(&self) -> Result<(), SailError> {
946 let mut stdin = self.shared.stdin.lock().await;
947 if stdin.eof_sent {
948 return Ok(());
949 }
950 if stdin.broken || stdin.write_in_flight {
951 stdin.eof_sent = true;
954 return Ok(());
955 }
956 match self.send_stdin(&mut stdin, &[], true).await {
960 Ok(()) => {
961 stdin.eof_sent = true;
962 Ok(())
963 }
964 Err(SailError::BrokenPipe { .. }) => {
966 stdin.eof_sent = true;
967 Ok(())
968 }
969 Err(err) => Err(err),
970 }
971 }
972
973 async fn send_stdin(
977 &self,
978 stdin: &mut StdinState,
979 payload: &[u8],
980 eof: bool,
981 ) -> Result<(), SailError> {
982 let endpoint = &self.shared.params.exec_endpoint;
983 let sailbox_id = &self.shared.params.sailbox_id;
984 let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
985 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
986 let mut sent = 0usize;
987 loop {
988 let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
989 let chunk = &payload[sent..end];
990 let last = end >= payload.len();
991 let message = pb::WriteSailboxExecStdinRequest {
992 sailbox_id: sailbox_id.clone(),
993 exec_request_id: self.exec_request_id.clone(),
994 offset: stdin.offset,
995 data: chunk.to_vec(),
996 eof: eof && last,
997 };
998 let request = match self.shared.worker.request_for(
1002 message,
1003 &[],
1004 Some(rpc_attempt_timeout(deadline)),
1005 ) {
1006 Ok(request) => request,
1007 Err(err) => return Err(err),
1008 };
1009 let result = match self.shared.worker.client_for(endpoint) {
1010 Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
1011 Err(err) => return Err(err),
1012 };
1013 match result {
1014 Ok(resp) => {
1015 let accepted_through = resp.into_inner().accepted_through;
1016 let accepted =
1021 ((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
1022 stdin.offset += accepted as i64;
1023 sent += accepted;
1024 if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
1025 return Ok(());
1026 }
1027 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
1029 if accepted > 0 {
1030 delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
1031 } else {
1032 delay = sleep_no_deadline(delay).await;
1034 }
1035 }
1036 Err(status) => {
1037 if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
1038 return Err(SailError::BrokenPipe {
1040 message: status.message().to_string(),
1041 });
1042 }
1043 if is_exec_not_ready(&status) {
1044 delay = sleep_no_deadline(delay).await;
1048 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
1049 continue;
1050 }
1051 if !should_retry_transient_exec_rpc(&status, deadline) {
1052 stdin.broken = true;
1055 return Err(SailError::from_exec_status(&status));
1056 }
1057 tracing::warn!(code = ?status.code(), "retrying exec stdin write");
1058 if should_invalidate_channel(&status) {
1059 self.shared.worker.channels().invalidate(endpoint);
1060 }
1061 delay = sleep_before_retry(delay, deadline).await;
1062 }
1063 }
1064 }
1065 }
1066}
1067
1068pub struct StreamReader {
1071 shared: Arc<ExecShared>,
1072 which: OutputStream,
1073 cursor: usize,
1074 dropped: bool,
1075 reset: bool,
1080 seen_front_clips: u64,
1083 seen_resets: u64,
1087}
1088
1089impl StreamReader {
1090 pub fn next(&mut self, timeout: Duration) -> ReadStep {
1094 let mut state = lock(&self.shared.state);
1095 loop {
1096 let ring = state.ring(self.which);
1097 if ring.resets != self.seen_resets {
1111 self.reset = true;
1112 self.dropped = ring.dropped;
1113 if self.cursor < ring.first_idx {
1114 self.cursor = ring.first_idx;
1115 }
1116 } else if self.cursor < ring.first_idx {
1117 self.cursor = ring.first_idx;
1118 self.dropped = true;
1119 } else if self.cursor == ring.first_idx && ring.front_clips != self.seen_front_clips {
1120 self.dropped = true;
1121 }
1122 self.seen_front_clips = ring.front_clips;
1123 self.seen_resets = ring.resets;
1124 let available = ring.first_idx + ring.pieces.len();
1125 if self.cursor < available {
1126 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
1127 self.cursor += 1;
1128 return ReadStep::Chunk(piece);
1129 }
1130 if state.ended {
1131 return ReadStep::Eof;
1132 }
1133 let (next_state, timed_out) = self
1134 .shared
1135 .cond
1136 .wait_timeout(state, timeout)
1137 .expect("exec state mutex poisoned");
1138 state = next_state;
1139 if timed_out.timed_out() {
1140 return ReadStep::Pending;
1141 }
1142 }
1143 }
1144
1145 pub fn try_next(&mut self) -> Option<Vec<u8>> {
1150 let state = lock(&self.shared.state);
1151 let ring = state.ring(self.which);
1152 if ring.resets != self.seen_resets {
1153 self.reset = true;
1154 self.dropped = ring.dropped;
1158 if self.cursor < ring.first_idx {
1159 self.cursor = ring.first_idx;
1160 }
1161 } else if self.cursor < ring.first_idx {
1162 self.cursor = ring.first_idx;
1163 self.dropped = true;
1164 } else if self.cursor == ring.first_idx && ring.front_clips != self.seen_front_clips {
1165 self.dropped = true;
1166 }
1167 self.seen_front_clips = ring.front_clips;
1168 self.seen_resets = ring.resets;
1169 if self.cursor < ring.first_idx + ring.pieces.len() {
1170 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
1171 self.cursor += 1;
1172 Some(piece)
1173 } else {
1174 None
1175 }
1176 }
1177
1178 pub fn took_drop(&mut self) -> bool {
1182 std::mem::take(&mut self.dropped)
1183 }
1184
1185 pub fn took_reset(&mut self) -> bool {
1191 std::mem::take(&mut self.reset)
1192 }
1193}
1194
1195pub struct AsyncStreamReader {
1199 shared: Arc<ExecShared>,
1200 which: OutputStream,
1201 cursor: usize,
1202}
1203
1204impl AsyncStreamReader {
1205 pub async fn next(&mut self) -> Option<Vec<u8>> {
1209 loop {
1210 let notified = self.shared.data_notify.notified();
1214 tokio::pin!(notified);
1215 notified.as_mut().enable();
1216 {
1217 let state = lock(&self.shared.state);
1218 let ring = state.ring(self.which);
1219 if self.cursor < ring.first_idx {
1220 self.cursor = ring.first_idx;
1221 }
1222 let available = ring.first_idx + ring.pieces.len();
1223 if self.cursor < available {
1224 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
1225 self.cursor += 1;
1226 return Some(piece);
1227 }
1228 if state.ended {
1229 return None;
1230 }
1231 }
1232 notified.await;
1233 }
1234 }
1235}
1236
1237async fn sleep_no_deadline(delay: f64) -> f64 {
1240 let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
1241 tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
1242 (delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
1243}
1244
1245fn is_exec_not_ready(status: &Status) -> bool {
1246 status.code() == Code::Unavailable && status.message().contains("; retry")
1247}
1248
1249fn terminal_status_error(status: i32) -> Option<SailError> {
1252 if status == pb::SailboxExecStatus::WorkerLost as i32 {
1253 return Some(SailError::WorkerLost {
1254 message: "the machine hosting your sailbox was lost before the command \
1255 finished; run exec again to retry"
1256 .to_string(),
1257 });
1258 }
1259 None
1260}
1261
1262async fn submit(
1265 worker: &Arc<WorkerProxy>,
1266 params: &ExecParams,
1267 stdout_resume_seq: i64,
1268 stderr_resume_seq: i64,
1269) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
1270 let deadline = retry_deadline(params.retry_timeout);
1271 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
1272 loop {
1273 let message = pb::StreamSailboxExecRequest {
1274 sailbox_id: params.sailbox_id.clone(),
1275 argv: params.argv.clone(),
1276 timeout_seconds: params.timeout_seconds,
1277 idempotency_key: params.idempotency_key.clone(),
1278 open_stdin: params.open_stdin,
1279 stdout_resume_seq,
1280 stderr_resume_seq,
1281 pty: params.pty,
1282 term_cols: params.cols,
1283 term_rows: params.rows,
1284 term: params.term.clone(),
1285 env: params.env.clone(),
1286 };
1287 let request =
1288 worker.request_for(message, ¶ms.extra_metadata, None)?;
1289 let status = match worker
1290 .client_for(¶ms.exec_endpoint)?
1291 .stream_sailbox_exec(request)
1292 .await
1293 {
1294 Ok(resp) => {
1295 let mut stream = resp.into_inner();
1296 match stream.message().await {
1297 Ok(Some(first)) => match first.frame {
1298 Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
1299 return Ok((started.exec_request_id, stream));
1300 }
1301 _ => {
1302 return Err(SailError::Execution {
1303 code: crate::error::GrpcCode::Internal,
1304 detail: "exec stream opened with a non-started frame".to_string(),
1305 });
1306 }
1307 },
1308 Ok(None) if stdout_resume_seq == 0 && stderr_resume_seq == 0 => {
1319 tonic::Status::unavailable(
1320 "exec stream ended before the server confirmed the launch",
1321 )
1322 }
1323 Ok(None) => {
1324 return Err(SailError::Execution {
1325 code: crate::error::GrpcCode::Internal,
1326 detail: "exec stream ended before the server confirmed the launch"
1327 .to_string(),
1328 });
1329 }
1330 Err(status) => status,
1331 }
1332 }
1333 Err(status) => status,
1334 };
1335 if !should_retry_transient_exec_rpc(&status, deadline) {
1336 return Err(SailError::from_exec_status(&status));
1337 }
1338 tracing::warn!(
1339 code = ?status.code(),
1340 stdout_resume_seq,
1341 stderr_resume_seq,
1342 "reconnecting exec stream"
1343 );
1344 if should_invalidate_channel(&status) {
1345 worker.channels().invalidate(¶ms.exec_endpoint);
1346 }
1347 delay = sleep_before_retry(delay, deadline).await;
1348 }
1349}
1350
1351struct Pump {
1356 shared: Arc<ExecShared>,
1357 stdout_seq: i64,
1358 stderr_seq: i64,
1359}
1360
1361impl Pump {
1362 fn new(shared: Arc<ExecShared>) -> Pump {
1363 Pump {
1364 shared,
1365 stdout_seq: 0,
1366 stderr_seq: 0,
1367 }
1368 }
1369
1370 fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
1376 match frame.frame {
1377 Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
1378 let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
1379 let which = if is_stderr {
1380 self.stderr_seq = self.stderr_seq.max(chunk.seq);
1381 OutputStream::Stderr
1382 } else {
1383 self.stdout_seq = self.stdout_seq.max(chunk.seq);
1384 OutputStream::Stdout
1385 };
1386 if !chunk.data.is_empty() {
1387 let mut state = lock(&self.shared.state);
1388 match which {
1389 OutputStream::Stdout => state.stdout.append(chunk.data),
1390 OutputStream::Stderr => state.stderr.append(chunk.data),
1391 }
1392 self.shared.cond.notify_all();
1393 self.shared.data_notify.notify_waiters();
1394 }
1395 false
1396 }
1397 Some(pb::stream_sailbox_exec_response::Frame::Snapshot(snap)) => {
1398 self.stdout_seq = snap.stdout_seq_basis;
1404 let mut state = lock(&self.shared.state);
1405 state.stdout.reset_to(snap.repaint);
1406 self.shared.cond.notify_all();
1407 self.shared.data_notify.notify_waiters();
1408 false
1409 }
1410 Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
1411 *lock(&self.shared.exit) = Some(ExitInfo {
1412 status: exit.status,
1413 exit_code: exit.return_code,
1414 timed_out: exit.timed_out,
1415 stdout_truncated: exit.stdout_truncated,
1416 stderr_truncated: exit.stderr_truncated,
1417 error_message: exit.error_message,
1418 stdout_seq: exit.stdout_seq,
1419 stderr_seq: exit.stderr_seq,
1420 stdout_total_bytes: exit.stdout_total_bytes,
1421 stderr_total_bytes: exit.stderr_total_bytes,
1422 });
1423 true
1424 }
1425 _ => false,
1426 }
1427 }
1428
1429 fn finalize(self) {
1432 {
1433 let mut state = lock(&self.shared.state);
1434 state.ended = true;
1435 self.shared.cond.notify_all();
1436 self.shared.data_notify.notify_waiters();
1437 }
1438 *lock(&self.shared.high_seq) = (self.stdout_seq, self.stderr_seq);
1439 self.shared.ended.store(true, Ordering::SeqCst);
1440 self.shared.ended_notify.notify_waiters();
1441 }
1442}
1443
1444async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
1447 let mut state = Pump::new(shared.clone());
1448 loop {
1449 if shared.closing.load(Ordering::SeqCst) {
1450 break;
1451 }
1452 let message = tokio::select! {
1453 biased;
1454 () = shared.close_notify.notified() => break,
1455 message = stream.message() => message,
1456 };
1457 match message {
1458 Ok(Some(frame)) => {
1459 if state.apply_frame(frame) {
1460 break;
1461 }
1462 }
1463 Ok(None) => break,
1465 Err(_status) => {
1466 if shared.closing.load(Ordering::SeqCst) {
1467 break;
1468 }
1469 match submit(
1472 &shared.worker,
1473 &shared.params,
1474 state.stdout_seq,
1475 state.stderr_seq,
1476 )
1477 .await
1478 {
1479 Ok((_id, fresh)) => {
1480 stream = fresh;
1481 continue;
1482 }
1483 Err(_) => break,
1484 }
1485 }
1486 }
1487 }
1488 state.finalize();
1489}
1490
1491#[cfg(any(test, feature = "fuzzing"))]
1497pub fn fuzz_exec_ring(data: &[u8]) {
1498 let mut chunked = Ring::default();
1502 let mut start = 0;
1503 for (i, byte) in data.iter().enumerate() {
1504 if byte & 1 == 1 {
1505 chunked.append(data[start..=i].to_vec());
1506 start = i + 1;
1507 }
1508 }
1509 if start < data.len() {
1510 chunked.append(data[start..].to_vec());
1511 }
1512 let mut whole = Ring::default();
1513 if !data.is_empty() {
1514 whole.append(data.to_vec());
1515 }
1516 assert_eq!(
1517 chunked.tail(),
1518 whole.tail(),
1519 "ring is not chunk-boundary invariant"
1520 );
1521
1522 let mut ring = Ring::default();
1526 ring.append(data.to_vec());
1527 let filler = vec![b'a'; STREAM_BUFFER_CAP_BYTES + 1 - data.len().min(STREAM_BUFFER_CAP_BYTES)];
1528 ring.append(filler.clone());
1529 assert!(
1530 ring.size <= STREAM_BUFFER_CAP_BYTES,
1531 "ring exceeded its cap"
1532 );
1533 let mut full = data.to_vec();
1534 full.extend_from_slice(&filler);
1535 assert!(
1536 full.ends_with(&ring.tail()),
1537 "ring tail is not a suffix of the appended bytes"
1538 );
1539
1540 let before_reset_pieces = ring.first_idx + ring.pieces.len();
1544 ring.reset_to(data.to_vec());
1545 assert!(
1546 ring.first_idx >= before_reset_pieces,
1547 "reset left old pieces reachable"
1548 );
1549 let expected = &data[data.len().saturating_sub(STREAM_BUFFER_CAP_BYTES)..];
1550 assert_eq!(
1551 ring.tail(),
1552 expected,
1553 "reset tail is not the capped repaint"
1554 );
1555}
1556
1557#[cfg(test)]
1558mod tests {
1559 use super::*;
1560
1561 #[test]
1562 fn shell_argv_wraps_cwd_and_background() {
1563 let plain = shell_argv("echo hi", &ExecOptions::default()).unwrap();
1564 assert_eq!(plain, ["/bin/sh", "-lc", "echo hi"]);
1565
1566 let cwd = shell_argv(
1567 "echo hi",
1568 &ExecOptions {
1569 cwd: Some("/app".to_string()),
1570 ..ExecOptions::default()
1571 },
1572 )
1573 .unwrap();
1574 assert_eq!(cwd[2], "cd '/app' && exec /bin/sh -lc 'echo hi'");
1575
1576 let background = shell_argv(
1577 "echo hi",
1578 &ExecOptions {
1579 cwd: Some("/app".to_string()),
1580 background: true,
1581 ..ExecOptions::default()
1582 },
1583 )
1584 .unwrap();
1585 assert_eq!(
1586 background[2],
1587 "nohup /bin/sh -lc 'cd '\\''/app'\\'' && exec /bin/sh -lc '\\''echo hi'\\''' </dev/null >/dev/null 2>&1 &"
1588 );
1589 }
1590
1591 #[test]
1592 fn shell_argv_rejects_invalid_combinations() {
1593 assert!(shell_argv("", &ExecOptions::default()).is_err());
1594 assert!(shell_argv(
1595 "x",
1596 &ExecOptions {
1597 cwd: Some(" ".to_string()),
1598 ..ExecOptions::default()
1599 }
1600 )
1601 .is_err());
1602 for (open_stdin, pty) in [(true, false), (false, true)] {
1603 assert!(shell_argv(
1604 "x",
1605 &ExecOptions {
1606 background: true,
1607 open_stdin,
1608 pty,
1609 ..ExecOptions::default()
1610 }
1611 )
1612 .is_err());
1613 }
1614 }
1615
1616 #[test]
1617 fn fuzz_exec_ring_holds_on_samples() {
1618 for sample in [
1621 &b""[..],
1622 b"hello",
1623 b"\xff\xfe\xfd",
1624 "€ µ é".as_bytes(),
1625 b"ab\xc3\xa9cd",
1626 &[0xC3, 0x28],
1627 ] {
1628 fuzz_exec_ring(sample);
1629 }
1630 }
1631
1632 #[test]
1633 fn ring_drops_oldest_past_cap_and_latches() {
1634 let mut ring = Ring::default();
1635 ring.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES]);
1636 assert!(!ring.dropped);
1637 ring.append(b"bbbb".to_vec());
1638 assert!(ring.dropped);
1639 assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
1640 assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
1641 assert!(ring.tail().ends_with(b"bbbb"));
1642 }
1643
1644 #[test]
1645 fn ring_clip_is_byte_exact() {
1646 let mut ring = Ring::default();
1650 let mut data = "é".as_bytes().to_vec();
1651 data.extend(std::iter::repeat_n(b'a', STREAM_BUFFER_CAP_BYTES - 1));
1652 ring.append(data);
1653 assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
1654 assert!(ring.dropped);
1655 let tail = ring.tail();
1656 assert_eq!(tail.len(), STREAM_BUFFER_CAP_BYTES);
1657 assert_eq!(tail[0], "é".as_bytes()[1]);
1659 }
1660
1661 #[test]
1662 fn ring_reset_to_supersedes_retained_pieces() {
1663 let mut ring = Ring::default();
1664 ring.append(b"old output".to_vec());
1665 ring.append(b"more".to_vec());
1666 let reachable_end = ring.first_idx + ring.pieces.len();
1667 ring.reset_to(b"\x1b[2J\x1b[Hrepaint".to_vec());
1668 assert!(ring.first_idx >= reachable_end);
1669 assert_eq!(ring.tail(), b"\x1b[2J\x1b[Hrepaint");
1670 assert!(!ring.dropped);
1672
1673 let mut empty = Ring::default();
1674 empty.append(b"x".to_vec());
1675 empty.reset_to(Vec::new());
1676 assert!(empty.tail().is_empty());
1677 }
1678
1679 #[test]
1683 fn pty_forward_env_filters_to_the_whitelist() {
1684 let vars = vec![
1685 ("COLORTERM".to_string(), "truecolor".to_string()),
1686 ("LC_ALL".to_string(), "en_US.UTF-8".to_string()),
1687 ("LANG".to_string(), "en_US.UTF-8".to_string()),
1688 ("TERM_PROGRAM".to_string(), "TestTerm".to_string()),
1689 ("PATH".to_string(), "/bin".to_string()),
1690 ("TERM".to_string(), "xterm".to_string()), ("SECRET_TOKEN".to_string(), "x".to_string()),
1692 ];
1693 let forwarded = pty_forward_env_from(vars.into_iter());
1694 let keys: Vec<&str> = forwarded.iter().map(|(k, _)| k.as_str()).collect();
1695 assert_eq!(keys, ["COLORTERM", "LC_ALL", "LANG", "TERM_PROGRAM"]);
1696 }
1697
1698 #[test]
1699 fn encode_env_rejects_malformed_keys() {
1700 for (key, value) in [
1701 ("", "v"),
1702 ("A=B", "v"),
1703 ("NUL\0KEY", "v"),
1704 ("K", "nul\0value"),
1705 ("1FOO", "v"),
1707 ("FO O", "v"),
1708 ("FOO-BAR", "v"),
1709 ("FOO.BAR", "v"),
1710 ] {
1711 let pairs = vec![(key.to_string(), value.to_string())];
1712 assert!(
1713 encode_env(&pairs).is_err(),
1714 "expected rejection for {key:?}={value:?}"
1715 );
1716 }
1717 let ok = encode_env(&[
1719 ("_FOO".to_string(), "bar=baz".to_string()),
1720 ("LC_ALL".to_string(), "C.UTF-8".to_string()),
1721 ])
1722 .unwrap();
1723 assert_eq!(ok.get("_FOO").map(String::as_str), Some("bar=baz"));
1724 }
1725
1726 #[test]
1727 fn ring_reset_to_clears_prior_dropped() {
1728 let mut ring = Ring::default();
1729 ring.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES + 1]);
1730 assert!(ring.dropped);
1731 ring.reset_to(b"repaint".to_vec());
1732 assert!(!ring.dropped);
1733 assert_eq!(ring.tail(), b"repaint");
1734
1735 let mut over = Ring::default();
1737 over.append(vec![b'b'; STREAM_BUFFER_CAP_BYTES + 1]);
1738 over.reset_to(vec![b'c'; STREAM_BUFFER_CAP_BYTES + 1]);
1739 assert!(over.dropped);
1740 assert_eq!(over.size, STREAM_BUFFER_CAP_BYTES);
1741 }
1742
1743 #[test]
1744 fn terminal_status_mapping() {
1745 assert!(matches!(
1746 terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
1747 Some(SailError::WorkerLost { .. })
1748 ));
1749 assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
1750 for retired in [9, 6, 7] {
1754 assert!(terminal_status_error(retired).is_none());
1755 }
1756 }
1757
1758 fn test_shared() -> Arc<ExecShared> {
1759 Arc::new(ExecShared {
1760 worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
1761 params: ExecParams {
1762 sailbox_id: "sb".into(),
1763 exec_endpoint: "endpoint".into(),
1764 argv: vec!["echo".into()],
1765 timeout_seconds: 0,
1766 idempotency_key: "idem".into(),
1767 open_stdin: false,
1768 pty: false,
1769 term: String::new(),
1770 cols: 0,
1771 rows: 0,
1772 env: std::collections::HashMap::default(),
1773 retry_timeout: 0.0,
1774 extra_metadata: vec![],
1775 },
1776 state: Mutex::new(State::default()),
1777 cond: Condvar::new(),
1778 data_notify: Notify::new(),
1779 exit: Mutex::new(None),
1780 high_seq: Mutex::new((0, 0)),
1781 stdin: AsyncMutex::new(StdinState::default()),
1782 ended: AtomicBool::new(false),
1783 ended_notify: Notify::new(),
1784 closing: AtomicBool::new(false),
1785 close_notify: Notify::new(),
1786 })
1787 }
1788
1789 fn chunk(which: OutputStream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
1790 let stream = match which {
1791 OutputStream::Stdout => pb::SailboxExecStream::Stdout,
1792 OutputStream::Stderr => pb::SailboxExecStream::Stderr,
1793 };
1794 pb::StreamSailboxExecResponse {
1795 frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
1796 pb::SailboxExecChunk {
1797 stream: stream as i32,
1798 data: data.to_vec(),
1799 seq,
1800 },
1801 )),
1802 }
1803 }
1804
1805 fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
1806 pb::StreamSailboxExecResponse {
1807 frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
1808 pb::SailboxExecExit {
1809 status: status as i32,
1810 return_code: 0,
1811 timed_out: false,
1812 stdout_truncated: false,
1813 stderr_truncated: false,
1814 error_message: String::new(),
1815 stdout_seq,
1816 stderr_seq: 0,
1817 ..Default::default()
1818 },
1819 )),
1820 }
1821 }
1822
1823 #[tokio::test]
1829 async fn exec_resume_carries_bytes_and_tracks_seq_across_break() {
1830 let shared = test_shared();
1831 let mut pump = Pump::new(shared.clone());
1832
1833 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"ab\xc3")));
1835 assert_eq!(shared.state.lock().unwrap().stdout.tail(), b"ab\xc3");
1836 assert_eq!(pump.stdout_seq, 1);
1838
1839 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"\xa9cd")));
1843 assert_eq!(
1844 shared.state.lock().unwrap().stdout.tail(),
1845 "abécd".as_bytes()
1846 );
1847 assert_eq!(lossy_tail(&shared.state.lock().unwrap().stdout), "abécd");
1848 assert_eq!(pump.stdout_seq, 2);
1849
1850 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"!")));
1852 assert_eq!(pump.stdout_seq, 2);
1853
1854 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
1855 pump.finalize();
1856 assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
1857 assert!(shared.ended.load(Ordering::SeqCst));
1858 }
1859
1860 fn snapshot_frame(repaint: &[u8], basis: i64) -> pb::StreamSailboxExecResponse {
1861 pb::StreamSailboxExecResponse {
1862 frame: Some(pb::stream_sailbox_exec_response::Frame::Snapshot(
1863 pb::SailboxExecSnapshot {
1864 repaint: repaint.to_vec(),
1865 stdout_seq_basis: basis,
1866 },
1867 )),
1868 }
1869 }
1870
1871 #[tokio::test]
1875 async fn snapshot_resets_ring_seq_basis_and_skips_readers() {
1876 let shared = test_shared();
1877 let mut pump = Pump::new(shared.clone());
1878 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"pre-disconnect ")));
1879 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"tail")));
1880
1881 assert!(!pump.apply_frame(snapshot_frame(b"\x1b[2J\x1b[Hscreen", 7)));
1883 assert_eq!(pump.stdout_seq, 7);
1884
1885 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 8, b" after")));
1887 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 8)));
1888 pump.finalize();
1889
1890 let mut reader = shared_reader(&shared, OutputStream::Stdout);
1891 let mut out = Vec::new();
1892 loop {
1893 match reader.next(Duration::from_millis(50)) {
1894 ReadStep::Chunk(piece) => out.extend_from_slice(&piece),
1895 ReadStep::Eof => break,
1896 ReadStep::Pending => panic!("ended ring should not return Pending"),
1897 }
1898 }
1899 assert_eq!(out, b"\x1b[2J\x1b[Hscreen after");
1900 assert_eq!(*shared.high_seq.lock().unwrap(), (8, 0));
1902 }
1903
1904 #[test]
1905 fn snapshot_with_empty_repaint_is_reset_only() {
1906 let shared = test_shared();
1907 let mut pump = Pump::new(shared.clone());
1908 pump.apply_frame(chunk(OutputStream::Stdout, 3, b"stale"));
1909 pump.apply_frame(snapshot_frame(b"", 3));
1910 assert_eq!(pump.stdout_seq, 3);
1911 let state = lock(&shared.state);
1912 assert!(state.stdout.tail().is_empty());
1913 assert!(!state.stdout.dropped);
1914 }
1915
1916 #[tokio::test]
1920 async fn exec_reader_follows_replayed_then_live() {
1921 let shared = test_shared();
1922 let mut reader = StreamReader {
1923 shared: shared.clone(),
1924 which: OutputStream::Stdout,
1925 cursor: 0,
1926 dropped: false,
1927 reset: false,
1928 seen_front_clips: 0,
1929 seen_resets: 0,
1930 };
1931 let collector = tokio::task::spawn_blocking(move || {
1932 let mut out = Vec::new();
1933 loop {
1934 match reader.next(Duration::from_millis(50)) {
1935 ReadStep::Chunk(piece) => out.extend_from_slice(&piece),
1936 ReadStep::Eof => return out,
1937 ReadStep::Pending => {}
1938 }
1939 }
1940 });
1941
1942 let mut pump = Pump::new(shared.clone());
1943 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"hello "));
1944 tokio::time::sleep(Duration::from_millis(10)).await;
1945 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"world"));
1947 tokio::time::sleep(Duration::from_millis(10)).await;
1948 pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
1949 pump.finalize();
1950
1951 assert_eq!(collector.await.unwrap(), b"hello world");
1952 }
1953
1954 #[test]
1957 fn exec_reader_started_late_replays_retained_tail() {
1958 let shared = test_shared();
1959 let mut pump = Pump::new(shared.clone());
1960 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"early "));
1961 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"output"));
1962 pump.finalize();
1963
1964 let mut reader = shared_reader(&shared, OutputStream::Stdout);
1966 let mut out = Vec::new();
1967 loop {
1968 match reader.next(Duration::from_millis(50)) {
1969 ReadStep::Chunk(piece) => out.extend_from_slice(&piece),
1970 ReadStep::Eof => break,
1971 ReadStep::Pending => panic!("ended ring should not return Pending"),
1972 }
1973 }
1974 assert_eq!(out, b"early output");
1975 }
1976
1977 #[test]
1981 fn reader_reports_drop_and_batch_drains() {
1982 let shared = test_shared();
1983 {
1984 let mut state = lock(&shared.state);
1987 state.stdout.append(b"HEAD".to_vec());
1988 state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
1989 state.ended = true;
1990 }
1991 let mut reader = shared_reader(&shared, OutputStream::Stdout);
1992
1993 let first = reader.next(Duration::from_millis(50));
1995 let ReadStep::Chunk(first) = first else {
1996 panic!("expected the retained chunk, got {first:?}");
1997 };
1998 assert!(
1999 reader.took_drop(),
2000 "the overflow evicted an unconsumed chunk"
2001 );
2002 assert!(!reader.took_drop(), "took_drop clears the latch");
2003
2004 assert_eq!(first, vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
2006 assert!(
2007 reader.try_next().is_none(),
2008 "a drained ring yields None without blocking"
2009 );
2010 }
2011
2012 #[test]
2013 fn reader_does_not_report_a_reset_repaint_as_a_drop() {
2014 let shared = test_shared();
2015 let mut reader = shared_reader(&shared, OutputStream::Stdout);
2016 {
2017 let mut state = lock(&shared.state);
2018 state.stdout.append(b"one".to_vec());
2019 }
2020 assert!(matches!(
2022 reader.next(Duration::from_millis(50)),
2023 ReadStep::Chunk(_)
2024 ));
2025 assert!(!reader.took_drop(), "a clean read latches no drop");
2026 assert!(!reader.took_reset(), "a clean read is not a reset");
2027
2028 {
2029 let mut state = lock(&shared.state);
2034 state.stdout.append(b"two".to_vec());
2035 state.stdout.append(b"three".to_vec());
2036 state.stdout.reset_to(b"REPAINT".to_vec());
2037 state.ended = true;
2038 }
2039 let repaint = reader.next(Duration::from_millis(50));
2040 let ReadStep::Chunk(repaint) = repaint else {
2041 panic!("expected the repaint chunk, got {repaint:?}");
2042 };
2043 assert_eq!(repaint, b"REPAINT");
2044 assert!(
2045 !reader.took_drop(),
2046 "reading a reset repaint is a heal, not a fall-behind drop"
2047 );
2048 assert!(
2049 reader.took_reset(),
2050 "the reset repaint is surfaced as a took_reset event so the pump can \
2051 drop stale terminal-local backlog buffered ahead of it"
2052 );
2053 assert!(!reader.took_reset(), "took_reset clears the latch");
2054 }
2055
2056 #[test]
2057 fn reader_reset_supersedes_a_previously_latched_drop() {
2058 let shared = test_shared();
2059 {
2060 let mut state = lock(&shared.state);
2063 state.stdout.append(b"HEAD".to_vec());
2064 state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
2065 }
2066 let mut reader = shared_reader(&shared, OutputStream::Stdout);
2067
2068 assert!(matches!(
2072 reader.next(Duration::from_millis(50)),
2073 ReadStep::Chunk(_)
2074 ));
2075
2076 {
2077 let mut state = lock(&shared.state);
2081 state.stdout.reset_to(b"REPAINT".to_vec());
2082 state.ended = true;
2083 }
2084 let repaint = reader.next(Duration::from_millis(50));
2085 let ReadStep::Chunk(repaint) = repaint else {
2086 panic!("expected the repaint chunk, got {repaint:?}");
2087 };
2088 assert_eq!(repaint, b"REPAINT");
2089 assert!(reader.took_reset(), "the repaint is surfaced as a reset");
2090 assert!(
2091 !reader.took_drop(),
2092 "the reset superseded the stale drop latch, so no redundant resync \
2093 clobbers the repaint"
2094 );
2095 }
2096
2097 #[test]
2098 fn reader_reports_a_drop_when_output_evicts_the_repaint_before_it_is_read() {
2099 let shared = test_shared();
2100 let mut reader = shared_reader(&shared, OutputStream::Stdout);
2101 {
2102 let mut state = lock(&shared.state);
2103 state.stdout.append(b"one".to_vec());
2104 }
2105 assert!(matches!(
2107 reader.next(Duration::from_millis(50)),
2108 ReadStep::Chunk(_)
2109 ));
2110 assert!(!reader.took_drop());
2111 assert!(!reader.took_reset());
2112
2113 {
2114 let mut state = lock(&shared.state);
2121 state.stdout.reset_to(b"REPAINT".to_vec());
2122 state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES + 1]);
2123 state.ended = true;
2124 }
2125 let chunk = reader.next(Duration::from_millis(50));
2126 let ReadStep::Chunk(chunk) = chunk else {
2127 panic!("expected the retained suffix, got {chunk:?}");
2128 };
2129 assert_ne!(chunk, b"REPAINT", "the repaint was evicted by the burst");
2130 assert!(reader.took_reset(), "a reset did occur");
2131 assert!(
2132 reader.took_drop(),
2133 "the repaint was evicted after the reset, so the reader reports a \
2134 drop and the pump resyncs instead of showing a torn suffix"
2135 );
2136 }
2137
2138 #[test]
2139 fn reader_try_next_batch_drains_buffered_chunks() {
2140 let shared = test_shared();
2141 {
2142 let mut state = lock(&shared.state);
2143 state.stdout.append(b"one".to_vec());
2144 state.stdout.append(b"two".to_vec());
2145 state.stdout.append(b"three".to_vec());
2146 state.ended = true;
2147 }
2148 let mut reader = shared_reader(&shared, OutputStream::Stdout);
2149
2150 let first = reader.next(Duration::from_millis(50));
2153 let ReadStep::Chunk(first) = first else {
2154 panic!("expected the first buffered chunk, got {first:?}");
2155 };
2156 assert_eq!(first, b"one");
2157 assert_eq!(reader.try_next(), Some(b"two".to_vec()));
2158 assert_eq!(reader.try_next(), Some(b"three".to_vec()));
2159 assert!(reader.try_next().is_none(), "the ring is drained");
2160 assert!(!reader.took_drop(), "no eviction, so no drop is latched");
2161 }
2162
2163 #[test]
2164 fn reader_reports_a_front_clip_as_a_drop() {
2165 let shared = test_shared();
2166 {
2167 let mut state = lock(&shared.state);
2169 state.stdout.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES]);
2170 }
2171 let mut reader = shared_reader(&shared, OutputStream::Stdout);
2172 {
2173 let mut state = lock(&shared.state);
2176 state.stdout.append(b"tail".to_vec());
2177 state.ended = true;
2178 assert_eq!(
2179 state.stdout.first_idx, 0,
2180 "the front piece was clipped, not evicted"
2181 );
2182 }
2183 let step = reader.next(Duration::from_millis(50));
2186 let ReadStep::Chunk(_) = step else {
2187 panic!("expected the clipped front piece, got {step:?}");
2188 };
2189 assert!(
2190 reader.took_drop(),
2191 "an in-place front clip of the parked piece is a drop"
2192 );
2193 }
2194
2195 fn shared_reader(shared: &Arc<ExecShared>, which: OutputStream) -> StreamReader {
2196 StreamReader {
2197 shared: shared.clone(),
2198 which,
2199 cursor: 0,
2200 dropped: false,
2201 reset: false,
2202 seen_front_clips: 0,
2203 seen_resets: 0,
2204 }
2205 }
2206
2207 use proptest::prelude::*;
2208
2209 proptest! {
2210 #[test]
2213 fn ring_without_overflow_is_lossless(
2214 pieces in proptest::collection::vec(proptest::collection::vec(any::<u8>(), 0..512), 0..32)
2215 ) {
2216 let concat: Vec<u8> = pieces.concat();
2217 prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
2218 let mut ring = Ring::default();
2219 for piece in &pieces {
2220 ring.append(piece.clone());
2221 }
2222 prop_assert!(!ring.dropped);
2223 prop_assert_eq!(ring.size, concat.len());
2224 prop_assert_eq!(ring.tail(), concat);
2225 }
2226 }
2227
2228 proptest! {
2229 #![proptest_config(ProptestConfig::with_cases(48))]
2231
2232 #[test]
2236 fn ring_eviction_keeps_byte_suffix_at_cap(
2237 overflow in 1usize..8192,
2238 tail_pieces in proptest::collection::vec(proptest::collection::vec(any::<u8>(), 0..64), 0..8),
2239 ) {
2240 let head = vec![b'h'; STREAM_BUFFER_CAP_BYTES + overflow];
2241 let mut full = head.clone();
2242 let mut ring = Ring::default();
2243 ring.append(head);
2244 for piece in &tail_pieces {
2245 ring.append(piece.clone());
2246 full.extend_from_slice(piece);
2247 }
2248 prop_assert!(ring.dropped);
2249 prop_assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
2250 let tail = ring.tail();
2251 prop_assert_eq!(tail.len(), STREAM_BUFFER_CAP_BYTES);
2252 prop_assert!(full.ends_with(&tail));
2253 }
2254 }
2255
2256 proptest! {
2257 #[test]
2261 fn pump_snapshot_assigns_basis_and_replaces_ring(
2262 pre_seqs in proptest::collection::vec(0i64..10_000, 0..16),
2263 basis in 0i64..10_000,
2264 ) {
2265 let shared = test_shared();
2266 let mut pump = Pump::new(shared.clone());
2267 for seq in pre_seqs {
2268 pump.apply_frame(chunk(OutputStream::Stdout, seq, b"x"));
2269 }
2270 pump.apply_frame(snapshot_frame(b"repaint", basis));
2271 prop_assert_eq!(pump.stdout_seq, basis);
2272 prop_assert_eq!(lock(&shared.state).stdout.tail(), b"repaint".to_vec());
2273 }
2274 }
2275
2276 proptest! {
2277 #[test]
2282 fn pump_seq_is_monotonic_running_max_per_stream(
2283 frames in proptest::collection::vec(
2284 (any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
2285 0..64,
2286 )
2287 ) {
2288 let mut pump = Pump::new(test_shared());
2289 let (mut expect_out, mut expect_err) = (0i64, 0i64);
2290 let (mut prev_out, mut prev_err) = (0i64, 0i64);
2291 for (is_stderr, seq, data) in frames {
2292 let which = if is_stderr { OutputStream::Stderr } else { OutputStream::Stdout };
2293 pump.apply_frame(chunk(which, seq, &data));
2294 if is_stderr {
2295 expect_err = expect_err.max(seq);
2296 } else {
2297 expect_out = expect_out.max(seq);
2298 }
2299 prop_assert_eq!(pump.stdout_seq, expect_out);
2300 prop_assert_eq!(pump.stderr_seq, expect_err);
2301 prop_assert!(pump.stdout_seq >= prev_out);
2302 prop_assert!(pump.stderr_seq >= prev_err);
2303 prev_out = pump.stdout_seq;
2304 prev_err = pump.stderr_seq;
2305 }
2306 }
2307 }
2308}