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
34pub const EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS: f64 = 30.0;
37const STREAM_BUFFER_CAP_BYTES: usize = 1024 * 1024;
43const STDIN_WRITE_CHUNK_BYTES: usize = 256 * 1024;
46
47fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
51 mutex
52 .lock()
53 .unwrap_or_else(std::sync::PoisonError::into_inner)
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum OutputStream {
59 Stdout,
61 Stderr,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum ReadStep {
68 Chunk(String),
70 Eof,
72 Pending,
75}
76
77#[derive(Debug, Clone, Serialize)]
79#[non_exhaustive]
80#[allow(clippy::struct_excessive_bools)]
81pub struct ExecResult {
82 pub stdout: String,
84 pub stderr: String,
86 pub exit_code: i32,
88 pub timed_out: bool,
90 pub stdout_truncated: bool,
93 pub stderr_truncated: bool,
96 pub stdout_complete: bool,
102 pub stderr_complete: bool,
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum CancelSignal {
110 Interrupt,
112 Kill,
114}
115
116impl CancelSignal {
117 fn is_force(self) -> bool {
119 matches!(self, CancelSignal::Kill)
120 }
121}
122
123#[derive(Debug, Clone, Copy, PartialEq)]
126pub enum RetryBudget {
127 None,
129 Within(Duration),
131 Forever,
133}
134
135impl RetryBudget {
136 pub fn as_secs_f64(self) -> f64 {
139 match self {
140 RetryBudget::None => 0.0,
141 RetryBudget::Within(d) => d.as_secs_f64(),
142 RetryBudget::Forever => f64::INFINITY,
143 }
144 }
145
146 pub fn from_secs_f64(secs: f64) -> RetryBudget {
149 if secs <= 0.0 {
150 RetryBudget::None
151 } else if secs.is_finite() {
152 RetryBudget::Within(Duration::from_secs_f64(secs))
153 } else {
154 RetryBudget::Forever
155 }
156 }
157}
158
159#[derive(Debug, Clone)]
165pub struct ExecOptions {
166 pub timeout: Option<Duration>,
170 pub open_stdin: bool,
172 pub pty: bool,
174 pub term: String,
176 pub cols: u32,
178 pub rows: u32,
180 pub idempotency_key: String,
183 pub retry_timeout: RetryBudget,
188 pub cwd: Option<String>,
191 pub background: bool,
196}
197
198impl Default for ExecOptions {
199 fn default() -> ExecOptions {
200 ExecOptions {
201 timeout: None,
202 open_stdin: false,
203 pty: false,
204 term: String::new(),
205 cols: 0,
206 rows: 0,
207 idempotency_key: String::new(),
208 retry_timeout: RetryBudget::Within(Duration::from_secs_f64(
209 EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
210 )),
211 cwd: None,
212 background: false,
213 }
214 }
215}
216
217pub(crate) fn sh_quote(value: &str) -> String {
219 format!("'{}'", value.replace('\'', "'\\''"))
220}
221
222pub fn shell_argv(command: &str, options: &ExecOptions) -> Result<Vec<String>, SailError> {
227 let invalid = |message: &str| {
228 Err(SailError::InvalidArgument {
229 message: message.to_string(),
230 })
231 };
232 if command.is_empty() {
233 return invalid("command must be non-empty");
234 }
235 if options.background && (options.open_stdin || options.pty) {
236 return invalid("background is not supported with open_stdin or pty");
237 }
238 let mut command = command.to_string();
239 if let Some(cwd) = &options.cwd {
240 let cwd = cwd.trim();
241 if cwd.is_empty() {
242 return invalid("cwd must be non-empty");
243 }
244 command = format!(
245 "cd {} && exec /bin/sh -lc {}",
246 sh_quote(cwd),
247 sh_quote(&command)
248 );
249 }
250 if options.background {
251 command = format!(
252 "nohup /bin/sh -lc {} </dev/null >/dev/null 2>&1 &",
253 sh_quote(&command)
254 );
255 }
256 Ok(vec!["/bin/sh".to_string(), "-lc".to_string(), command])
257}
258
259#[doc(hidden)]
261#[derive(Debug, Clone)]
262pub struct ExecParams {
263 pub sailbox_id: String,
265 pub exec_endpoint: String,
267 pub argv: Vec<String>,
269 pub timeout_seconds: u32,
272 pub idempotency_key: String,
275 pub open_stdin: bool,
277 pub pty: bool,
279 pub term: String,
281 pub cols: u32,
283 pub rows: u32,
285 pub retry_timeout: f64,
288 pub extra_metadata: Vec<(String, String)>,
290}
291
292#[derive(Default)]
296struct Utf8Decoder {
297 pending: Vec<u8>,
298}
299
300impl Utf8Decoder {
301 fn decode(&mut self, data: &[u8]) -> String {
302 self.pending.extend_from_slice(data);
303 let mut out = String::new();
304 loop {
305 match std::str::from_utf8(&self.pending) {
306 Ok(valid) => {
307 out.push_str(valid);
308 self.pending.clear();
309 break;
310 }
311 Err(err) => {
312 let valid_up_to = err.valid_up_to();
313 out.push_str(std::str::from_utf8(&self.pending[..valid_up_to]).unwrap());
315 if let Some(bad) = err.error_len() {
316 out.push('\u{FFFD}');
317 self.pending.drain(..valid_up_to + bad);
318 } else {
319 self.pending.drain(..valid_up_to);
321 break;
322 }
323 }
324 }
325 }
326 out
327 }
328
329 fn flush(&mut self) -> String {
330 if self.pending.is_empty() {
331 return String::new();
332 }
333 let out = String::from_utf8_lossy(&self.pending).into_owned();
334 self.pending.clear();
335 out
336 }
337}
338
339#[derive(Default)]
344struct Ring {
345 pieces: Vec<String>,
346 piece_bytes: Vec<usize>,
347 first_idx: usize,
348 size: usize,
349 dropped: bool,
350}
351
352impl Ring {
353 fn append(&mut self, text: String) {
354 let bytes = text.len();
355 self.pieces.push(text);
356 self.piece_bytes.push(bytes);
357 self.size += bytes;
358 while self.size > STREAM_BUFFER_CAP_BYTES {
359 let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
360 if self.piece_bytes[0] <= overflow {
361 self.pieces.remove(0);
362 self.size -= self.piece_bytes.remove(0);
363 self.first_idx += 1;
364 } else {
365 let piece = &self.pieces[0];
371 let mut cut = overflow;
372 while cut < piece.len() && !piece.is_char_boundary(cut) {
373 cut += 1;
374 }
375 let kept = piece[cut..].to_string();
376 self.size -= self.piece_bytes[0];
377 self.piece_bytes[0] = kept.len();
378 self.size += self.piece_bytes[0];
379 self.pieces[0] = kept;
380 }
381 self.dropped = true;
382 }
383 }
384
385 fn tail(&self) -> String {
386 self.pieces.concat()
387 }
388}
389
390#[derive(Default)]
391struct State {
392 stdout: Ring,
393 stderr: Ring,
394 ended: bool,
395}
396
397impl State {
398 fn ring(&self, which: OutputStream) -> &Ring {
399 match which {
400 OutputStream::Stdout => &self.stdout,
401 OutputStream::Stderr => &self.stderr,
402 }
403 }
404}
405
406#[derive(Clone)]
408struct ExitInfo {
409 status: i32,
410 exit_code: i32,
411 timed_out: bool,
412 stdout_truncated: bool,
413 stderr_truncated: bool,
414 error_message: String,
415 stdout_seq: i64,
416 stderr_seq: i64,
417}
418
419#[derive(Default)]
420struct StdinState {
421 offset: i64,
422 eof_sent: bool,
423 broken: bool,
424 write_in_flight: bool,
432}
433
434struct ExecShared {
435 worker: Arc<WorkerProxy>,
436 params: ExecParams,
437 state: Mutex<State>,
438 cond: Condvar,
441 data_notify: Notify,
444 exit: Mutex<Option<ExitInfo>>,
445 high_seq: Mutex<(i64, i64)>,
447 stdin: AsyncMutex<StdinState>,
448 ended: AtomicBool,
449 ended_notify: Notify,
450 closing: AtomicBool,
451 close_notify: Notify,
452}
453
454pub struct ExecProcess {
457 shared: Arc<ExecShared>,
458 exec_request_id: String,
459}
460
461impl Drop for ExecProcess {
462 fn drop(&mut self) {
463 self.close();
468 }
469}
470
471impl ExecProcess {
472 #[doc(hidden)]
480 pub async fn start(
481 worker: Arc<WorkerProxy>,
482 mut params: ExecParams,
483 ) -> Result<ExecProcess, SailError> {
484 let key = params.idempotency_key.trim();
488 params.idempotency_key = if key.is_empty() {
489 format!("exec_{}", uuid::Uuid::new_v4())
490 } else {
491 key.to_string()
492 };
493 let (exec_request_id, stream) = submit(&worker, ¶ms, 0, 0).await?;
494 let shared = Arc::new(ExecShared {
495 worker,
496 params,
497 state: Mutex::new(State::default()),
498 cond: Condvar::new(),
499 data_notify: Notify::new(),
500 exit: Mutex::new(None),
501 high_seq: Mutex::new((0, 0)),
502 stdin: AsyncMutex::new(StdinState::default()),
503 ended: AtomicBool::new(false),
504 ended_notify: Notify::new(),
505 closing: AtomicBool::new(false),
506 close_notify: Notify::new(),
507 });
508 let pump_shared = shared.clone();
509 tokio::spawn(async move { pump(pump_shared, stream).await });
510 Ok(ExecProcess {
511 shared,
512 exec_request_id,
513 })
514 }
515
516 pub fn exec_request_id(&self) -> &str {
519 &self.exec_request_id
520 }
521
522 pub fn idempotency_key(&self) -> &str {
525 &self.shared.params.idempotency_key
526 }
527
528 pub fn reader(&self, which: OutputStream) -> StreamReader {
531 StreamReader {
532 shared: self.shared.clone(),
533 which,
534 cursor: 0,
535 }
536 }
537
538 pub fn reader_async(&self, which: OutputStream) -> AsyncStreamReader {
541 AsyncStreamReader {
542 shared: self.shared.clone(),
543 which,
544 cursor: 0,
545 }
546 }
547
548 pub fn poll(&self) -> Option<Result<i32, SailError>> {
552 let exit = lock(&self.shared.exit);
553 exit.as_ref()
554 .map(|exit| match terminal_status_error(exit.status) {
555 Some(err) => Err(err),
556 None => Ok(exit.exit_code),
557 })
558 }
559
560 pub fn close(&self) {
562 self.shared.closing.store(true, Ordering::SeqCst);
563 self.shared.close_notify.notify_one();
568 }
569
570 pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
575 let _ = tokio::time::timeout(timeout, self.await_ended()).await;
576 self.shared.ended.load(Ordering::SeqCst)
577 }
578
579 async fn await_ended(&self) {
581 loop {
582 let notified = self.shared.ended_notify.notified();
583 if self.shared.ended.load(Ordering::SeqCst) {
584 return;
585 }
586 notified.await;
587 }
588 }
589
590 pub async fn wait(&self) -> Result<ExecResult, SailError> {
597 self.await_ended().await;
598 let exit = lock(&self.shared.exit).clone();
599 let (high_out, high_err) = *lock(&self.shared.high_seq);
600 let (out_dropped, err_dropped) = {
601 let state = lock(&self.shared.state);
602 (state.stdout.dropped, state.stderr.dropped)
603 };
604
605 let truncated = exit.as_ref().is_some_and(|exit| {
609 exit.stdout_truncated || exit.stderr_truncated || out_dropped || err_dropped
610 });
611 let incomplete = exit
612 .as_ref()
613 .is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
614
615 let stdout_complete = exit
620 .as_ref()
621 .is_some_and(|exit| exit.stdout_seq <= high_out);
622 let stderr_complete = exit
623 .as_ref()
624 .is_some_and(|exit| exit.stderr_seq <= high_err);
625
626 if exit.is_none() || truncated || incomplete {
627 let outcome = self
628 .shared
629 .worker
630 .wait_exec(
631 &self.shared.params.exec_endpoint,
632 &self.shared.params.sailbox_id,
633 &self.exec_request_id,
634 self.shared.params.retry_timeout,
635 )
636 .await?;
637 {
640 let mut exit_slot = lock(&self.shared.exit);
641 if exit_slot.is_none() {
642 *exit_slot = Some(ExitInfo {
643 status: outcome.status,
644 exit_code: outcome.exit_code,
645 timed_out: outcome.timed_out,
646 stdout_truncated: outcome.stdout_truncated,
647 stderr_truncated: outcome.stderr_truncated,
648 error_message: String::new(),
649 stdout_seq: 0,
650 stderr_seq: 0,
651 });
652 }
653 }
654 if let Some(witnessed) = exit.as_ref() {
660 if outcome.status == pb::SailboxExecStatus::WorkerLost as i32 {
661 let state = lock(&self.shared.state);
662 let mut stderr = state.stderr.tail();
663 if witnessed.status == pb::SailboxExecStatus::Failed as i32
664 && !witnessed.error_message.is_empty()
665 {
666 stderr = witnessed.error_message.clone();
667 }
668 return Ok(ExecResult {
669 stdout: state.stdout.tail(),
670 stderr,
671 exit_code: witnessed.exit_code,
672 timed_out: witnessed.timed_out,
673 stdout_truncated: witnessed.stdout_truncated
674 || out_dropped
675 || witnessed.stdout_seq > high_out,
676 stderr_truncated: witnessed.stderr_truncated
677 || err_dropped
678 || witnessed.stderr_seq > high_err,
679 stdout_complete,
680 stderr_complete,
681 });
682 }
683 }
684 if let Some(err) = terminal_status_error(outcome.status) {
685 return Err(err);
686 }
687 return Ok(ExecResult {
688 stdout: outcome.stdout,
689 stderr: outcome.stderr,
690 exit_code: outcome.exit_code,
691 timed_out: outcome.timed_out,
692 stdout_truncated: outcome.stdout_truncated,
693 stderr_truncated: outcome.stderr_truncated,
694 stdout_complete,
695 stderr_complete,
696 });
697 }
698
699 let exit = exit.expect("exit present on the clean path");
700 if let Some(err) = terminal_status_error(exit.status) {
701 return Err(err);
702 }
703 let state = lock(&self.shared.state);
704 let mut stderr = state.stderr.tail();
705 if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
706 stderr = exit.error_message.clone();
708 }
709 Ok(ExecResult {
710 stdout: state.stdout.tail(),
711 stderr,
712 exit_code: exit.exit_code,
713 timed_out: exit.timed_out,
714 stdout_truncated: false,
715 stderr_truncated: false,
716 stdout_complete,
717 stderr_complete,
718 })
719 }
720
721 pub async fn cancel(&self, signal: CancelSignal, retry: RetryBudget) -> Result<(), SailError> {
724 self.shared
725 .worker
726 .cancel_exec(
727 &self.shared.params.exec_endpoint,
728 &self.shared.params.sailbox_id,
729 &self.exec_request_id,
730 signal.is_force(),
731 retry.as_secs_f64(),
732 )
733 .await
734 }
735
736 pub async fn resize(&self, cols: u32, rows: u32) {
740 let message = pb::ResizeSailboxExecRequest {
741 sailbox_id: self.shared.params.sailbox_id.clone(),
742 exec_request_id: self.exec_request_id.clone(),
743 cols,
744 rows,
745 };
746 let Ok(request) =
747 self.shared
748 .worker
749 .request_for(message, &[], Some(Duration::from_secs(5)))
750 else {
751 return;
752 };
753 if let Ok(mut client) = self
754 .shared
755 .worker
756 .client_for(&self.shared.params.exec_endpoint)
757 {
758 let _ = client.resize_sailbox_exec(request).await;
759 }
760 }
761
762 pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
766 let mut stdin = self.shared.stdin.lock().await;
767 if stdin.eof_sent {
768 return Err(SailError::BrokenPipe {
769 message: "stdin is closed".to_string(),
770 });
771 }
772 if stdin.broken {
773 return Err(SailError::BrokenPipe {
774 message: "an earlier stdin write failed".to_string(),
775 });
776 }
777 if stdin.write_in_flight {
778 stdin.broken = true;
782 return Err(SailError::BrokenPipe {
783 message: "an earlier stdin write was interrupted".to_string(),
784 });
785 }
786 if data.is_empty() {
787 return Ok(());
788 }
789 stdin.write_in_flight = true;
790 let result = self.send_stdin(&mut stdin, data, false).await;
791 stdin.write_in_flight = false;
794 result
795 }
796
797 pub async fn close_stdin(&self) -> Result<(), SailError> {
799 let mut stdin = self.shared.stdin.lock().await;
800 if stdin.eof_sent {
801 return Ok(());
802 }
803 if stdin.broken || stdin.write_in_flight {
804 stdin.eof_sent = true;
807 return Ok(());
808 }
809 match self.send_stdin(&mut stdin, &[], true).await {
813 Ok(()) => {
814 stdin.eof_sent = true;
815 Ok(())
816 }
817 Err(SailError::BrokenPipe { .. }) => {
819 stdin.eof_sent = true;
820 Ok(())
821 }
822 Err(err) => Err(err),
823 }
824 }
825
826 async fn send_stdin(
830 &self,
831 stdin: &mut StdinState,
832 payload: &[u8],
833 eof: bool,
834 ) -> Result<(), SailError> {
835 let endpoint = &self.shared.params.exec_endpoint;
836 let sailbox_id = &self.shared.params.sailbox_id;
837 let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
838 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
839 let mut sent = 0usize;
840 loop {
841 let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
842 let chunk = &payload[sent..end];
843 let last = end >= payload.len();
844 let message = pb::WriteSailboxExecStdinRequest {
845 sailbox_id: sailbox_id.clone(),
846 exec_request_id: self.exec_request_id.clone(),
847 offset: stdin.offset,
848 data: chunk.to_vec(),
849 eof: eof && last,
850 };
851 let request = match self.shared.worker.request_for(
855 message,
856 &[],
857 Some(rpc_attempt_timeout(deadline)),
858 ) {
859 Ok(request) => request,
860 Err(err) => return Err(err),
861 };
862 let result = match self.shared.worker.client_for(endpoint) {
863 Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
864 Err(err) => return Err(err),
865 };
866 match result {
867 Ok(resp) => {
868 let accepted_through = resp.into_inner().accepted_through;
869 let accepted =
874 ((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
875 stdin.offset += accepted as i64;
876 sent += accepted;
877 if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
878 return Ok(());
879 }
880 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
882 if accepted > 0 {
883 delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
884 } else {
885 delay = sleep_no_deadline(delay).await;
887 }
888 }
889 Err(status) => {
890 if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
891 return Err(SailError::BrokenPipe {
893 message: status.message().to_string(),
894 });
895 }
896 if is_exec_not_ready(&status) {
897 delay = sleep_no_deadline(delay).await;
901 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
902 continue;
903 }
904 if !should_retry_transient_exec_rpc(&status, deadline) {
905 stdin.broken = true;
908 return Err(SailError::from_exec_status(&status));
909 }
910 tracing::warn!(code = ?status.code(), "retrying exec stdin write");
911 if should_invalidate_channel(&status) {
912 self.shared.worker.channels().invalidate(endpoint);
913 }
914 delay = sleep_before_retry(delay, deadline).await;
915 }
916 }
917 }
918 }
919}
920
921pub struct StreamReader {
924 shared: Arc<ExecShared>,
925 which: OutputStream,
926 cursor: usize,
927}
928
929impl StreamReader {
930 pub fn next(&mut self, timeout: Duration) -> ReadStep {
934 let mut state = lock(&self.shared.state);
935 loop {
936 let ring = state.ring(self.which);
937 if self.cursor < ring.first_idx {
938 self.cursor = ring.first_idx;
939 }
940 let available = ring.first_idx + ring.pieces.len();
941 if self.cursor < available {
942 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
943 self.cursor += 1;
944 return ReadStep::Chunk(piece);
945 }
946 if state.ended {
947 return ReadStep::Eof;
948 }
949 let (next_state, timed_out) = self
950 .shared
951 .cond
952 .wait_timeout(state, timeout)
953 .expect("exec state mutex poisoned");
954 state = next_state;
955 if timed_out.timed_out() {
956 return ReadStep::Pending;
957 }
958 }
959 }
960}
961
962pub struct AsyncStreamReader {
966 shared: Arc<ExecShared>,
967 which: OutputStream,
968 cursor: usize,
969}
970
971impl AsyncStreamReader {
972 pub async fn next(&mut self) -> Option<String> {
976 loop {
977 let notified = self.shared.data_notify.notified();
981 tokio::pin!(notified);
982 notified.as_mut().enable();
983 {
984 let state = lock(&self.shared.state);
985 let ring = state.ring(self.which);
986 if self.cursor < ring.first_idx {
987 self.cursor = ring.first_idx;
988 }
989 let available = ring.first_idx + ring.pieces.len();
990 if self.cursor < available {
991 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
992 self.cursor += 1;
993 return Some(piece);
994 }
995 if state.ended {
996 return None;
997 }
998 }
999 notified.await;
1000 }
1001 }
1002}
1003
1004async fn sleep_no_deadline(delay: f64) -> f64 {
1007 let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
1008 tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
1009 (delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
1010}
1011
1012fn is_exec_not_ready(status: &Status) -> bool {
1013 status.code() == Code::Unavailable && status.message().contains("; retry")
1014}
1015
1016fn terminal_status_error(status: i32) -> Option<SailError> {
1019 if status == pb::SailboxExecStatus::WorkerLost as i32 {
1020 return Some(SailError::WorkerLost {
1021 message: "the machine hosting your sailbox was lost before the command \
1022 finished; run exec again to retry"
1023 .to_string(),
1024 });
1025 }
1026 None
1027}
1028
1029async fn submit(
1032 worker: &Arc<WorkerProxy>,
1033 params: &ExecParams,
1034 stdout_resume_seq: i64,
1035 stderr_resume_seq: i64,
1036) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
1037 let deadline = retry_deadline(params.retry_timeout);
1038 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
1039 loop {
1040 let message = pb::StreamSailboxExecRequest {
1041 sailbox_id: params.sailbox_id.clone(),
1042 argv: params.argv.clone(),
1043 timeout_seconds: params.timeout_seconds,
1044 idempotency_key: params.idempotency_key.clone(),
1045 open_stdin: params.open_stdin,
1046 stdout_resume_seq,
1047 stderr_resume_seq,
1048 pty: params.pty,
1049 term_cols: params.cols,
1050 term_rows: params.rows,
1051 term: params.term.clone(),
1052 };
1053 let request = worker.request_for(message, ¶ms.extra_metadata, None)?;
1054 let status = match worker
1055 .client_for(¶ms.exec_endpoint)?
1056 .stream_sailbox_exec(request)
1057 .await
1058 {
1059 Ok(resp) => {
1060 let mut stream = resp.into_inner();
1061 match stream.message().await {
1062 Ok(Some(first)) => match first.frame {
1063 Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
1064 return Ok((started.exec_request_id, stream));
1065 }
1066 _ => {
1067 return Err(SailError::Execution {
1068 code: crate::error::GrpcCode::Internal,
1069 detail: "exec stream opened with a non-started frame".to_string(),
1070 });
1071 }
1072 },
1073 Ok(None) if stdout_resume_seq == 0 && stderr_resume_seq == 0 => {
1084 tonic::Status::unavailable(
1085 "exec stream ended before the server confirmed the launch",
1086 )
1087 }
1088 Ok(None) => {
1089 return Err(SailError::Execution {
1090 code: crate::error::GrpcCode::Internal,
1091 detail: "exec stream ended before the server confirmed the launch"
1092 .to_string(),
1093 });
1094 }
1095 Err(status) => status,
1096 }
1097 }
1098 Err(status) => status,
1099 };
1100 if !should_retry_transient_exec_rpc(&status, deadline) {
1101 return Err(SailError::from_exec_status(&status));
1102 }
1103 tracing::warn!(
1104 code = ?status.code(),
1105 stdout_resume_seq,
1106 stderr_resume_seq,
1107 "reconnecting exec stream"
1108 );
1109 if should_invalidate_channel(&status) {
1110 worker.channels().invalidate(¶ms.exec_endpoint);
1111 }
1112 delay = sleep_before_retry(delay, deadline).await;
1113 }
1114}
1115
1116struct Pump {
1122 shared: Arc<ExecShared>,
1123 stdout_dec: Utf8Decoder,
1124 stderr_dec: Utf8Decoder,
1125 stdout_seq: i64,
1126 stderr_seq: i64,
1127}
1128
1129impl Pump {
1130 fn new(shared: Arc<ExecShared>) -> Pump {
1131 Pump {
1132 shared,
1133 stdout_dec: Utf8Decoder::default(),
1134 stderr_dec: Utf8Decoder::default(),
1135 stdout_seq: 0,
1136 stderr_seq: 0,
1137 }
1138 }
1139
1140 fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
1145 match frame.frame {
1146 Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
1147 let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
1148 let (decoder, which) = if is_stderr {
1149 self.stderr_seq = self.stderr_seq.max(chunk.seq);
1150 (&mut self.stderr_dec, OutputStream::Stderr)
1151 } else {
1152 self.stdout_seq = self.stdout_seq.max(chunk.seq);
1153 (&mut self.stdout_dec, OutputStream::Stdout)
1154 };
1155 let text = decoder.decode(&chunk.data);
1156 if !text.is_empty() {
1157 let mut state = lock(&self.shared.state);
1158 match which {
1159 OutputStream::Stdout => state.stdout.append(text),
1160 OutputStream::Stderr => state.stderr.append(text),
1161 }
1162 self.shared.cond.notify_all();
1163 self.shared.data_notify.notify_waiters();
1164 }
1165 false
1166 }
1167 Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
1168 *lock(&self.shared.exit) = Some(ExitInfo {
1169 status: exit.status,
1170 exit_code: exit.return_code,
1171 timed_out: exit.timed_out,
1172 stdout_truncated: exit.stdout_truncated,
1173 stderr_truncated: exit.stderr_truncated,
1174 error_message: exit.error_message,
1175 stdout_seq: exit.stdout_seq,
1176 stderr_seq: exit.stderr_seq,
1177 });
1178 true
1179 }
1180 _ => false,
1181 }
1182 }
1183
1184 fn finalize(mut self) {
1187 let tail_out = self.stdout_dec.flush();
1188 let tail_err = self.stderr_dec.flush();
1189 {
1190 let mut state = lock(&self.shared.state);
1191 if !tail_out.is_empty() {
1192 state.stdout.append(tail_out);
1193 }
1194 if !tail_err.is_empty() {
1195 state.stderr.append(tail_err);
1196 }
1197 state.ended = true;
1198 self.shared.cond.notify_all();
1199 self.shared.data_notify.notify_waiters();
1200 }
1201 *lock(&self.shared.high_seq) = (self.stdout_seq, self.stderr_seq);
1202 self.shared.ended.store(true, Ordering::SeqCst);
1203 self.shared.ended_notify.notify_waiters();
1204 }
1205}
1206
1207async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
1210 let mut state = Pump::new(shared.clone());
1211 loop {
1212 if shared.closing.load(Ordering::SeqCst) {
1213 break;
1214 }
1215 let message = tokio::select! {
1216 biased;
1217 () = shared.close_notify.notified() => break,
1218 message = stream.message() => message,
1219 };
1220 match message {
1221 Ok(Some(frame)) => {
1222 if state.apply_frame(frame) {
1223 break;
1224 }
1225 }
1226 Ok(None) => break,
1228 Err(_status) => {
1229 if shared.closing.load(Ordering::SeqCst) {
1230 break;
1231 }
1232 match submit(
1235 &shared.worker,
1236 &shared.params,
1237 state.stdout_seq,
1238 state.stderr_seq,
1239 )
1240 .await
1241 {
1242 Ok((_id, fresh)) => {
1243 stream = fresh;
1244 continue;
1245 }
1246 Err(_) => break,
1247 }
1248 }
1249 }
1250 }
1251 state.finalize();
1252}
1253
1254#[cfg(any(test, feature = "fuzzing"))]
1260pub fn fuzz_decode_ring(data: &[u8]) {
1261 let mut whole = Utf8Decoder::default();
1263 let mut whole_out = whole.decode(data);
1264 whole_out.push_str(&whole.flush());
1265
1266 let mut chunked = Utf8Decoder::default();
1269 let mut chunked_out = String::new();
1270 let mut start = 0;
1271 for (i, byte) in data.iter().enumerate() {
1272 if byte & 1 == 1 {
1273 chunked_out.push_str(&chunked.decode(&data[start..=i]));
1274 start = i + 1;
1275 }
1276 }
1277 chunked_out.push_str(&chunked.decode(&data[start..]));
1278 chunked_out.push_str(&chunked.flush());
1279
1280 assert_eq!(
1282 chunked_out, whole_out,
1283 "decoder is not chunk-boundary invariant"
1284 );
1285 assert_eq!(
1286 whole_out,
1287 String::from_utf8_lossy(data),
1288 "decoder disagrees with a lossy decode of the whole input"
1289 );
1290
1291 let mut ring = Ring::default();
1296 ring.append(whole_out.clone());
1297 let filler = STREAM_BUFFER_CAP_BYTES + 1 - whole_out.len().min(STREAM_BUFFER_CAP_BYTES);
1298 let filler = "a".repeat(filler);
1299 ring.append(filler.clone());
1300 assert!(
1301 ring.size <= STREAM_BUFFER_CAP_BYTES,
1302 "ring exceeded its cap"
1303 );
1304 let full = format!("{whole_out}{filler}");
1305 let tail = ring.tail();
1306 assert!(
1307 full.as_bytes().ends_with(tail.as_bytes()),
1308 "ring tail is not a suffix of the appended bytes"
1309 );
1310}
1311
1312#[cfg(test)]
1313mod tests {
1314 use super::*;
1315
1316 #[test]
1317 fn shell_argv_wraps_cwd_and_background() {
1318 let plain = shell_argv("echo hi", &ExecOptions::default()).unwrap();
1319 assert_eq!(plain, ["/bin/sh", "-lc", "echo hi"]);
1320
1321 let cwd = shell_argv(
1322 "echo hi",
1323 &ExecOptions {
1324 cwd: Some("/app".to_string()),
1325 ..ExecOptions::default()
1326 },
1327 )
1328 .unwrap();
1329 assert_eq!(cwd[2], "cd '/app' && exec /bin/sh -lc 'echo hi'");
1330
1331 let background = shell_argv(
1332 "echo hi",
1333 &ExecOptions {
1334 cwd: Some("/app".to_string()),
1335 background: true,
1336 ..ExecOptions::default()
1337 },
1338 )
1339 .unwrap();
1340 assert_eq!(
1341 background[2],
1342 "nohup /bin/sh -lc 'cd '\\''/app'\\'' && exec /bin/sh -lc '\\''echo hi'\\''' </dev/null >/dev/null 2>&1 &"
1343 );
1344 }
1345
1346 #[test]
1347 fn shell_argv_rejects_invalid_combinations() {
1348 assert!(shell_argv("", &ExecOptions::default()).is_err());
1349 assert!(shell_argv(
1350 "x",
1351 &ExecOptions {
1352 cwd: Some(" ".to_string()),
1353 ..ExecOptions::default()
1354 }
1355 )
1356 .is_err());
1357 for (open_stdin, pty) in [(true, false), (false, true)] {
1358 assert!(shell_argv(
1359 "x",
1360 &ExecOptions {
1361 background: true,
1362 open_stdin,
1363 pty,
1364 ..ExecOptions::default()
1365 }
1366 )
1367 .is_err());
1368 }
1369 }
1370
1371 #[test]
1372 fn fuzz_decode_ring_holds_on_samples() {
1373 for sample in [
1376 &b""[..],
1377 b"hello",
1378 b"\xff\xfe\xfd",
1379 "€ µ é".as_bytes(),
1380 b"ab\xc3\xa9cd",
1381 &[0xC3, 0x28],
1382 ] {
1383 fuzz_decode_ring(sample);
1384 }
1385 }
1386
1387 #[test]
1388 fn ring_drops_oldest_past_cap_and_latches() {
1389 let mut ring = Ring::default();
1390 ring.append("a".repeat(STREAM_BUFFER_CAP_BYTES));
1391 assert!(!ring.dropped);
1392 ring.append("bbbb".to_string());
1393 assert!(ring.dropped);
1394 assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
1395 assert!(ring.tail().ends_with("bbbb"));
1396 }
1397
1398 #[test]
1399 fn ring_clip_on_split_multibyte_char_terminates() {
1400 let mut ring = Ring::default();
1405 ring.append(format!("é{}", "a".repeat(STREAM_BUFFER_CAP_BYTES - 1)));
1406 assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1407 assert!(ring.dropped);
1408 let tail = ring.tail();
1409 assert!(!tail.contains('é'));
1410 assert!(!tail.contains('\u{FFFD}'));
1411 }
1412
1413 #[test]
1414 fn decoder_carries_split_multibyte_char() {
1415 let mut dec = Utf8Decoder::default();
1416 let euro = "€".as_bytes(); assert_eq!(dec.decode(&euro[..2]), "");
1418 assert_eq!(dec.decode(&euro[2..]), "€");
1419 }
1420
1421 #[test]
1422 fn decoder_replaces_invalid_bytes() {
1423 let mut dec = Utf8Decoder::default();
1424 assert_eq!(dec.decode(&[0xff, b'x']), "\u{FFFD}x");
1425 }
1426
1427 #[test]
1428 fn terminal_status_mapping() {
1429 assert!(matches!(
1430 terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
1431 Some(SailError::WorkerLost { .. })
1432 ));
1433 assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
1434 for retired in [9, 6, 7] {
1438 assert!(terminal_status_error(retired).is_none());
1439 }
1440 }
1441
1442 fn test_shared() -> Arc<ExecShared> {
1443 Arc::new(ExecShared {
1444 worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
1445 params: ExecParams {
1446 sailbox_id: "sb".into(),
1447 exec_endpoint: "endpoint".into(),
1448 argv: vec!["echo".into()],
1449 timeout_seconds: 0,
1450 idempotency_key: "idem".into(),
1451 open_stdin: false,
1452 pty: false,
1453 term: String::new(),
1454 cols: 0,
1455 rows: 0,
1456 retry_timeout: 0.0,
1457 extra_metadata: vec![],
1458 },
1459 state: Mutex::new(State::default()),
1460 cond: Condvar::new(),
1461 data_notify: Notify::new(),
1462 exit: Mutex::new(None),
1463 high_seq: Mutex::new((0, 0)),
1464 stdin: AsyncMutex::new(StdinState::default()),
1465 ended: AtomicBool::new(false),
1466 ended_notify: Notify::new(),
1467 closing: AtomicBool::new(false),
1468 close_notify: Notify::new(),
1469 })
1470 }
1471
1472 fn chunk(which: OutputStream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
1473 let stream = match which {
1474 OutputStream::Stdout => pb::SailboxExecStream::Stdout,
1475 OutputStream::Stderr => pb::SailboxExecStream::Stderr,
1476 };
1477 pb::StreamSailboxExecResponse {
1478 frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
1479 pb::SailboxExecChunk {
1480 stream: stream as i32,
1481 data: data.to_vec(),
1482 seq,
1483 },
1484 )),
1485 }
1486 }
1487
1488 fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
1489 pb::StreamSailboxExecResponse {
1490 frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
1491 pb::SailboxExecExit {
1492 status: status as i32,
1493 return_code: 0,
1494 timed_out: false,
1495 stdout_truncated: false,
1496 stderr_truncated: false,
1497 error_message: String::new(),
1498 stdout_seq,
1499 stderr_seq: 0,
1500 },
1501 )),
1502 }
1503 }
1504
1505 #[tokio::test]
1510 async fn exec_resume_carries_utf8_and_tracks_seq_across_break() {
1511 let shared = test_shared();
1512 let mut pump = Pump::new(shared.clone());
1513
1514 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"ab\xc3")));
1516 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "ab");
1517 assert_eq!(pump.stdout_seq, 1);
1519
1520 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"\xa9cd")));
1523 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "abécd");
1524 assert_eq!(pump.stdout_seq, 2);
1525
1526 assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"!")));
1528 assert_eq!(pump.stdout_seq, 2);
1529
1530 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
1531 pump.finalize();
1532 assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
1533 assert!(shared.ended.load(Ordering::SeqCst));
1534 }
1535
1536 #[tokio::test]
1540 async fn exec_reader_follows_replayed_then_live() {
1541 let shared = test_shared();
1542 let mut reader = StreamReader {
1543 shared: shared.clone(),
1544 which: OutputStream::Stdout,
1545 cursor: 0,
1546 };
1547 let collector = tokio::task::spawn_blocking(move || {
1548 let mut out = String::new();
1549 loop {
1550 match reader.next(Duration::from_millis(50)) {
1551 ReadStep::Chunk(piece) => out.push_str(&piece),
1552 ReadStep::Eof => return out,
1553 ReadStep::Pending => {}
1554 }
1555 }
1556 });
1557
1558 let mut pump = Pump::new(shared.clone());
1559 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"hello "));
1560 tokio::time::sleep(Duration::from_millis(10)).await;
1561 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"world"));
1563 tokio::time::sleep(Duration::from_millis(10)).await;
1564 pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
1565 pump.finalize();
1566
1567 assert_eq!(collector.await.unwrap(), "hello world");
1568 }
1569
1570 #[test]
1573 fn exec_reader_started_late_replays_retained_tail() {
1574 let shared = test_shared();
1575 let mut pump = Pump::new(shared.clone());
1576 pump.apply_frame(chunk(OutputStream::Stdout, 1, b"early "));
1577 pump.apply_frame(chunk(OutputStream::Stdout, 2, b"output"));
1578 pump.finalize();
1579
1580 let mut reader = shared_reader(&shared, OutputStream::Stdout);
1582 let mut out = String::new();
1583 loop {
1584 match reader.next(Duration::from_millis(50)) {
1585 ReadStep::Chunk(piece) => out.push_str(&piece),
1586 ReadStep::Eof => break,
1587 ReadStep::Pending => panic!("ended ring should not return Pending"),
1588 }
1589 }
1590 assert_eq!(out, "early output");
1591 }
1592
1593 fn shared_reader(shared: &Arc<ExecShared>, which: OutputStream) -> StreamReader {
1594 StreamReader {
1595 shared: shared.clone(),
1596 which,
1597 cursor: 0,
1598 }
1599 }
1600
1601 use proptest::prelude::*;
1602
1603 fn utf8ish_bytes() -> impl Strategy<Value = Vec<u8>> {
1607 prop_oneof![
1608 proptest::collection::vec(any::<u8>(), 0..256),
1609 any::<String>().prop_map(String::into_bytes),
1610 ]
1611 }
1612
1613 proptest! {
1614 #[test]
1617 fn ring_without_overflow_is_lossless(
1618 pieces in proptest::collection::vec(any::<String>(), 0..32)
1619 ) {
1620 let concat: String = pieces.concat();
1621 prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
1622 let mut ring = Ring::default();
1623 for piece in &pieces {
1624 ring.append(piece.clone());
1625 }
1626 prop_assert!(!ring.dropped);
1627 prop_assert_eq!(ring.size, concat.len());
1628 prop_assert_eq!(ring.tail(), concat);
1629 }
1630 }
1631
1632 proptest! {
1633 #![proptest_config(ProptestConfig::with_cases(48))]
1635
1636 #[test]
1641 fn ring_eviction_keeps_valid_suffix_under_cap(
1642 overflow in 1usize..8192,
1643 tail_pieces in proptest::collection::vec("[a-z]{0,64}", 0..8),
1644 ) {
1645 let head = "é".repeat(STREAM_BUFFER_CAP_BYTES / 2 + overflow);
1648 let mut full = head.clone();
1649 let mut ring = Ring::default();
1650 ring.append(head);
1651 for piece in &tail_pieces {
1652 ring.append(piece.clone());
1653 full.push_str(piece);
1654 }
1655 prop_assert!(ring.dropped);
1656 prop_assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1657 let tail = ring.tail();
1658 let has_replacement_char = tail.contains('\u{FFFD}');
1661 prop_assert!(!has_replacement_char);
1662 prop_assert!(full.as_bytes().ends_with(tail.as_bytes()));
1663 }
1664 }
1665
1666 proptest! {
1667 #[test]
1671 fn decoder_is_chunk_boundary_invariant(
1672 data in utf8ish_bytes(),
1673 cuts in proptest::collection::vec(any::<prop::sample::Index>(), 0..16),
1674 ) {
1675 let mut whole = Utf8Decoder::default();
1676 let mut whole_out = whole.decode(&data);
1677 whole_out.push_str(&whole.flush());
1678
1679 let mut bounds: Vec<usize> = cuts.iter().map(|c| c.index(data.len() + 1)).collect();
1680 bounds.sort_unstable();
1681 let mut dec = Utf8Decoder::default();
1682 let mut chunked_out = String::new();
1683 let mut prev = 0;
1684 for bound in bounds.into_iter().chain(std::iter::once(data.len())) {
1685 let bound = bound.clamp(prev, data.len());
1686 chunked_out.push_str(&dec.decode(&data[prev..bound]));
1687 prev = bound;
1688 }
1689 chunked_out.push_str(&dec.flush());
1690
1691 prop_assert_eq!(&chunked_out, &whole_out);
1692 prop_assert_eq!(whole_out, String::from_utf8_lossy(&data).into_owned());
1693 }
1694 }
1695
1696 proptest! {
1697 #[test]
1702 fn pump_seq_is_monotonic_running_max_per_stream(
1703 frames in proptest::collection::vec(
1704 (any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
1705 0..64,
1706 )
1707 ) {
1708 let mut pump = Pump::new(test_shared());
1709 let (mut expect_out, mut expect_err) = (0i64, 0i64);
1710 let (mut prev_out, mut prev_err) = (0i64, 0i64);
1711 for (is_stderr, seq, data) in frames {
1712 let which = if is_stderr { OutputStream::Stderr } else { OutputStream::Stdout };
1713 pump.apply_frame(chunk(which, seq, &data));
1714 if is_stderr {
1715 expect_err = expect_err.max(seq);
1716 } else {
1717 expect_out = expect_out.max(seq);
1718 }
1719 prop_assert_eq!(pump.stdout_seq, expect_out);
1720 prop_assert_eq!(pump.stderr_seq, expect_err);
1721 prop_assert!(pump.stdout_seq >= prev_out);
1722 prop_assert!(pump.stderr_seq >= prev_err);
1723 prev_out = pump.stdout_seq;
1724 prev_err = pump.stderr_seq;
1725 }
1726 }
1727 }
1728}