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;
36const STREAM_BUFFER_CAP_BYTES: usize = 1024 * 1024;
40const STDIN_WRITE_CHUNK_BYTES: usize = 256 * 1024;
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum Stream {
47 Stdout,
49 Stderr,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum ReadStep {
56 Chunk(String),
58 Eof,
60 Pending,
63}
64
65#[derive(Debug, Clone, Serialize)]
67#[non_exhaustive]
68#[allow(clippy::struct_excessive_bools)]
69pub struct ExecResult {
70 pub stdout: String,
72 pub stderr: String,
74 pub return_code: i32,
76 pub timed_out: bool,
78 pub stdout_truncated: bool,
81 pub stderr_truncated: bool,
84 pub stdout_complete: bool,
90 pub stderr_complete: bool,
93}
94
95#[derive(Debug, Clone)]
100pub struct ExecOptions {
101 pub timeout_seconds: u32,
104 pub open_stdin: bool,
106 pub pty: bool,
108 pub term: String,
110 pub cols: u32,
112 pub rows: u32,
114 pub idempotency_key: String,
117 pub retry_timeout: f64,
119}
120
121impl Default for ExecOptions {
122 fn default() -> ExecOptions {
123 ExecOptions {
124 timeout_seconds: 0,
125 open_stdin: false,
126 pty: false,
127 term: String::new(),
128 cols: 0,
129 rows: 0,
130 idempotency_key: String::new(),
131 retry_timeout: EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
132 }
133 }
134}
135
136#[doc(hidden)]
138#[derive(Debug, Clone)]
139pub struct ExecParams {
140 pub sailbox_id: String,
142 pub exec_endpoint: String,
144 pub argv: Vec<String>,
146 pub timeout_seconds: u32,
149 pub idempotency_key: String,
152 pub open_stdin: bool,
154 pub pty: bool,
156 pub term: String,
158 pub cols: u32,
160 pub rows: u32,
162 pub retry_timeout: f64,
165 pub extra_metadata: Vec<(String, String)>,
167}
168
169#[derive(Default)]
173struct Utf8Decoder {
174 pending: Vec<u8>,
175}
176
177impl Utf8Decoder {
178 fn decode(&mut self, data: &[u8]) -> String {
179 self.pending.extend_from_slice(data);
180 let mut out = String::new();
181 loop {
182 match std::str::from_utf8(&self.pending) {
183 Ok(valid) => {
184 out.push_str(valid);
185 self.pending.clear();
186 break;
187 }
188 Err(err) => {
189 let valid_up_to = err.valid_up_to();
190 out.push_str(std::str::from_utf8(&self.pending[..valid_up_to]).unwrap());
192 if let Some(bad) = err.error_len() {
193 out.push('\u{FFFD}');
194 self.pending.drain(..valid_up_to + bad);
195 } else {
196 self.pending.drain(..valid_up_to);
198 break;
199 }
200 }
201 }
202 }
203 out
204 }
205
206 fn flush(&mut self) -> String {
207 if self.pending.is_empty() {
208 return String::new();
209 }
210 let out = String::from_utf8_lossy(&self.pending).into_owned();
211 self.pending.clear();
212 out
213 }
214}
215
216#[derive(Default)]
221struct Ring {
222 pieces: Vec<String>,
223 piece_bytes: Vec<usize>,
224 first_idx: usize,
225 size: usize,
226 dropped: bool,
227}
228
229impl Ring {
230 fn append(&mut self, text: String) {
231 let bytes = text.len();
232 self.pieces.push(text);
233 self.piece_bytes.push(bytes);
234 self.size += bytes;
235 while self.size > STREAM_BUFFER_CAP_BYTES {
236 let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
237 if self.piece_bytes[0] <= overflow {
238 self.pieces.remove(0);
239 self.size -= self.piece_bytes.remove(0);
240 self.first_idx += 1;
241 } else {
242 let piece = &self.pieces[0];
248 let mut cut = overflow;
249 while cut < piece.len() && !piece.is_char_boundary(cut) {
250 cut += 1;
251 }
252 let kept = piece[cut..].to_string();
253 self.size -= self.piece_bytes[0];
254 self.piece_bytes[0] = kept.len();
255 self.size += self.piece_bytes[0];
256 self.pieces[0] = kept;
257 }
258 self.dropped = true;
259 }
260 }
261
262 fn tail(&self) -> String {
263 self.pieces.concat()
264 }
265}
266
267#[derive(Default)]
268struct State {
269 stdout: Ring,
270 stderr: Ring,
271 ended: bool,
272}
273
274impl State {
275 fn ring(&self, which: Stream) -> &Ring {
276 match which {
277 Stream::Stdout => &self.stdout,
278 Stream::Stderr => &self.stderr,
279 }
280 }
281}
282
283#[derive(Clone)]
285struct ExitInfo {
286 status: i32,
287 return_code: i32,
288 timed_out: bool,
289 stdout_truncated: bool,
290 stderr_truncated: bool,
291 error_message: String,
292 stdout_seq: i64,
293 stderr_seq: i64,
294}
295
296#[derive(Default)]
297struct StdinState {
298 offset: i64,
299 eof_sent: bool,
300 broken: bool,
301 write_in_flight: bool,
309}
310
311struct ExecShared {
312 worker: Arc<WorkerProxy>,
313 params: ExecParams,
314 state: Mutex<State>,
315 cond: Condvar,
318 data_notify: Notify,
321 exit: Mutex<Option<ExitInfo>>,
322 high_seq: Mutex<(i64, i64)>,
324 stdin: AsyncMutex<StdinState>,
325 ended: AtomicBool,
326 ended_notify: Notify,
327 closing: AtomicBool,
328 close_notify: Notify,
329}
330
331pub struct ExecProcess {
334 shared: Arc<ExecShared>,
335 exec_request_id: String,
336}
337
338impl Drop for ExecProcess {
339 fn drop(&mut self) {
340 self.close();
345 }
346}
347
348impl ExecProcess {
349 #[doc(hidden)]
357 pub async fn start(
358 worker: Arc<WorkerProxy>,
359 mut params: ExecParams,
360 ) -> Result<ExecProcess, SailError> {
361 let key = params.idempotency_key.trim();
365 params.idempotency_key = if key.is_empty() {
366 format!("exec_{}", uuid::Uuid::new_v4())
367 } else {
368 key.to_string()
369 };
370 let (exec_request_id, stream) = submit(&worker, ¶ms, 0, 0).await?;
371 let shared = Arc::new(ExecShared {
372 worker,
373 params,
374 state: Mutex::new(State::default()),
375 cond: Condvar::new(),
376 data_notify: Notify::new(),
377 exit: Mutex::new(None),
378 high_seq: Mutex::new((0, 0)),
379 stdin: AsyncMutex::new(StdinState::default()),
380 ended: AtomicBool::new(false),
381 ended_notify: Notify::new(),
382 closing: AtomicBool::new(false),
383 close_notify: Notify::new(),
384 });
385 let pump_shared = shared.clone();
386 tokio::spawn(async move { pump(pump_shared, stream).await });
387 Ok(ExecProcess {
388 shared,
389 exec_request_id,
390 })
391 }
392
393 pub fn exec_request_id(&self) -> &str {
396 &self.exec_request_id
397 }
398
399 pub fn idempotency_key(&self) -> &str {
402 &self.shared.params.idempotency_key
403 }
404
405 pub fn reader(&self, which: Stream) -> StreamReader {
408 StreamReader {
409 shared: self.shared.clone(),
410 which,
411 cursor: 0,
412 }
413 }
414
415 pub fn reader_async(&self, which: Stream) -> AsyncStreamReader {
418 AsyncStreamReader {
419 shared: self.shared.clone(),
420 which,
421 cursor: 0,
422 }
423 }
424
425 pub fn poll(&self) -> Option<Result<i32, SailError>> {
429 let exit = self.shared.exit.lock().unwrap();
430 exit.as_ref()
431 .map(|exit| match terminal_status_error(exit.status) {
432 Some(err) => Err(err),
433 None => Ok(exit.return_code),
434 })
435 }
436
437 pub fn close(&self) {
439 self.shared.closing.store(true, Ordering::SeqCst);
440 self.shared.close_notify.notify_one();
445 }
446
447 pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
452 let _ = tokio::time::timeout(timeout, self.await_ended()).await;
453 self.shared.ended.load(Ordering::SeqCst)
454 }
455
456 async fn await_ended(&self) {
458 loop {
459 let notified = self.shared.ended_notify.notified();
460 if self.shared.ended.load(Ordering::SeqCst) {
461 return;
462 }
463 notified.await;
464 }
465 }
466
467 pub async fn wait(&self, retry_timeout: f64) -> Result<ExecResult, SailError> {
473 self.await_ended().await;
474 let exit = self.shared.exit.lock().unwrap().clone();
475 let (high_out, high_err) = *self.shared.high_seq.lock().unwrap();
476 let (out_dropped, err_dropped) = {
477 let state = self.shared.state.lock().unwrap();
478 (state.stdout.dropped, state.stderr.dropped)
479 };
480
481 let truncated = exit.as_ref().is_some_and(|exit| {
484 exit.stdout_truncated || exit.stderr_truncated || out_dropped || err_dropped
485 });
486 let incomplete = exit
487 .as_ref()
488 .is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
489
490 let stdout_complete = exit
495 .as_ref()
496 .is_some_and(|exit| exit.stdout_seq <= high_out);
497 let stderr_complete = exit
498 .as_ref()
499 .is_some_and(|exit| exit.stderr_seq <= high_err);
500
501 if exit.is_none() || truncated || incomplete {
502 let outcome = self
503 .shared
504 .worker
505 .wait_exec(
506 &self.shared.params.exec_endpoint,
507 &self.shared.params.sailbox_id,
508 &self.exec_request_id,
509 retry_timeout,
510 )
511 .await?;
512 {
515 let mut exit_slot = self.shared.exit.lock().unwrap();
516 if exit_slot.is_none() {
517 *exit_slot = Some(ExitInfo {
518 status: outcome.status,
519 return_code: outcome.return_code,
520 timed_out: outcome.timed_out,
521 stdout_truncated: outcome.stdout_truncated,
522 stderr_truncated: outcome.stderr_truncated,
523 error_message: String::new(),
524 stdout_seq: 0,
525 stderr_seq: 0,
526 });
527 }
528 }
529 if let Some(err) = terminal_status_error(outcome.status) {
530 return Err(err);
531 }
532 return Ok(ExecResult {
533 stdout: outcome.stdout,
534 stderr: outcome.stderr,
535 return_code: outcome.return_code,
536 timed_out: outcome.timed_out,
537 stdout_truncated: outcome.stdout_truncated,
538 stderr_truncated: outcome.stderr_truncated,
539 stdout_complete,
540 stderr_complete,
541 });
542 }
543
544 let exit = exit.expect("exit present on the clean path");
545 if let Some(err) = terminal_status_error(exit.status) {
546 return Err(err);
547 }
548 let state = self.shared.state.lock().unwrap();
549 let mut stderr = state.stderr.tail();
550 if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
551 stderr = exit.error_message.clone();
553 }
554 Ok(ExecResult {
555 stdout: state.stdout.tail(),
556 stderr,
557 return_code: exit.return_code,
558 timed_out: exit.timed_out,
559 stdout_truncated: false,
560 stderr_truncated: false,
561 stdout_complete,
562 stderr_complete,
563 })
564 }
565
566 pub async fn cancel(&self, force: bool, retry_timeout: f64) -> Result<(), SailError> {
568 self.shared
569 .worker
570 .cancel_exec(
571 &self.shared.params.exec_endpoint,
572 &self.shared.params.sailbox_id,
573 &self.exec_request_id,
574 force,
575 retry_timeout,
576 )
577 .await
578 }
579
580 pub async fn resize(&self, cols: u32, rows: u32) {
584 let message = pb::ResizeSailboxExecRequest {
585 sailbox_id: self.shared.params.sailbox_id.clone(),
586 exec_request_id: self.exec_request_id.clone(),
587 cols,
588 rows,
589 };
590 let Ok(request) =
591 self.shared
592 .worker
593 .request_for(message, &[], Some(Duration::from_secs(5)))
594 else {
595 return;
596 };
597 if let Ok(mut client) = self
598 .shared
599 .worker
600 .client_for(&self.shared.params.exec_endpoint)
601 {
602 let _ = client.resize_sailbox_exec(request).await;
603 }
604 }
605
606 pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
610 let mut stdin = self.shared.stdin.lock().await;
611 if stdin.eof_sent {
612 return Err(SailError::BrokenPipe {
613 message: "stdin is closed".to_string(),
614 });
615 }
616 if stdin.broken {
617 return Err(SailError::BrokenPipe {
618 message: "an earlier stdin write failed".to_string(),
619 });
620 }
621 if stdin.write_in_flight {
622 stdin.broken = true;
626 return Err(SailError::BrokenPipe {
627 message: "an earlier stdin write was interrupted".to_string(),
628 });
629 }
630 if data.is_empty() {
631 return Ok(());
632 }
633 stdin.write_in_flight = true;
634 let result = self.send_stdin(&mut stdin, data, false).await;
635 stdin.write_in_flight = false;
638 result
639 }
640
641 pub async fn eof_stdin(&self) -> Result<(), SailError> {
643 let mut stdin = self.shared.stdin.lock().await;
644 if stdin.eof_sent {
645 return Ok(());
646 }
647 if stdin.broken || stdin.write_in_flight {
648 stdin.eof_sent = true;
651 return Ok(());
652 }
653 match self.send_stdin(&mut stdin, &[], true).await {
657 Ok(()) => {
658 stdin.eof_sent = true;
659 Ok(())
660 }
661 Err(SailError::BrokenPipe { .. }) => {
663 stdin.eof_sent = true;
664 Ok(())
665 }
666 Err(err) => Err(err),
667 }
668 }
669
670 async fn send_stdin(
674 &self,
675 stdin: &mut StdinState,
676 payload: &[u8],
677 eof: bool,
678 ) -> Result<(), SailError> {
679 let endpoint = &self.shared.params.exec_endpoint;
680 let sailbox_id = &self.shared.params.sailbox_id;
681 let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
682 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
683 let mut sent = 0usize;
684 loop {
685 let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
686 let chunk = &payload[sent..end];
687 let last = end >= payload.len();
688 let message = pb::WriteSailboxExecStdinRequest {
689 sailbox_id: sailbox_id.clone(),
690 exec_request_id: self.exec_request_id.clone(),
691 offset: stdin.offset,
692 data: chunk.to_vec(),
693 eof: eof && last,
694 };
695 let request = match self.shared.worker.request_for(
699 message,
700 &[],
701 Some(rpc_attempt_timeout(deadline)),
702 ) {
703 Ok(request) => request,
704 Err(err) => return Err(err),
705 };
706 let result = match self.shared.worker.client_for(endpoint) {
707 Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
708 Err(err) => return Err(err),
709 };
710 match result {
711 Ok(resp) => {
712 let accepted_through = resp.into_inner().accepted_through;
713 let accepted =
718 ((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
719 stdin.offset += accepted as i64;
720 sent += accepted;
721 if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
722 return Ok(());
723 }
724 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
726 if accepted > 0 {
727 delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
728 } else {
729 delay = sleep_no_deadline(delay).await;
731 }
732 }
733 Err(status) => {
734 if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
735 return Err(SailError::BrokenPipe {
737 message: status.message().to_string(),
738 });
739 }
740 if is_exec_not_ready(&status) {
741 delay = sleep_no_deadline(delay).await;
745 deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
746 continue;
747 }
748 if !should_retry_transient_exec_rpc(&status, deadline) {
749 stdin.broken = true;
752 return Err(SailError::from_exec_status(&status));
753 }
754 tracing::warn!(code = ?status.code(), "retrying exec stdin write");
755 if should_invalidate_channel(&status) {
756 self.shared.worker.channels().invalidate(endpoint);
757 }
758 delay = sleep_before_retry(delay, deadline).await;
759 }
760 }
761 }
762 }
763}
764
765pub struct StreamReader {
768 shared: Arc<ExecShared>,
769 which: Stream,
770 cursor: usize,
771}
772
773impl StreamReader {
774 pub fn next(&mut self, timeout: Duration) -> ReadStep {
778 let mut state = self.shared.state.lock().unwrap();
779 loop {
780 let ring = state.ring(self.which);
781 if self.cursor < ring.first_idx {
782 self.cursor = ring.first_idx;
783 }
784 let available = ring.first_idx + ring.pieces.len();
785 if self.cursor < available {
786 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
787 self.cursor += 1;
788 return ReadStep::Chunk(piece);
789 }
790 if state.ended {
791 return ReadStep::Eof;
792 }
793 let (next_state, timed_out) = self
794 .shared
795 .cond
796 .wait_timeout(state, timeout)
797 .expect("exec state mutex poisoned");
798 state = next_state;
799 if timed_out.timed_out() {
800 return ReadStep::Pending;
801 }
802 }
803 }
804}
805
806pub struct AsyncStreamReader {
810 shared: Arc<ExecShared>,
811 which: Stream,
812 cursor: usize,
813}
814
815impl AsyncStreamReader {
816 pub async fn next(&mut self) -> Option<String> {
820 loop {
821 let notified = self.shared.data_notify.notified();
825 tokio::pin!(notified);
826 notified.as_mut().enable();
827 {
828 let state = self.shared.state.lock().unwrap();
829 let ring = state.ring(self.which);
830 if self.cursor < ring.first_idx {
831 self.cursor = ring.first_idx;
832 }
833 let available = ring.first_idx + ring.pieces.len();
834 if self.cursor < available {
835 let piece = ring.pieces[self.cursor - ring.first_idx].clone();
836 self.cursor += 1;
837 return Some(piece);
838 }
839 if state.ended {
840 return None;
841 }
842 }
843 notified.await;
844 }
845 }
846}
847
848async fn sleep_no_deadline(delay: f64) -> f64 {
851 let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
852 tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
853 (delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
854}
855
856fn is_exec_not_ready(status: &Status) -> bool {
857 status.code() == Code::Unavailable && status.message().contains("; retry")
858}
859
860fn terminal_status_error(status: i32) -> Option<SailError> {
863 if status == pb::SailboxExecStatus::WorkerLost as i32 {
864 return Some(SailError::WorkerLost {
865 message: "the machine hosting your sailbox was lost before the command \
866 finished; run exec again to retry"
867 .to_string(),
868 });
869 }
870 if status == pb::SailboxExecStatus::InterruptedRetryable as i32 {
871 return Some(SailError::ExecInterrupted {
872 retryable: true,
873 message: "the command was interrupted before it finished; run exec again to retry"
874 .to_string(),
875 });
876 }
877 if status == pb::SailboxExecStatus::InterruptedUnsafeToRetry as i32 {
878 return Some(SailError::ExecInterrupted {
879 retryable: false,
880 message: "the command was interrupted before it finished and may have partially \
881 run, so it was not retried automatically"
882 .to_string(),
883 });
884 }
885 if status == pb::SailboxExecStatus::Canceled as i32 {
886 return Some(SailError::ExecCanceled {
887 message: "the command was canceled before it finished".to_string(),
888 });
889 }
890 None
891}
892
893async fn submit(
896 worker: &Arc<WorkerProxy>,
897 params: &ExecParams,
898 stdout_resume_seq: i64,
899 stderr_resume_seq: i64,
900) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
901 let deadline = retry_deadline(params.retry_timeout);
902 let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
903 loop {
904 let message = pb::StreamSailboxExecRequest {
905 sailbox_id: params.sailbox_id.clone(),
906 argv: params.argv.clone(),
907 timeout_seconds: params.timeout_seconds,
908 idempotency_key: params.idempotency_key.clone(),
909 open_stdin: params.open_stdin,
910 stdout_resume_seq,
911 stderr_resume_seq,
912 pty: params.pty,
913 term_cols: params.cols,
914 term_rows: params.rows,
915 term: params.term.clone(),
916 };
917 let request = worker.request_for(message, ¶ms.extra_metadata, None)?;
918 let status = match worker
919 .client_for(¶ms.exec_endpoint)?
920 .stream_sailbox_exec(request)
921 .await
922 {
923 Ok(resp) => {
924 let mut stream = resp.into_inner();
925 match stream.message().await {
926 Ok(Some(first)) => match first.frame {
927 Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
928 return Ok((started.exec_request_id, stream));
929 }
930 _ => {
931 return Err(SailError::Execution {
932 code: crate::error::GrpcCode::Internal,
933 detail: "exec stream opened with a non-started frame".to_string(),
934 });
935 }
936 },
937 Ok(None) => {
938 return Err(SailError::Execution {
939 code: crate::error::GrpcCode::Internal,
940 detail: "exec stream ended before the server confirmed the launch"
941 .to_string(),
942 });
943 }
944 Err(status) => status,
945 }
946 }
947 Err(status) => status,
948 };
949 if !should_retry_transient_exec_rpc(&status, deadline) {
950 return Err(SailError::from_exec_status(&status));
951 }
952 tracing::warn!(
953 code = ?status.code(),
954 stdout_resume_seq,
955 stderr_resume_seq,
956 "reconnecting exec stream"
957 );
958 if should_invalidate_channel(&status) {
959 worker.channels().invalidate(¶ms.exec_endpoint);
960 }
961 delay = sleep_before_retry(delay, deadline).await;
962 }
963}
964
965struct Pump {
971 shared: Arc<ExecShared>,
972 stdout_dec: Utf8Decoder,
973 stderr_dec: Utf8Decoder,
974 stdout_seq: i64,
975 stderr_seq: i64,
976}
977
978impl Pump {
979 fn new(shared: Arc<ExecShared>) -> Pump {
980 Pump {
981 shared,
982 stdout_dec: Utf8Decoder::default(),
983 stderr_dec: Utf8Decoder::default(),
984 stdout_seq: 0,
985 stderr_seq: 0,
986 }
987 }
988
989 fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
994 match frame.frame {
995 Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
996 let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
997 let (decoder, which) = if is_stderr {
998 self.stderr_seq = self.stderr_seq.max(chunk.seq);
999 (&mut self.stderr_dec, Stream::Stderr)
1000 } else {
1001 self.stdout_seq = self.stdout_seq.max(chunk.seq);
1002 (&mut self.stdout_dec, Stream::Stdout)
1003 };
1004 let text = decoder.decode(&chunk.data);
1005 if !text.is_empty() {
1006 let mut state = self.shared.state.lock().unwrap();
1007 match which {
1008 Stream::Stdout => state.stdout.append(text),
1009 Stream::Stderr => state.stderr.append(text),
1010 }
1011 self.shared.cond.notify_all();
1012 self.shared.data_notify.notify_waiters();
1013 }
1014 false
1015 }
1016 Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
1017 *self.shared.exit.lock().unwrap() = Some(ExitInfo {
1018 status: exit.status,
1019 return_code: exit.return_code,
1020 timed_out: exit.timed_out,
1021 stdout_truncated: exit.stdout_truncated,
1022 stderr_truncated: exit.stderr_truncated,
1023 error_message: exit.error_message,
1024 stdout_seq: exit.stdout_seq,
1025 stderr_seq: exit.stderr_seq,
1026 });
1027 true
1028 }
1029 _ => false,
1030 }
1031 }
1032
1033 fn finalize(mut self) {
1036 let tail_out = self.stdout_dec.flush();
1037 let tail_err = self.stderr_dec.flush();
1038 {
1039 let mut state = self.shared.state.lock().unwrap();
1040 if !tail_out.is_empty() {
1041 state.stdout.append(tail_out);
1042 }
1043 if !tail_err.is_empty() {
1044 state.stderr.append(tail_err);
1045 }
1046 state.ended = true;
1047 self.shared.cond.notify_all();
1048 self.shared.data_notify.notify_waiters();
1049 }
1050 *self.shared.high_seq.lock().unwrap() = (self.stdout_seq, self.stderr_seq);
1051 self.shared.ended.store(true, Ordering::SeqCst);
1052 self.shared.ended_notify.notify_waiters();
1053 }
1054}
1055
1056async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
1059 let mut state = Pump::new(shared.clone());
1060 loop {
1061 if shared.closing.load(Ordering::SeqCst) {
1062 break;
1063 }
1064 let message = tokio::select! {
1065 biased;
1066 () = shared.close_notify.notified() => break,
1067 message = stream.message() => message,
1068 };
1069 match message {
1070 Ok(Some(frame)) => {
1071 if state.apply_frame(frame) {
1072 break;
1073 }
1074 }
1075 Ok(None) => break,
1077 Err(_status) => {
1078 if shared.closing.load(Ordering::SeqCst) {
1079 break;
1080 }
1081 match submit(
1084 &shared.worker,
1085 &shared.params,
1086 state.stdout_seq,
1087 state.stderr_seq,
1088 )
1089 .await
1090 {
1091 Ok((_id, fresh)) => {
1092 stream = fresh;
1093 continue;
1094 }
1095 Err(_) => break,
1096 }
1097 }
1098 }
1099 }
1100 state.finalize();
1101}
1102
1103#[cfg(any(test, feature = "fuzzing"))]
1109pub fn fuzz_decode_ring(data: &[u8]) {
1110 let mut whole = Utf8Decoder::default();
1112 let mut whole_out = whole.decode(data);
1113 whole_out.push_str(&whole.flush());
1114
1115 let mut chunked = Utf8Decoder::default();
1118 let mut chunked_out = String::new();
1119 let mut start = 0;
1120 for (i, byte) in data.iter().enumerate() {
1121 if byte & 1 == 1 {
1122 chunked_out.push_str(&chunked.decode(&data[start..=i]));
1123 start = i + 1;
1124 }
1125 }
1126 chunked_out.push_str(&chunked.decode(&data[start..]));
1127 chunked_out.push_str(&chunked.flush());
1128
1129 assert_eq!(
1131 chunked_out, whole_out,
1132 "decoder is not chunk-boundary invariant"
1133 );
1134 assert_eq!(
1135 whole_out,
1136 String::from_utf8_lossy(data),
1137 "decoder disagrees with a lossy decode of the whole input"
1138 );
1139
1140 let mut ring = Ring::default();
1145 ring.append(whole_out.clone());
1146 let filler = STREAM_BUFFER_CAP_BYTES + 1 - whole_out.len().min(STREAM_BUFFER_CAP_BYTES);
1147 let filler = "a".repeat(filler);
1148 ring.append(filler.clone());
1149 assert!(
1150 ring.size <= STREAM_BUFFER_CAP_BYTES,
1151 "ring exceeded its cap"
1152 );
1153 let full = format!("{whole_out}{filler}");
1154 let tail = ring.tail();
1155 assert!(
1156 full.as_bytes().ends_with(tail.as_bytes()),
1157 "ring tail is not a suffix of the appended bytes"
1158 );
1159}
1160
1161#[cfg(test)]
1162mod tests {
1163 use super::*;
1164
1165 #[test]
1166 fn fuzz_decode_ring_holds_on_samples() {
1167 for sample in [
1170 &b""[..],
1171 b"hello",
1172 b"\xff\xfe\xfd",
1173 "€ µ é".as_bytes(),
1174 b"ab\xc3\xa9cd",
1175 &[0xC3, 0x28],
1176 ] {
1177 fuzz_decode_ring(sample);
1178 }
1179 }
1180
1181 #[test]
1182 fn ring_drops_oldest_past_cap_and_latches() {
1183 let mut ring = Ring::default();
1184 ring.append("a".repeat(STREAM_BUFFER_CAP_BYTES));
1185 assert!(!ring.dropped);
1186 ring.append("bbbb".to_string());
1187 assert!(ring.dropped);
1188 assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
1189 assert!(ring.tail().ends_with("bbbb"));
1190 }
1191
1192 #[test]
1193 fn ring_clip_on_split_multibyte_char_terminates() {
1194 let mut ring = Ring::default();
1199 ring.append(format!("é{}", "a".repeat(STREAM_BUFFER_CAP_BYTES - 1)));
1200 assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1201 assert!(ring.dropped);
1202 let tail = ring.tail();
1203 assert!(!tail.contains('é'));
1204 assert!(!tail.contains('\u{FFFD}'));
1205 }
1206
1207 #[test]
1208 fn decoder_carries_split_multibyte_char() {
1209 let mut dec = Utf8Decoder::default();
1210 let euro = "€".as_bytes(); assert_eq!(dec.decode(&euro[..2]), "");
1212 assert_eq!(dec.decode(&euro[2..]), "€");
1213 }
1214
1215 #[test]
1216 fn decoder_replaces_invalid_bytes() {
1217 let mut dec = Utf8Decoder::default();
1218 assert_eq!(dec.decode(&[0xff, b'x']), "\u{FFFD}x");
1219 }
1220
1221 #[test]
1222 fn terminal_status_mapping() {
1223 assert!(matches!(
1224 terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
1225 Some(SailError::WorkerLost { .. })
1226 ));
1227 assert!(matches!(
1228 terminal_status_error(pb::SailboxExecStatus::Canceled as i32),
1229 Some(SailError::ExecCanceled { .. })
1230 ));
1231 assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
1232 }
1233
1234 fn test_shared() -> Arc<ExecShared> {
1235 Arc::new(ExecShared {
1236 worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
1237 params: ExecParams {
1238 sailbox_id: "sb".into(),
1239 exec_endpoint: "endpoint".into(),
1240 argv: vec!["echo".into()],
1241 timeout_seconds: 0,
1242 idempotency_key: "idem".into(),
1243 open_stdin: false,
1244 pty: false,
1245 term: String::new(),
1246 cols: 0,
1247 rows: 0,
1248 retry_timeout: 0.0,
1249 extra_metadata: vec![],
1250 },
1251 state: Mutex::new(State::default()),
1252 cond: Condvar::new(),
1253 data_notify: Notify::new(),
1254 exit: Mutex::new(None),
1255 high_seq: Mutex::new((0, 0)),
1256 stdin: AsyncMutex::new(StdinState::default()),
1257 ended: AtomicBool::new(false),
1258 ended_notify: Notify::new(),
1259 closing: AtomicBool::new(false),
1260 close_notify: Notify::new(),
1261 })
1262 }
1263
1264 fn chunk(which: Stream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
1265 let stream = match which {
1266 Stream::Stdout => pb::SailboxExecStream::Stdout,
1267 Stream::Stderr => pb::SailboxExecStream::Stderr,
1268 };
1269 pb::StreamSailboxExecResponse {
1270 frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
1271 pb::SailboxExecChunk {
1272 stream: stream as i32,
1273 data: data.to_vec(),
1274 seq,
1275 },
1276 )),
1277 }
1278 }
1279
1280 fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
1281 pb::StreamSailboxExecResponse {
1282 frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
1283 pb::SailboxExecExit {
1284 status: status as i32,
1285 return_code: 0,
1286 timed_out: false,
1287 stdout_truncated: false,
1288 stderr_truncated: false,
1289 error_message: String::new(),
1290 stdout_seq,
1291 stderr_seq: 0,
1292 },
1293 )),
1294 }
1295 }
1296
1297 #[tokio::test]
1302 async fn exec_resume_carries_utf8_and_tracks_seq_across_break() {
1303 let shared = test_shared();
1304 let mut pump = Pump::new(shared.clone());
1305
1306 assert!(!pump.apply_frame(chunk(Stream::Stdout, 1, b"ab\xc3")));
1308 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "ab");
1309 assert_eq!(pump.stdout_seq, 1);
1311
1312 assert!(!pump.apply_frame(chunk(Stream::Stdout, 2, b"\xa9cd")));
1315 assert_eq!(shared.state.lock().unwrap().stdout.tail(), "abécd");
1316 assert_eq!(pump.stdout_seq, 2);
1317
1318 assert!(!pump.apply_frame(chunk(Stream::Stdout, 1, b"!")));
1320 assert_eq!(pump.stdout_seq, 2);
1321
1322 assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
1323 pump.finalize();
1324 assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
1325 assert!(shared.ended.load(Ordering::SeqCst));
1326 }
1327
1328 #[tokio::test]
1332 async fn exec_reader_follows_replayed_then_live() {
1333 let shared = test_shared();
1334 let mut reader = StreamReader {
1335 shared: shared.clone(),
1336 which: Stream::Stdout,
1337 cursor: 0,
1338 };
1339 let collector = tokio::task::spawn_blocking(move || {
1340 let mut out = String::new();
1341 loop {
1342 match reader.next(Duration::from_millis(50)) {
1343 ReadStep::Chunk(piece) => out.push_str(&piece),
1344 ReadStep::Eof => return out,
1345 ReadStep::Pending => {}
1346 }
1347 }
1348 });
1349
1350 let mut pump = Pump::new(shared.clone());
1351 pump.apply_frame(chunk(Stream::Stdout, 1, b"hello "));
1352 tokio::time::sleep(Duration::from_millis(10)).await;
1353 pump.apply_frame(chunk(Stream::Stdout, 2, b"world"));
1355 tokio::time::sleep(Duration::from_millis(10)).await;
1356 pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
1357 pump.finalize();
1358
1359 assert_eq!(collector.await.unwrap(), "hello world");
1360 }
1361
1362 #[test]
1365 fn exec_reader_started_late_replays_retained_tail() {
1366 let shared = test_shared();
1367 let mut pump = Pump::new(shared.clone());
1368 pump.apply_frame(chunk(Stream::Stdout, 1, b"early "));
1369 pump.apply_frame(chunk(Stream::Stdout, 2, b"output"));
1370 pump.finalize();
1371
1372 let mut reader = shared_reader(&shared, Stream::Stdout);
1374 let mut out = String::new();
1375 loop {
1376 match reader.next(Duration::from_millis(50)) {
1377 ReadStep::Chunk(piece) => out.push_str(&piece),
1378 ReadStep::Eof => break,
1379 ReadStep::Pending => panic!("ended ring should not return Pending"),
1380 }
1381 }
1382 assert_eq!(out, "early output");
1383 }
1384
1385 fn shared_reader(shared: &Arc<ExecShared>, which: Stream) -> StreamReader {
1386 StreamReader {
1387 shared: shared.clone(),
1388 which,
1389 cursor: 0,
1390 }
1391 }
1392
1393 use proptest::prelude::*;
1394
1395 fn utf8ish_bytes() -> impl Strategy<Value = Vec<u8>> {
1399 prop_oneof![
1400 proptest::collection::vec(any::<u8>(), 0..256),
1401 any::<String>().prop_map(String::into_bytes),
1402 ]
1403 }
1404
1405 proptest! {
1406 #[test]
1409 fn ring_without_overflow_is_lossless(
1410 pieces in proptest::collection::vec(any::<String>(), 0..32)
1411 ) {
1412 let concat: String = pieces.concat();
1413 prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
1414 let mut ring = Ring::default();
1415 for piece in &pieces {
1416 ring.append(piece.clone());
1417 }
1418 prop_assert!(!ring.dropped);
1419 prop_assert_eq!(ring.size, concat.len());
1420 prop_assert_eq!(ring.tail(), concat);
1421 }
1422 }
1423
1424 proptest! {
1425 #![proptest_config(ProptestConfig::with_cases(48))]
1427
1428 #[test]
1433 fn ring_eviction_keeps_valid_suffix_under_cap(
1434 overflow in 1usize..8192,
1435 tail_pieces in proptest::collection::vec("[a-z]{0,64}", 0..8),
1436 ) {
1437 let head = "é".repeat(STREAM_BUFFER_CAP_BYTES / 2 + overflow);
1440 let mut full = head.clone();
1441 let mut ring = Ring::default();
1442 ring.append(head);
1443 for piece in &tail_pieces {
1444 ring.append(piece.clone());
1445 full.push_str(piece);
1446 }
1447 prop_assert!(ring.dropped);
1448 prop_assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
1449 let tail = ring.tail();
1450 let has_replacement_char = tail.contains('\u{FFFD}');
1453 prop_assert!(!has_replacement_char);
1454 prop_assert!(full.as_bytes().ends_with(tail.as_bytes()));
1455 }
1456 }
1457
1458 proptest! {
1459 #[test]
1463 fn decoder_is_chunk_boundary_invariant(
1464 data in utf8ish_bytes(),
1465 cuts in proptest::collection::vec(any::<prop::sample::Index>(), 0..16),
1466 ) {
1467 let mut whole = Utf8Decoder::default();
1468 let mut whole_out = whole.decode(&data);
1469 whole_out.push_str(&whole.flush());
1470
1471 let mut bounds: Vec<usize> = cuts.iter().map(|c| c.index(data.len() + 1)).collect();
1472 bounds.sort_unstable();
1473 let mut dec = Utf8Decoder::default();
1474 let mut chunked_out = String::new();
1475 let mut prev = 0;
1476 for bound in bounds.into_iter().chain(std::iter::once(data.len())) {
1477 let bound = bound.clamp(prev, data.len());
1478 chunked_out.push_str(&dec.decode(&data[prev..bound]));
1479 prev = bound;
1480 }
1481 chunked_out.push_str(&dec.flush());
1482
1483 prop_assert_eq!(&chunked_out, &whole_out);
1484 prop_assert_eq!(whole_out, String::from_utf8_lossy(&data).into_owned());
1485 }
1486 }
1487
1488 proptest! {
1489 #[test]
1494 fn pump_seq_is_monotonic_running_max_per_stream(
1495 frames in proptest::collection::vec(
1496 (any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
1497 0..64,
1498 )
1499 ) {
1500 let mut pump = Pump::new(test_shared());
1501 let (mut expect_out, mut expect_err) = (0i64, 0i64);
1502 let (mut prev_out, mut prev_err) = (0i64, 0i64);
1503 for (is_stderr, seq, data) in frames {
1504 let which = if is_stderr { Stream::Stderr } else { Stream::Stdout };
1505 pump.apply_frame(chunk(which, seq, &data));
1506 if is_stderr {
1507 expect_err = expect_err.max(seq);
1508 } else {
1509 expect_out = expect_out.max(seq);
1510 }
1511 prop_assert_eq!(pump.stdout_seq, expect_out);
1512 prop_assert_eq!(pump.stderr_seq, expect_err);
1513 prop_assert!(pump.stdout_seq >= prev_out);
1514 prop_assert!(pump.stderr_seq >= prev_err);
1515 prev_out = pump.stdout_seq;
1516 prev_err = pump.stderr_seq;
1517 }
1518 }
1519 }
1520}