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,
};
#[doc(hidden)]
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(Vec<u8>),
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,
pub stdout_total_bytes: i64,
pub stderr_total_bytes: i64,
}
#[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 {
#[doc(hidden)]
pub fn as_secs_f64(self) -> f64 {
match self {
RetryBudget::None => 0.0,
RetryBudget::Within(d) => d.as_secs_f64(),
RetryBudget::Forever => f64::INFINITY,
}
}
#[doc(hidden)]
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 env: Vec<(String, String)>,
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,
env: Vec::new(),
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('\'', "'\\''"))
}
const PTY_ENV_WHITELIST: [&str; 3] = ["COLORTERM", "LANG", "TERM_PROGRAM"];
fn pty_env_whitelisted(key: &str) -> bool {
PTY_ENV_WHITELIST.contains(&key) || key.starts_with("LC_")
}
pub(crate) fn pty_forward_env() -> Vec<(String, String)> {
pty_forward_env_from(
std::env::vars_os()
.filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))),
)
}
fn pty_forward_env_from(vars: impl Iterator<Item = (String, String)>) -> Vec<(String, String)> {
vars.filter(|(key, _)| pty_env_whitelisted(key)).collect()
}
#[doc(hidden)]
pub fn encode_env(
pairs: &[(String, String)],
) -> Result<std::collections::HashMap<String, String>, SailError> {
let mut env = std::collections::HashMap::with_capacity(pairs.len());
for (key, value) in pairs {
if !is_portable_env_name(key) || value.contains('\0') {
return Err(SailError::InvalidArgument {
message: format!("invalid env entry {key:?}"),
});
}
env.insert(key.clone(), value.clone());
}
Ok(env)
}
fn is_portable_env_name(name: &str) -> bool {
let mut chars = name.chars();
match chars.next() {
Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
_ => return false,
}
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
#[doc(hidden)]
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 env: std::collections::HashMap<String, String>,
pub retry_timeout: f64,
pub extra_metadata: Vec<(String, String)>,
}
#[derive(Default)]
struct Ring {
pieces: Vec<Vec<u8>>,
first_idx: usize,
size: usize,
dropped: bool,
front_clips: u64,
resets: u64,
}
impl Ring {
fn append(&mut self, data: Vec<u8>) {
self.size += data.len();
self.pieces.push(data);
while self.size > STREAM_BUFFER_CAP_BYTES {
let overflow = self.size - STREAM_BUFFER_CAP_BYTES;
if self.pieces[0].len() <= overflow {
self.size -= self.pieces[0].len();
self.pieces.remove(0);
self.first_idx += 1;
} else {
self.pieces[0].drain(..overflow);
self.size -= overflow;
self.front_clips += 1;
}
self.dropped = true;
}
}
fn reset_to(&mut self, repaint: Vec<u8>) {
self.first_idx += self.pieces.len();
self.pieces.clear();
self.size = 0;
self.dropped = false;
self.resets += 1;
if !repaint.is_empty() {
self.append(repaint);
}
}
fn tail(&self) -> Vec<u8> {
self.pieces.concat()
}
}
fn lossy_tail(ring: &Ring) -> String {
String::from_utf8_lossy(&ring.tail()).replace('\0', "\u{FFFD}")
}
#[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,
stdout_total_bytes: i64,
stderr_total_bytes: 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()
};
if params.pty {
for (key, value) in pty_forward_env() {
params.env.entry(key).or_insert(value);
}
}
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,
dropped: false,
reset: false,
seen_front_clips: 0,
seen_resets: 0,
}
}
pub fn reader_async(&self, which: OutputStream) -> AsyncStreamReader {
AsyncStreamReader {
shared: self.shared.clone(),
which,
cursor: 0,
}
}
#[doc(hidden)]
pub fn buffered_output(&self, which: OutputStream) -> Vec<u8> {
lock(&self.shared.state).ring(which).tail()
}
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();
}
#[doc(hidden)]
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 incomplete = exit
.as_ref()
.is_some_and(|exit| exit.stdout_seq > high_out || exit.stderr_seq > high_err);
let stdout_truncated_flag = exit
.as_ref()
.is_some_and(|exit| exit.stdout_truncated || out_dropped);
let stderr_truncated_flag = exit
.as_ref()
.is_some_and(|exit| exit.stderr_truncated || err_dropped);
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);
let stdout_total_bytes = exit.as_ref().map_or(0, |exit| exit.stdout_total_bytes);
let stderr_total_bytes = exit.as_ref().map_or(0, |exit| exit.stderr_total_bytes);
if exit.is_none() || 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,
stdout_total_bytes: 0,
stderr_total_bytes: 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 = lossy_tail(&state.stderr);
if witnessed.status == pb::SailboxExecStatus::Failed as i32
&& !witnessed.error_message.is_empty()
{
stderr = witnessed.error_message.clone();
}
return Ok(ExecResult {
stdout: lossy_tail(&state.stdout),
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,
stdout_total_bytes,
stderr_total_bytes,
});
}
}
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,
stdout_total_bytes,
stderr_total_bytes,
});
}
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 = lossy_tail(&state.stderr);
if exit.status == pb::SailboxExecStatus::Failed as i32 && !exit.error_message.is_empty() {
stderr = exit.error_message.clone();
}
Ok(ExecResult {
stdout: lossy_tail(&state.stdout),
stderr,
exit_code: exit.exit_code,
timed_out: exit.timed_out,
stdout_truncated: stdout_truncated_flag,
stderr_truncated: stderr_truncated_flag,
stdout_complete,
stderr_complete,
stdout_total_bytes,
stderr_total_bytes,
})
}
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 resync(&self) {
let message = pb::ResyncSailboxExecRequest {
sailbox_id: self.shared.params.sailbox_id.clone(),
exec_request_id: self.exec_request_id.clone(),
};
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.resync_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,
dropped: bool,
reset: bool,
seen_front_clips: u64,
seen_resets: u64,
}
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 ring.resets != self.seen_resets {
self.reset = true;
self.dropped = ring.dropped;
if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
}
} else if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
self.dropped = true;
} else if self.cursor == ring.first_idx && ring.front_clips != self.seen_front_clips {
self.dropped = true;
}
self.seen_front_clips = ring.front_clips;
self.seen_resets = ring.resets;
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 fn try_next(&mut self) -> Option<Vec<u8>> {
let state = lock(&self.shared.state);
let ring = state.ring(self.which);
if ring.resets != self.seen_resets {
self.reset = true;
self.dropped = ring.dropped;
if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
}
} else if self.cursor < ring.first_idx {
self.cursor = ring.first_idx;
self.dropped = true;
} else if self.cursor == ring.first_idx && ring.front_clips != self.seen_front_clips {
self.dropped = true;
}
self.seen_front_clips = ring.front_clips;
self.seen_resets = ring.resets;
if self.cursor < ring.first_idx + ring.pieces.len() {
let piece = ring.pieces[self.cursor - ring.first_idx].clone();
self.cursor += 1;
Some(piece)
} else {
None
}
}
pub fn took_drop(&mut self) -> bool {
std::mem::take(&mut self.dropped)
}
pub fn took_reset(&mut self) -> bool {
std::mem::take(&mut self.reset)
}
}
pub struct AsyncStreamReader {
shared: Arc<ExecShared>,
which: OutputStream,
cursor: usize,
}
impl AsyncStreamReader {
pub async fn next(&mut self) -> Option<Vec<u8>> {
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(),
env: params.env.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_seq: i64,
stderr_seq: i64,
}
impl Pump {
fn new(shared: Arc<ExecShared>) -> Pump {
Pump {
shared,
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 which = if is_stderr {
self.stderr_seq = self.stderr_seq.max(chunk.seq);
OutputStream::Stderr
} else {
self.stdout_seq = self.stdout_seq.max(chunk.seq);
OutputStream::Stdout
};
if !chunk.data.is_empty() {
let mut state = lock(&self.shared.state);
match which {
OutputStream::Stdout => state.stdout.append(chunk.data),
OutputStream::Stderr => state.stderr.append(chunk.data),
}
self.shared.cond.notify_all();
self.shared.data_notify.notify_waiters();
}
false
}
Some(pb::stream_sailbox_exec_response::Frame::Snapshot(snap)) => {
self.stdout_seq = snap.stdout_seq_basis;
let mut state = lock(&self.shared.state);
state.stdout.reset_to(snap.repaint);
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,
stdout_total_bytes: exit.stdout_total_bytes,
stderr_total_bytes: exit.stderr_total_bytes,
});
true
}
_ => false,
}
}
fn finalize(self) {
{
let mut state = lock(&self.shared.state);
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_exec_ring(data: &[u8]) {
let mut chunked = Ring::default();
let mut start = 0;
for (i, byte) in data.iter().enumerate() {
if byte & 1 == 1 {
chunked.append(data[start..=i].to_vec());
start = i + 1;
}
}
if start < data.len() {
chunked.append(data[start..].to_vec());
}
let mut whole = Ring::default();
if !data.is_empty() {
whole.append(data.to_vec());
}
assert_eq!(
chunked.tail(),
whole.tail(),
"ring is not chunk-boundary invariant"
);
let mut ring = Ring::default();
ring.append(data.to_vec());
let filler = vec![b'a'; STREAM_BUFFER_CAP_BYTES + 1 - data.len().min(STREAM_BUFFER_CAP_BYTES)];
ring.append(filler.clone());
assert!(
ring.size <= STREAM_BUFFER_CAP_BYTES,
"ring exceeded its cap"
);
let mut full = data.to_vec();
full.extend_from_slice(&filler);
assert!(
full.ends_with(&ring.tail()),
"ring tail is not a suffix of the appended bytes"
);
let before_reset_pieces = ring.first_idx + ring.pieces.len();
ring.reset_to(data.to_vec());
assert!(
ring.first_idx >= before_reset_pieces,
"reset left old pieces reachable"
);
let expected = &data[data.len().saturating_sub(STREAM_BUFFER_CAP_BYTES)..];
assert_eq!(
ring.tail(),
expected,
"reset tail is not the capped repaint"
);
}
#[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_exec_ring_holds_on_samples() {
for sample in [
&b""[..],
b"hello",
b"\xff\xfe\xfd",
"€ µ é".as_bytes(),
b"ab\xc3\xa9cd",
&[0xC3, 0x28],
] {
fuzz_exec_ring(sample);
}
}
#[test]
fn ring_drops_oldest_past_cap_and_latches() {
let mut ring = Ring::default();
ring.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES]);
assert!(!ring.dropped);
ring.append(b"bbbb".to_vec());
assert!(ring.dropped);
assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
assert_eq!(ring.tail().len(), STREAM_BUFFER_CAP_BYTES);
assert!(ring.tail().ends_with(b"bbbb"));
}
#[test]
fn ring_clip_is_byte_exact() {
let mut ring = Ring::default();
let mut data = "é".as_bytes().to_vec();
data.extend(std::iter::repeat_n(b'a', STREAM_BUFFER_CAP_BYTES - 1));
ring.append(data);
assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
assert!(ring.dropped);
let tail = ring.tail();
assert_eq!(tail.len(), STREAM_BUFFER_CAP_BYTES);
assert_eq!(tail[0], "é".as_bytes()[1]);
}
#[test]
fn ring_reset_to_supersedes_retained_pieces() {
let mut ring = Ring::default();
ring.append(b"old output".to_vec());
ring.append(b"more".to_vec());
let reachable_end = ring.first_idx + ring.pieces.len();
ring.reset_to(b"\x1b[2J\x1b[Hrepaint".to_vec());
assert!(ring.first_idx >= reachable_end);
assert_eq!(ring.tail(), b"\x1b[2J\x1b[Hrepaint");
assert!(!ring.dropped);
let mut empty = Ring::default();
empty.append(b"x".to_vec());
empty.reset_to(Vec::new());
assert!(empty.tail().is_empty());
}
#[test]
fn pty_forward_env_filters_to_the_whitelist() {
let vars = vec![
("COLORTERM".to_string(), "truecolor".to_string()),
("LC_ALL".to_string(), "en_US.UTF-8".to_string()),
("LANG".to_string(), "en_US.UTF-8".to_string()),
("TERM_PROGRAM".to_string(), "TestTerm".to_string()),
("PATH".to_string(), "/bin".to_string()),
("TERM".to_string(), "xterm".to_string()), ("SECRET_TOKEN".to_string(), "x".to_string()),
];
let forwarded = pty_forward_env_from(vars.into_iter());
let keys: Vec<&str> = forwarded.iter().map(|(k, _)| k.as_str()).collect();
assert_eq!(keys, ["COLORTERM", "LC_ALL", "LANG", "TERM_PROGRAM"]);
}
#[test]
fn encode_env_rejects_malformed_keys() {
for (key, value) in [
("", "v"),
("A=B", "v"),
("NUL\0KEY", "v"),
("K", "nul\0value"),
("1FOO", "v"),
("FO O", "v"),
("FOO-BAR", "v"),
("FOO.BAR", "v"),
] {
let pairs = vec![(key.to_string(), value.to_string())];
assert!(
encode_env(&pairs).is_err(),
"expected rejection for {key:?}={value:?}"
);
}
let ok = encode_env(&[
("_FOO".to_string(), "bar=baz".to_string()),
("LC_ALL".to_string(), "C.UTF-8".to_string()),
])
.unwrap();
assert_eq!(ok.get("_FOO").map(String::as_str), Some("bar=baz"));
}
#[test]
fn ring_reset_to_clears_prior_dropped() {
let mut ring = Ring::default();
ring.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES + 1]);
assert!(ring.dropped);
ring.reset_to(b"repaint".to_vec());
assert!(!ring.dropped);
assert_eq!(ring.tail(), b"repaint");
let mut over = Ring::default();
over.append(vec![b'b'; STREAM_BUFFER_CAP_BYTES + 1]);
over.reset_to(vec![b'c'; STREAM_BUFFER_CAP_BYTES + 1]);
assert!(over.dropped);
assert_eq!(over.size, STREAM_BUFFER_CAP_BYTES);
}
#[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,
env: std::collections::HashMap::default(),
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,
..Default::default()
},
)),
}
}
#[tokio::test]
async fn exec_resume_carries_bytes_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(), b"ab\xc3");
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".as_bytes()
);
assert_eq!(lossy_tail(&shared.state.lock().unwrap().stdout), "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));
}
fn snapshot_frame(repaint: &[u8], basis: i64) -> pb::StreamSailboxExecResponse {
pb::StreamSailboxExecResponse {
frame: Some(pb::stream_sailbox_exec_response::Frame::Snapshot(
pb::SailboxExecSnapshot {
repaint: repaint.to_vec(),
stdout_seq_basis: basis,
},
)),
}
}
#[tokio::test]
async fn snapshot_resets_ring_seq_basis_and_skips_readers() {
let shared = test_shared();
let mut pump = Pump::new(shared.clone());
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 1, b"pre-disconnect ")));
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 2, b"tail")));
assert!(!pump.apply_frame(snapshot_frame(b"\x1b[2J\x1b[Hscreen", 7)));
assert_eq!(pump.stdout_seq, 7);
assert!(!pump.apply_frame(chunk(OutputStream::Stdout, 8, b" after")));
assert!(pump.apply_frame(exit_frame(pb::SailboxExecStatus::Succeeded, 8)));
pump.finalize();
let mut reader = shared_reader(&shared, OutputStream::Stdout);
let mut out = Vec::new();
loop {
match reader.next(Duration::from_millis(50)) {
ReadStep::Chunk(piece) => out.extend_from_slice(&piece),
ReadStep::Eof => break,
ReadStep::Pending => panic!("ended ring should not return Pending"),
}
}
assert_eq!(out, b"\x1b[2J\x1b[Hscreen after");
assert_eq!(*shared.high_seq.lock().unwrap(), (8, 0));
}
#[test]
fn snapshot_with_empty_repaint_is_reset_only() {
let shared = test_shared();
let mut pump = Pump::new(shared.clone());
pump.apply_frame(chunk(OutputStream::Stdout, 3, b"stale"));
pump.apply_frame(snapshot_frame(b"", 3));
assert_eq!(pump.stdout_seq, 3);
let state = lock(&shared.state);
assert!(state.stdout.tail().is_empty());
assert!(!state.stdout.dropped);
}
#[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,
dropped: false,
reset: false,
seen_front_clips: 0,
seen_resets: 0,
};
let collector = tokio::task::spawn_blocking(move || {
let mut out = Vec::new();
loop {
match reader.next(Duration::from_millis(50)) {
ReadStep::Chunk(piece) => out.extend_from_slice(&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(), b"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 = Vec::new();
loop {
match reader.next(Duration::from_millis(50)) {
ReadStep::Chunk(piece) => out.extend_from_slice(&piece),
ReadStep::Eof => break,
ReadStep::Pending => panic!("ended ring should not return Pending"),
}
}
assert_eq!(out, b"early output");
}
#[test]
fn reader_reports_drop_and_batch_drains() {
let shared = test_shared();
{
let mut state = lock(&shared.state);
state.stdout.append(b"HEAD".to_vec());
state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
state.ended = true;
}
let mut reader = shared_reader(&shared, OutputStream::Stdout);
let first = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(first) = first else {
panic!("expected the retained chunk, got {first:?}");
};
assert!(
reader.took_drop(),
"the overflow evicted an unconsumed chunk"
);
assert!(!reader.took_drop(), "took_drop clears the latch");
assert_eq!(first, vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
assert!(
reader.try_next().is_none(),
"a drained ring yields None without blocking"
);
}
#[test]
fn reader_does_not_report_a_reset_repaint_as_a_drop() {
let shared = test_shared();
let mut reader = shared_reader(&shared, OutputStream::Stdout);
{
let mut state = lock(&shared.state);
state.stdout.append(b"one".to_vec());
}
assert!(matches!(
reader.next(Duration::from_millis(50)),
ReadStep::Chunk(_)
));
assert!(!reader.took_drop(), "a clean read latches no drop");
assert!(!reader.took_reset(), "a clean read is not a reset");
{
let mut state = lock(&shared.state);
state.stdout.append(b"two".to_vec());
state.stdout.append(b"three".to_vec());
state.stdout.reset_to(b"REPAINT".to_vec());
state.ended = true;
}
let repaint = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(repaint) = repaint else {
panic!("expected the repaint chunk, got {repaint:?}");
};
assert_eq!(repaint, b"REPAINT");
assert!(
!reader.took_drop(),
"reading a reset repaint is a heal, not a fall-behind drop"
);
assert!(
reader.took_reset(),
"the reset repaint is surfaced as a took_reset event so the pump can \
drop stale terminal-local backlog buffered ahead of it"
);
assert!(!reader.took_reset(), "took_reset clears the latch");
}
#[test]
fn reader_reset_supersedes_a_previously_latched_drop() {
let shared = test_shared();
{
let mut state = lock(&shared.state);
state.stdout.append(b"HEAD".to_vec());
state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES]);
}
let mut reader = shared_reader(&shared, OutputStream::Stdout);
assert!(matches!(
reader.next(Duration::from_millis(50)),
ReadStep::Chunk(_)
));
{
let mut state = lock(&shared.state);
state.stdout.reset_to(b"REPAINT".to_vec());
state.ended = true;
}
let repaint = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(repaint) = repaint else {
panic!("expected the repaint chunk, got {repaint:?}");
};
assert_eq!(repaint, b"REPAINT");
assert!(reader.took_reset(), "the repaint is surfaced as a reset");
assert!(
!reader.took_drop(),
"the reset superseded the stale drop latch, so no redundant resync \
clobbers the repaint"
);
}
#[test]
fn reader_reports_a_drop_when_output_evicts_the_repaint_before_it_is_read() {
let shared = test_shared();
let mut reader = shared_reader(&shared, OutputStream::Stdout);
{
let mut state = lock(&shared.state);
state.stdout.append(b"one".to_vec());
}
assert!(matches!(
reader.next(Duration::from_millis(50)),
ReadStep::Chunk(_)
));
assert!(!reader.took_drop());
assert!(!reader.took_reset());
{
let mut state = lock(&shared.state);
state.stdout.reset_to(b"REPAINT".to_vec());
state.stdout.append(vec![b'x'; STREAM_BUFFER_CAP_BYTES + 1]);
state.ended = true;
}
let chunk = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(chunk) = chunk else {
panic!("expected the retained suffix, got {chunk:?}");
};
assert_ne!(chunk, b"REPAINT", "the repaint was evicted by the burst");
assert!(reader.took_reset(), "a reset did occur");
assert!(
reader.took_drop(),
"the repaint was evicted after the reset, so the reader reports a \
drop and the pump resyncs instead of showing a torn suffix"
);
}
#[test]
fn reader_try_next_batch_drains_buffered_chunks() {
let shared = test_shared();
{
let mut state = lock(&shared.state);
state.stdout.append(b"one".to_vec());
state.stdout.append(b"two".to_vec());
state.stdout.append(b"three".to_vec());
state.ended = true;
}
let mut reader = shared_reader(&shared, OutputStream::Stdout);
let first = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(first) = first else {
panic!("expected the first buffered chunk, got {first:?}");
};
assert_eq!(first, b"one");
assert_eq!(reader.try_next(), Some(b"two".to_vec()));
assert_eq!(reader.try_next(), Some(b"three".to_vec()));
assert!(reader.try_next().is_none(), "the ring is drained");
assert!(!reader.took_drop(), "no eviction, so no drop is latched");
}
#[test]
fn reader_reports_a_front_clip_as_a_drop() {
let shared = test_shared();
{
let mut state = lock(&shared.state);
state.stdout.append(vec![b'a'; STREAM_BUFFER_CAP_BYTES]);
}
let mut reader = shared_reader(&shared, OutputStream::Stdout);
{
let mut state = lock(&shared.state);
state.stdout.append(b"tail".to_vec());
state.ended = true;
assert_eq!(
state.stdout.first_idx, 0,
"the front piece was clipped, not evicted"
);
}
let step = reader.next(Duration::from_millis(50));
let ReadStep::Chunk(_) = step else {
panic!("expected the clipped front piece, got {step:?}");
};
assert!(
reader.took_drop(),
"an in-place front clip of the parked piece is a drop"
);
}
fn shared_reader(shared: &Arc<ExecShared>, which: OutputStream) -> StreamReader {
StreamReader {
shared: shared.clone(),
which,
cursor: 0,
dropped: false,
reset: false,
seen_front_clips: 0,
seen_resets: 0,
}
}
use proptest::prelude::*;
proptest! {
#[test]
fn ring_without_overflow_is_lossless(
pieces in proptest::collection::vec(proptest::collection::vec(any::<u8>(), 0..512), 0..32)
) {
let concat: Vec<u8> = 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_byte_suffix_at_cap(
overflow in 1usize..8192,
tail_pieces in proptest::collection::vec(proptest::collection::vec(any::<u8>(), 0..64), 0..8),
) {
let head = vec![b'h'; STREAM_BUFFER_CAP_BYTES + overflow];
let mut full = head.clone();
let mut ring = Ring::default();
ring.append(head);
for piece in &tail_pieces {
ring.append(piece.clone());
full.extend_from_slice(piece);
}
prop_assert!(ring.dropped);
prop_assert_eq!(ring.size, STREAM_BUFFER_CAP_BYTES);
let tail = ring.tail();
prop_assert_eq!(tail.len(), STREAM_BUFFER_CAP_BYTES);
prop_assert!(full.ends_with(&tail));
}
}
proptest! {
#[test]
fn pump_snapshot_assigns_basis_and_replaces_ring(
pre_seqs in proptest::collection::vec(0i64..10_000, 0..16),
basis in 0i64..10_000,
) {
let shared = test_shared();
let mut pump = Pump::new(shared.clone());
for seq in pre_seqs {
pump.apply_frame(chunk(OutputStream::Stdout, seq, b"x"));
}
pump.apply_frame(snapshot_frame(b"repaint", basis));
prop_assert_eq!(pump.stdout_seq, basis);
prop_assert_eq!(lock(&shared.state).stdout.tail(), b"repaint".to_vec());
}
}
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;
}
}
}
}