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 Stream {
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 return_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)]
164pub struct ExecOptions {
165 pub timeout: Option<Duration>,
169 pub open_stdin: bool,
171 pub pty: bool,
173 pub term: String,
175 pub cols: u32,
177 pub rows: u32,
179 pub idempotency_key: String,
182 pub retry_timeout: RetryBudget,
186}
187
188impl Default for ExecOptions {
189 fn default() -> ExecOptions {
190 ExecOptions {
191 timeout: None,
192 open_stdin: false,
193 pty: false,
194 term: String::new(),
195 cols: 0,
196 rows: 0,
197 idempotency_key: String::new(),
198 retry_timeout: RetryBudget::Within(Duration::from_secs_f64(
199 EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
200 )),
201 }
202 }
203}
204
205#[doc(hidden)]
207#[derive(Debug, Clone)]
208pub struct ExecParams {
209 pub sailbox_id: String,
211 pub exec_endpoint: String,
213 pub argv: Vec<String>,
215 pub timeout_seconds: u32,
218 pub idempotency_key: String,
221 pub open_stdin: bool,
223 pub pty: bool,
225 pub term: String,
227 pub cols: u32,
229 pub rows: u32,
231 pub retry_timeout: f64,
234 pub extra_metadata: Vec<(String, String)>,
236}
237
238#[derive(Default)]
242struct Utf8Decoder {
243 pending: Vec<u8>,
244}
245
246impl Utf8Decoder {
247 fn decode(&mut self, data: &[u8]) -> String {
248 self.pending.extend_from_slice(data);
249 let mut out = String::new();
250 loop {
251 match std::str::from_utf8(&self.pending) {
252 Ok(valid) => {
253 out.push_str(valid);
254 self.pending.clear();
255 break;
256 }
257 Err(err) => {
258 let valid_up_to = err.valid_up_to();
259 out.push_str(std::str::from_utf8(&self.pending[..valid_up_to]).unwrap());
261 if let Some(bad) = err.error_len() {
262 out.push('\u{FFFD}');
263 self.pending.drain(..valid_up_to + bad);
264 } else {
265 self.pending.drain(..valid_up_to);
267 break;
268 }
269 }
270 }
271 }
272 out
273 }
274
275 fn flush(&mut self) -> String {
276 if self.pending.is_empty() {
277 return String::new();
278 }
279 let out = String::from_utf8_lossy(&self.pending).into_owned();
280 self.pending.clear();
281 out
282 }
283}
284
285#[derive(Default)]
290struct Ring {
291 pieces: Vec<String>,
292 piece_bytes: Vec<usize>,
293 first_idx: usize,
294 size: usize,
295 dropped: bool,
296}
297
298impl Ring {
299 fn append(&mut self, text: String) {
300 let bytes = text.len();
301 self.pieces.push(text);
302 self.piece_bytes.push(bytes);
303 self.size += bytes;
304 while self.size > STREAM_BUFFER_CAP_BYTES {
305 let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
306 if self.piece_bytes[0] <= overflow {
307 self.pieces.remove(0);
308 self.size -= self.piece_bytes.remove(0);
309 self.first_idx += 1;
310 } else {
311 let piece = &self.pieces[0];
317 let mut cut = overflow;
318 while cut < piece.len() && !piece.is_char_boundary(cut) {
319 cut += 1;
320 }
321 let kept = piece[cut..].to_string();
322 self.size -= self.piece_bytes[0];
323 self.piece_bytes[0] = kept.len();
324 self.size += self.piece_bytes[0];
325 self.pieces[0] = kept;
326 }
327 self.dropped = true;
328 }
329 }
330
331 fn tail(&self) -> String {
332 self.pieces.concat()
333 }
334}
335
336#[derive(Default)]
337struct State {
338 stdout: Ring,
339 stderr: Ring,
340 ended: bool,
341}
342
343impl State {
344 fn ring(&self, which: Stream) -> &Ring {
345 match which {
346 Stream::Stdout => &self.stdout,
347 Stream::Stderr => &self.stderr,
348 }
349 }
350}
351
352#[derive(Clone)]
354struct ExitInfo {
355 status: i32,
356 return_code: i32,
357 timed_out: bool,
358 stdout_truncated: bool,
359 stderr_truncated: bool,
360 error_message: String,
361 stdout_seq: i64,
362 stderr_seq: i64,
363}
364
365#[derive(Default)]
366struct StdinState {
367 offset: i64,
368 eof_sent: bool,
369 broken: bool,
370 write_in_flight: bool,
378}
379
380struct ExecShared {
381 worker: Arc<WorkerProxy>,
382 params: ExecParams,
383 state: Mutex<State>,
384 cond: Condvar,
387 data_notify: Notify,
390 exit: Mutex<Option<ExitInfo>>,
391 high_seq: Mutex<(i64, i64)>,
393 stdin: AsyncMutex<StdinState>,
394 ended: AtomicBool,
395 ended_notify: Notify,
396 closing: AtomicBool,
397 close_notify: Notify,
398}
399
400pub struct ExecProcess {
403 shared: Arc<ExecShared>,
404 exec_request_id: String,
405}
406
407impl Drop for ExecProcess {
408 fn drop(&mut self) {
409 self.close();
414 }
415}
416
417impl ExecProcess {
418 #[doc(hidden)]
426 pub async fn start(
427 worker: Arc<WorkerProxy>,
428 mut params: ExecParams,
429 ) -> Result<ExecProcess, SailError> {
430 let key = params.idempotency_key.trim();
434 params.idempotency_key = if key.is_empty() {
435 format!("exec_{}", uuid::Uuid::new_v4())
436 } else {
437 key.to_string()
438 };
439 let (exec_request_id, stream) = submit(&worker, ¶ms, 0, 0).await?;
440 let shared = Arc::new(ExecShared {
441 worker,
442 params,
443 state: Mutex::new(State::default()),
444 cond: Condvar::new(),
445 data_notify: Notify::new(),
446 exit: Mutex::new(None),
447 high_seq: Mutex::new((0, 0)),
448 stdin: AsyncMutex::new(StdinState::default()),
449 ended: AtomicBool::new(false),
450 ended_notify: Notify::new(),
451 closing: AtomicBool::new(false),
452 close_notify: Notify::new(),
453 });
454 let pump_shared = shared.clone();
455 tokio::spawn(async move { pump(pump_shared, stream).await });
456 Ok(ExecProcess {
457 shared,
458 exec_request_id,
459 })
460 }
461
462 pub fn exec_request_id(&self) -> &str {
465 &self.exec_request_id
466 }
467
468 pub fn idempotency_key(&self) -> &str {
471 &self.shared.params.idempotency_key
472 }
473
474 pub fn reader(&self, which: Stream) -> StreamReader {
477 StreamReader {
478 shared: self.shared.clone(),
479 which,
480 cursor: 0,
481 }
482 }
483
484 pub fn reader_async(&self, which: Stream) -> AsyncStreamReader {
487 AsyncStreamReader {
488 shared: self.shared.clone(),
489 which,
490 cursor: 0,
491 }
492 }
493
494 pub fn poll(&self) -> Option<Result<i32, SailError>> {
498 let exit = lock(&self.shared.exit);
499 exit.as_ref()
500 .map(|exit| match terminal_status_error(exit.status) {
501 Some(err) => Err(err),
502 None => Ok(exit.return_code),
503 })
504 }
505
506 pub fn close(&self) {
508 self.shared.closing.store(true, Ordering::SeqCst);
509 self.shared.close_notify.notify_one();
514 }
515
516 pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
521 let _ = tokio::time::timeout(timeout, self.await_ended()).await;
522 self.shared.ended.load(Ordering::SeqCst)
523 }
524
525 async fn await_ended(&self) {
527 loop {
528 let notified = self.shared.ended_notify.notified();
529 if self.shared.ended.load(Ordering::SeqCst) {
530 return;
531 }
532 notified.await;
533 }
534 }
535
536 pub async fn wait(&self) -> Result<ExecResult, SailError> {
543 self.await_ended().await;
544 let exit = lock(&self.shared.exit).clone();
545 let (high_out, high_err) = *lock(&self.shared.high_seq);
546 let (out_dropped, err_dropped) = {
547 let state = lock(&self.shared.state);
548 (state.stdout.dropped, state.stderr.dropped)
549 };
550
551 let truncated = exit.as_ref().is_some_and(|exit| {
554 exit.stdout_truncated || exit.stderr_truncated || out_dropped || err_dropped
555 });
556 let incomplete = exit
557 .as_ref()
558 .is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
559
560 let stdout_complete = exit
565 .as_ref()
566 .is_some_and(|exit| exit.stdout_seq <= high_out);
567 let stderr_complete = exit
568 .as_ref()
569 .is_some_and(|exit| exit.stderr_seq <= high_err);
570
571 if exit.is_none() || truncated || incomplete {
572 let outcome = self
573 .shared
574 .worker
575 .wait_exec(
576 &self.shared.params.exec_endpoint,
577 &self.shared.params.sailbox_id,
578 &self.exec_request_id,
579 self.shared.params.retry_timeout,
580 )
581 .await?;
582 {
585 let mut exit_slot = lock(&self.shared.exit);
586 if exit_slot.is_none() {
587 *exit_slot = Some(ExitInfo {
588 status: outcome.status,
589 return_code: outcome.return_code,
590 timed_out: outcome.timed_out,
591 stdout_truncated: outcome.stdout_truncated,
592 stderr_truncated: outcome.stderr_truncated,
593 error_message: String::new(),
594 stdout_seq: 0,
595 stderr_seq: 0,
596 });
597 }
598 }
599 if let Some(err) = terminal_status_error(outcome.status) {
600 return Err(err);
601 }
602 return Ok(ExecResult {
603 stdout: outcome.stdout,
604 stderr: outcome.stderr,
605 return_code: outcome.return_code,
606 timed_out: outcome.timed_out,
607 stdout_truncated: outcome.stdout_truncated,
608 stderr_truncated: outcome.stderr_truncated,
609 stdout_complete,
610 stderr_complete,
611 });
612 }
613
614 let exit = exit.expect("exit present on the clean path");
615 if let Some(err) = terminal_status_error(exit.status) {
616 return Err(err);
617 }
618 let state = lock(&self.shared.state);
619 let mut stderr = state.stderr.tail();
620 if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
621 stderr = exit.error_message.clone();
623 }
624 Ok(ExecResult {
625 stdout: state.stdout.tail(),
626 stderr,
627 return_code: exit.return_code,
628 timed_out: exit.timed_out,
629 stdout_truncated: false,
630 stderr_truncated: false,
631 stdout_complete,
632 stderr_complete,
633 })
634 }
635
636 pub async fn cancel(&self, signal: CancelSignal, retry: RetryBudget) -> Result<(), SailError> {
639 self.shared
640 .worker
641 .cancel_exec(
642 &self.shared.params.exec_endpoint,
643 &self.shared.params.sailbox_id,
644 &self.exec_request_id,
645 signal.is_force(),
646 retry.as_secs_f64(),
647 )
648 .await
649 }
650
651 pub async fn resize(&self, cols: u32, rows: u32) {
655 let message = pb::ResizeSailboxExecRequest {
656 sailbox_id: self.shared.params.sailbox_id.clone(),
657 exec_request_id: self.exec_request_id.clone(),
658 cols,
659 rows,
660 };
661 let Ok(request) =
662 self.shared
663 .worker
664 .request_for(message, &[], Some(Duration::from_secs(5)))
665 else {
666 return;
667 };
668 if let Ok(mut client) = self
669 .shared
670 .worker
671 .client_for(&self.shared.params.exec_endpoint)
672 {
673 let _ = client.resize_sailbox_exec(request).await;
674 }
675 }
676
677 pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
681 let mut stdin = self.shared.stdin.lock().await;
682 if stdin.eof_sent {
683 return Err(SailError::BrokenPipe {
684 message: "stdin is closed".to_string(),
685 });
686 }
687 if stdin.broken {
688 return Err(SailError::BrokenPipe {
689 message: "an earlier stdin write failed".to_string(),
690 });
691 }
692 if stdin.write_in_flight {
693 stdin.broken = true;
697 return Err(SailError::BrokenPipe {
698 message: "an earlier stdin write was interrupted".to_string(),
699 });
700 }
701 if data.is_empty() {
702 return Ok(());
703 }
704 stdin.write_in_flight = true;
705 let result = self.send_stdin(&mut stdin, data, false).await;
706 stdin.write_in_flight = false;
709 result
710 }
711
712 pub async fn eof_stdin(&self) -> Result<(), SailError> {
714 let mut stdin = self.shared.stdin.lock().await;
715 if stdin.eof_sent {
716 return Ok(());
717 }
718 if stdin.broken || stdin.write_in_flight {
719 stdin.eof_sent = true;
722 return Ok(());
723 }
724 match self.send_stdin(&mut stdin, &[], true).await {
728 Ok(()) => {
729 stdin.eof_sent = true;
730 Ok(())
731 }
732 Err(SailError::BrokenPipe { .. }) => {
734 stdin.eof_sent = true;
735 Ok(())
736 }
737 Err(err) => Err(err),
738 }
739 }
740
741 async fn send_stdin(
745 &self,
746 stdin: &mut StdinState,
747 payload: &[u8],
748 eof: bool,
749 ) -> Result<(), SailError> {
750 let endpoint = &self.shared.params.exec_endpoint;
751 let sailbox_id = &self.shared.params.sailbox_id;
752 let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
753 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
754 let mut sent = 0usize;
755 loop {
756 let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
757 let chunk = &payload[sent..end];
758 let last = end >= payload.len();
759 let message = pb::WriteSailboxExecStdinRequest {
760 sailbox_id: sailbox_id.clone(),
761 exec_request_id: self.exec_request_id.clone(),
762 offset: stdin.offset,
763 data: chunk.to_vec(),
764 eof: eof && last,
765 };
766 let request = match self.shared.worker.request_for(
770 message,
771 &[],
772 Some(rpc_attempt_timeout(deadline)),
773 ) {
774 Ok(request) => request,
775 Err(err) => return Err(err),
776 };
777 let result = match self.shared.worker.client_for(endpoint) {
778 Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
779 Err(err) => return Err(err),
780 };
781 match result {
782 Ok(resp) => {
783 let accepted_through = resp.into_inner().accepted_through;
784 let accepted =
789 ((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
790 stdin.offset += accepted as i64;
791 sent += accepted;
792 if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
793 return Ok(());
794 }
795 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
797 if accepted > 0 {
798 delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
799 } else {
800 delay = sleep_no_deadline(delay).await;
802 }
803 }
804 Err(status) => {
805 if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
806 return Err(SailError::BrokenPipe {
808 message: status.message().to_string(),
809 });
810 }
811 if is_exec_not_ready(&status) {
812 delay = sleep_no_deadline(delay).await;
816 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
817 continue;
818 }
819 if !should_retry_transient_exec_rpc(&status, deadline) {
820 stdin.broken = true;
823 return Err(SailError::from_exec_status(&status));
824 }
825 tracing::warn!(code = ?status.code(), "retrying exec stdin write");
826 if should_invalidate_channel(&status) {
827 self.shared.worker.channels().invalidate(endpoint);
828 }
829 delay = sleep_before_retry(delay, deadline).await;
830 }
831 }
832 }
833 }
834}
835
836pub struct StreamReader {
839 shared: Arc<ExecShared>,
840 which: Stream,
841 cursor: usize,
842}
843
844impl StreamReader {
845 pub fn next(&mut self, timeout: Duration) -> ReadStep {
849 let mut state = lock(&self.shared.state);
850 loop {
851 let ring = state.ring(self.which);
852 if self.cursor < ring.first_idx {
853 self.cursor = ring.first_idx;
854 }
855 let available = ring.first_idx + ring.pieces.len();
856 if self.cursor < available {
857 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
858 self.cursor += 1;
859 return ReadStep::Chunk(piece);
860 }
861 if state.ended {
862 return ReadStep::Eof;
863 }
864 let (next_state, timed_out) = self
865 .shared
866 .cond
867 .wait_timeout(state, timeout)
868 .expect("exec state mutex poisoned");
869 state = next_state;
870 if timed_out.timed_out() {
871 return ReadStep::Pending;
872 }
873 }
874 }
875}
876
877pub struct AsyncStreamReader {
881 shared: Arc<ExecShared>,
882 which: Stream,
883 cursor: usize,
884}
885
886impl AsyncStreamReader {
887 pub async fn next(&mut self) -> Option<String> {
891 loop {
892 let notified = self.shared.data_notify.notified();
896 tokio::pin!(notified);
897 notified.as_mut().enable();
898 {
899 let state = lock(&self.shared.state);
900 let ring = state.ring(self.which);
901 if self.cursor < ring.first_idx {
902 self.cursor = ring.first_idx;
903 }
904 let available = ring.first_idx + ring.pieces.len();
905 if self.cursor < available {
906 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
907 self.cursor += 1;
908 return Some(piece);
909 }
910 if state.ended {
911 return None;
912 }
913 }
914 notified.await;
915 }
916 }
917}
918
919async fn sleep_no_deadline(delay: f64) -> f64 {
922 let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
923 tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
924 (delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
925}
926
927fn is_exec_not_ready(status: &Status) -> bool {
928 status.code() == Code::Unavailable && status.message().contains("; retry")
929}
930
931fn terminal_status_error(status: i32) -> Option<SailError> {
934 if status == pb::SailboxExecStatus::WorkerLost as i32 {
935 return Some(SailError::WorkerLost {
936 message: "the machine hosting your sailbox was lost before the command \
937 finished; run exec again to retry"
938 .to_string(),
939 });
940 }
941 if status == pb::SailboxExecStatus::InterruptedRetryable as i32 {
942 return Some(SailError::ExecInterrupted {
943 retryable: true,
944 message: "the command was interrupted before it finished; run exec again to retry"
945 .to_string(),
946 });
947 }
948 if status == pb::SailboxExecStatus::InterruptedUnsafeToRetry as i32 {
949 return Some(SailError::ExecInterrupted {
950 retryable: false,
951 message: "the command was interrupted before it finished and may have partially \
952 run, so it was not retried automatically"
953 .to_string(),
954 });
955 }
956 if status == pb::SailboxExecStatus::Canceled as i32 {
957 return Some(SailError::ExecCanceled {
958 message: "the command was canceled before it finished".to_string(),
959 });
960 }
961 None
962}
963
964async fn submit(
967 worker: &Arc<WorkerProxy>,
968 params: &ExecParams,
969 stdout_resume_seq: i64,
970 stderr_resume_seq: i64,
971) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
972 let deadline = retry_deadline(params.retry_timeout);
973 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
974 loop {
975 let message = pb::StreamSailboxExecRequest {
976 sailbox_id: params.sailbox_id.clone(),
977 argv: params.argv.clone(),
978 timeout_seconds: params.timeout_seconds,
979 idempotency_key: params.idempotency_key.clone(),
980 open_stdin: params.open_stdin,
981 stdout_resume_seq,
982 stderr_resume_seq,
983 pty: params.pty,
984 term_cols: params.cols,
985 term_rows: params.rows,
986 term: params.term.clone(),
987 };
988 let request = worker.request_for(message, ¶ms.extra_metadata, None)?;
989 let status = match worker
990 .client_for(¶ms.exec_endpoint)?
991 .stream_sailbox_exec(request)
992 .await
993 {
994 Ok(resp) => {
995 let mut stream = resp.into_inner();
996 match stream.message().await {
997 Ok(Some(first)) => match first.frame {
998 Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
999 return Ok((started.exec_request_id, stream));
1000 }
1001 _ => {
1002 return Err(SailError::Execution {
1003 code: crate::error::GrpcCode::Internal,
1004 detail: "exec stream opened with a non-started frame".to_string(),
1005 });
1006 }
1007 },
1008 Ok(None) => {
1009 return Err(SailError::Execution {
1010 code: crate::error::GrpcCode::Internal,
1011 detail: "exec stream ended before the server confirmed the launch"
1012 .to_string(),
1013 });
1014 }
1015 Err(status) => status,
1016 }
1017 }
1018 Err(status) => status,
1019 };
1020 if !should_retry_transient_exec_rpc(&status, deadline) {
1021 return Err(SailError::from_exec_status(&status));
1022 }
1023 tracing::warn!(
1024 code = ?status.code(),
1025 stdout_resume_seq,
1026 stderr_resume_seq,
1027 "reconnecting exec stream"
1028 );
1029 if should_invalidate_channel(&status) {
1030 worker.channels().invalidate(¶ms.exec_endpoint);
1031 }
1032 delay = sleep_before_retry(delay, deadline).await;
1033 }
1034}
1035
1036struct Pump {
1042 shared: Arc<ExecShared>,
1043 stdout_dec: Utf8Decoder,
1044 stderr_dec: Utf8Decoder,
1045 stdout_seq: i64,
1046 stderr_seq: i64,
1047}
1048
1049impl Pump {
1050 fn new(shared: Arc<ExecShared>) -> Pump {
1051 Pump {
1052 shared,
1053 stdout_dec: Utf8Decoder::default(),
1054 stderr_dec: Utf8Decoder::default(),
1055 stdout_seq: 0,
1056 stderr_seq: 0,
1057 }
1058 }
1059
1060 fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
1065 match frame.frame {
1066 Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
1067 let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
1068 let (decoder, which) = if is_stderr {
1069 self.stderr_seq = self.stderr_seq.max(chunk.seq);
1070 (&mut self.stderr_dec, Stream::Stderr)
1071 } else {
1072 self.stdout_seq = self.stdout_seq.max(chunk.seq);
1073 (&mut self.stdout_dec, Stream::Stdout)
1074 };
1075 let text = decoder.decode(&chunk.data);
1076 if !text.is_empty() {
1077 let mut state = lock(&self.shared.state);
1078 match which {
1079 Stream::Stdout => state.stdout.append(text),
1080 Stream::Stderr => state.stderr.append(text),
1081 }
1082 self.shared.cond.notify_all();
1083 self.shared.data_notify.notify_waiters();
1084 }
1085 false
1086 }
1087 Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
1088 *lock(&self.shared.exit) = Some(ExitInfo {
1089 status: exit.status,
1090 return_code: exit.return_code,
1091 timed_out: exit.timed_out,
1092 stdout_truncated: exit.stdout_truncated,
1093 stderr_truncated: exit.stderr_truncated,
1094 error_message: exit.error_message,
1095 stdout_seq: exit.stdout_seq,
1096 stderr_seq: exit.stderr_seq,
1097 });
1098 true
1099 }
1100 _ => false,
1101 }
1102 }
1103
1104 fn finalize(mut self) {
1107 let tail_out = self.stdout_dec.flush();
1108 let tail_err = self.stderr_dec.flush();
1109 {
1110 let mut state = lock(&self.shared.state);
1111 if !tail_out.is_empty() {
1112 state.stdout.append(tail_out);
1113 }
1114 if !tail_err.is_empty() {
1115 state.stderr.append(tail_err);
1116 }
1117 state.ended = true;
1118 self.shared.cond.notify_all();
1119 self.shared.data_notify.notify_waiters();
1120 }
1121 *lock(&self.shared.high_seq) = (self.stdout_seq, self.stderr_seq);
1122 self.shared.ended.store(true, Ordering::SeqCst);
1123 self.shared.ended_notify.notify_waiters();
1124 }
1125}
1126
1127async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
1130 let mut state = Pump::new(shared.clone());
1131 loop {
1132 if shared.closing.load(Ordering::SeqCst) {
1133 break;
1134 }
1135 let message = tokio::select! {
1136 biased;
1137 () = shared.close_notify.notified() => break,
1138 message = stream.message() => message,
1139 };
1140 match message {
1141 Ok(Some(frame)) => {
1142 if state.apply_frame(frame) {
1143 break;
1144 }
1145 }
1146 Ok(None) => break,
1148 Err(_status) => {
1149 if shared.closing.load(Ordering::SeqCst) {
1150 break;
1151 }
1152 match submit(
1155 &shared.worker,
1156 &shared.params,
1157 state.stdout_seq,
1158 state.stderr_seq,
1159 )
1160 .await
1161 {
1162 Ok((_id, fresh)) => {
1163 stream = fresh;
1164 continue;
1165 }
1166 Err(_) => break,
1167 }
1168 }
1169 }
1170 }
1171 state.finalize();
1172}
1173
1174#[cfg(any(test, feature = "fuzzing"))]
1180pub fn fuzz_decode_ring(data: &[u8]) {
1181 let mut whole = Utf8Decoder::default();
1183 let mut whole_out = whole.decode(data);
1184 whole_out.push_str(&whole.flush());
1185
1186 let mut chunked = Utf8Decoder::default();
1189 let mut chunked_out = String::new();
1190 let mut start = 0;
1191 for (i, byte) in data.iter().enumerate() {
1192 if byte & 1 == 1 {
1193 chunked_out.push_str(&chunked.decode(&data[start..=i]));
1194 start = i + 1;
1195 }
1196 }
1197 chunked_out.push_str(&chunked.decode(&data[start..]));
1198 chunked_out.push_str(&chunked.flush());
1199
1200 assert_eq!(
1202 chunked_out, whole_out,
1203 "decoder is not chunk-boundary invariant"
1204 );
1205 assert_eq!(
1206 whole_out,
1207 String::from_utf8_lossy(data),
1208 "decoder disagrees with a lossy decode of the whole input"
1209 );
1210
1211 let mut ring = Ring::default();
1216 ring.append(whole_out.clone());
1217 let filler = STREAM_BUFFER_CAP_BYTES + 1 - whole_out.len().min(STREAM_BUFFER_CAP_BYTES);
1218 let filler = "a".repeat(filler);
1219 ring.append(filler.clone());
1220 assert!(
1221 ring.size <= STREAM_BUFFER_CAP_BYTES,
1222 "ring exceeded its cap"
1223 );
1224 let full = format!("{whole_out}{filler}");
1225 let tail = ring.tail();
1226 assert!(
1227 full.as_bytes().ends_with(tail.as_bytes()),
1228 "ring tail is not a suffix of the appended bytes"
1229 );
1230}
1231
1232#[cfg(test)]
1233mod tests {
1234 use super::*;
1235
1236 #[test]
1237 fn fuzz_decode_ring_holds_on_samples() {
1238 for sample in [
1241 &b""[..],
1242 b"hello",
1243 b"\xff\xfe\xfd",
1244 "€ µ é".as_bytes(),
1245 b"ab\xc3\xa9cd",
1246 &[0xC3, 0x28],
1247 ] {
1248 fuzz_decode_ring(sample);
1249 }
1250 }
1251
1252 #[test]
1253 fn ring_drops_oldest_past_cap_and_latches() {
1254 let mut ring = Ring::default();
1255 ring.append("a".repeat(STREAM_BUFFER_CAP_BYTES));
1256 assert!(!ring.dropped);
1257 ring.append("bbbb".to_string());
1258 assert!(ring.dropped);
1259 assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
1260 assert!(ring.tail().ends_with("bbbb"));
1261 }
1262
1263 #[test]
1264 fn ring_clip_on_split_multibyte_char_terminates() {
1265 let mut ring = Ring::default();
1270 ring.append(format!("é{}", "a".repeat(STREAM_BUFFER_CAP_BYTES - 1)));
1271 assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1272 assert!(ring.dropped);
1273 let tail = ring.tail();
1274 assert!(!tail.contains('é'));
1275 assert!(!tail.contains('\u{FFFD}'));
1276 }
1277
1278 #[test]
1279 fn decoder_carries_split_multibyte_char() {
1280 let mut dec = Utf8Decoder::default();
1281 let euro = "€".as_bytes(); assert_eq!(dec.decode(&euro[..2]), "");
1283 assert_eq!(dec.decode(&euro[2..]), "€");
1284 }
1285
1286 #[test]
1287 fn decoder_replaces_invalid_bytes() {
1288 let mut dec = Utf8Decoder::default();
1289 assert_eq!(dec.decode(&[0xff, b'x']), "\u{FFFD}x");
1290 }
1291
1292 #[test]
1293 fn terminal_status_mapping() {
1294 assert!(matches!(
1295 terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
1296 Some(SailError::WorkerLost { .. })
1297 ));
1298 assert!(matches!(
1299 terminal_status_error(pb::SailboxExecStatus::Canceled as i32),
1300 Some(SailError::ExecCanceled { .. })
1301 ));
1302 assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
1303 }
1304
1305 fn test_shared() -> Arc<ExecShared> {
1306 Arc::new(ExecShared {
1307 worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
1308 params: ExecParams {
1309 sailbox_id: "sb".into(),
1310 exec_endpoint: "endpoint".into(),
1311 argv: vec!["echo".into()],
1312 timeout_seconds: 0,
1313 idempotency_key: "idem".into(),
1314 open_stdin: false,
1315 pty: false,
1316 term: String::new(),
1317 cols: 0,
1318 rows: 0,
1319 retry_timeout: 0.0,
1320 extra_metadata: vec![],
1321 },
1322 state: Mutex::new(State::default()),
1323 cond: Condvar::new(),
1324 data_notify: Notify::new(),
1325 exit: Mutex::new(None),
1326 high_seq: Mutex::new((0, 0)),
1327 stdin: AsyncMutex::new(StdinState::default()),
1328 ended: AtomicBool::new(false),
1329 ended_notify: Notify::new(),
1330 closing: AtomicBool::new(false),
1331 close_notify: Notify::new(),
1332 })
1333 }
1334
1335 fn chunk(which: Stream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
1336 let stream = match which {
1337 Stream::Stdout => pb::SailboxExecStream::Stdout,
1338 Stream::Stderr => pb::SailboxExecStream::Stderr,
1339 };
1340 pb::StreamSailboxExecResponse {
1341 frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
1342 pb::SailboxExecChunk {
1343 stream: stream as i32,
1344 data: data.to_vec(),
1345 seq,
1346 },
1347 )),
1348 }
1349 }
1350
1351 fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
1352 pb::StreamSailboxExecResponse {
1353 frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
1354 pb::SailboxExecExit {
1355 status: status as i32,
1356 return_code: 0,
1357 timed_out: false,
1358 stdout_truncated: false,
1359 stderr_truncated: false,
1360 error_message: String::new(),
1361 stdout_seq,
1362 stderr_seq: 0,
1363 },
1364 )),
1365 }
1366 }
1367
1368 #[tokio::test]
1373 async fn exec_resume_carries_utf8_and_tracks_seq_across_break() {
1374 let shared = test_shared();
1375 let mut pump = Pump::new(shared.clone());
1376
1377 assert!(!pump.apply_frame(chunk(Stream::Stdout, 1, b"ab\xc3")));
1379 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "ab");
1380 assert_eq!(pump.stdout_seq, 1);
1382
1383 assert!(!pump.apply_frame(chunk(Stream::Stdout, 2, b"\xa9cd")));
1386 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "abécd");
1387 assert_eq!(pump.stdout_seq, 2);
1388
1389 assert!(!pump.apply_frame(chunk(Stream::Stdout, 1, b"!")));
1391 assert_eq!(pump.stdout_seq, 2);
1392
1393 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
1394 pump.finalize();
1395 assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
1396 assert!(shared.ended.load(Ordering::SeqCst));
1397 }
1398
1399 #[tokio::test]
1403 async fn exec_reader_follows_replayed_then_live() {
1404 let shared = test_shared();
1405 let mut reader = StreamReader {
1406 shared: shared.clone(),
1407 which: Stream::Stdout,
1408 cursor: 0,
1409 };
1410 let collector = tokio::task::spawn_blocking(move || {
1411 let mut out = String::new();
1412 loop {
1413 match reader.next(Duration::from_millis(50)) {
1414 ReadStep::Chunk(piece) => out.push_str(&piece),
1415 ReadStep::Eof => return out,
1416 ReadStep::Pending => {}
1417 }
1418 }
1419 });
1420
1421 let mut pump = Pump::new(shared.clone());
1422 pump.apply_frame(chunk(Stream::Stdout, 1, b"hello "));
1423 tokio::time::sleep(Duration::from_millis(10)).await;
1424 pump.apply_frame(chunk(Stream::Stdout, 2, b"world"));
1426 tokio::time::sleep(Duration::from_millis(10)).await;
1427 pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
1428 pump.finalize();
1429
1430 assert_eq!(collector.await.unwrap(), "hello world");
1431 }
1432
1433 #[test]
1436 fn exec_reader_started_late_replays_retained_tail() {
1437 let shared = test_shared();
1438 let mut pump = Pump::new(shared.clone());
1439 pump.apply_frame(chunk(Stream::Stdout, 1, b"early "));
1440 pump.apply_frame(chunk(Stream::Stdout, 2, b"output"));
1441 pump.finalize();
1442
1443 let mut reader = shared_reader(&shared, Stream::Stdout);
1445 let mut out = String::new();
1446 loop {
1447 match reader.next(Duration::from_millis(50)) {
1448 ReadStep::Chunk(piece) => out.push_str(&piece),
1449 ReadStep::Eof => break,
1450 ReadStep::Pending => panic!("ended ring should not return Pending"),
1451 }
1452 }
1453 assert_eq!(out, "early output");
1454 }
1455
1456 fn shared_reader(shared: &Arc<ExecShared>, which: Stream) -> StreamReader {
1457 StreamReader {
1458 shared: shared.clone(),
1459 which,
1460 cursor: 0,
1461 }
1462 }
1463
1464 use proptest::prelude::*;
1465
1466 fn utf8ish_bytes() -> impl Strategy<Value = Vec<u8>> {
1470 prop_oneof![
1471 proptest::collection::vec(any::<u8>(), 0..256),
1472 any::<String>().prop_map(String::into_bytes),
1473 ]
1474 }
1475
1476 proptest! {
1477 #[test]
1480 fn ring_without_overflow_is_lossless(
1481 pieces in proptest::collection::vec(any::<String>(), 0..32)
1482 ) {
1483 let concat: String = pieces.concat();
1484 prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
1485 let mut ring = Ring::default();
1486 for piece in &pieces {
1487 ring.append(piece.clone());
1488 }
1489 prop_assert!(!ring.dropped);
1490 prop_assert_eq!(ring.size, concat.len());
1491 prop_assert_eq!(ring.tail(), concat);
1492 }
1493 }
1494
1495 proptest! {
1496 #![proptest_config(ProptestConfig::with_cases(48))]
1498
1499 #[test]
1504 fn ring_eviction_keeps_valid_suffix_under_cap(
1505 overflow in 1usize..8192,
1506 tail_pieces in proptest::collection::vec("[a-z]{0,64}", 0..8),
1507 ) {
1508 let head = "é".repeat(STREAM_BUFFER_CAP_BYTES / 2 + overflow);
1511 let mut full = head.clone();
1512 let mut ring = Ring::default();
1513 ring.append(head);
1514 for piece in &tail_pieces {
1515 ring.append(piece.clone());
1516 full.push_str(piece);
1517 }
1518 prop_assert!(ring.dropped);
1519 prop_assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1520 let tail = ring.tail();
1521 let has_replacement_char = tail.contains('\u{FFFD}');
1524 prop_assert!(!has_replacement_char);
1525 prop_assert!(full.as_bytes().ends_with(tail.as_bytes()));
1526 }
1527 }
1528
1529 proptest! {
1530 #[test]
1534 fn decoder_is_chunk_boundary_invariant(
1535 data in utf8ish_bytes(),
1536 cuts in proptest::collection::vec(any::<prop::sample::Index>(), 0..16),
1537 ) {
1538 let mut whole = Utf8Decoder::default();
1539 let mut whole_out = whole.decode(&data);
1540 whole_out.push_str(&whole.flush());
1541
1542 let mut bounds: Vec<usize> = cuts.iter().map(|c| c.index(data.len() + 1)).collect();
1543 bounds.sort_unstable();
1544 let mut dec = Utf8Decoder::default();
1545 let mut chunked_out = String::new();
1546 let mut prev = 0;
1547 for bound in bounds.into_iter().chain(std::iter::once(data.len())) {
1548 let bound = bound.clamp(prev, data.len());
1549 chunked_out.push_str(&dec.decode(&data[prev..bound]));
1550 prev = bound;
1551 }
1552 chunked_out.push_str(&dec.flush());
1553
1554 prop_assert_eq!(&chunked_out, &whole_out);
1555 prop_assert_eq!(whole_out, String::from_utf8_lossy(&data).into_owned());
1556 }
1557 }
1558
1559 proptest! {
1560 #[test]
1565 fn pump_seq_is_monotonic_running_max_per_stream(
1566 frames in proptest::collection::vec(
1567 (any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
1568 0..64,
1569 )
1570 ) {
1571 let mut pump = Pump::new(test_shared());
1572 let (mut expect_out, mut expect_err) = (0i64, 0i64);
1573 let (mut prev_out, mut prev_err) = (0i64, 0i64);
1574 for (is_stderr, seq, data) in frames {
1575 let which = if is_stderr { Stream::Stderr } else { Stream::Stdout };
1576 pump.apply_frame(chunk(which, seq, &data));
1577 if is_stderr {
1578 expect_err = expect_err.max(seq);
1579 } else {
1580 expect_out = expect_out.max(seq);
1581 }
1582 prop_assert_eq!(pump.stdout_seq, expect_out);
1583 prop_assert_eq!(pump.stderr_seq, expect_err);
1584 prop_assert!(pump.stdout_seq >= prev_out);
1585 prop_assert!(pump.stderr_seq >= prev_err);
1586 prev_out = pump.stdout_seq;
1587 prev_err = pump.stderr_seq;
1588 }
1589 }
1590 }
1591}