diff --git a/src/boxlite/src/portal/interfaces/exec.rs b/src/boxlite/src/portal/interfaces/exec.rs index e9c9e0e37..6c4fea036 100644 --- a/src/boxlite/src/portal/interfaces/exec.rs +++ b/src/boxlite/src/portal/interfaces/exec.rs @@ -409,6 +409,23 @@ impl ExecProtocol { tracing::trace!(len = chunk.data.len(), "Received exec stderr"); stderr.send_bytes(chunk.data); } + Some(exec_output::Event::Dropped(dropped)) => { + tracing::warn!( + stdout_bytes = dropped.stdout_bytes, + stderr_bytes = dropped.stderr_bytes, + "Guest output buffer dropped older output" + ); + if dropped.stdout_bytes > 0 { + stdout.flush(); + } + if dropped.stderr_bytes > 0 { + stderr.flush(); + } + let _ = stderr.tx.send(format!( + "[boxlite] output dropped (stdout: {} bytes, stderr: {} bytes)\n", + dropped.stdout_bytes, dropped.stderr_bytes + )); + } None => {} } } @@ -1162,6 +1179,45 @@ mod tests { assert!(stderr_rx.try_recv().is_err()); } + #[test] + fn output_drop_only_flushes_the_affected_utf8_decoder() { + use boxlite_shared::{OutputDropped, Stderr as StderrMsg, exec_output}; + + let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); + let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); + let mut stdout = DecodedStream::new(stdout_tx); + let mut stderr = DecodedStream::new(stderr_tx); + + let output = |event| ExecOutput { event: Some(event) }; + ExecProtocol::route_output( + output(exec_output::Event::Stderr(StderrMsg { data: vec![0xE2] })), + &mut stdout, + &mut stderr, + ); + ExecProtocol::route_output( + output(exec_output::Event::Dropped(OutputDropped { + stdout_bytes: 1, + stderr_bytes: 0, + })), + &mut stdout, + &mut stderr, + ); + ExecProtocol::route_output( + output(exec_output::Event::Stderr(StderrMsg { + data: vec![0x94, 0x80], + })), + &mut stdout, + &mut stderr, + ); + + assert!(stdout_rx.try_recv().is_err()); + assert_eq!( + stderr_rx.try_recv().ok(), + Some("[boxlite] output dropped (stdout: 1 bytes, stderr: 0 bytes)\n".to_string()) + ); + assert_eq!(stderr_rx.try_recv().ok(), Some("─".to_string())); + } + /// Flushing a DecodedStream must drain held-over bytes, leave the /// decoder in a valid drained state, and be idempotent. The attach loop /// flushes both streams on every exit path (clean EOF, transport error, diff --git a/src/boxlite/tests/run_main_command.rs b/src/boxlite/tests/run_main_command.rs index 1c637c427..21893b7b0 100644 --- a/src/boxlite/tests/run_main_command.rs +++ b/src/boxlite/tests/run_main_command.rs @@ -63,6 +63,92 @@ async fn attached_stdout(opts: BoxOptions) -> String { stdout } +#[tokio::test] +async fn main_command_exits_after_large_output_without_attach() { + let home = boxlite_test_utils::home::PerTestBoxHome::new(); + let runtime = boxlite::BoxliteRuntime::new(boxlite::runtime::options::BoxliteOptions { + home_dir: home.path.clone(), + image_registries: common::test_registries(), + }) + .expect("create runtime"); + + let handle = runtime + .create( + main_command_opts(&["sh", "-c", "head -c 1048576 /dev/zero; exit 23"], false), + None, + ) + .await + .expect("create box"); + + let completed = tokio::time::timeout(std::time::Duration::from_secs(30), async { + handle.start().await.expect("start box"); + loop { + if handle.info().status == boxlite::BoxStatus::Stopped { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + }) + .await + .unwrap_or(false); + + let _ = handle.stop().await; + let _ = runtime.remove(handle.id().as_str(), true).await; + let _ = runtime.shutdown(Some(common::TEST_SHUTDOWN_TIMEOUT)).await; + + assert!( + completed, + "a main command that has no Attach consumer must still drain and exit" + ); +} + +#[tokio::test] +async fn late_attach_reports_output_dropped() { + let home = boxlite_test_utils::home::PerTestBoxHome::new(); + let runtime = boxlite::BoxliteRuntime::new(boxlite::runtime::options::BoxliteOptions { + home_dir: home.path.clone(), + image_registries: common::test_registries(), + }) + .expect("create runtime"); + + let handle = runtime + .create( + main_command_opts(&["sh", "-c", "head -c 2097152 /dev/zero; sleep 30"], false), + None, + ) + .await + .expect("create box"); + + handle.start().await.expect("start box"); + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + + let mut execution = handle + .attach(None) + .await + .expect("attach to the main command"); + let mut stderr = execution.stderr().expect("stderr stream"); + let dropped = tokio::time::timeout(std::time::Duration::from_secs(5), async { + while let Some(chunk) = stderr.next().await { + if chunk.contains("[boxlite] output dropped") { + return Some(chunk); + } + } + None + }) + .await + .unwrap_or(None); + + let _ = handle.stop().await; + let _ = runtime.remove(handle.id().as_str(), true).await; + let _ = runtime.shutdown(Some(common::TEST_SHUTDOWN_TIMEOUT)).await; + + let dropped = dropped.expect("late attach must report overwritten output"); + assert!( + dropped.contains("stdout:") && dropped.contains("stderr: 0 bytes"), + "stdout-only loss must preserve the stderr decoder: {dropped:?}" + ); +} + /// `tty: true` must give the *main command* a real terminal. /// /// `test -t 0` asks the kernel, so this cannot pass unless init's fd 0 really diff --git a/src/guest/src/service/exec/mod.rs b/src/guest/src/service/exec/mod.rs index e46f9c432..3ce6321e3 100644 --- a/src/guest/src/service/exec/mod.rs +++ b/src/guest/src/service/exec/mod.rs @@ -19,6 +19,7 @@ #[cfg(target_os = "linux")] pub mod exec_handle; pub(in crate::service) mod executor; +mod output; pub(in crate::service) mod registry; pub(in crate::service) mod state; mod timeout; @@ -35,7 +36,6 @@ use boxlite_shared::{ }; use futures::stream::Stream; use std::pin::Pin; -use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; use tracing::{debug, info, warn}; @@ -82,11 +82,9 @@ impl Execution for GuestServer { .ok_or_else(|| Status::not_found(format!("Execution not found: {}", exec_id)))?; // Call state directly - let rx = state.attach(&exec_id).await?; + let output = state.attach().await?; - Ok(Response::new( - Box::pin(ReceiverStream::new(rx)) as Self::AttachStream - )) + Ok(Response::new(output)) } async fn send_input( diff --git a/src/guest/src/service/exec/output.rs b/src/guest/src/service/exec/output.rs new file mode 100644 index 000000000..2de51035c --- /dev/null +++ b/src/guest/src/service/exec/output.rs @@ -0,0 +1,225 @@ +use crate::service::exec::exec_handle::{ExecStderr, ExecStdout}; +use async_stream::stream; +use boxlite_shared::{exec_output, ExecOutput, OutputDropped, Stderr, Stdout}; +use futures::{Stream, StreamExt}; +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use tokio::sync::{watch, Mutex}; +use tonic::Status; + +const BUFFER_CAPACITY_BYTES: usize = 1024 * 1024; + +pub(crate) type OutputStream = Pin> + Send>>; + +#[derive(Clone)] +pub(crate) struct OutputManager { + inner: Arc>, + updated: watch::Sender<()>, +} + +struct OutputState { + entries: VecDeque, + buffered_bytes: usize, + oldest_sequence: u64, + next_sequence: u64, + pending_dropped: DroppedBytes, + open_readers: usize, + attached: bool, +} + +struct OutputEntry { + sequence: u64, + output: ExecOutput, + source: OutputSource, + byte_len: usize, +} + +#[derive(Clone, Copy)] +enum OutputSource { + Stdout, + Stderr, +} + +#[derive(Default)] +struct DroppedBytes { + stdout: u64, + stderr: u64, +} + +impl DroppedBytes { + fn record(&mut self, source: OutputSource, byte_len: usize) { + let byte_len = byte_len as u64; + match source { + OutputSource::Stdout => self.stdout += byte_len, + OutputSource::Stderr => self.stderr += byte_len, + } + } + + fn take(&mut self) -> Self { + std::mem::take(self) + } +} + +impl OutputManager { + pub(crate) fn new(stdout: Option, stderr: Option) -> Self { + let open_readers = usize::from(stdout.is_some()) + usize::from(stderr.is_some()); + let (updated, _) = watch::channel(()); + let manager = Self { + inner: Arc::new(Mutex::new(OutputState { + entries: VecDeque::new(), + buffered_bytes: 0, + oldest_sequence: 0, + next_sequence: 0, + pending_dropped: DroppedBytes::default(), + open_readers, + attached: false, + })), + updated, + }; + + if let Some(stdout) = stdout { + manager.spawn_stdout(stdout); + } + if let Some(stderr) = stderr { + manager.spawn_stderr(stderr); + } + + manager + } + + pub(crate) async fn attach(&self) -> Result { + { + let mut state = self.inner.lock().await; + if state.attached { + return Err(Status::already_exists("Already attached")); + } + state.attached = true; + } + + let manager = self.clone(); + let output = stream! { + let mut next_sequence = 0; + let mut updates = manager.updated.subscribe(); + + loop { + enum Next { + Item(ExecOutput), + Wait, + Done, + } + + let next = { + let mut state = manager.inner.lock().await; + + if next_sequence < state.oldest_sequence { + next_sequence = state.oldest_sequence; + Next::Item(dropped_output(state.pending_dropped.take())) + } else if next_sequence < state.next_sequence { + let index = (next_sequence - state.oldest_sequence) as usize; + let entry = state.entries.get(index).expect("ring sequence must exist"); + next_sequence += 1; + Next::Item(entry.output.clone()) + } else if state.open_readers == 0 { + Next::Done + } else { + Next::Wait + } + }; + + match next { + Next::Item(item) => yield Ok(item), + Next::Done => break, + Next::Wait => { + if updates.changed().await.is_err() { + break; + } + } + } + } + }; + + Ok(Box::pin(output)) + } + + fn spawn_stdout(&self, stdout: ExecStdout) { + let manager = self.clone(); + tokio::spawn(async move { + manager.drain(stdout, OutputSource::Stdout).await; + }); + } + + fn spawn_stderr(&self, stderr: ExecStderr) { + let manager = self.clone(); + tokio::spawn(async move { + manager.drain(stderr, OutputSource::Stderr).await; + }); + } + + async fn drain(&self, mut stream: S, source: OutputSource) + where + S: Stream> + Unpin, + { + while let Some(data) = stream.next().await { + self.push(source, data).await; + } + self.reader_finished().await; + } + + async fn push(&self, source: OutputSource, data: Vec) { + let byte_len = data.len(); + let output = match source { + OutputSource::Stdout => ExecOutput { + event: Some(exec_output::Event::Stdout(Stdout { data })), + }, + OutputSource::Stderr => ExecOutput { + event: Some(exec_output::Event::Stderr(Stderr { data })), + }, + }; + + let mut state = self.inner.lock().await; + let sequence = state.next_sequence; + state.next_sequence += 1; + + while state.buffered_bytes + byte_len > BUFFER_CAPACITY_BYTES { + let Some(removed) = state.entries.pop_front() else { + state.pending_dropped.record(source, byte_len); + state.oldest_sequence = state.next_sequence; + break; + }; + state.buffered_bytes -= removed.byte_len; + state + .pending_dropped + .record(removed.source, removed.byte_len); + state.oldest_sequence = removed.sequence + 1; + } + + if byte_len <= BUFFER_CAPACITY_BYTES { + state.buffered_bytes += byte_len; + state.entries.push_back(OutputEntry { + sequence, + output, + source, + byte_len, + }); + } + drop(state); + self.updated.send_replace(()); + } + + async fn reader_finished(&self) { + let mut state = self.inner.lock().await; + state.open_readers -= 1; + drop(state); + self.updated.send_replace(()); + } +} + +fn dropped_output(dropped: DroppedBytes) -> ExecOutput { + ExecOutput { + event: Some(exec_output::Event::Dropped(OutputDropped { + stdout_bytes: dropped.stdout, + stderr_bytes: dropped.stderr, + })), + } +} diff --git a/src/guest/src/service/exec/state.rs b/src/guest/src/service/exec/state.rs index af93b9015..89078dc7c 100644 --- a/src/guest/src/service/exec/state.rs +++ b/src/guest/src/service/exec/state.rs @@ -1,11 +1,10 @@ use crate::service::exec::exec_handle::ExecHandle; -use boxlite_shared::ExecOutput; +use crate::service::exec::output::{OutputManager, OutputStream}; use std::os::unix::io::AsRawFd; use std::sync::Arc; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::Mutex; use tokio::task::JoinHandle; use tonic::Status; -use tracing::info; /// Abstraction for checking container init health. /// @@ -25,8 +24,7 @@ pub(crate) trait InitHealthCheck: Send + Sync { struct Inner { /// The process handle (owns pid, pty_controller, stdin, stdout, stderr) handle: Option, - /// Stdout/stderr forwarding tasks (set on attach) - output_tasks: Vec>, + output: OutputManager, /// Timeout flag #[allow(dead_code)] // Will be used for timeout handling timed_out: bool, @@ -35,10 +33,6 @@ struct Inner { init_health: Option>>, } -/// Execution state. -/// -/// Handle owns pid, pty_controller, stdin, stdout, stderr. -/// stdin is taken on send_input(), stdout/stderr are taken on attach(). #[derive(Clone)] pub(crate) struct ExecutionState { inner: Arc>, @@ -49,24 +43,25 @@ pub(crate) struct ExecutionState { } impl ExecutionState { - fn from_inner(inner: Inner, exit: crate::reaper::ExitSlot) -> Self { + fn from_handle( + mut handle: ExecHandle, + init_health: Option>>, + exit: crate::reaper::ExitSlot, + ) -> Self { Self { - inner: Arc::new(Mutex::new(inner)), + inner: Arc::new(Mutex::new(Inner { + output: OutputManager::new(handle.stdout(), handle.stderr()), + handle: Some(handle), + timed_out: false, + init_health, + })), exit, } } /// Create new execution state for a guest-side process. pub(super) fn new(handle: ExecHandle, exit: crate::reaper::ExitSlot) -> Self { - Self::from_inner( - Inner { - handle: Some(handle), - output_tasks: Vec::new(), - timed_out: false, - init_health: None, - }, - exit, - ) + Self::from_handle(handle, None, exit) } /// Create execution state with an init health checker. @@ -78,15 +73,7 @@ impl ExecutionState { init_health: Arc>, exit: crate::reaper::ExitSlot, ) -> Self { - Self::from_inner( - Inner { - handle: Some(handle), - output_tasks: Vec::new(), - timed_out: false, - init_health: Some(init_health), - }, - exit, - ) + Self::from_handle(handle, Some(init_health), exit) } /// Create execution state for the container's init process itself. @@ -95,15 +82,7 @@ impl ExecutionState { /// reparents to guest main (the boxlite-guest agent process), which owns /// `waitpid(-1)`. See `wait_process`. pub(crate) fn new_init_session(handle: ExecHandle, exit: crate::reaper::ExitSlot) -> Self { - Self::from_inner( - Inner { - handle: Some(handle), - output_tasks: Vec::new(), - timed_out: false, - init_health: None, - }, - exit, - ) + Self::from_handle(handle, None, exit) } /// Check if the container init process died. @@ -192,84 +171,9 @@ impl ExecutionState { self.exit.get().await } - /// Attach to execution output. - /// - /// Takes stdout/stderr from handle and starts forwarding tasks. - /// Returns stream of output chunks. - pub async fn attach( - &self, - exec_id: &str, - ) -> Result>, Status> { - use boxlite_shared::{exec_output, Stderr, Stdout}; - use futures::StreamExt; - - let (tx, rx) = mpsc::channel(100); - - // Take stdout/stderr from handle - let (stdout, stderr) = { - let mut inner = self.inner.lock().await; - - if !inner.output_tasks.is_empty() { - return Err(Status::already_exists("Already attached")); - } - - let handle = inner - .handle - .as_mut() - .ok_or_else(|| Status::failed_precondition("Handle not available"))?; - - let stdout = handle.stdout(); - let stderr = handle.stderr(); - - (stdout, stderr) - }; - - // Spawn forwarding tasks - let mut tasks = Vec::new(); - - // Spawn stdout forwarding task - let exec_id_string = exec_id.to_string(); - if let Some(mut stdout) = stdout { - let tx = tx.clone(); - let handle = tokio::spawn(async move { - while let Some(chunk) = stdout.next().await { - let msg = ExecOutput { - event: Some(exec_output::Event::Stdout(Stdout { data: chunk })), - }; - if tx.send(Ok(msg)).await.is_err() { - break; - } - } - info!(execution = ?exec_id_string, "Stdout forwarding task ended"); - }); - tasks.push(handle); - } - - // Spawn stderr forwarding task - let exec_id_string = exec_id.to_string(); - if let Some(mut stderr) = stderr { - let tx = tx.clone(); - let handle = tokio::spawn(async move { - while let Some(chunk) = stderr.next().await { - let msg = ExecOutput { - event: Some(exec_output::Event::Stderr(Stderr { data: chunk })), - }; - if tx.send(Ok(msg)).await.is_err() { - break; - } - } - info!(execution = ?exec_id_string, "Stderr forwarding task ended"); - }); - tasks.push(handle); - } - - // Store tasks - { - let mut inner = self.inner.lock().await; - inner.output_tasks = tasks; - } - - Ok(rx) + pub async fn attach(&self) -> Result { + let output = self.inner.lock().await.output.clone(); + output.attach().await } /// Kill process with signal. diff --git a/src/shared/proto/boxlite/v1/service.proto b/src/shared/proto/boxlite/v1/service.proto index 4069a9483..4fb6495a6 100644 --- a/src/shared/proto/boxlite/v1/service.proto +++ b/src/shared/proto/boxlite/v1/service.proto @@ -354,9 +354,15 @@ message ExecOutput { oneof event { Stdout stdout = 1; Stderr stderr = 2; + OutputDropped dropped = 3; } } +message OutputDropped { + uint64 stdout_bytes = 1; + uint64 stderr_bytes = 2; +} + message Stdout { bytes data = 1; }