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(String),
71 Eof,
73 Pending,
76}
77
78#[derive(Debug, Clone, Serialize)]
80#[non_exhaustive]
81#[allow(clippy::struct_excessive_bools)]
82pub struct ExecResult {
83 pub stdout: String,
85 pub stderr: String,
87 pub exit_code: i32,
89 pub timed_out: bool,
91 pub stdout_truncated: bool,
94 pub stderr_truncated: bool,
97 pub stdout_complete: bool,
103 pub stderr_complete: bool,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum CancelSignal {
111 Interrupt,
113 Kill,
115}
116
117impl CancelSignal {
118 fn is_force(self) -> bool {
120 matches!(self, CancelSignal::Kill)
121 }
122}
123
124#[derive(Debug, Clone, Copy, PartialEq)]
127pub enum RetryBudget {
128 None,
130 Within(Duration),
132 Forever,
134}
135
136impl RetryBudget {
137 #[doc(hidden)]
140 pub fn as_secs_f64(self) -> f64 {
141 match self {
142 RetryBudget::None => 0.0,
143 RetryBudget::Within(d) => d.as_secs_f64(),
144 RetryBudget::Forever => f64::INFINITY,
145 }
146 }
147
148 #[doc(hidden)]
151 pub fn from_secs_f64(secs: f64) -> RetryBudget {
152 if secs <= 0.0 {
153 RetryBudget::None
154 } else if secs.is_finite() {
155 RetryBudget::Within(Duration::from_secs_f64(secs))
156 } else {
157 RetryBudget::Forever
158 }
159 }
160}
161
162#[derive(Debug, Clone)]
168pub struct ExecOptions {
169 pub timeout: Option<Duration>,
173 pub open_stdin: bool,
175 pub pty: bool,
177 pub term: String,
179 pub cols: u32,
181 pub rows: u32,
183 pub idempotency_key: String,
186 pub retry_timeout: RetryBudget,
191 pub cwd: Option<String>,
194 pub background: bool,
199}
200
201impl Default for ExecOptions {
202 fn default() -> ExecOptions {
203 ExecOptions {
204 timeout: None,
205 open_stdin: false,
206 pty: false,
207 term: String::new(),
208 cols: 0,
209 rows: 0,
210 idempotency_key: String::new(),
211 retry_timeout: RetryBudget::Within(Duration::from_secs_f64(
212 EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
213 )),
214 cwd: None,
215 background: false,
216 }
217 }
218}
219
220pub(crate) fn sh_quote(value: &str) -> String {
222 format!("'{}'", value.replace('\'', "'\\''"))
223}
224
225#[doc(hidden)]
230pub fn shell_argv(command: &str, options: &ExecOptions) -> Result<Vec<String>, SailError> {
231 let invalid = |message: &str| {
232 Err(SailError::InvalidArgument {
233 message: message.to_string(),
234 })
235 };
236 if command.is_empty() {
237 return invalid("command must be non-empty");
238 }
239 if options.background && (options.open_stdin || options.pty) {
240 return invalid("background is not supported with open_stdin or pty");
241 }
242 let mut command = command.to_string();
243 if let Some(cwd) = &options.cwd {
244 let cwd = cwd.trim();
245 if cwd.is_empty() {
246 return invalid("cwd must be non-empty");
247 }
248 command = format!(
249 "cd {} && exec /bin/sh -lc {}",
250 sh_quote(cwd),
251 sh_quote(&command)
252 );
253 }
254 if options.background {
255 command = format!(
256 "nohup /bin/sh -lc {} </dev/null >/dev/null 2>&1 &",
257 sh_quote(&command)
258 );
259 }
260 Ok(vec!["/bin/sh".to_string(), "-lc".to_string(), command])
261}
262
263#[doc(hidden)]
265#[derive(Debug, Clone)]
266pub struct ExecParams {
267 pub sailbox_id: String,
269 pub exec_endpoint: String,
271 pub argv: Vec<String>,
273 pub timeout_seconds: u32,
276 pub idempotency_key: String,
279 pub open_stdin: bool,
281 pub pty: bool,
283 pub term: String,
285 pub cols: u32,
287 pub rows: u32,
289 pub retry_timeout: f64,
292 pub extra_metadata: Vec<(String, String)>,
294}
295
296#[derive(Default)]
300struct Utf8Decoder {
301 pending: Vec<u8>,
302}
303
304impl Utf8Decoder {
305 fn decode(&mut self, data: &[u8]) -> String {
306 self.pending.extend_from_slice(data);
307 let mut out = String::new();
308 loop {
309 match std::str::from_utf8(&self.pending) {
310 Ok(valid) => {
311 out.push_str(valid);
312 self.pending.clear();
313 break;
314 }
315 Err(err) => {
316 let valid_up_to = err.valid_up_to();
317 out.push_str(std::str::from_utf8(&self.pending[..valid_up_to]).unwrap());
319 if let Some(bad) = err.error_len() {
320 out.push('\u{FFFD}');
321 self.pending.drain(..valid_up_to + bad);
322 } else {
323 self.pending.drain(..valid_up_to);
325 break;
326 }
327 }
328 }
329 }
330 out
331 }
332
333 fn flush(&mut self) -> String {
334 if self.pending.is_empty() {
335 return String::new();
336 }
337 let out = String::from_utf8_lossy(&self.pending).into_owned();
338 self.pending.clear();
339 out
340 }
341}
342
343#[derive(Default)]
348struct Ring {
349 pieces: Vec<String>,
350 piece_bytes: Vec<usize>,
351 first_idx: usize,
352 size: usize,
353 dropped: bool,
354}
355
356impl Ring {
357 fn append(&mut self, text: String) {
358 let bytes = text.len();
359 self.pieces.push(text);
360 self.piece_bytes.push(bytes);
361 self.size += bytes;
362 while self.size > STREAM_BUFFER_CAP_BYTES {
363 let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
364 if self.piece_bytes[0] <= overflow {
365 self.pieces.remove(0);
366 self.size -= self.piece_bytes.remove(0);
367 self.first_idx += 1;
368 } else {
369 let piece = &self.pieces[0];
375 let mut cut = overflow;
376 while cut < piece.len() && !piece.is_char_boundary(cut) {
377 cut += 1;
378 }
379 let kept = piece[cut..].to_string();
380 self.size -= self.piece_bytes[0];
381 self.piece_bytes[0] = kept.len();
382 self.size += self.piece_bytes[0];
383 self.pieces[0] = kept;
384 }
385 self.dropped = true;
386 }
387 }
388
389 fn tail(&self) -> String {
390 self.pieces.concat()
391 }
392}
393
394#[derive(Default)]
395struct State {
396 stdout: Ring,
397 stderr: Ring,
398 ended: bool,
399}
400
401impl State {
402 fn ring(&self, which: OutputStream) -> &Ring {
403 match which {
404 OutputStream::Stdout => &self.stdout,
405 OutputStream::Stderr => &self.stderr,
406 }
407 }
408}
409
410#[derive(Clone)]
412struct ExitInfo {
413 status: i32,
414 exit_code: i32,
415 timed_out: bool,
416 stdout_truncated: bool,
417 stderr_truncated: bool,
418 error_message: String,
419 stdout_seq: i64,
420 stderr_seq: i64,
421}
422
423#[derive(Default)]
424struct StdinState {
425 offset: i64,
426 eof_sent: bool,
427 broken: bool,
428 write_in_flight: bool,
436}
437
438struct ExecShared {
439 worker: Arc<WorkerProxy>,
440 params: ExecParams,
441 state: Mutex<State>,
442 cond: Condvar,
445 data_notify: Notify,
448 exit: Mutex<Option<ExitInfo>>,
449 high_seq: Mutex<(i64, i64)>,
451 stdin: AsyncMutex<StdinState>,
452 ended: AtomicBool,
453 ended_notify: Notify,
454 closing: AtomicBool,
455 close_notify: Notify,
456}
457
458pub struct ExecProcess {
461 shared: Arc<ExecShared>,
462 exec_request_id: String,
463}
464
465impl Drop for ExecProcess {
466 fn drop(&mut self) {
467 self.close();
472 }
473}
474
475impl ExecProcess {
476 #[doc(hidden)]
484 pub async fn start(
485 worker: Arc<WorkerProxy>,
486 mut params: ExecParams,
487 ) -> Result<ExecProcess, SailError> {
488 let key = params.idempotency_key.trim();
492 params.idempotency_key = if key.is_empty() {
493 format!("exec_{}", uuid::Uuid::new_v4())
494 } else {
495 key.to_string()
496 };
497 let (exec_request_id, stream) = submit(
498 &worker, ¶ms, 0, 0,
499 )
500 .await?;
501 let shared = Arc::new(ExecShared {
502 worker,
503 params,
504 state: Mutex::new(State::default()),
505 cond: Condvar::new(),
506 data_notify: Notify::new(),
507 exit: Mutex::new(None),
508 high_seq: Mutex::new((0, 0)),
509 stdin: AsyncMutex::new(StdinState::default()),
510 ended: AtomicBool::new(false),
511 ended_notify: Notify::new(),
512 closing: AtomicBool::new(false),
513 close_notify: Notify::new(),
514 });
515 let pump_shared = shared.clone();
516 tokio::spawn(async move { pump(pump_shared, stream).await });
517 Ok(ExecProcess {
518 shared,
519 exec_request_id,
520 })
521 }
522
523 pub fn exec_request_id(&self) -> &str {
526 &self.exec_request_id
527 }
528
529 pub fn idempotency_key(&self) -> &str {
532 &self.shared.params.idempotency_key
533 }
534
535 pub fn reader(&self, which: OutputStream) -> StreamReader {
538 StreamReader {
539 shared: self.shared.clone(),
540 which,
541 cursor: 0,
542 }
543 }
544
545 pub fn reader_async(&self, which: OutputStream) -> AsyncStreamReader {
548 AsyncStreamReader {
549 shared: self.shared.clone(),
550 which,
551 cursor: 0,
552 }
553 }
554
555 pub fn poll(&self) -> Option<Result<i32, SailError>> {
559 let exit = lock(&self.shared.exit);
560 exit.as_ref()
561 .map(|exit| match terminal_status_error(exit.status) {
562 Some(err) => Err(err),
563 None => Ok(exit.exit_code),
564 })
565 }
566
567 pub fn close(&self) {
569 self.shared.closing.store(true, Ordering::SeqCst);
570 self.shared.close_notify.notify_one();
575 }
576
577 #[doc(hidden)]
582 pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
583 let _ = tokio::time::timeout(timeout, self.await_ended()).await;
584 self.shared.ended.load(Ordering::SeqCst)
585 }
586
587 async fn await_ended(&self) {
589 loop {
590 let notified = self.shared.ended_notify.notified();
591 if self.shared.ended.load(Ordering::SeqCst) {
592 return;
593 }
594 notified.await;
595 }
596 }
597
598 pub async fn wait(&self) -> Result<ExecResult, SailError> {
605 self.await_ended().await;
606 let exit = lock(&self.shared.exit).clone();
607 let (high_out, high_err) = *lock(&self.shared.high_seq);
608 let (out_dropped, err_dropped) = {
609 let state = lock(&self.shared.state);
610 (state.stdout.dropped, state.stderr.dropped)
611 };
612
613 let truncated = exit.as_ref().is_some_and(|exit| {
617 exit.stdout_truncated || exit.stderr_truncated || out_dropped || err_dropped
618 });
619 let incomplete = exit
620 .as_ref()
621 .is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
622
623 let stdout_complete = exit
628 .as_ref()
629 .is_some_and(|exit| exit.stdout_seq <= high_out);
630 let stderr_complete = exit
631 .as_ref()
632 .is_some_and(|exit| exit.stderr_seq <= high_err);
633
634 if exit.is_none() || truncated || incomplete {
635 let outcome = self
636 .shared
637 .worker
638 .wait_exec(
639 &self.shared.params.exec_endpoint,
640 &self.shared.params.sailbox_id,
641 &self.exec_request_id,
642 self.shared.params.retry_timeout,
643 )
644 .await?;
645 {
648 let mut exit_slot = lock(&self.shared.exit);
649 if exit_slot.is_none() {
650 *exit_slot = Some(ExitInfo {
651 status: outcome.status,
652 exit_code: outcome.exit_code,
653 timed_out: outcome.timed_out,
654 stdout_truncated: outcome.stdout_truncated,
655 stderr_truncated: outcome.stderr_truncated,
656 error_message: String::new(),
657 stdout_seq: 0,
658 stderr_seq: 0,
659 });
660 }
661 }
662 if let Some(witnessed) = exit.as_ref() {
668 if outcome.status == pb::SailboxExecStatus::WorkerLost as i32 {
669 let state = lock(&self.shared.state);
670 let mut stderr = state.stderr.tail();
671 if witnessed.status == pb::SailboxExecStatus::Failed as i32
672 && !witnessed.error_message.is_empty()
673 {
674 stderr = witnessed.error_message.clone();
675 }
676 return Ok(ExecResult {
677 stdout: state.stdout.tail(),
678 stderr,
679 exit_code: witnessed.exit_code,
680 timed_out: witnessed.timed_out,
681 stdout_truncated: witnessed.stdout_truncated
682 || out_dropped
683 || witnessed.stdout_seq > high_out,
684 stderr_truncated: witnessed.stderr_truncated
685 || err_dropped
686 || witnessed.stderr_seq > high_err,
687 stdout_complete,
688 stderr_complete,
689 });
690 }
691 }
692 if let Some(err) = terminal_status_error(outcome.status) {
693 return Err(err);
694 }
695 return Ok(ExecResult {
696 stdout: outcome.stdout,
697 stderr: outcome.stderr,
698 exit_code: outcome.exit_code,
699 timed_out: outcome.timed_out,
700 stdout_truncated: outcome.stdout_truncated,
701 stderr_truncated: outcome.stderr_truncated,
702 stdout_complete,
703 stderr_complete,
704 });
705 }
706
707 let exit = exit.expect("exit present on the clean path");
708 if let Some(err) = terminal_status_error(exit.status) {
709 return Err(err);
710 }
711 let state = lock(&self.shared.state);
712 let mut stderr = state.stderr.tail();
713 if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
714 stderr = exit.error_message.clone();
716 }
717 Ok(ExecResult {
718 stdout: state.stdout.tail(),
719 stderr,
720 exit_code: exit.exit_code,
721 timed_out: exit.timed_out,
722 stdout_truncated: false,
723 stderr_truncated: false,
724 stdout_complete,
725 stderr_complete,
726 })
727 }
728
729 pub async fn cancel(&self, signal: CancelSignal, retry: RetryBudget) -> Result<(), SailError> {
732 self.shared
733 .worker
734 .cancel_exec(
735 &self.shared.params.exec_endpoint,
736 &self.shared.params.sailbox_id,
737 &self.exec_request_id,
738 signal.is_force(),
739 retry.as_secs_f64(),
740 )
741 .await
742 }
743
744 pub async fn resize(&self, cols: u32, rows: u32) {
748 let message = pb::ResizeSailboxExecRequest {
749 sailbox_id: self.shared.params.sailbox_id.clone(),
750 exec_request_id: self.exec_request_id.clone(),
751 cols,
752 rows,
753 };
754 let Ok(request) =
755 self.shared
756 .worker
757 .request_for(message, &[], Some(Duration::from_secs(5)))
758 else {
759 return;
760 };
761 if let Ok(mut client) = self
762 .shared
763 .worker
764 .client_for(&self.shared.params.exec_endpoint)
765 {
766 let _ = client.resize_sailbox_exec(request).await;
767 }
768 }
769
770 pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
774 let mut stdin = self.shared.stdin.lock().await;
775 if stdin.eof_sent {
776 return Err(SailError::BrokenPipe {
777 message: "stdin is closed".to_string(),
778 });
779 }
780 if stdin.broken {
781 return Err(SailError::BrokenPipe {
782 message: "an earlier stdin write failed".to_string(),
783 });
784 }
785 if stdin.write_in_flight {
786 stdin.broken = true;
790 return Err(SailError::BrokenPipe {
791 message: "an earlier stdin write was interrupted".to_string(),
792 });
793 }
794 if data.is_empty() {
795 return Ok(());
796 }
797 stdin.write_in_flight = true;
798 let result = self.send_stdin(&mut stdin, data, false).await;
799 stdin.write_in_flight = false;
802 result
803 }
804
805 pub async fn close_stdin(&self) -> Result<(), SailError> {
807 let mut stdin = self.shared.stdin.lock().await;
808 if stdin.eof_sent {
809 return Ok(());
810 }
811 if stdin.broken || stdin.write_in_flight {
812 stdin.eof_sent = true;
815 return Ok(());
816 }
817 match self.send_stdin(&mut stdin, &[], true).await {
821 Ok(()) => {
822 stdin.eof_sent = true;
823 Ok(())
824 }
825 Err(SailError::BrokenPipe { .. }) => {
827 stdin.eof_sent = true;
828 Ok(())
829 }
830 Err(err) => Err(err),
831 }
832 }
833
834 async fn send_stdin(
838 &self,
839 stdin: &mut StdinState,
840 payload: &[u8],
841 eof: bool,
842 ) -> Result<(), SailError> {
843 let endpoint = &self.shared.params.exec_endpoint;
844 let sailbox_id = &self.shared.params.sailbox_id;
845 let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
846 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
847 let mut sent = 0usize;
848 loop {
849 let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
850 let chunk = &payload[sent..end];
851 let last = end >= payload.len();
852 let message = pb::WriteSailboxExecStdinRequest {
853 sailbox_id: sailbox_id.clone(),
854 exec_request_id: self.exec_request_id.clone(),
855 offset: stdin.offset,
856 data: chunk.to_vec(),
857 eof: eof && last,
858 };
859 let request = match self.shared.worker.request_for(
863 message,
864 &[],
865 Some(rpc_attempt_timeout(deadline)),
866 ) {
867 Ok(request) => request,
868 Err(err) => return Err(err),
869 };
870 let result = match self.shared.worker.client_for(endpoint) {
871 Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
872 Err(err) => return Err(err),
873 };
874 match result {
875 Ok(resp) => {
876 let accepted_through = resp.into_inner().accepted_through;
877 let accepted =
882 ((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
883 stdin.offset += accepted as i64;
884 sent += accepted;
885 if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
886 return Ok(());
887 }
888 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
890 if accepted > 0 {
891 delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
892 } else {
893 delay = sleep_no_deadline(delay).await;
895 }
896 }
897 Err(status) => {
898 if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
899 return Err(SailError::BrokenPipe {
901 message: status.message().to_string(),
902 });
903 }
904 if is_exec_not_ready(&status) {
905 delay = sleep_no_deadline(delay).await;
909 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
910 continue;
911 }
912 if !should_retry_transient_exec_rpc(&status, deadline) {
913 stdin.broken = true;
916 return Err(SailError::from_exec_status(&status));
917 }
918 tracing::warn!(code = ?status.code(), "retrying exec stdin write");
919 if should_invalidate_channel(&status) {
920 self.shared.worker.channels().invalidate(endpoint);
921 }
922 delay = sleep_before_retry(delay, deadline).await;
923 }
924 }
925 }
926 }
927}
928
929pub struct StreamReader {
932 shared: Arc<ExecShared>,
933 which: OutputStream,
934 cursor: usize,
935}
936
937impl StreamReader {
938 pub fn next(&mut self, timeout: Duration) -> ReadStep {
942 let mut state = lock(&self.shared.state);
943 loop {
944 let ring = state.ring(self.which);
945 if self.cursor < ring.first_idx {
946 self.cursor = ring.first_idx;
947 }
948 let available = ring.first_idx + ring.pieces.len();
949 if self.cursor < available {
950 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
951 self.cursor += 1;
952 return ReadStep::Chunk(piece);
953 }
954 if state.ended {
955 return ReadStep::Eof;
956 }
957 let (next_state, timed_out) = self
958 .shared
959 .cond
960 .wait_timeout(state, timeout)
961 .expect("exec state mutex poisoned");
962 state = next_state;
963 if timed_out.timed_out() {
964 return ReadStep::Pending;
965 }
966 }
967 }
968}
969
970pub struct AsyncStreamReader {
974 shared: Arc<ExecShared>,
975 which: OutputStream,
976 cursor: usize,
977}
978
979impl AsyncStreamReader {
980 pub async fn next(&mut self) -> Option<String> {
984 loop {
985 let notified = self.shared.data_notify.notified();
989 tokio::pin!(notified);
990 notified.as_mut().enable();
991 {
992 let state = lock(&self.shared.state);
993 let ring = state.ring(self.which);
994 if self.cursor < ring.first_idx {
995 self.cursor = ring.first_idx;
996 }
997 let available = ring.first_idx + ring.pieces.len();
998 if self.cursor < available {
999 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
1000 self.cursor += 1;
1001 return Some(piece);
1002 }
1003 if state.ended {
1004 return None;
1005 }
1006 }
1007 notified.await;
1008 }
1009 }
1010}
1011
1012async fn sleep_no_deadline(delay: f64) -> f64 {
1015 let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
1016 tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
1017 (delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
1018}
1019
1020fn is_exec_not_ready(status: &Status) -> bool {
1021 status.code() == Code::Unavailable && status.message().contains("; retry")
1022}
1023
1024fn terminal_status_error(status: i32) -> Option<SailError> {
1027 if status == pb::SailboxExecStatus::WorkerLost as i32 {
1028 return Some(SailError::WorkerLost {
1029 message: "the machine hosting your sailbox was lost before the command \
1030 finished; run exec again to retry"
1031 .to_string(),
1032 });
1033 }
1034 None
1035}
1036
1037async fn submit(
1040 worker: &Arc<WorkerProxy>,
1041 params: &ExecParams,
1042 stdout_resume_seq: i64,
1043 stderr_resume_seq: i64,
1044) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
1045 let deadline = retry_deadline(params.retry_timeout);
1046 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
1047 loop {
1048 let message = pb::StreamSailboxExecRequest {
1049 sailbox_id: params.sailbox_id.clone(),
1050 argv: params.argv.clone(),
1051 timeout_seconds: params.timeout_seconds,
1052 idempotency_key: params.idempotency_key.clone(),
1053 open_stdin: params.open_stdin,
1054 stdout_resume_seq,
1055 stderr_resume_seq,
1056 pty: params.pty,
1057 term_cols: params.cols,
1058 term_rows: params.rows,
1059 term: params.term.clone(),
1060 };
1061 let request =
1062 worker.request_for(message, ¶ms.extra_metadata, None)?;
1063 let status = match worker
1064 .client_for(¶ms.exec_endpoint)?
1065 .stream_sailbox_exec(request)
1066 .await
1067 {
1068 Ok(resp) => {
1069 let mut stream = resp.into_inner();
1070 match stream.message().await {
1071 Ok(Some(first)) => match first.frame {
1072 Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
1073 return Ok((started.exec_request_id, stream));
1074 }
1075 _ => {
1076 return Err(SailError::Execution {
1077 code: crate::error::GrpcCode::Internal,
1078 detail: "exec stream opened with a non-started frame".to_string(),
1079 });
1080 }
1081 },
1082 Ok(None) if stdout_resume_seq == 0 && stderr_resume_seq == 0 => {
1093 tonic::Status::unavailable(
1094 "exec stream ended before the server confirmed the launch",
1095 )
1096 }
1097 Ok(None) => {
1098 return Err(SailError::Execution {
1099 code: crate::error::GrpcCode::Internal,
1100 detail: "exec stream ended before the server confirmed the launch"
1101 .to_string(),
1102 });
1103 }
1104 Err(status) => status,
1105 }
1106 }
1107 Err(status) => status,
1108 };
1109 if !should_retry_transient_exec_rpc(&status, deadline) {
1110 return Err(SailError::from_exec_status(&status));
1111 }
1112 tracing::warn!(
1113 code = ?status.code(),
1114 stdout_resume_seq,
1115 stderr_resume_seq,
1116 "reconnecting exec stream"
1117 );
1118 if should_invalidate_channel(&status) {
1119 worker.channels().invalidate(¶ms.exec_endpoint);
1120 }
1121 delay = sleep_before_retry(delay, deadline).await;
1122 }
1123}
1124
1125struct Pump {
1131 shared: Arc<ExecShared>,
1132 stdout_dec: Utf8Decoder,
1133 stderr_dec: Utf8Decoder,
1134 stdout_seq: i64,
1135 stderr_seq: i64,
1136}
1137
1138impl Pump {
1139 fn new(shared: Arc<ExecShared>) -> Pump {
1140 Pump {
1141 shared,
1142 stdout_dec: Utf8Decoder::default(),
1143 stderr_dec: Utf8Decoder::default(),
1144 stdout_seq: 0,
1145 stderr_seq: 0,
1146 }
1147 }
1148
1149 fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
1154 match frame.frame {
1155 Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
1156 let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
1157 let (decoder, which) = if is_stderr {
1158 self.stderr_seq = self.stderr_seq.max(chunk.seq);
1159 (&mut self.stderr_dec, OutputStream::Stderr)
1160 } else {
1161 self.stdout_seq = self.stdout_seq.max(chunk.seq);
1162 (&mut self.stdout_dec, OutputStream::Stdout)
1163 };
1164 let text = decoder.decode(&chunk.data);
1165 if !text.is_empty() {
1166 let mut state = lock(&self.shared.state);
1167 match which {
1168 OutputStream::Stdout => state.stdout.append(text),
1169 OutputStream::Stderr => state.stderr.append(text),
1170 }
1171 self.shared.cond.notify_all();
1172 self.shared.data_notify.notify_waiters();
1173 }
1174 false
1175 }
1176 Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
1177 *lock(&self.shared.exit) = Some(ExitInfo {
1178 status: exit.status,
1179 exit_code: exit.return_code,
1180 timed_out: exit.timed_out,
1181 stdout_truncated: exit.stdout_truncated,
1182 stderr_truncated: exit.stderr_truncated,
1183 error_message: exit.error_message,
1184 stdout_seq: exit.stdout_seq,
1185 stderr_seq: exit.stderr_seq,
1186 });
1187 true
1188 }
1189 _ => false,
1190 }
1191 }
1192
1193 fn finalize(mut self) {
1196 let tail_out = self.stdout_dec.flush();
1197 let tail_err = self.stderr_dec.flush();
1198 {
1199 let mut state = lock(&self.shared.state);
1200 if !tail_out.is_empty() {
1201 state.stdout.append(tail_out);
1202 }
1203 if !tail_err.is_empty() {
1204 state.stderr.append(tail_err);
1205 }
1206 state.ended = true;
1207 self.shared.cond.notify_all();
1208 self.shared.data_notify.notify_waiters();
1209 }
1210 *lock(&self.shared.high_seq) = (self.stdout_seq, self.stderr_seq);
1211 self.shared.ended.store(true, Ordering::SeqCst);
1212 self.shared.ended_notify.notify_waiters();
1213 }
1214}
1215
1216async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
1219 let mut state = Pump::new(shared.clone());
1220 loop {
1221 if shared.closing.load(Ordering::SeqCst) {
1222 break;
1223 }
1224 let message = tokio::select! {
1225 biased;
1226 () = shared.close_notify.notified() => break,
1227 message = stream.message() => message,
1228 };
1229 match message {
1230 Ok(Some(frame)) => {
1231 if state.apply_frame(frame) {
1232 break;
1233 }
1234 }
1235 Ok(None) => break,
1237 Err(_status) => {
1238 if shared.closing.load(Ordering::SeqCst) {
1239 break;
1240 }
1241 match submit(
1244 &shared.worker,
1245 &shared.params,
1246 state.stdout_seq,
1247 state.stderr_seq,
1248 )
1249 .await
1250 {
1251 Ok((_id, fresh)) => {
1252 stream = fresh;
1253 continue;
1254 }
1255 Err(_) => break,
1256 }
1257 }
1258 }
1259 }
1260 state.finalize();
1261}
1262
1263#[cfg(any(test, feature = "fuzzing"))]
1269pub fn fuzz_decode_ring(data: &[u8]) {
1270 let mut whole = Utf8Decoder::default();
1272 let mut whole_out = whole.decode(data);
1273 whole_out.push_str(&whole.flush());
1274
1275 let mut chunked = Utf8Decoder::default();
1278 let mut chunked_out = String::new();
1279 let mut start = 0;
1280 for (i, byte) in data.iter().enumerate() {
1281 if byte & 1 == 1 {
1282 chunked_out.push_str(&chunked.decode(&data[start..=i]));
1283 start = i + 1;
1284 }
1285 }
1286 chunked_out.push_str(&chunked.decode(&data[start..]));
1287 chunked_out.push_str(&chunked.flush());
1288
1289 assert_eq!(
1291 chunked_out, whole_out,
1292 "decoder is not chunk-boundary invariant"
1293 );
1294 assert_eq!(
1295 whole_out,
1296 String::from_utf8_lossy(data),
1297 "decoder disagrees with a lossy decode of the whole input"
1298 );
1299
1300 let mut ring = Ring::default();
1305 ring.append(whole_out.clone());
1306 let filler = STREAM_BUFFER_CAP_BYTES + 1 - whole_out.len().min(STREAM_BUFFER_CAP_BYTES);
1307 let filler = "a".repeat(filler);
1308 ring.append(filler.clone());
1309 assert!(
1310 ring.size <= STREAM_BUFFER_CAP_BYTES,
1311 "ring exceeded its cap"
1312 );
1313 let full = format!("{whole_out}{filler}");
1314 let tail = ring.tail();
1315 assert!(
1316 full.as_bytes().ends_with(tail.as_bytes()),
1317 "ring tail is not a suffix of the appended bytes"
1318 );
1319}
1320
1321#[cfg(test)]
1322mod tests {
1323 use super::*;
1324
1325 #[test]
1326 fn shell_argv_wraps_cwd_and_background() {
1327 let plain = shell_argv("echo hi", &ExecOptions::default()).unwrap();
1328 assert_eq!(plain, ["/bin/sh", "-lc", "echo hi"]);
1329
1330 let cwd = shell_argv(
1331 "echo hi",
1332 &ExecOptions {
1333 cwd: Some("/app".to_string()),
1334 ..ExecOptions::default()
1335 },
1336 )
1337 .unwrap();
1338 assert_eq!(cwd[2], "cd '/app' && exec /bin/sh -lc 'echo hi'");
1339
1340 let background = shell_argv(
1341 "echo hi",
1342 &ExecOptions {
1343 cwd: Some("/app".to_string()),
1344 background: true,
1345 ..ExecOptions::default()
1346 },
1347 )
1348 .unwrap();
1349 assert_eq!(
1350 background[2],
1351 "nohup /bin/sh -lc 'cd '\\''/app'\\'' && exec /bin/sh -lc '\\''echo hi'\\''' </dev/null >/dev/null 2>&1 &"
1352 );
1353 }
1354
1355 #[test]
1356 fn shell_argv_rejects_invalid_combinations() {
1357 assert!(shell_argv("", &ExecOptions::default()).is_err());
1358 assert!(shell_argv(
1359 "x",
1360 &ExecOptions {
1361 cwd: Some(" ".to_string()),
1362 ..ExecOptions::default()
1363 }
1364 )
1365 .is_err());
1366 for (open_stdin, pty) in [(true, false), (false, true)] {
1367 assert!(shell_argv(
1368 "x",
1369 &ExecOptions {
1370 background: true,
1371 open_stdin,
1372 pty,
1373 ..ExecOptions::default()
1374 }
1375 )
1376 .is_err());
1377 }
1378 }
1379
1380 #[test]
1381 fn fuzz_decode_ring_holds_on_samples() {
1382 for sample in [
1385 &b""[..],
1386 b"hello",
1387 b"\xff\xfe\xfd",
1388 "€ µ é".as_bytes(),
1389 b"ab\xc3\xa9cd",
1390 &[0xC3, 0x28],
1391 ] {
1392 fuzz_decode_ring(sample);
1393 }
1394 }
1395
1396 #[test]
1397 fn ring_drops_oldest_past_cap_and_latches() {
1398 let mut ring = Ring::default();
1399 ring.append("a".repeat(STREAM_BUFFER_CAP_BYTES));
1400 assert!(!ring.dropped);
1401 ring.append("bbbb".to_string());
1402 assert!(ring.dropped);
1403 assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
1404 assert!(ring.tail().ends_with("bbbb"));
1405 }
1406
1407 #[test]
1408 fn ring_clip_on_split_multibyte_char_terminates() {
1409 let mut ring = Ring::default();
1414 ring.append(format!("é{}", "a".repeat(STREAM_BUFFER_CAP_BYTES - 1)));
1415 assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1416 assert!(ring.dropped);
1417 let tail = ring.tail();
1418 assert!(!tail.contains('é'));
1419 assert!(!tail.contains('\u{FFFD}'));
1420 }
1421
1422 #[test]
1423 fn decoder_carries_split_multibyte_char() {
1424 let mut dec = Utf8Decoder::default();
1425 let euro = "€".as_bytes(); assert_eq!(dec.decode(&euro[..2]), "");
1427 assert_eq!(dec.decode(&euro[2..]), "€");
1428 }
1429
1430 #[test]
1431 fn decoder_replaces_invalid_bytes() {
1432 let mut dec = Utf8Decoder::default();
1433 assert_eq!(dec.decode(&[0xff, b'x']), "\u{FFFD}x");
1434 }
1435
1436 #[test]
1437 fn terminal_status_mapping() {
1438 assert!(matches!(
1439 terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
1440 Some(SailError::WorkerLost { .. })
1441 ));
1442 assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
1443 for retired in [9, 6, 7] {
1447 assert!(terminal_status_error(retired).is_none());
1448 }
1449 }
1450
1451 fn test_shared() -> Arc<ExecShared> {
1452 Arc::new(ExecShared {
1453 worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
1454 params: ExecParams {
1455 sailbox_id: "sb".into(),
1456 exec_endpoint: "endpoint".into(),
1457 argv: vec!["echo".into()],
1458 timeout_seconds: 0,
1459 idempotency_key: "idem".into(),
1460 open_stdin: false,
1461 pty: false,
1462 term: String::new(),
1463 cols: 0,
1464 rows: 0,
1465 retry_timeout: 0.0,
1466 extra_metadata: vec![],
1467 },
1468 state: Mutex::new(State::default()),
1469 cond: Condvar::new(),
1470 data_notify: Notify::new(),
1471 exit: Mutex::new(None),
1472 high_seq: Mutex::new((0, 0)),
1473 stdin: AsyncMutex::new(StdinState::default()),
1474 ended: AtomicBool::new(false),
1475 ended_notify: Notify::new(),
1476 closing: AtomicBool::new(false),
1477 close_notify: Notify::new(),
1478 })
1479 }
1480
1481 fn chunk(which: OutputStream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
1482 let stream = match which {
1483 OutputStream::Stdout => pb::SailboxExecStream::Stdout,
1484 OutputStream::Stderr => pb::SailboxExecStream::Stderr,
1485 };
1486 pb::StreamSailboxExecResponse {
1487 frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
1488 pb::SailboxExecChunk {
1489 stream: stream as i32,
1490 data: data.to_vec(),
1491 seq,
1492 },
1493 )),
1494 }
1495 }
1496
1497 fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
1498 pb::StreamSailboxExecResponse {
1499 frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
1500 pb::SailboxExecExit {
1501 status: status as i32,
1502 return_code: 0,
1503 timed_out: false,
1504 stdout_truncated: false,
1505 stderr_truncated: false,
1506 error_message: String::new(),
1507 stdout_seq,
1508 stderr_seq: 0,
1509 },
1510 )),
1511 }
1512 }
1513
1514 #[tokio::test]
1519 async fn exec_resume_carries_utf8_and_tracks_seq_across_break() {
1520 let shared = test_shared();
1521 let mut pump = Pump::new(shared.clone());
1522
1523 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"ab\xc3")));
1525 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "ab");
1526 assert_eq!(pump.stdout_seq, 1);
1528
1529 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"\xa9cd")));
1532 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "abécd");
1533 assert_eq!(pump.stdout_seq, 2);
1534
1535 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"!")));
1537 assert_eq!(pump.stdout_seq, 2);
1538
1539 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
1540 pump.finalize();
1541 assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
1542 assert!(shared.ended.load(Ordering::SeqCst));
1543 }
1544
1545 #[tokio::test]
1549 async fn exec_reader_follows_replayed_then_live() {
1550 let shared = test_shared();
1551 let mut reader = StreamReader {
1552 shared: shared.clone(),
1553 which: OutputStream::Stdout,
1554 cursor: 0,
1555 };
1556 let collector = tokio::task::spawn_blocking(move || {
1557 let mut out = String::new();
1558 loop {
1559 match reader.next(Duration::from_millis(50)) {
1560 ReadStep::Chunk(piece) => out.push_str(&piece),
1561 ReadStep::Eof => return out,
1562 ReadStep::Pending => {}
1563 }
1564 }
1565 });
1566
1567 let mut pump = Pump::new(shared.clone());
1568 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"hello "));
1569 tokio::time::sleep(Duration::from_millis(10)).await;
1570 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"world"));
1572 tokio::time::sleep(Duration::from_millis(10)).await;
1573 pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
1574 pump.finalize();
1575
1576 assert_eq!(collector.await.unwrap(), "hello world");
1577 }
1578
1579 #[test]
1582 fn exec_reader_started_late_replays_retained_tail() {
1583 let shared = test_shared();
1584 let mut pump = Pump::new(shared.clone());
1585 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"early "));
1586 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"output"));
1587 pump.finalize();
1588
1589 let mut reader = shared_reader(&shared, OutputStream::Stdout);
1591 let mut out = String::new();
1592 loop {
1593 match reader.next(Duration::from_millis(50)) {
1594 ReadStep::Chunk(piece) => out.push_str(&piece),
1595 ReadStep::Eof => break,
1596 ReadStep::Pending => panic!("ended ring should not return Pending"),
1597 }
1598 }
1599 assert_eq!(out, "early output");
1600 }
1601
1602 fn shared_reader(shared: &Arc<ExecShared>, which: OutputStream) -> StreamReader {
1603 StreamReader {
1604 shared: shared.clone(),
1605 which,
1606 cursor: 0,
1607 }
1608 }
1609
1610 use proptest::prelude::*;
1611
1612 fn utf8ish_bytes() -> impl Strategy<Value = Vec<u8>> {
1616 prop_oneof![
1617 proptest::collection::vec(any::<u8>(), 0..256),
1618 any::<String>().prop_map(String::into_bytes),
1619 ]
1620 }
1621
1622 proptest! {
1623 #[test]
1626 fn ring_without_overflow_is_lossless(
1627 pieces in proptest::collection::vec(any::<String>(), 0..32)
1628 ) {
1629 let concat: String = pieces.concat();
1630 prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
1631 let mut ring = Ring::default();
1632 for piece in &pieces {
1633 ring.append(piece.clone());
1634 }
1635 prop_assert!(!ring.dropped);
1636 prop_assert_eq!(ring.size, concat.len());
1637 prop_assert_eq!(ring.tail(), concat);
1638 }
1639 }
1640
1641 proptest! {
1642 #![proptest_config(ProptestConfig::with_cases(48))]
1644
1645 #[test]
1650 fn ring_eviction_keeps_valid_suffix_under_cap(
1651 overflow in 1usize..8192,
1652 tail_pieces in proptest::collection::vec("[a-z]{0,64}", 0..8),
1653 ) {
1654 let head = "é".repeat(STREAM_BUFFER_CAP_BYTES / 2 + overflow);
1657 let mut full = head.clone();
1658 let mut ring = Ring::default();
1659 ring.append(head);
1660 for piece in &tail_pieces {
1661 ring.append(piece.clone());
1662 full.push_str(piece);
1663 }
1664 prop_assert!(ring.dropped);
1665 prop_assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1666 let tail = ring.tail();
1667 let has_replacement_char = tail.contains('\u{FFFD}');
1670 prop_assert!(!has_replacement_char);
1671 prop_assert!(full.as_bytes().ends_with(tail.as_bytes()));
1672 }
1673 }
1674
1675 proptest! {
1676 #[test]
1680 fn decoder_is_chunk_boundary_invariant(
1681 data in utf8ish_bytes(),
1682 cuts in proptest::collection::vec(any::<prop::sample::Index>(), 0..16),
1683 ) {
1684 let mut whole = Utf8Decoder::default();
1685 let mut whole_out = whole.decode(&data);
1686 whole_out.push_str(&whole.flush());
1687
1688 let mut bounds: Vec<usize> = cuts.iter().map(|c| c.index(data.len() + 1)).collect();
1689 bounds.sort_unstable();
1690 let mut dec = Utf8Decoder::default();
1691 let mut chunked_out = String::new();
1692 let mut prev = 0;
1693 for bound in bounds.into_iter().chain(std::iter::once(data.len())) {
1694 let bound = bound.clamp(prev, data.len());
1695 chunked_out.push_str(&dec.decode(&data[prev..bound]));
1696 prev = bound;
1697 }
1698 chunked_out.push_str(&dec.flush());
1699
1700 prop_assert_eq!(&chunked_out, &whole_out);
1701 prop_assert_eq!(whole_out, String::from_utf8_lossy(&data).into_owned());
1702 }
1703 }
1704
1705 proptest! {
1706 #[test]
1711 fn pump_seq_is_monotonic_running_max_per_stream(
1712 frames in proptest::collection::vec(
1713 (any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
1714 0..64,
1715 )
1716 ) {
1717 let mut pump = Pump::new(test_shared());
1718 let (mut expect_out, mut expect_err) = (0i64, 0i64);
1719 let (mut prev_out, mut prev_err) = (0i64, 0i64);
1720 for (is_stderr, seq, data) in frames {
1721 let which = if is_stderr { OutputStream::Stderr } else { OutputStream::Stdout };
1722 pump.apply_frame(chunk(which, seq, &data));
1723 if is_stderr {
1724 expect_err = expect_err.max(seq);
1725 } else {
1726 expect_out = expect_out.max(seq);
1727 }
1728 prop_assert_eq!(pump.stdout_seq, expect_out);
1729 prop_assert_eq!(pump.stderr_seq, expect_err);
1730 prop_assert!(pump.stdout_seq >= prev_out);
1731 prop_assert!(pump.stderr_seq >= prev_err);
1732 prev_out = pump.stdout_seq;
1733 prev_err = pump.stderr_seq;
1734 }
1735 }
1736 }
1737}