use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::time::Duration;
use serde::Serialize;
use tokio::sync::Mutex as AsyncMutex;
use tokio::sync::Notify;
use tonic::{Code, Status, Streaming};
use crate::error::SailError;
use crate::pb::workerproxy::v1 as pb;
use crate::worker::{
retry_deadline, rpc_attempt_timeout, should_invalidate_channel,
should_retry_transient_exec_rpc, sleep_before_retry, WorkerProxy,
EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS,
};
pub const EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS: f64 = 30.0;
const STREAM_BUFFER_CAP_BYTES: usize = 1024 * 1024;
const STDIN_WRITE_CHUNK_BYTES: usize = 256 * 1024;
fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputStream {
Stdout,
Stderr,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReadStep {
Chunk(String),
Eof,
Pending,
}
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
#[allow(clippy::struct_excessive_bools)]
pub struct ExecResult {
pub stdout: String,
pub stderr: String,
pub exit_code: i32,
pub timed_out: bool,
pub stdout_truncated: bool,
pub stderr_truncated: bool,
pub stdout_complete: bool,
pub stderr_complete: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CancelSignal {
Interrupt,
Kill,
}
impl CancelSignal {
fn is_force(self) -> bool {
matches!(self, CancelSignal::Kill)
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RetryBudget {
None,
Within(Duration),
Forever,
}
impl RetryBudget {
pub fn as_secs_f64(self) -> f64 {
match self {
RetryBudget::None => 0.0,
RetryBudget::Within(d) => d.as_secs_f64(),
RetryBudget::Forever => f64::INFINITY,
}
}
pub fn from_secs_f64(secs: f64) -> RetryBudget {
if secs <= 0.0 {
RetryBudget::None
} else if secs.is_finite() {
RetryBudget::Within(Duration::from_secs_f64(secs))
} else {
RetryBudget::Forever
}
}
}
#[derive(Debug, Clone)]
pub struct ExecOptions {
pub timeout: Option<Duration>,
pub open_stdin: bool,
pub pty: bool,
pub term: String,
pub cols: u32,
pub rows: u32,
pub idempotency_key: String,
pub retry_timeout: RetryBudget,
pub cwd: Option<String>,
pub background: bool,
}
impl Default for ExecOptions {
fn default() -> ExecOptions {
ExecOptions {
timeout: None,
open_stdin: false,
pty: false,
term: String::new(),
cols: 0,
rows: 0,
idempotency_key: String::new(),
retry_timeout: RetryBudget::Within(Duration::from_secs_f64(
EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
)),
cwd: None,
background: false,
}
}
}
pub(crate) fn sh_quote(value: &str) -> String {
format!("'{}'", value.replace('\'', "'\\''"))
}
pub fn shell_argv(command: &str, options: &ExecOptions) -> Result<Vec<String>, SailError> {
let invalid = |message: &str| {
Err(SailError::InvalidArgument {
message: message.to_string(),
})
};
if command.is_empty() {
return invalid("command must be non-empty");
}
if options.background && (options.open_stdin || options.pty) {
return invalid("background is not supported with open_stdin or pty");
}
let mut command = command.to_string();
if let Some(cwd) = &options.cwd {
let cwd = cwd.trim();
if cwd.is_empty() {
return invalid("cwd must be non-empty");
}
command = format!(
"cd {} && exec /bin/sh -lc {}",
sh_quote(cwd),
sh_quote(&command)
);
}
if options.background {
command = format!(
"nohup /bin/sh -lc {} </dev/null >/dev/null 2>&1 &",
sh_quote(&command)
);
}
Ok(vec!["/bin/sh".to_string(), "-lc".to_string(), command])
}
#[doc(hidden)]
#[derive(Debug, Clone)]
pub struct ExecParams {
pub sailbox_id: String,
pub exec_endpoint: String,
pub argv: Vec<String>,
pub timeout_seconds: u32,
pub idempotency_key: String,
pub open_stdin: bool,
pub pty: bool,
pub term: String,
pub cols: u32,
pub rows: u32,
pub retry_timeout: f64,
pub extra_metadata: Vec<(String, String)>,
}
#[derive(Default)]
struct Utf8Decoder {
pending: Vec<u8>,
}
impl Utf8Decoder {
fn decode(&mut self, data: &[u8]) -> String {
self.pending.extend_from_slice(data);
let mut out = String::new();
loop {
match std::str::from_utf8(&self.pending) {
Ok(valid) => {
out.push_str(valid);
self.pending.clear();
break;
}
Err(err) => {
let valid_up_to = err.valid_up_to();
out.push_str(std::str::from_utf8(&self.pending[..valid_up_to]).unwrap());
if let Some(bad) = err.error_len() {
out.push('\u{FFFD}');
self.pending.drain(..valid_up_to + bad);
} else {
self.pending.drain(..valid_up_to);
break;
}
}
}
}
out
}
fn flush(&mut self) -> String {
if self.pending.is_empty() {
return String::new();
}
let out = String::from_utf8_lossy(&self.pending).into_owned();
self.pending.clear();
out
}
}
#[derive(Default)]
struct Ring {
pieces: Vec<String>,
piece_bytes: Vec<usize>,
first_idx: usize,
size: usize,
dropped: bool,
}
impl Ring {
fn append(&mut self, text: String) {
let bytes = text.len();
self.pieces.push(text);
self.piece_bytes.push(bytes);
self.size += bytes;
while self.size > STREAM_BUFFER_CAP_BYTES {
let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
if self.piece_bytes[0] <= overflow {
self.pieces.remove(0);
self.size -= self.piece_bytes.remove(0);
self.first_idx += 1;
} else {
let piece = &self.pieces[0];
let mut cut = overflow;
while cut < piece.len() && !piece.is_char_boundary(cut) {
cut += 1;
}
let kept = piece[cut..].to_string();
self.size -= self.piece_bytes[0];
self.piece_bytes[0] = kept.len();
self.size += self.piece_bytes[0];
self.pieces[0] = kept;
}
self.dropped = true;
}
}
fn tail(&self) -> String {
self.pieces.concat()
}
}
#[derive(Default)]
struct State {
stdout: Ring,
stderr: Ring,
ended: bool,
}
impl State {
fn ring(&self, which: OutputStream) -> &Ring {
match which {
OutputStream::Stdout => &self.stdout,
OutputStream::Stderr => &self.stderr,
}
}
}
#[derive(Clone)]
struct ExitInfo {
status: i32,
exit_code: i32,
timed_out: bool,
stdout_truncated: bool,
stderr_truncated: bool,
error_message: String,
stdout_seq: i64,
stderr_seq: i64,
}
#[derive(Default)]
struct StdinState {
offset: i64,
eof_sent: bool,
broken: bool,
write_in_flight: bool,
}
struct ExecShared {
worker: Arc<WorkerProxy>,
params: ExecParams,
state: Mutex<State>,
cond: Condvar,
data_notify: Notify,
exit: Mutex<Option<ExitInfo>>,
high_seq: Mutex<(i64, i64)>,
stdin: AsyncMutex<StdinState>,
ended: AtomicBool,
ended_notify: Notify,
closing: AtomicBool,
close_notify: Notify,
}
pub struct ExecProcess {
shared: Arc<ExecShared>,
exec_request_id: String,
}
impl Drop for ExecProcess {
fn drop(&mut self) {
self.close();
}
}
impl ExecProcess {
#[doc(hidden)]
pub async fn start(
worker: Arc<WorkerProxy>,
mut params: ExecParams,
) -> Result<ExecProcess, SailError> {
let key = params.idempotency_key.trim();
params.idempotency_key = if key.is_empty() {
format!("exec_{}", uuid::Uuid::new_v4())
} else {
key.to_string()
};
let (exec_request_id, stream) = submit(&worker, ¶ms, 0, 0).await?;
let shared = Arc::new(ExecShared {
worker,
params,
state: Mutex::new(State::default()),
cond: Condvar::new(),
data_notify: Notify::new(),
exit: Mutex::new(None),
high_seq: Mutex::new((0, 0)),
stdin: AsyncMutex::new(StdinState::default()),
ended: AtomicBool::new(false),
ended_notify: Notify::new(),
closing: AtomicBool::new(false),
close_notify: Notify::new(),
});
let pump_shared = shared.clone();
tokio::spawn(async move { pump(pump_shared, stream).await });
Ok(ExecProcess {
shared,
exec_request_id,
})
}
pub fn exec_request_id(&self) -> &str {
&self.exec_request_id
}
pub fn idempotency_key(&self) -> &str {
&self.shared.params.idempotency_key
}
pub fn reader(&self, which: OutputStream) -> StreamReader {
StreamReader {
shared: self.shared.clone(),
which,
cursor: 0,
}
}
pub fn reader_async(&self, which: OutputStream) -> AsyncStreamReader {
AsyncStreamReader {
shared: self.shared.clone(),
which,
cursor: 0,
}
}
pub fn poll(&self) -> Option<Result<i32, SailError>> {
let exit = lock(&self.shared.exit);
exit.as_ref()
.map(|exit| match terminal_status_error(exit.status) {
Some(err) => Err(err),
None => Ok(exit.exit_code),
})
}
pub fn close(&self) {
self.shared.closing.store(true, Ordering::SeqCst);
self.shared.close_notify.notify_one();
}
pub async fn wait_stream_ended(&self, timeout: Duration) -> bool {
let _ = tokio::time::timeout(timeout, self.await_ended()).await;
self.shared.ended.load(Ordering::SeqCst)
}
async fn await_ended(&self) {
loop {
let notified = self.shared.ended_notify.notified();
if self.shared.ended.load(Ordering::SeqCst) {
return;
}
notified.await;
}
}
pub async fn wait(&self) -> Result<ExecResult, SailError> {
self.await_ended().await;
let exit = lock(&self.shared.exit).clone();
let (high_out, high_err) = *lock(&self.shared.high_seq);
let (out_dropped, err_dropped) = {
let state = lock(&self.shared.state);
(state.stdout.dropped, state.stderr.dropped)
};
let truncated = exit.as_ref().is_some_and(|exit| {
exit.stdout_truncated || exit.stderr_truncated || out_dropped || err_dropped
});
let incomplete = exit
.as_ref()
.is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
let stdout_complete = exit
.as_ref()
.is_some_and(|exit| exit.stdout_seq <= high_out);
let stderr_complete = exit
.as_ref()
.is_some_and(|exit| exit.stderr_seq <= high_err);
if exit.is_none() || truncated || incomplete {
let outcome = self
.shared
.worker
.wait_exec(
&self.shared.params.exec_endpoint,
&self.shared.params.sailbox_id,
&self.exec_request_id,
self.shared.params.retry_timeout,
)
.await?;
{
let mut exit_slot = lock(&self.shared.exit);
if exit_slot.is_none() {
*exit_slot = Some(ExitInfo {
status: outcome.status,
exit_code: outcome.exit_code,
timed_out: outcome.timed_out,
stdout_truncated: outcome.stdout_truncated,
stderr_truncated: outcome.stderr_truncated,
error_message: String::new(),
stdout_seq: 0,
stderr_seq: 0,
});
}
}
if let Some(witnessed) = exit.as_ref() {
if outcome.status == pb::SailboxExecStatus::WorkerLost as i32 {
let state = lock(&self.shared.state);
let mut stderr = state.stderr.tail();
if witnessed.status == pb::SailboxExecStatus::Failed as i32
&& !witnessed.error_message.is_empty()
{
stderr = witnessed.error_message.clone();
}
return Ok(ExecResult {
stdout: state.stdout.tail(),
stderr,
exit_code: witnessed.exit_code,
timed_out: witnessed.timed_out,
stdout_truncated: witnessed.stdout_truncated
|| out_dropped
|| witnessed.stdout_seq > high_out,
stderr_truncated: witnessed.stderr_truncated
|| err_dropped
|| witnessed.stderr_seq > high_err,
stdout_complete,
stderr_complete,
});
}
}
if let Some(err) = terminal_status_error(outcome.status) {
return Err(err);
}
return Ok(ExecResult {
stdout: outcome.stdout,
stderr: outcome.stderr,
exit_code: outcome.exit_code,
timed_out: outcome.timed_out,
stdout_truncated: outcome.stdout_truncated,
stderr_truncated: outcome.stderr_truncated,
stdout_complete,
stderr_complete,
});
}
let exit = exit.expect("exit present on the clean path");
if let Some(err) = terminal_status_error(exit.status) {
return Err(err);
}
let state = lock(&self.shared.state);
let mut stderr = state.stderr.tail();
if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
stderr = exit.error_message.clone();
}
Ok(ExecResult {
stdout: state.stdout.tail(),
stderr,
exit_code: exit.exit_code,
timed_out: exit.timed_out,
stdout_truncated: false,
stderr_truncated: false,
stdout_complete,
stderr_complete,
})
}
pub async fn cancel(&self, signal: CancelSignal, retry: RetryBudget) -> Result<(), SailError> {
self.shared
.worker
.cancel_exec(
&self.shared.params.exec_endpoint,
&self.shared.params.sailbox_id,
&self.exec_request_id,
signal.is_force(),
retry.as_secs_f64(),
)
.await
}
pub async fn resize(&self, cols: u32, rows: u32) {
let message = pb::ResizeSailboxExecRequest {
sailbox_id: self.shared.params.sailbox_id.clone(),
exec_request_id: self.exec_request_id.clone(),
cols,
rows,
};
let Ok(request) =
self.shared
.worker
.request_for(message, &[], Some(Duration::from_secs(5)))
else {
return;
};
if let Ok(mut client) = self
.shared
.worker
.client_for(&self.shared.params.exec_endpoint)
{
let _ = client.resize_sailbox_exec(request).await;
}
}
pub async fn write_stdin(&self, data: &[u8]) -> Result<(), SailError> {
let mut stdin = self.shared.stdin.lock().await;
if stdin.eof_sent {
return Err(SailError::BrokenPipe {
message: "stdin is closed".to_string(),
});
}
if stdin.broken {
return Err(SailError::BrokenPipe {
message: "an earlier stdin write failed".to_string(),
});
}
if stdin.write_in_flight {
stdin.broken = true;
return Err(SailError::BrokenPipe {
message: "an earlier stdin write was interrupted".to_string(),
});
}
if data.is_empty() {
return Ok(());
}
stdin.write_in_flight = true;
let result = self.send_stdin(&mut stdin, data, false).await;
stdin.write_in_flight = false;
result
}
pub async fn close_stdin(&self) -> Result<(), SailError> {
let mut stdin = self.shared.stdin.lock().await;
if stdin.eof_sent {
return Ok(());
}
if stdin.broken || stdin.write_in_flight {
stdin.eof_sent = true;
return Ok(());
}
match self.send_stdin(&mut stdin, &[], true).await {
Ok(()) => {
stdin.eof_sent = true;
Ok(())
}
Err(SailError::BrokenPipe { .. }) => {
stdin.eof_sent = true;
Ok(())
}
Err(err) => Err(err),
}
}
async fn send_stdin(
&self,
stdin: &mut StdinState,
payload: &[u8],
eof: bool,
) -> Result<(), SailError> {
let endpoint = &self.shared.params.exec_endpoint;
let sailbox_id = &self.shared.params.sailbox_id;
let mut deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
let mut sent = 0usize;
loop {
let end = (sent + STDIN_WRITE_CHUNK_BYTES).min(payload.len());
let chunk = &payload[sent..end];
let last = end >= payload.len();
let message = pb::WriteSailboxExecStdinRequest {
sailbox_id: sailbox_id.clone(),
exec_request_id: self.exec_request_id.clone(),
offset: stdin.offset,
data: chunk.to_vec(),
eof: eof && last,
};
let request = match self.shared.worker.request_for(
message,
&[],
Some(rpc_attempt_timeout(deadline)),
) {
Ok(request) => request,
Err(err) => return Err(err),
};
let result = match self.shared.worker.client_for(endpoint) {
Ok(mut client) => client.write_sailbox_exec_stdin(request).await,
Err(err) => return Err(err),
};
match result {
Ok(resp) => {
let accepted_through = resp.into_inner().accepted_through;
let accepted =
((accepted_through - stdin.offset).max(0) as usize).min(chunk.len());
stdin.offset += accepted as i64;
sent += accepted;
if sent >= payload.len() && (!eof || (last && accepted == chunk.len())) {
return Ok(());
}
deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
if accepted > 0 {
delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
} else {
delay = sleep_no_deadline(delay).await;
}
}
Err(status) => {
if matches!(status.code(), Code::NotFound | Code::FailedPrecondition) {
return Err(SailError::BrokenPipe {
message: status.message().to_string(),
});
}
if is_exec_not_ready(&status) {
delay = sleep_no_deadline(delay).await;
deadline = retry_deadline(EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS);
continue;
}
if !should_retry_transient_exec_rpc(&status, deadline) {
stdin.broken = true;
return Err(SailError::from_exec_status(&status));
}
tracing::warn!(code = ?status.code(), "retrying exec stdin write");
if should_invalidate_channel(&status) {
self.shared.worker.channels().invalidate(endpoint);
}
delay = sleep_before_retry(delay, deadline).await;
}
}
}
}
}
pub struct StreamReader {
shared: Arc<ExecShared>,
which: OutputStream,
cursor: usize,
}
impl StreamReader {
pub fn next(&mut self, timeout: Duration) -> ReadStep {
let mut state = lock(&self.shared.state);
loop {
let ring = state.ring(self.which);
if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
}
let available = ring.first_idx + ring.pieces.len();
if self.cursor < available {
let piece = ring.pieces[self.cursor - ring.first_idx].clone();
self.cursor += 1;
return ReadStep::Chunk(piece);
}
if state.ended {
return ReadStep::Eof;
}
let (next_state, timed_out) = self
.shared
.cond
.wait_timeout(state, timeout)
.expect("exec state mutex poisoned");
state = next_state;
if timed_out.timed_out() {
return ReadStep::Pending;
}
}
}
}
pub struct AsyncStreamReader {
shared: Arc<ExecShared>,
which: OutputStream,
cursor: usize,
}
impl AsyncStreamReader {
pub async fn next(&mut self) -> Option<String> {
loop {
let notified = self.shared.data_notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
{
let state = lock(&self.shared.state);
let ring = state.ring(self.which);
if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
}
let available = ring.first_idx + ring.pieces.len();
if self.cursor < available {
let piece = ring.pieces[self.cursor - ring.first_idx].clone();
self.cursor += 1;
return Some(piece);
}
if state.ended {
return None;
}
}
notified.await;
}
}
}
async fn sleep_no_deadline(delay: f64) -> f64 {
let sleep_for = delay.min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS);
tokio::time::sleep(Duration::from_secs_f64(sleep_for.max(0.0))).await;
(delay * 2.0).min(crate::worker::EXEC_TRANSIENT_RETRY_MAX_DELAY_SECONDS)
}
fn is_exec_not_ready(status: &Status) -> bool {
status.code() == Code::Unavailable && status.message().contains("; retry")
}
fn terminal_status_error(status: i32) -> Option<SailError> {
if status == pb::SailboxExecStatus::WorkerLost as i32 {
return Some(SailError::WorkerLost {
message: "the machine hosting your sailbox was lost before the command \
finished; run exec again to retry"
.to_string(),
});
}
None
}
async fn submit(
worker: &Arc<WorkerProxy>,
params: &ExecParams,
stdout_resume_seq: i64,
stderr_resume_seq: i64,
) -> Result<(String, Streaming<pb::StreamSailboxExecResponse>), SailError> {
let deadline = retry_deadline(params.retry_timeout);
let mut delay = EXEC_TRANSIENT_RETRY_INITIAL_DELAY_SECONDS;
loop {
let message = pb::StreamSailboxExecRequest {
sailbox_id: params.sailbox_id.clone(),
argv: params.argv.clone(),
timeout_seconds: params.timeout_seconds,
idempotency_key: params.idempotency_key.clone(),
open_stdin: params.open_stdin,
stdout_resume_seq,
stderr_resume_seq,
pty: params.pty,
term_cols: params.cols,
term_rows: params.rows,
term: params.term.clone(),
};
let request = worker.request_for(message, ¶ms.extra_metadata, None)?;
let status = match worker
.client_for(¶ms.exec_endpoint)?
.stream_sailbox_exec(request)
.await
{
Ok(resp) => {
let mut stream = resp.into_inner();
match stream.message().await {
Ok(Some(first)) => match first.frame {
Some(pb::stream_sailbox_exec_response::Frame::Started(started)) => {
return Ok((started.exec_request_id, stream));
}
_ => {
return Err(SailError::Execution {
code: crate::error::GrpcCode::Internal,
detail: "exec stream opened with a non-started frame".to_string(),
});
}
},
Ok(None) if stdout_resume_seq == 0 && stderr_resume_seq == 0 => {
tonic::Status::unavailable(
"exec stream ended before the server confirmed the launch",
)
}
Ok(None) => {
return Err(SailError::Execution {
code: crate::error::GrpcCode::Internal,
detail: "exec stream ended before the server confirmed the launch"
.to_string(),
});
}
Err(status) => status,
}
}
Err(status) => status,
};
if !should_retry_transient_exec_rpc(&status, deadline) {
return Err(SailError::from_exec_status(&status));
}
tracing::warn!(
code = ?status.code(),
stdout_resume_seq,
stderr_resume_seq,
"reconnecting exec stream"
);
if should_invalidate_channel(&status) {
worker.channels().invalidate(¶ms.exec_endpoint);
}
delay = sleep_before_retry(delay, deadline).await;
}
}
struct Pump {
shared: Arc<ExecShared>,
stdout_dec: Utf8Decoder,
stderr_dec: Utf8Decoder,
stdout_seq: i64,
stderr_seq: i64,
}
impl Pump {
fn new(shared: Arc<ExecShared>) -> Pump {
Pump {
shared,
stdout_dec: Utf8Decoder::default(),
stderr_dec: Utf8Decoder::default(),
stdout_seq: 0,
stderr_seq: 0,
}
}
fn apply_frame(&mut self, frame: pb::StreamSailboxExecResponse) -> bool {
match frame.frame {
Some(pb::stream_sailbox_exec_response::Frame::Chunk(chunk)) => {
let is_stderr = chunk.stream == pb::SailboxExecStream::Stderr as i32;
let (decoder, which) = if is_stderr {
self.stderr_seq = self.stderr_seq.max(chunk.seq);
(&mut self.stderr_dec, OutputStream::Stderr)
} else {
self.stdout_seq = self.stdout_seq.max(chunk.seq);
(&mut self.stdout_dec, OutputStream::Stdout)
};
let text = decoder.decode(&chunk.data);
if !text.is_empty() {
let mut state = lock(&self.shared.state);
match which {
OutputStream::Stdout => state.stdout.append(text),
OutputStream::Stderr => state.stderr.append(text),
}
self.shared.cond.notify_all();
self.shared.data_notify.notify_waiters();
}
false
}
Some(pb::stream_sailbox_exec_response::Frame::Exit(exit)) => {
*lock(&self.shared.exit) = Some(ExitInfo {
status: exit.status,
exit_code: exit.return_code,
timed_out: exit.timed_out,
stdout_truncated: exit.stdout_truncated,
stderr_truncated: exit.stderr_truncated,
error_message: exit.error_message,
stdout_seq: exit.stdout_seq,
stderr_seq: exit.stderr_seq,
});
true
}
_ => false,
}
}
fn finalize(mut self) {
let tail_out = self.stdout_dec.flush();
let tail_err = self.stderr_dec.flush();
{
let mut state = lock(&self.shared.state);
if !tail_out.is_empty() {
state.stdout.append(tail_out);
}
if !tail_err.is_empty() {
state.stderr.append(tail_err);
}
state.ended = true;
self.shared.cond.notify_all();
self.shared.data_notify.notify_waiters();
}
*lock(&self.shared.high_seq) = (self.stdout_seq, self.stderr_seq);
self.shared.ended.store(true, Ordering::SeqCst);
self.shared.ended_notify.notify_waiters();
}
}
async fn pump(shared: Arc<ExecShared>, mut stream: Streaming<pb::StreamSailboxExecResponse>) {
let mut state = Pump::new(shared.clone());
loop {
if shared.closing.load(Ordering::SeqCst) {
break;
}
let message = tokio::select! {
biased;
() = shared.close_notify.notified() => break,
message = stream.message() => message,
};
match message {
Ok(Some(frame)) => {
if state.apply_frame(frame) {
break;
}
}
Ok(None) => break,
Err(_status) => {
if shared.closing.load(Ordering::SeqCst) {
break;
}
match submit(
&shared.worker,
&shared.params,
state.stdout_seq,
state.stderr_seq,
)
.await
{
Ok((_id, fresh)) => {
stream = fresh;
continue;
}
Err(_) => break,
}
}
}
}
state.finalize();
}
#[cfg(any(test, feature = "fuzzing"))]
pub fn fuzz_decode_ring(data: &[u8]) {
let mut whole = Utf8Decoder::default();
let mut whole_out = whole.decode(data);
whole_out.push_str(&whole.flush());
let mut chunked = Utf8Decoder::default();
let mut chunked_out = String::new();
let mut start = 0;
for (i, byte) in data.iter().enumerate() {
if byte & 1 == 1 {
chunked_out.push_str(&chunked.decode(&data[start..=i]));
start = i + 1;
}
}
chunked_out.push_str(&chunked.decode(&data[start..]));
chunked_out.push_str(&chunked.flush());
assert_eq!(
chunked_out, whole_out,
"decoder is not chunk-boundary invariant"
);
assert_eq!(
whole_out,
String::from_utf8_lossy(data),
"decoder disagrees with a lossy decode of the whole input"
);
let mut ring = Ring::default();
ring.append(whole_out.clone());
let filler = STREAM_BUFFER_CAP_BYTES + 1 - whole_out.len().min(STREAM_BUFFER_CAP_BYTES);
let filler = "a".repeat(filler);
ring.append(filler.clone());
assert!(
ring.size <= STREAM_BUFFER_CAP_BYTES,
"ring exceeded its cap"
);
let full = format!("{whole_out}{filler}");
let tail = ring.tail();
assert!(
full.as_bytes().ends_with(tail.as_bytes()),
"ring tail is not a suffix of the appended bytes"
);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shell_argv_wraps_cwd_and_background() {
let plain = shell_argv("echo hi", &ExecOptions::default()).unwrap();
assert_eq!(plain, ["/bin/sh", "-lc", "echo hi"]);
let cwd = shell_argv(
"echo hi",
&ExecOptions {
cwd: Some("/app".to_string()),
..ExecOptions::default()
},
)
.unwrap();
assert_eq!(cwd[2], "cd '/app' && exec /bin/sh -lc 'echo hi'");
let background = shell_argv(
"echo hi",
&ExecOptions {
cwd: Some("/app".to_string()),
background: true,
..ExecOptions::default()
},
)
.unwrap();
assert_eq!(
background[2],
"nohup /bin/sh -lc 'cd '\\''/app'\\'' && exec /bin/sh -lc '\\''echo hi'\\''' </dev/null >/dev/null 2>&1 &"
);
}
#[test]
fn shell_argv_rejects_invalid_combinations() {
assert!(shell_argv("", &ExecOptions::default()).is_err());
assert!(shell_argv(
"x",
&ExecOptions {
cwd: Some(" ".to_string()),
..ExecOptions::default()
}
)
.is_err());
for (open_stdin, pty) in [(true, false), (false, true)] {
assert!(shell_argv(
"x",
&ExecOptions {
background: true,
open_stdin,
pty,
..ExecOptions::default()
}
)
.is_err());
}
}
#[test]
fn fuzz_decode_ring_holds_on_samples() {
for sample in [
&b""[..],
b"hello",
b"\xff\xfe\xfd",
"€ µ é".as_bytes(),
b"ab\xc3\xa9cd",
&[0xC3, 0x28],
] {
fuzz_decode_ring(sample);
}
}
#[test]
fn ring_drops_oldest_past_cap_and_latches() {
let mut ring = Ring::default();
ring.append("a".repeat(STREAM_BUFFER_CAP_BYTES));
assert!(!ring.dropped);
ring.append("bbbb".to_string());
assert!(ring.dropped);
assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
assert!(ring.tail().ends_with("bbbb"));
}
#[test]
fn ring_clip_on_split_multibyte_char_terminates() {
let mut ring = Ring::default();
ring.append(format!("é{}", "a".repeat(STREAM_BUFFER_CAP_BYTES - 1)));
assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
assert!(ring.dropped);
let tail = ring.tail();
assert!(!tail.contains('é'));
assert!(!tail.contains('\u{FFFD}'));
}
#[test]
fn decoder_carries_split_multibyte_char() {
let mut dec = Utf8Decoder::default();
let euro = "€".as_bytes(); assert_eq!(dec.decode(&euro[..2]), "");
assert_eq!(dec.decode(&euro[2..]), "€");
}
#[test]
fn decoder_replaces_invalid_bytes() {
let mut dec = Utf8Decoder::default();
assert_eq!(dec.decode(&[0xff, b'x']), "\u{FFFD}x");
}
#[test]
fn terminal_status_mapping() {
assert!(matches!(
terminal_status_error(pb::SailboxExecStatus::WorkerLost as i32),
Some(SailError::WorkerLost { .. })
));
assert!(terminal_status_error(pb::SailboxExecStatus::Succeeded as i32).is_none());
for retired in [9, 6, 7] {
assert!(terminal_status_error(retired).is_none());
}
}
fn test_shared() -> Arc<ExecShared> {
Arc::new(ExecShared {
worker: Arc::new(WorkerProxy::new("test-key").unwrap()),
params: ExecParams {
sailbox_id: "sb".into(),
exec_endpoint: "endpoint".into(),
argv: vec!["echo".into()],
timeout_seconds: 0,
idempotency_key: "idem".into(),
open_stdin: false,
pty: false,
term: String::new(),
cols: 0,
rows: 0,
retry_timeout: 0.0,
extra_metadata: vec![],
},
state: Mutex::new(State::default()),
cond: Condvar::new(),
data_notify: Notify::new(),
exit: Mutex::new(None),
high_seq: Mutex::new((0, 0)),
stdin: AsyncMutex::new(StdinState::default()),
ended: AtomicBool::new(false),
ended_notify: Notify::new(),
closing: AtomicBool::new(false),
close_notify: Notify::new(),
})
}
fn chunk(which: OutputStream, seq: i64, data: &[u8]) -> pb::StreamSailboxExecResponse {
let stream = match which {
OutputStream::Stdout => pb::SailboxExecStream::Stdout,
OutputStream::Stderr => pb::SailboxExecStream::Stderr,
};
pb::StreamSailboxExecResponse {
frame: Some(pb::stream_sailbox_exec_response::Frame::Chunk(
pb::SailboxExecChunk {
stream: stream as i32,
data: data.to_vec(),
seq,
},
)),
}
}
fn exit_frame(status: pb::SailboxExecStatus, stdout_seq: i64) -> pb::StreamSailboxExecResponse {
pb::StreamSailboxExecResponse {
frame: Some(pb::stream_sailbox_exec_response::Frame::Exit(
pb::SailboxExecExit {
status: status as i32,
return_code: 0,
timed_out: false,
stdout_truncated: false,
stderr_truncated: false,
error_message: String::new(),
stdout_seq,
stderr_seq: 0,
},
)),
}
}
#[tokio::test]
async fn exec_resume_carries_utf8_and_tracks_seq_across_break() {
let shared = test_shared();
let mut pump = Pump::new(shared.clone());
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"ab\xc3")));
assert_eq!(shared.state.lock().unwrap().stdout.tail(), "ab");
assert_eq!(pump.stdout_seq, 1);
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"\xa9cd")));
assert_eq!(shared.state.lock().unwrap().stdout.tail(), "abécd");
assert_eq!(pump.stdout_seq, 2);
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"!")));
assert_eq!(pump.stdout_seq, 2);
assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2)));
pump.finalize();
assert_eq!(*shared.high_seq.lock().unwrap(), (2, 0));
assert!(shared.ended.load(Ordering::SeqCst));
}
#[tokio::test]
async fn exec_reader_follows_replayed_then_live() {
let shared = test_shared();
let mut reader = StreamReader {
shared: shared.clone(),
which: OutputStream::Stdout,
cursor: 0,
};
let collector = tokio::task::spawn_blocking(move || {
let mut out = String::new();
loop {
match reader.next(Duration::from_millis(50)) {
ReadStep::Chunk(piece) => out.push_str(&piece),
ReadStep::Eof => return out,
ReadStep::Pending => {}
}
}
});
let mut pump = Pump::new(shared.clone());
pump.apply_frame(chunk(OutputStream::Stdout, 1, b"hello "));
tokio::time::sleep(Duration::from_millis(10)).await;
pump.apply_frame(chunk(OutputStream::Stdout, 2, b"world"));
tokio::time::sleep(Duration::from_millis(10)).await;
pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 2));
pump.finalize();
assert_eq!(collector.await.unwrap(), "hello world");
}
#[test]
fn exec_reader_started_late_replays_retained_tail() {
let shared = test_shared();
let mut pump = Pump::new(shared.clone());
pump.apply_frame(chunk(OutputStream::Stdout, 1, b"early "));
pump.apply_frame(chunk(OutputStream::Stdout, 2, b"output"));
pump.finalize();
let mut reader = shared_reader(&shared, OutputStream::Stdout);
let mut out = String::new();
loop {
match reader.next(Duration::from_millis(50)) {
ReadStep::Chunk(piece) => out.push_str(&piece),
ReadStep::Eof => break,
ReadStep::Pending => panic!("ended ring should not return Pending"),
}
}
assert_eq!(out, "early output");
}
fn shared_reader(shared: &Arc<ExecShared>, which: OutputStream) -> StreamReader {
StreamReader {
shared: shared.clone(),
which,
cursor: 0,
}
}
use proptest::prelude::*;
fn utf8ish_bytes() -> impl Strategy<Value = Vec<u8>> {
prop_oneof![
proptest::collection::vec(any::<u8>(), 0..256),
any::<String>().prop_map(String::into_bytes),
]
}
proptest! {
#[test]
fn ring_without_overflow_is_lossless(
pieces in proptest::collection::vec(any::<String>(), 0..32)
) {
let concat: String = pieces.concat();
prop_assume!(concat.len() <= STREAM_BUFFER_CAP_BYTES);
let mut ring = Ring::default();
for piece in &pieces {
ring.append(piece.clone());
}
prop_assert!(!ring.dropped);
prop_assert_eq!(ring.size, concat.len());
prop_assert_eq!(ring.tail(), concat);
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(48))]
#[test]
fn ring_eviction_keeps_valid_suffix_under_cap(
overflow in 1usize..8192,
tail_pieces in proptest::collection::vec("[a-z]{0,64}", 0..8),
) {
let head = "é".repeat(STREAM_BUFFER_CAP_BYTES / 2 + overflow);
let mut full = head.clone();
let mut ring = Ring::default();
ring.append(head);
for piece in &tail_pieces {
ring.append(piece.clone());
full.push_str(piece);
}
prop_assert!(ring.dropped);
prop_assert!(ring.size <= STREAM_BUFFER_CAP_BYTES);
let tail = ring.tail();
let has_replacement_char = tail.contains('\u{FFFD}');
prop_assert!(!has_replacement_char);
prop_assert!(full.as_bytes().ends_with(tail.as_bytes()));
}
}
proptest! {
#[test]
fn decoder_is_chunk_boundary_invariant(
data in utf8ish_bytes(),
cuts in proptest::collection::vec(any::<prop::sample::Index>(), 0..16),
) {
let mut whole = Utf8Decoder::default();
let mut whole_out = whole.decode(&data);
whole_out.push_str(&whole.flush());
let mut bounds: Vec<usize> = cuts.iter().map(|c| c.index(data.len() + 1)).collect();
bounds.sort_unstable();
let mut dec = Utf8Decoder::default();
let mut chunked_out = String::new();
let mut prev = 0;
for bound in bounds.into_iter().chain(std::iter::once(data.len())) {
let bound = bound.clamp(prev, data.len());
chunked_out.push_str(&dec.decode(&data[prev..bound]));
prev = bound;
}
chunked_out.push_str(&dec.flush());
prop_assert_eq!(&chunked_out, &whole_out);
prop_assert_eq!(whole_out, String::from_utf8_lossy(&data).into_owned());
}
}
proptest! {
#[test]
fn pump_seq_is_monotonic_running_max_per_stream(
frames in proptest::collection::vec(
(any::<bool>(), 0i64..10_000, proptest::collection::vec(any::<u8>(), 0..8)),
0..64,
)
) {
let mut pump = Pump::new(test_shared());
let (mut expect_out, mut expect_err) = (0i64, 0i64);
let (mut prev_out, mut prev_err) = (0i64, 0i64);
for (is_stderr, seq, data) in frames {
let which = if is_stderr { OutputStream::Stderr } else { OutputStream::Stdout };
pump.apply_frame(chunk(which, seq, &data));
if is_stderr {
expect_err = expect_err.max(seq);
} else {
expect_out = expect_out.max(seq);
}
prop_assert_eq!(pump.stdout_seq, expect_out);
prop_assert_eq!(pump.stderr_seq, expect_err);
prop_assert!(pump.stdout_seq >= prev_out);
prop_assert!(pump.stderr_seq >= prev_err);
prev_out = pump.stdout_seq;
prev_err = pump.stderr_seq;
}
}
}
}