diff --git a/src/guest/src/service/container.rs b/src/guest/src/service/container.rs index 0f02de58c..bc91fc3ec 100644 --- a/src/guest/src/service/container.rs +++ b/src/guest/src/service/container.rs @@ -483,9 +483,24 @@ impl ContainerService for GuestServer { exit.clone(), process, ); - self.registry - .register(init_execution_id.clone(), state) - .await; + if !self + .registry + .register(init_execution_id.clone(), state.clone()) + .await + { + state.abort_unpublished().await; + // Init's session is not shutdown-managed, so the + // teardown above leaves the slot — and the power-off + // action parked on it — in place. + reaper.release_slot(&exit); + return Ok(Response::new(ContainerInitResponse { + result: Some(container_init_response::Result::Error( + internal_init_error( + "Execution registry is shutting down or already owns the init session", + ), + )), + })); + } exit } Ok(None) => { diff --git a/src/guest/src/service/exec/mod.rs b/src/guest/src/service/exec/mod.rs index 6a010322b..bdc80971a 100644 --- a/src/guest/src/service/exec/mod.rs +++ b/src/guest/src/service/exec/mod.rs @@ -32,7 +32,8 @@ pub(crate) use state::InitHealthCheck; use crate::service::exec::error::ExecutionError; use crate::service::exec::executor::{ContainerExecutor, GuestExecutor}; -use crate::service::exec::state::{ExecutionExit, ExecutionState}; +use crate::service::exec::registry::ExecutionLookup; +use crate::service::exec::state::ExecutionExit; use crate::service::server::GuestServer; use boxlite_shared::{ constants::executor as executor_const, AttachRequest, ExecError, ExecOutput, ExecRequest, @@ -55,9 +56,9 @@ use tracing::{debug, info, warn}; /// privileged, and neither can drift from the other, because the behaviour /// lives here rather than in either adapter. impl GuestServer { - async fn execution(&self, exec_id: &str) -> Result { + async fn execution_lookup(&self, exec_id: &str) -> Result { self.registry - .get(exec_id) + .lookup(exec_id) .await .ok_or_else(|| ExecutionError::NotFound(exec_id.to_string())) } @@ -68,7 +69,28 @@ impl GuestServer { exec_id: &str, ) -> Result>, ExecutionError> { info!(execution_id = %exec_id, "attach request"); - self.execution(exec_id).await?.attach(exec_id).await + match self.execution_lookup(exec_id).await? { + ExecutionLookup::Live(state) => match state.attach(exec_id).await { + Ok(output) => Ok(output), + Err(ExecutionError::AlreadyAttached) => state.attach_retained(exec_id).await, + Err(error @ ExecutionError::HandleUnavailable) => state + .sealed_terminal_output_summary() + .await + .map(terminal_output_receiver) + .ok_or(error), + Err(error) => Err(error), + }, + ExecutionLookup::Retained { state, snapshot } => { + match state.attach_retained(exec_id).await { + Ok(output) => Ok(output), + Err(ExecutionError::HandleUnavailable) => { + Ok(terminal_output_receiver(snapshot.output)) + } + Err(error) => Err(error), + } + } + ExecutionLookup::Tombstone(snapshot) => Ok(terminal_output_receiver(snapshot.output)), + } } /// Forward `first` and then everything `stream` yields to the execution's @@ -84,10 +106,12 @@ impl GuestServer { )); } let exec_id = first.execution_id.clone(); - self.execution(&exec_id) - .await? - .send_input(first, stream) - .await + match self.execution_lookup(&exec_id).await? { + ExecutionLookup::Live(state) => state.send_input(first, stream).await, + ExecutionLookup::Retained { .. } | ExecutionLookup::Tombstone(_) => { + Err(ExecutionError::HandleUnavailable) + } + } } /// Wait for exit, already classified — see [`ExecutionState::wait_exit`]. @@ -96,7 +120,12 @@ impl GuestServer { exec_id: &str, ) -> Result { debug!(execution_id = %exec_id, "wait request"); - Ok(self.execution(exec_id).await?.wait_exit(exec_id).await) + match self.execution_lookup(exec_id).await? { + ExecutionLookup::Live(state) => Ok(state.wait_exit(exec_id).await), + ExecutionLookup::Retained { snapshot, .. } | ExecutionLookup::Tombstone(snapshot) => { + Ok(snapshot.exit) + } + } } /// Signal the execution. `false` means the process had already exited. @@ -110,10 +139,13 @@ impl GuestServer { // Resolve before parsing: signal 0 is the conventional liveness probe and // is not a `Signal`, so parsing first would report an unknown execution // as an invalid argument. - let state = self.execution(exec_id).await?; + let state = self.execution_lookup(exec_id).await?; let parsed = nix::sys::signal::Signal::try_from(signal) .map_err(|_| ExecutionError::InvalidArgument(format!("signal number {signal}")))?; - let sent = state.kill(parsed, process_group).await; + let sent = match state { + ExecutionLookup::Live(state) => state.kill(parsed, process_group).await, + ExecutionLookup::Retained { .. } | ExecutionLookup::Tombstone(_) => false, + }; if sent { info!(execution_id = %exec_id, signal, "signal sent"); } else { @@ -152,23 +184,22 @@ impl GuestServer { let y_pixels = dimension("y_pixels", y_pixels)?; info!(execution_id = %exec_id, rows, cols, "resize_tty request"); - Ok( - match self - .execution(exec_id) - .await? - .resize_pty(rows, cols, x_pixels, y_pixels) - .await - { - Ok(()) => { - info!(execution_id = %exec_id, rows, cols, "tty resized"); - TtyResize::Resized - } - Err(error) => { - info!(execution_id = %exec_id, %error, "failed to resize tty"); - TtyResize::Rejected(error) - } - }, - ) + let outcome = match self.execution_lookup(exec_id).await? { + ExecutionLookup::Live(state) => state.resize_pty(rows, cols, x_pixels, y_pixels).await, + ExecutionLookup::Retained { .. } | ExecutionLookup::Tombstone(_) => { + Err(ExecutionError::HandleUnavailable) + } + }; + Ok(match outcome { + Ok(()) => { + info!(execution_id = %exec_id, rows, cols, "tty resized"); + TtyResize::Resized + } + Err(error) => { + info!(execution_id = %exec_id, %error, "failed to resize tty"); + TtyResize::Rejected(error) + } + }) } } @@ -273,6 +304,22 @@ impl Execution for GuestServer { } } +fn terminal_output_receiver( + summary: output::OutputTerminalSummary, +) -> mpsc::Receiver> { + let (tx, rx) = mpsc::channel(2); + if let Some(failure) = summary.reader_failure { + tx.try_send(Err(Status::internal(failure))) + .expect("terminal attach receiver must be live"); + } else { + for event in output::terminal_events(&summary) { + tx.try_send(Ok(event)) + .expect("terminal attach receiver must be live"); + } + } + rx +} + /// Start a typed workload selected by the in-process SSH server. This entry /// point is intentionally outside the public Execution RPC contract. pub(crate) async fn start_ssh_execution( @@ -288,17 +335,60 @@ async fn start_execution( req: ExecRequest, ssh_workload: Option, ) -> ExecResponse { - let execution_id = req - .execution_id - .clone() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let execution_id = + match execution_id_for_request(req.execution_id.as_deref(), ssh_workload.is_some()) { + Ok(id) => id, + Err(detail) => { + return error_response( + req.execution_id.clone().unwrap_or_default(), + "invalid_argument", + detail, + ); + } + }; - if server.registry.exists(&execution_id).await { - return error_response(execution_id, "execution_exists", "Execution already exists"); + let reservation = if ssh_workload.is_some() { + if server.registry.exists(&execution_id).await { + return error_response(execution_id, "execution_exists", "Execution already exists"); + } + None + } else { + match server.registry.reserve(execution_id.clone()).await { + Some(reservation) => Some(reservation), + None => { + return error_response( + execution_id, + "execution_exists", + "Execution already exists", + ); + } + } + }; + + match spawn_execution( + server, + execution_id, + req, + ssh_workload, + reservation.as_ref(), + ) + .await + { + Ok(response) => response, + Err(response) => { + if let Some(reservation) = reservation { + server.registry.release_reservation(&reservation).await; + } + response + } } +} - match spawn_execution(server, execution_id, req, ssh_workload).await { - Ok(response) | Err(response) => response, +fn execution_id_for_request(requested: Option<&str>, is_ssh: bool) -> Result { + match (requested, is_ssh) { + (Some(_), false) => Err("execution_id is assigned by the guest"), + (Some(id), true) => Ok(id.to_owned()), + (None, _) => Ok(uuid::Uuid::new_v4().to_string()), } } @@ -308,6 +398,7 @@ async fn spawn_execution( execution_id: String, req: ExecRequest, ssh_workload: Option, + reservation: Option<®istry::ExecutionReservation>, ) -> Result { let started_at_ms = now_ms(); @@ -328,6 +419,9 @@ async fn spawn_execution( ); } + // Register this pid's exit slot at the spawn, not at the wait: a detached + // exec sends no Wait until its caller chooses to, and in between the exit + // would age out as an ownerless stray. let reaper = crate::reaper::REAPER .get() .expect("reaper installed at startup"); @@ -342,18 +436,42 @@ async fn spawn_execution( } None => state::ExecutionState::new(child, exit, process), }; - server + if let Some(reservation) = reservation { + if !server.registry.publish(reservation, state.clone()).await { + state.abort_unpublished().await; + return Err(error_response( + execution_id, + "execution_cancelled", + "Execution reservation was released before spawn completed", + )); + } + } else if !server .registry .register(execution_id.clone(), state.clone()) - .await; + .await + { + state.abort_unpublished().await; + return Err(error_response( + execution_id, + "execution_cancelled", + "Execution registry is shutting down", + )); + } // Step 3: Start timeout watcher (if requested) if req.timeout_ms > 0 { - timeout::start_timeout_watcher( + let timeout_task = timeout::start_timeout_watcher( timeout::TimeoutTarget::new(process), execution_id.clone(), std::time::Duration::from_millis(req.timeout_ms), ); + state.set_timeout_task(timeout_task).await; + } + + if reservation.is_some() { + server + .registry + .observe_terminal(execution_id.clone(), state.clone()); } Ok(ExecResponse { @@ -551,7 +669,15 @@ fn is_single_path_component(id: &str) -> bool { #[cfg(test)] mod container_id_path_tests { - use super::is_single_path_component; + use super::{execution_id_for_request, is_single_path_component, GuestServer, TtyResize}; + use crate::layout::GuestLayout; + use crate::reaper::ExitSlot; + use crate::service::exec::error::ExecutionError; + use crate::service::exec::exec_handle::{ExecHandle, ExitStatus}; + use crate::service::exec::output::{OutputStreamSummary, OutputTerminalSummary}; + use crate::service::exec::state::{ExecutionExit, ExecutionState, TerminalSnapshot}; + use boxlite_shared::{exec_output, ExecStdin}; + use nix::unistd::{pipe, write, Pid}; #[test] fn rejects_ids_that_would_escape_the_containers_dir() { @@ -562,4 +688,264 @@ mod container_id_path_tests { "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" )); } + + #[test] + fn public_exec_rejects_a_caller_supplied_execution_id() { + assert!(execution_id_for_request(Some("caller-id"), false).is_err()); + assert_eq!( + execution_id_for_request(Some("ssh-id"), true).unwrap(), + "ssh-id" + ); + } + + #[tokio::test] + async fn tombstone_routes_terminal_rpcs_without_reopening_the_process() { + let server = GuestServer::new(GuestLayout::new()); + server + .registry + .store_tombstone( + "terminal-exec".into(), + TerminalSnapshot { + exit: ExecutionExit { + exit_code: 9, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 4, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }, + ) + .await; + + assert_eq!( + server + .wait_execution("terminal-exec") + .await + .unwrap() + .exit_code, + 9 + ); + assert_eq!( + server + .wait_execution("terminal-exec") + .await + .unwrap() + .exit_code, + 9 + ); + + let mut output = server.attach_execution("terminal-exec").await.unwrap(); + let event = output.recv().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = event.event else { + panic!("tombstone attach must return stdout terminal event"); + }; + assert!(stdout.data.is_empty()); + assert_eq!(stdout.offset, Some(4)); + assert_eq!(stdout.total_bytes, Some(4)); + + let input = ExecStdin { + execution_id: "terminal-exec".into(), + data: Vec::new(), + close: true, + }; + assert!(matches!( + server + .send_execution_input( + input, + futures::stream::empty::>() + ) + .await, + Err(ExecutionError::HandleUnavailable) + )); + assert!(!server + .kill_execution("terminal-exec", 15, false) + .await + .unwrap()); + assert!(matches!( + server + .resize_execution_tty("terminal-exec", 24, 80, 0, 0) + .await, + Ok(TtyResize::Rejected(ExecutionError::HandleUnavailable)) + )); + } + + #[tokio::test] + async fn reader_failure_tombstone_attach_returns_the_stored_internal_error() { + let server = GuestServer::new(GuestLayout::new()); + server + .registry + .store_tombstone( + "reader-failure".into(), + TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 7, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: Some("stdout reader failed".into()), + }, + }, + ) + .await; + + let mut output = server.attach_execution("reader-failure").await.unwrap(); + let error = output.recv().await.unwrap().unwrap_err(); + assert_eq!(error.code(), tonic::Code::Internal); + assert_eq!(error.message(), "stdout reader failed"); + assert!(output.recv().await.is_none()); + } + + #[tokio::test] + async fn retained_attach_returns_terminal_output_after_the_ring_is_sealed() { + let server = GuestServer::new(GuestLayout::new()); + let (_stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + drop(stdout_peer); + drop(stderr_peer); + let handle = ExecHandle::new(Pid::from_raw(13_001), stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"); + let state = + ExecutionState::new_for_test(handle, ExitSlot::settled_for_test(ExitStatus::Code(0))); + let snapshot = state.wait_terminal_snapshot("retained-exec").await; + server + .registry + .register("retained-exec".into(), state) + .await; + assert!(server.registry.retain("retained-exec", snapshot, 0).await); + + let mut output = server.attach_execution("retained-exec").await.unwrap(); + let event = output.recv().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = event.event else { + panic!("retained attach must return stdout terminal event"); + }; + assert!(stdout.data.is_empty()); + assert_eq!(stdout.offset, Some(0)); + assert_eq!(stdout.total_bytes, Some(0)); + } + + #[tokio::test] + async fn retained_attach_replays_buffered_output_before_its_terminal_event() { + let server = GuestServer::new(GuestLayout::new()); + let (_stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + write(&stdout_peer, b"snapshot-data").unwrap(); + drop(stdout_peer); + drop(stderr_peer); + let handle = ExecHandle::new(Pid::from_raw(13_004), stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"); + let state = + ExecutionState::new_for_test(handle, ExitSlot::settled_for_test(ExitStatus::Code(0))); + let snapshot = state.wait_terminal_snapshot("retained-output-exec").await; + server + .registry + .register("retained-output-exec".into(), state) + .await; + assert!( + server + .registry + .retain("retained-output-exec", snapshot, b"snapshot-data".len()) + .await + ); + + let mut output = server + .attach_execution("retained-output-exec") + .await + .unwrap(); + let mut saw_data = false; + while let Some(event) = output.recv().await { + let Some(exec_output::Event::Stdout(stdout)) = event.unwrap().event else { + continue; + }; + if !saw_data { + assert_eq!(stdout.data, b"snapshot-data"); + assert_eq!(stdout.offset, Some(0)); + assert_eq!(stdout.total_bytes, None); + saw_data = true; + continue; + } + assert!(stdout.data.is_empty()); + assert_eq!(stdout.offset, Some(b"snapshot-data".len() as u64)); + assert_eq!(stdout.total_bytes, Some(b"snapshot-data".len() as u64)); + return; + } + panic!("retained attach must replay stdout before its terminal event"); + } + + #[tokio::test] + async fn sealed_live_attach_returns_terminal_output_while_retention_is_pending() { + let server = GuestServer::new(GuestLayout::new()); + let (_stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + drop(stdout_peer); + drop(stderr_peer); + let handle = ExecHandle::new(Pid::from_raw(13_002), stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"); + let state = + ExecutionState::new_for_test(handle, ExitSlot::settled_for_test(ExitStatus::Code(0))); + state.wait_terminal_output_summary().await; + server + .registry + .register("sealed-live-exec".into(), state) + .await; + + let mut output = server.attach_execution("sealed-live-exec").await.unwrap(); + let event = output.recv().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = event.event else { + panic!("sealed live attach must return stdout terminal event"); + }; + assert!(stdout.data.is_empty()); + assert_eq!(stdout.total_bytes, Some(0)); + } + + #[tokio::test] + async fn sealed_live_attach_returns_terminal_output_after_resource_release() { + let server = GuestServer::new(GuestLayout::new()); + let (_stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + drop(stdout_peer); + drop(stderr_peer); + let handle = ExecHandle::new(Pid::from_raw(13_003), stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"); + let state = + ExecutionState::new_for_test(handle, ExitSlot::settled_for_test(ExitStatus::Code(0))); + state.wait_terminal_output_summary().await; + server + .registry + .register("released-sealed-live-exec".into(), state.clone()) + .await; + assert!(state.release_resources().await); + + let mut output = server + .attach_execution("released-sealed-live-exec") + .await + .expect("sealed live state must retain its terminal output"); + let event = output.recv().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = event.event else { + panic!("released sealed live attach must return stdout terminal event"); + }; + assert!(stdout.data.is_empty()); + assert_eq!(stdout.total_bytes, Some(0)); + } } diff --git a/src/guest/src/service/exec/output.rs b/src/guest/src/service/exec/output.rs index 47c2ac870..acd113cef 100644 --- a/src/guest/src/service/exec/output.rs +++ b/src/guest/src/service/exec/output.rs @@ -4,6 +4,7 @@ use boxlite_shared::{exec_output, ExecOutput, Stderr, Stdout}; use futures::{Stream, StreamExt}; use std::collections::VecDeque; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use tokio::sync::{watch, Mutex}; use tokio::task::JoinHandle; @@ -14,10 +15,24 @@ const BUFFER_CAPACITY_BYTES: usize = 1024 * 1024; pub(crate) type AttachStream = Pin> + Send>>; +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct OutputTerminalSummary { + pub(crate) stdout: OutputStreamSummary, + pub(crate) stderr: OutputStreamSummary, + pub(crate) reader_failure: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct OutputStreamSummary { + pub(crate) enabled: bool, + pub(crate) total_bytes: u64, +} + #[derive(Clone)] pub(crate) struct OutputManager { inner: Arc>, updated: watch::Sender<()>, + consumer_lease: Arc, drain_tasks: Arc>>>, } @@ -27,11 +42,23 @@ struct OutputState { oldest_sequence: u64, next_sequence: u64, failure: Option, - attached: bool, + sealed: bool, stdout: StreamState, stderr: StreamState, } +struct ConsumerLease { + claimed: Arc, + updated: watch::Sender<()>, +} + +impl Drop for ConsumerLease { + fn drop(&mut self) { + self.claimed.store(false, Ordering::Release); + self.updated.send_replace(()); + } +} + struct ReaderFailure { sequence: u64, message: String, @@ -68,7 +95,7 @@ impl OutputManager { oldest_sequence: 0, next_sequence: 0, failure: None, - attached: false, + sealed: false, stdout: StreamState { enabled: stdout_enabled, finished: !stdout_enabled, @@ -83,6 +110,7 @@ impl OutputManager { }, })), updated, + consumer_lease: Arc::new(AtomicBool::new(false)), drain_tasks: Arc::new(StdMutex::new(Vec::new())), }; @@ -112,16 +140,45 @@ impl OutputManager { } pub(crate) async fn attach(&self) -> Result { + let lease = self.claim_consumer().await?; + if self.inner.lock().await.sealed { + return Err(Status::failed_precondition( + "execution output is finalizing", + )); + } + Ok(self.attach_stream(lease)) + } + + pub(crate) async fn attach_retained(&self) -> Result { + let lease = self.claim_consumer().await?; + Ok(self.attach_stream(lease)) + } + + /// The lease flag itself, so a holder can test it without taking any lock. + pub(crate) fn consumer_lease_flag(&self) -> Arc { + Arc::clone(&self.consumer_lease) + } + + async fn claim_consumer(&self) -> Result { + if self + .consumer_lease + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() { - let mut state = self.inner.lock().await; - if state.attached { - return Err(Status::already_exists("Already attached")); - } - state.attached = true; + return Err(Status::already_exists("Already attached")); } + let lease = ConsumerLease { + claimed: Arc::clone(&self.consumer_lease), + updated: self.updated.clone(), + }; + Ok(lease) + } + + fn attach_stream(&self, lease: ConsumerLease) -> AttachStream { let manager = self.clone(); let output = stream! { + let _lease = lease; let mut next_sequence = 0; let mut stdout_end_sent = false; let mut stderr_end_sent = false; @@ -184,7 +241,78 @@ impl OutputManager { } }; - Ok(Box::pin(output)) + Box::pin(output) + } + + pub(crate) async fn seal(&self) -> bool { + let mut state = self.inner.lock().await; + if !state.stdout.finished || !state.stderr.finished { + return false; + } + state.sealed = true; + drop(state); + self.updated.send_replace(()); + true + } + + pub(crate) async fn terminal_summary(&self) -> Option { + let state = self.inner.lock().await; + if !state.stdout.finished || !state.stderr.finished { + return None; + } + + Some(OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: state.stdout.enabled, + total_bytes: state.stdout.total_bytes, + }, + stderr: OutputStreamSummary { + enabled: state.stderr.enabled, + total_bytes: state.stderr.total_bytes, + }, + reader_failure: state + .failure + .as_ref() + .map(|failure| failure.message.clone()), + }) + } + + pub(crate) async fn sealed_terminal_summary(&self) -> Option { + let state = self.inner.lock().await; + if !state.sealed { + return None; + } + Some(OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: state.stdout.enabled, + total_bytes: state.stdout.total_bytes, + }, + stderr: OutputStreamSummary { + enabled: state.stderr.enabled, + total_bytes: state.stderr.total_bytes, + }, + reader_failure: state + .failure + .as_ref() + .map(|failure| failure.message.clone()), + }) + } + + pub(crate) async fn wait_terminal_summary(&self) -> OutputTerminalSummary { + let mut updates = self.updated.subscribe(); + loop { + if let Some(summary) = self.terminal_summary().await { + return summary; + } + updates + .changed() + .await + .expect("output manager must outlive its waiters"); + } + } + + pub(crate) async fn retained_bytes(&self) -> usize { + self.inner.lock().await.buffered_bytes } fn spawn(&self, stream: S, source: OutputSource) @@ -324,6 +452,17 @@ fn end_output(source: OutputSource, total_bytes: u64) -> ExecOutput { ExecOutput { event: Some(event) } } +pub(crate) fn terminal_events(summary: &OutputTerminalSummary) -> Vec { + let mut events = Vec::new(); + if summary.stdout.enabled { + events.push(end_output(OutputSource::Stdout, summary.stdout.total_bytes)); + } + if summary.stderr.enabled { + events.push(end_output(OutputSource::Stderr, summary.stderr.total_bytes)); + } + events +} + #[cfg(test)] mod tests { use super::*; @@ -353,6 +492,23 @@ mod tests { assert_eq!(end.total_bytes, Some(b"already sent".len() as u64)); } + #[tokio::test] + async fn waiting_for_a_terminal_summary_returns_after_both_streams_finish() { + let manager = OutputManager::new(None, None); + + let summary = manager.wait_terminal_summary().await; + assert_eq!(summary.stdout.total_bytes, 0); + assert_eq!(summary.stderr.total_bytes, 0); + } + + #[tokio::test] + async fn retained_bytes_reports_the_shared_ring_size() { + let manager = OutputManager::new(None, None); + manager.push(OutputSource::Stdout, b"ring".to_vec()).await; + + assert_eq!(manager.retained_bytes().await, 4); + } + #[tokio::test] async fn closed_stdout_emits_its_end_before_stderr_closes() { let (stdout_read, stdout_write) = nix::unistd::pipe().unwrap(); @@ -473,4 +629,98 @@ mod tests { let error = output.next().await.unwrap().unwrap_err(); assert!(error.message().contains("stdout failure")); } + + #[tokio::test] + async fn dropping_attach_stream_releases_consumer_lease() { + let manager = OutputManager::new(None, None); + + let first = manager.attach().await.unwrap(); + let error = manager + .attach() + .await + .err() + .expect("the second Attach must be rejected"); + assert_eq!(error.code(), tonic::Code::AlreadyExists); + + drop(first); + + assert!(manager.attach().await.is_ok()); + } + + #[tokio::test] + async fn terminal_summary_records_enabled_stream_totals() { + let (stdout_read, stdout_write) = nix::unistd::pipe().unwrap(); + let manager = OutputManager::new(Some(ExecStdout::new(stdout_read).unwrap()), None); + assert!(manager.terminal_summary().await.is_none()); + + nix::unistd::write(&stdout_write, b"stdout").unwrap(); + drop(stdout_write); + + let summary = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if let Some(summary) = manager.terminal_summary().await { + break summary; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("stdout EOF must produce a terminal summary"); + assert!(summary.stdout.enabled); + assert_eq!(summary.stdout.total_bytes, 6); + assert!(!summary.stderr.enabled); + assert_eq!(summary.stderr.total_bytes, 0); + } + + #[tokio::test] + async fn seal_requires_terminal_output_and_rejects_new_attach() { + let (stdout_read, stdout_write) = nix::unistd::pipe().unwrap(); + let manager = OutputManager::new(Some(ExecStdout::new(stdout_read).unwrap()), None); + assert!(!manager.seal().await); + + drop(stdout_write); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if manager.terminal_summary().await.is_some() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("stdout EOF must make sealing possible"); + assert!(manager.seal().await); + + let error = manager + .attach() + .await + .err() + .expect("sealed output must reject Attach"); + assert_eq!(error.code(), tonic::Code::FailedPrecondition); + } + + #[tokio::test] + async fn sealing_keeps_an_existing_attach_stream_until_it_emits_terminal_output() { + let (stdout_read, stdout_write) = nix::unistd::pipe().unwrap(); + let manager = OutputManager::new(Some(ExecStdout::new(stdout_read).unwrap()), None); + let mut output = manager.attach().await.unwrap(); + + nix::unistd::write(&stdout_write, b"buffered").unwrap(); + drop(stdout_write); + manager.wait_terminal_summary().await; + assert!(manager.seal().await); + + let data = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = data.event else { + panic!("sealed stream must retain buffered stdout"); + }; + assert_eq!(stdout.data, b"buffered"); + + let end = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout)) = end.event else { + panic!("sealed stream must emit stdout terminal event"); + }; + assert!(stdout.data.is_empty()); + assert_eq!(stdout.total_bytes, Some(8)); + } } diff --git a/src/guest/src/service/exec/registry.rs b/src/guest/src/service/exec/registry.rs index 654e1c2e5..40a0831c1 100644 --- a/src/guest/src/service/exec/registry.rs +++ b/src/guest/src/service/exec/registry.rs @@ -3,7 +3,7 @@ //! Manages the state of all active executions, providing thread-safe access //! to execution metadata, I/O channels, and completion status. -use crate::service::exec::state::ExecutionState; +use crate::service::exec::state::{ExecutionState, TerminalSnapshot}; use nix::sys::signal::Signal; /// How long [`ExecutionRegistry::shutdown_all`] gives execs to die between @@ -13,41 +13,542 @@ use nix::sys::signal::Signal; /// host-driven Shutdown RPC and the guest's own power-off when the main /// command exits — drain the same execs and should wait the same amount. pub(crate) const SHUTDOWN_TIMEOUT_MS: u64 = 1000; +const RETAIN_GRACE: Duration = Duration::from_secs(5 * 60); +const TOMBSTONE_TTL: Duration = Duration::from_secs(15 * 60); +const MAX_RETAINED_ENTRIES: usize = 64; +const MAX_RETAINED_BYTES: usize = 8 * 1024 * 1024; +const MAX_TOMBSTONE_ENTRIES: usize = 1024; +const MAX_TOMBSTONE_METADATA_BYTES: usize = 5 * 1024 * 1024; +const MAX_TOMBSTONE_DIAGNOSTIC_BYTES: usize = 4 * 1024; +const TOMBSTONE_TRUNCATION_SUFFIX: &str = "…[truncated]"; use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use tokio::sync::Mutex; +use std::time::{Duration, Instant}; +use tokio::sync::{watch, Mutex}; +use tokio::task::JoinHandle; use tracing::{info, warn}; +/// Holds an execution id from before its process is spawned until the state is +/// published. +/// +/// Releases on drop, not on the error path alone: tonic drops the whole request +/// future when a client disconnects, and nothing prunes a `Reserved` entry, so a +/// cancelled Exec would otherwise hold its id for the life of the guest. +pub(crate) struct ExecutionReservation { + registry: ExecutionRegistry, + execution_id: String, + ticket: u64, + released: AtomicBool, +} + +impl Drop for ExecutionReservation { + fn drop(&mut self) { + if self.released.load(Ordering::Acquire) { + return; + } + // The registry is behind an async mutex, so the release cannot run here. + // Same shape as `ChannelBridge::drop`, which spawns its teardown. + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let registry = self.registry.clone(); + let execution_id = std::mem::take(&mut self.execution_id); + let ticket = self.ticket; + runtime.spawn(async move { + registry.release_reserved_id(&execution_id, ticket).await; + }); + } +} + +enum ExecutionEntry { + Reserved { + ticket: u64, + }, + Live(ExecutionState), + Retained { + state: ExecutionState, + snapshot: TerminalSnapshot, + retained_bytes: usize, + expires_at: Instant, + last_access: u64, + }, + Tombstone { + snapshot: TerminalSnapshot, + expires_at: Instant, + metadata_bytes: usize, + last_access: u64, + }, +} + +#[derive(Clone)] +pub(crate) enum ExecutionLookup { + Live(ExecutionState), + Retained { + state: ExecutionState, + snapshot: TerminalSnapshot, + }, + Tombstone(TerminalSnapshot), +} + +struct RegistryInner { + entries: HashMap, + next_ticket: u64, + next_access: u64, + lifecycle_manager: Option>, + /// Signals the lifecycle manager to leave its loop. Shutdown waits for that + /// instead of aborting, because `prune_inner` tombstones entries under the + /// lock and releases their resources after dropping it — an abort between + /// those two steps strands resources no later caller can reach. + lifecycle_shutdown: watch::Sender, + is_shutting_down: bool, + lifecycle_ended_cleanly: bool, +} + +impl Default for RegistryInner { + fn default() -> Self { + Self { + entries: HashMap::new(), + next_ticket: 0, + next_access: 0, + lifecycle_manager: None, + lifecycle_shutdown: watch::channel(false).0, + is_shutting_down: false, + lifecycle_ended_cleanly: false, + } + } +} + +fn next_access(inner: &mut RegistryInner) -> u64 { + inner.next_access = inner.next_access.wrapping_add(1).max(1); + inner.next_access +} + +fn truncate_diagnostic(message: &str) -> String { + if message.len() <= MAX_TOMBSTONE_DIAGNOSTIC_BYTES { + return message.to_owned(); + } + let mut end = MAX_TOMBSTONE_DIAGNOSTIC_BYTES - TOMBSTONE_TRUNCATION_SUFFIX.len(); + while !message.is_char_boundary(end) { + end -= 1; + } + format!("{}{}", &message[..end], TOMBSTONE_TRUNCATION_SUFFIX) +} + +fn tombstone_entry( + snapshot: TerminalSnapshot, + expires_at: Instant, + last_access: u64, +) -> ExecutionEntry { + let mut snapshot = snapshot; + snapshot.exit.error_message = truncate_diagnostic(&snapshot.exit.error_message); + snapshot.output.reader_failure = snapshot + .output + .reader_failure + .as_deref() + .map(truncate_diagnostic); + let metadata_bytes = snapshot.exit.error_message.len() + + snapshot + .output + .reader_failure + .as_ref() + .map_or(0, String::len); + ExecutionEntry::Tombstone { + snapshot, + expires_at, + metadata_bytes, + last_access, + } +} + +fn enforce_tombstone_limits(inner: &mut RegistryInner) { + let mut tombstones: Vec<_> = inner + .entries + .iter() + .filter_map(|(id, entry)| match entry { + ExecutionEntry::Tombstone { + metadata_bytes, + last_access, + .. + } => Some((id.clone(), *metadata_bytes, *last_access)), + _ => None, + }) + .collect(); + tombstones.sort_by_key(|(_, _, last_access)| *last_access); + let mut total_bytes: usize = tombstones.iter().map(|(_, bytes, _)| *bytes).sum(); + while tombstones.len() > MAX_TOMBSTONE_ENTRIES || total_bytes > MAX_TOMBSTONE_METADATA_BYTES { + let (id, bytes, _) = tombstones.remove(0); + total_bytes -= bytes; + inner.entries.remove(&id); + } +} + /// Registry of active executions. /// /// Thread-safe registry that stores execution state and provides /// methods for registration, lookup, and lifecycle management. #[derive(Clone)] pub(crate) struct ExecutionRegistry { - executions: Arc>>, + inner: Arc>, } impl ExecutionRegistry { /// Create new registry. pub fn new() -> Self { Self { - executions: Arc::new(Mutex::new(HashMap::new())), + inner: Arc::new(Mutex::new(RegistryInner::default())), } } /// Check if execution exists. pub async fn exists(&self, exec_id: &str) -> bool { - self.executions.lock().await.contains_key(exec_id) + self.inner.lock().await.entries.contains_key(exec_id) } /// Get execution state. pub async fn get(&self, exec_id: &str) -> Option { - self.executions.lock().await.get(exec_id).cloned() + match self.lookup(exec_id).await { + Some(ExecutionLookup::Live(state) | ExecutionLookup::Retained { state, .. }) => { + Some(state) + } + Some(ExecutionLookup::Tombstone(_)) | None => None, + } + } + + pub async fn lookup(&self, exec_id: &str) -> Option { + let mut inner = self.inner.lock().await; + let access = next_access(&mut inner); + match inner.entries.get_mut(exec_id) { + Some(ExecutionEntry::Live(state)) => Some(ExecutionLookup::Live(state.clone())), + Some(ExecutionEntry::Retained { + state, + snapshot, + last_access, + .. + }) => { + *last_access = access; + Some(ExecutionLookup::Retained { + state: state.clone(), + snapshot: snapshot.clone(), + }) + } + Some(ExecutionEntry::Tombstone { + snapshot, + last_access, + .. + }) => { + *last_access = access; + Some(ExecutionLookup::Tombstone(snapshot.clone())) + } + Some(ExecutionEntry::Reserved { .. }) | None => None, + } } /// Register new execution state. - pub async fn register(&self, exec_id: String, state: ExecutionState) { - self.executions.lock().await.insert(exec_id, state); + pub async fn register(&self, exec_id: String, state: ExecutionState) -> bool { + let mut inner = self.inner.lock().await; + if inner.is_shutting_down || inner.entries.contains_key(&exec_id) { + return false; + } + inner.entries.insert(exec_id, ExecutionEntry::Live(state)); + true + } + + pub async fn reserve(&self, execution_id: String) -> Option { + let mut inner = self.inner.lock().await; + if inner.is_shutting_down || inner.entries.contains_key(&execution_id) { + return None; + } + inner.next_ticket = inner.next_ticket.wrapping_add(1).max(1); + let ticket = inner.next_ticket; + inner + .entries + .insert(execution_id.clone(), ExecutionEntry::Reserved { ticket }); + Some(ExecutionReservation { + registry: self.clone(), + execution_id, + ticket, + released: AtomicBool::new(false), + }) + } + + pub async fn release_reservation(&self, reservation: &ExecutionReservation) -> bool { + reservation.released.store(true, Ordering::Release); + self.release_reserved_id(&reservation.execution_id, reservation.ticket) + .await + } + + /// Remove a reservation only while this exact ticket still owns the id, so a + /// late release cannot evict the entry that replaced it. + async fn release_reserved_id(&self, execution_id: &str, ticket: u64) -> bool { + let mut inner = self.inner.lock().await; + if inner.entries.get(execution_id).is_some_and( + |entry| matches!(entry, ExecutionEntry::Reserved { ticket: held } if *held == ticket), + ) { + inner.entries.remove(execution_id); + true + } else { + false + } + } + + pub async fn publish(&self, reservation: &ExecutionReservation, state: ExecutionState) -> bool { + let mut inner = self.inner.lock().await; + if inner.is_shutting_down { + return false; + } + let Some(entry) = inner.entries.get_mut(&reservation.execution_id) else { + return false; + }; + if !matches!(entry, ExecutionEntry::Reserved { ticket } if *ticket == reservation.ticket) { + return false; + } + *entry = ExecutionEntry::Live(state); + reservation.released.store(true, Ordering::Release); + true + } + + #[cfg(test)] + pub async fn store_tombstone(&self, execution_id: String, snapshot: TerminalSnapshot) { + let mut inner = self.inner.lock().await; + let access = next_access(&mut inner); + inner.entries.insert( + execution_id, + tombstone_entry(snapshot, Instant::now() + TOMBSTONE_TTL, access), + ); + enforce_tombstone_limits(&mut inner); + } + + #[cfg(test)] + pub async fn terminal_snapshot(&self, execution_id: &str) -> Option { + let mut inner = self.inner.lock().await; + let access = next_access(&mut inner); + match inner.entries.get_mut(execution_id) { + Some(ExecutionEntry::Retained { + snapshot, + last_access, + .. + }) + | Some(ExecutionEntry::Tombstone { + snapshot, + last_access, + .. + }) => { + *last_access = access; + Some(snapshot.clone()) + } + _ => None, + } + } + + pub async fn retain( + &self, + execution_id: &str, + snapshot: TerminalSnapshot, + retained_bytes: usize, + ) -> bool { + let evicted = { + let mut inner = self.inner.lock().await; + if inner.is_shutting_down { + return false; + } + let Some(ExecutionEntry::Live(state)) = inner.entries.get(execution_id) else { + return false; + }; + let state = state.clone(); + if snapshot.output.reader_failure.is_some() { + let access = next_access(&mut inner); + inner.entries.insert( + execution_id.to_string(), + tombstone_entry(snapshot, Instant::now() + TOMBSTONE_TTL, access), + ); + enforce_tombstone_limits(&mut inner); + vec![state] + } else { + let access = next_access(&mut inner); + inner.entries.insert( + execution_id.to_string(), + ExecutionEntry::Retained { + state: state.clone(), + snapshot, + retained_bytes, + expires_at: Instant::now() + RETAIN_GRACE, + last_access: access, + }, + ); + let mut retained: Vec<_> = inner + .entries + .iter() + .filter_map(|(id, entry)| match entry { + ExecutionEntry::Retained { + state, + snapshot, + retained_bytes, + last_access, + .. + } => Some(( + id.clone(), + state.clone(), + snapshot.clone(), + *retained_bytes, + *last_access, + )), + _ => None, + }) + .collect(); + retained.sort_by_key(|(_, _, _, _, last_access)| *last_access); + let mut total_bytes: usize = + retained.iter().map(|(_, _, _, bytes, _)| *bytes).sum(); + let mut evicted = Vec::new(); + while retained.len() > MAX_RETAINED_ENTRIES || total_bytes > MAX_RETAINED_BYTES { + // A reader mid-replay would see its stream cut with no way to + // tell truncation from a normal end, so skip it and take the + // next oldest. When every candidate is being read the caps + // stay breached until those readers finish. + let Some(oldest_idle) = retained + .iter() + .position(|(_, state, _, _, _)| !state.has_active_reader()) + else { + warn!( + retained = retained.len(), + retained_bytes = total_bytes, + "retention over budget: every retained session still has a reader" + ); + break; + }; + let (id, state, snapshot, bytes, _) = retained.remove(oldest_idle); + total_bytes -= bytes; + let access = next_access(&mut inner); + inner.entries.insert( + id, + tombstone_entry(snapshot, Instant::now() + TOMBSTONE_TTL, access), + ); + evicted.push(state); + } + enforce_tombstone_limits(&mut inner); + evicted + } + }; + for state in evicted { + state.release_resources().await; + } + true + } + + #[cfg(test)] + async fn prune_at(&self, now: Instant) { + Self::prune_inner(&self.inner, now).await; + } + + async fn prune_inner(inner: &Arc>, now: Instant) { + let states = { + let mut inner = inner.lock().await; + let expired: Vec<_> = inner + .entries + .iter() + .filter_map(|(id, entry)| match entry { + // A reader outliving the grace keeps its session: tombstoning + // it would abort the forwarder mid-replay, which the client + // cannot tell from a normal end of output. The next tick + // retires it once the reader detaches. + ExecutionEntry::Retained { + state, + snapshot, + expires_at, + .. + } if *expires_at <= now && !state.has_active_reader() => { + Some((id.clone(), state.clone(), snapshot.clone())) + } + _ => None, + }) + .collect(); + for (id, _, snapshot) in &expired { + let access = next_access(&mut inner); + inner.entries.insert( + id.clone(), + tombstone_entry(snapshot.clone(), now + TOMBSTONE_TTL, access), + ); + } + inner.entries.retain(|_, entry| { + !matches!(entry, ExecutionEntry::Tombstone { expires_at, .. } if *expires_at <= now) + }); + enforce_tombstone_limits(&mut inner); + expired + .into_iter() + .map(|(_, state, _)| state) + .collect::>() + }; + for state in states { + state.release_resources().await; + } + } + + #[cfg(test)] + async fn prune_for_test(&self, now: Instant) { + self.prune_at(now).await; + } + + #[cfg(test)] + async fn lifecycle_manager_is_running(&self) -> bool { + self.inner.lock().await.lifecycle_manager.is_some() + } + + async fn ensure_lifecycle_manager(&self) { + let weak_inner = Arc::downgrade(&self.inner); + let mut inner = self.inner.lock().await; + if inner.is_shutting_down || inner.lifecycle_manager.is_some() { + return; + } + let mut shutdown = inner.lifecycle_shutdown.subscribe(); + inner.lifecycle_manager = Some(tokio::spawn(async move { + loop { + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(1)) => {} + _ = shutdown.changed() => break, + } + let Some(inner) = weak_inner.upgrade() else { + break; + }; + ExecutionRegistry::prune_inner(&inner, Instant::now()).await; + } + })); + } + + async fn stop_lifecycle_manager(&self) { + let lifecycle_manager = { + let mut inner = self.inner.lock().await; + inner.is_shutting_down = true; + inner.lifecycle_shutdown.send_replace(true); + inner.lifecycle_manager.take() + }; + let Some(lifecycle_manager) = lifecycle_manager else { + return; + }; + let ended_cleanly = lifecycle_manager.await.is_ok(); + if !ended_cleanly { + warn!("execution lifecycle manager did not stop cleanly"); + } + self.inner.lock().await.lifecycle_ended_cleanly = ended_cleanly; + } + + #[cfg(test)] + async fn lifecycle_manager_ended_cleanly(&self) -> bool { + self.inner.lock().await.lifecycle_ended_cleanly + } + + pub fn observe_terminal(&self, execution_id: String, state: ExecutionState) { + let registry = self.clone(); + tokio::spawn(async move { + registry.ensure_lifecycle_manager().await; + let exit = state.wait_exit(&execution_id).await; + state.cancel_timeout_task().await; + let output = state.wait_terminal_output_summary().await; + let snapshot = TerminalSnapshot { exit, output }; + let retained_bytes = state.retained_output_bytes().await; + registry + .retain(&execution_id, snapshot, retained_bytes) + .await; + }); } /// Release one explicitly ephemeral execution. @@ -57,8 +558,18 @@ impl ExecutionRegistry { /// their normal wait paths, preserving repeatable waits for those callers. pub async fn release_ephemeral(&self, exec_id: &str) -> bool { let state = { - let mut executions = self.executions.lock().await; - executions.remove(exec_id) + let mut inner = self.inner.lock().await; + match inner.entries.remove(exec_id) { + Some(ExecutionEntry::Live(state)) + | Some(ExecutionEntry::Retained { state, .. }) => Some(state), + Some( + entry @ (ExecutionEntry::Reserved { .. } | ExecutionEntry::Tombstone { .. }), + ) => { + inner.entries.insert(exec_id.to_string(), entry); + None + } + None => None, + } }; let Some(state) = state else { return false; @@ -71,15 +582,22 @@ impl ExecutionRegistry { /// /// Sends SIGTERM first, waits for exit with timeout, then SIGKILL if needed. pub async fn shutdown_all(&self, timeout_ms: u64) { - // Step 1: SIGTERM every execution whose process identity is current. + self.stop_lifecycle_manager().await; + let mut states_to_wait = Vec::new(); let states: Vec<_> = self - .executions + .inner .lock() .await + .entries .iter() - .map(|(exec_id, state)| (exec_id.clone(), state.clone())) + .filter_map(|(exec_id, entry)| match entry { + ExecutionEntry::Live(state) => Some((exec_id.clone(), state.clone())), + ExecutionEntry::Reserved { .. } + | ExecutionEntry::Retained { .. } + | ExecutionEntry::Tombstone { .. } => None, + }) .collect(); for (exec_id, state) in states { match state.signal_owned_process_if_current(Signal::SIGTERM).await { @@ -94,37 +612,54 @@ impl ExecutionRegistry { if states_to_wait.is_empty() { info!("No running executions to shutdown"); - return; - } + } else { + let start = std::time::Instant::now(); + while start.elapsed().as_millis() < timeout_ms as u128 { + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - // Step 2: Wait for graceful exit with timeout - let start = std::time::Instant::now(); - while start.elapsed().as_millis() < timeout_ms as u128 { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + let mut still_running = Vec::new(); + for (exec_id, state) in states_to_wait { + if state.owned_process_is_current().await { + still_running.push((exec_id, state)); + } + } + states_to_wait = still_running; - let mut still_running = Vec::new(); - for (exec_id, state) in states_to_wait { - if state.owned_process_is_current().await { - still_running.push((exec_id, state)); + if states_to_wait.is_empty() { + info!("All executions exited gracefully"); + break; } } - states_to_wait = still_running; - if states_to_wait.is_empty() { - info!("All executions exited gracefully"); - return; + for (exec_id, state) in &states_to_wait { + match state.signal_owned_process_if_current(Signal::SIGKILL).await { + Ok(true) => { + warn!(exec_id = %exec_id, "Execution didn't exit gracefully, sending SIGKILL"); + } + Ok(false) => {} + Err(error) => warn!(exec_id = %exec_id, %error, "shutdown SIGKILL failed"), + } } } - // Step 3: SIGKILL remaining executions - for (exec_id, state) in &states_to_wait { - match state.signal_owned_process_if_current(Signal::SIGKILL).await { - Ok(true) => { - warn!(exec_id = %exec_id, "Execution didn't exit gracefully, sending SIGKILL"); - } - Ok(false) => {} - Err(error) => warn!(exec_id = %exec_id, %error, "shutdown SIGKILL failed"), - } + self.release_remaining_states().await; + } + + async fn release_remaining_states(&self) { + let states = { + let mut inner = self.inner.lock().await; + std::mem::take(&mut inner.entries) + .into_values() + .filter_map(|entry| match entry { + ExecutionEntry::Live(state) | ExecutionEntry::Retained { state, .. } => { + Some(state) + } + ExecutionEntry::Reserved { .. } | ExecutionEntry::Tombstone { .. } => None, + }) + .collect::>() + }; + for state in states { + state.release_resources().await; } } } @@ -134,7 +669,11 @@ mod release_tests { use super::*; use crate::reaper::ExitSlot; use crate::service::exec::exec_handle::{ExecHandle, ExitStatus}; + use crate::service::exec::output::{OutputStreamSummary, OutputTerminalSummary}; + use crate::service::exec::state::{ExecutionExit, TerminalSnapshot}; use nix::unistd::{pipe, Pid}; + use std::os::unix::process::ExitStatusExt; + use std::time::Instant; fn settled_state(pid: i32, is_init: bool) -> ExecutionState { let (_stdin_peer, stdin) = pipe().unwrap(); @@ -150,6 +689,406 @@ mod release_tests { } } + #[tokio::test] + async fn releasing_a_reservation_returns_its_id_to_absent() { + let registry = ExecutionRegistry::new(); + let first = registry + .reserve("reserved-exec".into()) + .await + .expect("first reservation must succeed"); + + assert!(registry.reserve("reserved-exec".into()).await.is_none()); + assert!(registry.release_reservation(&first).await); + assert!(registry.reserve("reserved-exec".into()).await.is_some()); + } + + #[tokio::test] + async fn shutdown_rejects_a_late_ssh_registration() { + let registry = ExecutionRegistry::new(); + registry.shutdown_all(0).await; + + assert!( + !registry + .register("late-ssh-exec".into(), settled_state(11_010, false)) + .await + ); + + assert!(!registry.exists("late-ssh-exec").await); + } + + #[tokio::test] + async fn only_the_matching_reservation_can_publish_a_live_state() { + let registry = ExecutionRegistry::new(); + let reservation = registry + .reserve("reserved-exec".into()) + .await + .expect("reservation must succeed"); + + assert!( + registry + .publish(&reservation, settled_state(11_004, false)) + .await + ); + assert!(registry.get("reserved-exec").await.is_some()); + assert!(!registry.release_reservation(&reservation).await); + } + + #[tokio::test] + async fn a_tombstone_keeps_its_terminal_snapshot_repeatable() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 7, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 12, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + + registry + .store_tombstone("terminal-exec".into(), snapshot.clone()) + .await; + assert_eq!( + registry.terminal_snapshot("terminal-exec").await, + Some(snapshot.clone()) + ); + assert_eq!( + registry.terminal_snapshot("terminal-exec").await, + Some(snapshot) + ); + } + + #[tokio::test] + async fn retained_entry_exposes_its_terminal_snapshot_before_tombstoning() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 3, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + registry + .register("retained-exec".into(), settled_state(11_005, false)) + .await; + + assert!(registry.retain("retained-exec", snapshot.clone(), 0).await); + assert_eq!( + registry.terminal_snapshot("retained-exec").await, + Some(snapshot) + ); + } + + #[tokio::test] + async fn terminal_observer_retains_a_completed_live_execution() { + let registry = ExecutionRegistry::new(); + let state = settled_state(11_006, false); + registry + .register("completed-exec".into(), state.clone()) + .await; + + registry.observe_terminal("completed-exec".into(), state); + let snapshot = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if let Some(snapshot) = registry.terminal_snapshot("completed-exec").await { + break snapshot; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("completed execution must become retained"); + assert_eq!(snapshot.exit.exit_code, 7); + } + + #[tokio::test] + async fn terminal_observer_cancels_timeout_when_leader_exits_before_pipe_eof() { + let registry = ExecutionRegistry::new(); + let (_stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + drop(stderr_peer); + let state = ExecutionState::new_for_test( + ExecHandle::new(Pid::from_raw(11_009), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"), + ExitSlot::settled_for_test(ExitStatus::Code(0)), + ); + let timeout_task = tokio::spawn(std::future::pending()); + let timeout_abort = timeout_task.abort_handle(); + state.set_timeout_task(timeout_task).await; + + registry + .register("leader-exited-pipe-open".into(), state) + .await; + registry.observe_terminal( + "leader-exited-pipe-open".into(), + registry.get("leader-exited-pipe-open").await.unwrap(), + ); + + tokio::time::timeout(Duration::from_secs(1), async { + while !timeout_abort.is_finished() { + tokio::task::yield_now().await; + } + }) + .await + .expect("leader exit must cancel its timeout before output EOF"); + assert!(matches!( + registry.lookup("leader-exited-pipe-open").await, + Some(ExecutionLookup::Live(_)) + )); + drop(stdout_peer); + } + + #[tokio::test] + async fn expired_retained_entries_become_tombstones_then_expire() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + registry + .register("expiring-exec".into(), settled_state(11_007, false)) + .await; + assert!(registry.retain("expiring-exec", snapshot, 0).await); + + let now = Instant::now(); + registry.prune_for_test(now + RETAIN_GRACE).await; + assert!(registry.terminal_snapshot("expiring-exec").await.is_some()); + assert!(registry.get("expiring-exec").await.is_none()); + + registry + .prune_for_test(now + RETAIN_GRACE + TOMBSTONE_TTL) + .await; + assert!(registry.terminal_snapshot("expiring-exec").await.is_none()); + } + + #[tokio::test] + async fn retained_entry_limit_evicts_the_oldest_to_a_tombstone() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + for index in 0..=MAX_RETAINED_ENTRIES { + let id = format!("retained-{index}"); + registry + .register(id.clone(), settled_state(12_000 + index as i32, false)) + .await; + assert!(registry.retain(&id, snapshot.clone(), 0).await); + } + + assert!(registry.get("retained-0").await.is_none()); + assert_eq!( + registry.terminal_snapshot("retained-0").await, + Some(snapshot) + ); + } + + #[tokio::test] + async fn retained_byte_limit_evicts_the_oldest_to_a_tombstone() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + registry + .register("first".into(), settled_state(12_100, false)) + .await; + registry + .register("second".into(), settled_state(12_101, false)) + .await; + + assert!( + registry + .retain("first", snapshot.clone(), MAX_RETAINED_BYTES) + .await + ); + assert!(registry.retain("second", snapshot.clone(), 1).await); + + assert!(registry.get("first").await.is_none()); + assert_eq!(registry.terminal_snapshot("first").await, Some(snapshot)); + } + + #[tokio::test] + async fn tombstone_entry_limit_evicts_the_oldest_snapshot() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + for index in 0..=1024 { + registry + .store_tombstone(format!("tombstone-{index}"), snapshot.clone()) + .await; + } + + assert!(registry.terminal_snapshot("tombstone-0").await.is_none()); + assert_eq!( + registry.terminal_snapshot("tombstone-1024").await, + Some(snapshot) + ); + } + + #[tokio::test] + async fn tombstone_diagnostics_are_utf8_bounded_and_count_toward_eviction() { + let registry = ExecutionRegistry::new(); + let diagnostic = "€".repeat(2_000); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: diagnostic.clone(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: Some(diagnostic), + }, + }; + registry + .store_tombstone("bounded".into(), snapshot.clone()) + .await; + let bounded = registry.terminal_snapshot("bounded").await.unwrap(); + assert!(bounded.exit.error_message.len() <= MAX_TOMBSTONE_DIAGNOSTIC_BYTES); + assert!(bounded.exit.error_message.ends_with("…[truncated]")); + assert!(bounded + .output + .reader_failure + .as_ref() + .is_some_and(|message| { + message.len() <= MAX_TOMBSTONE_DIAGNOSTIC_BYTES && message.ends_with("…[truncated]") + })); + + for index in 0..700 { + registry + .store_tombstone(format!("metadata-{index}"), snapshot.clone()) + .await; + } + + assert!(registry.terminal_snapshot("metadata-0").await.is_none()); + assert!(registry.terminal_snapshot("metadata-699").await.is_some()); + } + + #[tokio::test] + async fn reader_failure_is_tombstoned_without_a_retained_grace_period() { + let registry = ExecutionRegistry::new(); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 10, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: Some("stdout reader failed".into()), + }, + }; + registry + .register("reader-failure".into(), settled_state(12_200, false)) + .await; + + assert!( + registry + .retain("reader-failure", snapshot.clone(), 10) + .await + ); + assert!(registry.get("reader-failure").await.is_none()); + assert_eq!( + registry.terminal_snapshot("reader-failure").await, + Some(snapshot) + ); + } + #[tokio::test] async fn release_is_scoped_and_preserves_retained_waits_and_init_sessions() { let registry = ExecutionRegistry::new(); @@ -194,9 +1133,64 @@ mod release_tests { } #[tokio::test] - async fn shutdown_does_not_signal_an_execution_without_identity() { - use std::os::unix::process::ExitStatusExt; + async fn shutdown_stops_the_lifecycle_manager() { + let registry = ExecutionRegistry::new(); + registry.ensure_lifecycle_manager().await; + assert!(registry.lifecycle_manager_is_running().await); + + registry.shutdown_all(0).await; + + assert!(!registry.lifecycle_manager_is_running().await); + } + #[tokio::test] + async fn dropping_registry_does_not_keep_its_lifecycle_manager_alive() { + let registry = ExecutionRegistry::new(); + registry.ensure_lifecycle_manager().await; + let inner = Arc::downgrade(®istry.inner); + + drop(registry); + + assert!(inner.upgrade().is_none()); + } + + #[tokio::test] + async fn shutdown_releases_live_and_retained_state_resources() { + let registry = ExecutionRegistry::new(); + let live = settled_state(12_300, false); + let retained = settled_state(12_301, false); + let snapshot = TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + }; + registry.register("live".into(), live.clone()).await; + registry.register("retained".into(), retained.clone()).await; + assert!(registry.retain("retained", snapshot, 0).await); + + registry.shutdown_all(0).await; + + assert!(live.get_pid().await.is_none()); + assert!(retained.get_pid().await.is_none()); + assert!(!registry.exists("live").await); + assert!(!registry.exists("retained").await); + } + + #[tokio::test] + async fn shutdown_does_not_signal_an_execution_without_identity() { let _test_guard = crate::reaper::reap_test_guard().await; let mut child = std::process::Command::new("/bin/sleep") .arg("30") @@ -223,9 +1217,188 @@ mod release_tests { .expect("wait task must not panic"); assert_eq!( status.signal(), - Some(nix::sys::signal::Signal::SIGKILL as i32), - "child must die from this test's SIGKILL; SIGTERM means shutdown_all \ - signalled a state that has no process identity" + Some(nix::sys::signal::Signal::SIGKILL as i32) + ); + } + + #[tokio::test] + async fn shutdown_does_not_signal_an_init_target() { + let _test_guard = crate::reaper::reap_test_guard().await; + let mut child = std::process::Command::new("/bin/sleep") + .arg("30") + .spawn() + .expect("spawn sleep"); + let leader = Pid::from_raw(child.id() as i32); + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let init = ExecutionState::new_init_session( + ExecHandle::new(leader, stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"), + ExitSlot::settled_for_test(ExitStatus::Code(0)), + crate::service::exec::process_instance::ProcessInstance::capture(leader), + ); + let registry = ExecutionRegistry::new(); + registry.register("init".into(), init).await; + + registry.shutdown_all(0).await; + + assert!( + child.try_wait().expect("check child status").is_none(), + "registry must leave init shutdown to the container lifecycle" + ); + child.kill().expect("kill test child"); + let status = tokio::task::spawn_blocking(move || { + let _fence = crate::reaper::reap_fence(); + child.wait().expect("wait for test child") + }) + .await + .expect("wait task must not panic"); + assert_eq!( + status.signal(), + Some(nix::sys::signal::Signal::SIGKILL as i32) + ); + drop((stdin_peer, stdout_peer, stderr_peer)); + } + + fn exit_snapshot() -> TerminalSnapshot { + TerminalSnapshot { + exit: ExecutionExit { + exit_code: 0, + signal: 0, + error_message: String::new(), + }, + output: OutputTerminalSummary { + stdout: OutputStreamSummary { + enabled: true, + total_bytes: 0, + }, + stderr: OutputStreamSummary { + enabled: false, + total_bytes: 0, + }, + reader_failure: None, + }, + } + } + + /// A live state whose pipe peers stay open, so a forwarder attached to it + /// keeps holding the output lease instead of hitting EOF. + fn state_with_open_pipes(pid: i32) -> (ExecutionState, Vec) { + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let handle = ExecHandle::new(Pid::from_raw(pid), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); + let state = ExecutionState::new( + handle, + ExitSlot::settled_for_test(ExitStatus::Code(0)), + None, + ); + (state, vec![stdin_peer, stdout_peer, stderr_peer]) + } + + /// A cancelled Exec RPC drops its reservation without ever publishing. The + /// id must not stay reserved for the life of the guest: `prune_inner` never + /// touches `Reserved`, so nothing else would ever free it. + #[tokio::test] + async fn dropping_an_unpublished_reservation_frees_its_execution_id() { + let registry = ExecutionRegistry::new(); + drop( + registry + .reserve("cancelled-exec".into()) + .await + .expect("reservation must succeed"), + ); + + tokio::time::timeout(Duration::from_secs(1), async { + while registry.exists("cancelled-exec").await { + tokio::task::yield_now().await; + } + }) + .await + .expect("a dropped reservation must free its execution id"); + } + + /// A reader attached to a retained session holds its output lease. Evicting + /// it mid-replay aborts the forwarder, and the client cannot tell that + /// truncation apart from a normal end of output. + #[tokio::test] + async fn eviction_spares_a_retained_session_with_an_active_reader() { + let registry = ExecutionRegistry::new(); + let snapshot = exit_snapshot(); + + let (attached, _peers) = state_with_open_pipes(12_400); + registry.register("attached".into(), attached.clone()).await; + let _reader = attached + .attach_retained("attached") + .await + .expect("attach must claim the output lease"); + assert!( + registry + .retain("attached", snapshot.clone(), MAX_RETAINED_BYTES) + .await + ); + + registry + .register("newcomer".into(), settled_state(12_401, false)) + .await; + assert!(registry.retain("newcomer", snapshot, 1).await); + + assert!( + registry.get("attached").await.is_some(), + "a retained session with an active reader must not be evicted" + ); + + // The forwarder parks on pipes this test keeps open, so join it here + // rather than leaving a detached task behind for the next test. + attached.release_resources().await; + drop((_reader, _peers)); + } + + /// The retain grace elapsing is the other way a reader loses its stream, and + /// unlike the caps it cannot be deferred by attaching: `lookup` refreshes + /// `last_access` for LRU but never extends `expires_at`. + #[tokio::test] + async fn grace_expiry_spares_a_retained_session_with_an_active_reader() { + let registry = ExecutionRegistry::new(); + let (attached, _peers) = state_with_open_pipes(12_402); + registry.register("attached".into(), attached.clone()).await; + let _reader = attached + .attach_retained("attached") + .await + .expect("attach must claim the output lease"); + assert!(registry.retain("attached", exit_snapshot(), 0).await); + + registry + .prune_for_test(Instant::now() + RETAIN_GRACE + Duration::from_secs(1)) + .await; + + assert!( + registry.get("attached").await.is_some(), + "grace expiry must not tombstone a session with an active reader" + ); + + attached.release_resources().await; + drop((_reader, _peers)); + } + + /// `prune_inner` tombstones entries under the lock and releases their + /// resources after dropping it, so aborting the manager between those two + /// steps strands resources no later caller can reach. Shutdown therefore + /// signals the manager and waits for it to leave the loop itself. + #[tokio::test] + async fn stopping_the_lifecycle_manager_lets_it_exit_on_its_own() { + let registry = ExecutionRegistry::new(); + registry.ensure_lifecycle_manager().await; + assert!(registry.lifecycle_manager_is_running().await); + + registry.shutdown_all(0).await; + + assert!(!registry.lifecycle_manager_is_running().await); + assert!( + registry.lifecycle_manager_ended_cleanly().await, + "the manager must finish its own loop rather than be aborted" ); } } diff --git a/src/guest/src/service/exec/state.rs b/src/guest/src/service/exec/state.rs index f9abf2494..67f823c48 100644 --- a/src/guest/src/service/exec/state.rs +++ b/src/guest/src/service/exec/state.rs @@ -1,12 +1,12 @@ use crate::service::exec::error::ExecutionError; use crate::service::exec::exec_handle::ExecHandle; -use crate::service::exec::output::OutputManager; +use crate::service::exec::output::{OutputManager, OutputTerminalSummary}; use crate::service::exec::process_instance::ProcessInstance; use boxlite_shared::ExecOutput; use futures::{Stream, StreamExt as _}; use std::os::unix::io::AsRawFd; use std::sync::Arc; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::{mpsc, Mutex, OnceCell}; use tokio::task::{AbortHandle, JoinHandle}; use tonic::Status; @@ -31,18 +31,20 @@ struct Inner { output: OutputManager, /// Abort handles for stdin forwarding tasks that may own the taken stdin FD. input_tasks: Vec, - /// Stdout/stderr forwarding tasks (set on attach) - output_tasks: Vec>, + output_task: Option>, + timeout_task: Option>, /// Set once ephemeral callers explicitly release this execution's resources. released: bool, - /// Timeout flag - #[allow(dead_code)] // Will be used for timeout handling - timed_out: bool, /// Optional init health checker for the container this exec runs in. /// Used to detect container init death when exec gets SIGKILL. init_health: Option>>, } +enum OutputAttachMode { + Live, + Retained, +} + /// How an execution ended, already classified. /// /// `error_message` carries the container-death diagnosis when pid-namespace @@ -54,6 +56,12 @@ pub(crate) struct ExecutionExit { pub error_message: String, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TerminalSnapshot { + pub(crate) exit: ExecutionExit, + pub(crate) output: OutputTerminalSummary, +} + /// Execution state. /// /// Handle owns pid, pty_controller, stdin, stdout, stderr. @@ -66,6 +74,10 @@ pub(crate) struct ExecutionState { exit: crate::reaper::ExitSlot, process: Option, shutdown_managed: bool, + consumer_lease: Arc, + /// Classified once: the container-death diagnosis drains init's pipes, so a + /// second reader would get a different answer. + terminal_exit: Arc>, } impl ExecutionState { @@ -76,22 +88,35 @@ impl ExecutionState { process: Option, shutdown_managed: bool, ) -> Self { + let output = OutputManager::new(handle.stdout(), handle.stderr()); + let consumer_lease = output.consumer_lease_flag(); Self { inner: Arc::new(Mutex::new(Inner { - output: OutputManager::new(handle.stdout(), handle.stderr()), + output, handle: Some(handle), input_tasks: Vec::new(), - output_tasks: Vec::new(), + output_task: None, + timeout_task: None, released: false, - timed_out: false, init_health, })), exit, process, shutdown_managed, + consumer_lease, + terminal_exit: Arc::new(OnceCell::new()), } } + /// Whether a reader is currently streaming this execution's output. + /// + /// Kept outside `inner` so eviction can ask while holding the registry lock, + /// which must not await on a state's own mutex. + pub(crate) fn has_active_reader(&self) -> bool { + self.consumer_lease + .load(std::sync::atomic::Ordering::Acquire) + } + /// Create new execution state for a guest-side process. pub(super) fn new( handle: ExecHandle, @@ -243,6 +268,63 @@ impl ExecutionState { self.exit.get().await } + pub(crate) async fn wait_terminal_output_summary(&self) -> OutputTerminalSummary { + let output = self.inner.lock().await.output.clone(); + let summary = output.wait_terminal_summary().await; + let sealed = output.seal().await; + debug_assert!(sealed, "terminal output must be sealable after its summary"); + summary + } + + pub(crate) async fn sealed_terminal_output_summary(&self) -> Option { + let output = self.inner.lock().await.output.clone(); + output.sealed_terminal_summary().await + } + + pub(crate) async fn retained_output_bytes(&self) -> usize { + let output = self.inner.lock().await.output.clone(); + output.retained_bytes().await + } + + pub(crate) async fn set_timeout_task(&self, task: JoinHandle<()>) { + let task_to_abort = { + let mut inner = self.inner.lock().await; + if inner.released { + Some(task) + } else { + debug_assert!(inner.timeout_task.is_none()); + inner.timeout_task = Some(task); + None + } + }; + if let Some(task) = task_to_abort { + task.abort(); + let _ = task.await; + } + } + + pub(crate) async fn cancel_timeout_task(&self) { + let task = self.inner.lock().await.timeout_task.take(); + if let Some(task) = task { + task.abort(); + let _ = task.await; + } + } + + pub(crate) async fn abort_unpublished(&self) { + let _ = self + .signal_owned_process_if_current(nix::sys::signal::Signal::SIGKILL) + .await; + self.release_resources().await; + } + + #[cfg(test)] + pub(crate) async fn wait_terminal_snapshot(&self, exec_id: &str) -> TerminalSnapshot { + let (exit, output) = + tokio::join!(self.wait_exit(exec_id), self.wait_terminal_output_summary(),); + TerminalSnapshot { exit, output } + } + /// Wait for exit and classify it. /// /// The SIGKILL diagnosis lives here rather than in a caller because @@ -251,6 +333,13 @@ impl ExecutionState { /// the same answer, including the reason a tenant was killed by pid-namespace /// teardown rather than by its own exit. pub(crate) async fn wait_exit(&self, exec_id: &str) -> ExecutionExit { + self.terminal_exit + .get_or_init(|| self.classify_exit(exec_id)) + .await + .clone() + } + + async fn classify_exit(&self, exec_id: &str) -> ExecutionExit { use crate::service::exec::exec_handle::ExitStatus; match self.wait_process().await { @@ -294,6 +383,24 @@ impl ExecutionState { &self, exec_id: &str, ) -> Result>, ExecutionError> { + self.attach_output(exec_id, OutputAttachMode::Live).await + } + + pub(crate) async fn attach_retained( + &self, + exec_id: &str, + ) -> Result>, ExecutionError> { + self.attach_output(exec_id, OutputAttachMode::Retained) + .await + } + + async fn attach_output( + &self, + exec_id: &str, + mode: OutputAttachMode, + ) -> Result>, ExecutionError> { + self.join_finished_output_task().await; + let output = { let inner = self.inner.lock().await; if inner.released { @@ -301,10 +408,11 @@ impl ExecutionState { } inner.output.clone() }; - let mut output = output - .attach() - .await - .map_err(|_| ExecutionError::AlreadyAttached)?; + let mut output = match mode { + OutputAttachMode::Live => output.attach().await, + OutputAttachMode::Retained => output.attach_retained().await, + } + .map_err(|_| ExecutionError::AlreadyAttached)?; let (tx, rx) = mpsc::channel(100); let execution_id = exec_id.to_owned(); let task = tokio::spawn(async move { @@ -316,15 +424,55 @@ impl ExecutionState { tracing::info!(%execution_id, "execution output forwarding ended"); }); - let mut inner = self.inner.lock().await; - if inner.released { + let (task_to_abort, displaced) = { + let mut inner = self.inner.lock().await; + if inner.released { + (Some(task), None) + } else { + (None, inner.output_task.replace(task)) + } + }; + // A forwarder releases the consumer lease when its stream drops, which + // happens before its handle reports finished. So a displaced forwarder is + // already ending, and this join cannot wait on live output. + if let Some(displaced) = displaced { + let _ = displaced.await; + } + if let Some(task) = task_to_abort { task.abort(); + let _ = task.await; return Err(ExecutionError::HandleUnavailable); } - inner.output_tasks.push(task); Ok(rx) } + async fn join_finished_output_task(&self) { + let task = { + let mut inner = self.inner.lock().await; + if inner + .output_task + .as_ref() + .is_some_and(JoinHandle::is_finished) + { + inner.output_task.take() + } else { + None + } + }; + if let Some(task) = task { + let _ = task.await; + } + } + + #[cfg(test)] + async fn output_task_count(&self) -> usize { + if self.inner.lock().await.output_task.is_some() { + 1 + } else { + 0 + } + } + /// Drop every resource owned by an explicitly ephemeral execution. /// /// Ordinary SDK executions never call this method, so their repeatable wait @@ -333,7 +481,7 @@ impl ExecutionState { /// clone still exists. Forwarders are aborted because stdin may already have /// moved out of the handle into one of those tasks. pub(super) async fn release_resources(&self) -> bool { - let (output, input_tasks, output_tasks) = { + let (output, input_tasks, output_task, timeout_task) = { let mut inner = self.inner.lock().await; if inner.released { return false; @@ -341,17 +489,23 @@ impl ExecutionState { inner.released = true; let output = inner.output.clone(); let input_tasks = std::mem::take(&mut inner.input_tasks); - let output_tasks = std::mem::take(&mut inner.output_tasks); + let output_task = inner.output_task.take(); + let timeout_task = inner.timeout_task.take(); drop(inner.handle.take()); drop(inner.init_health.take()); - (output, input_tasks, output_tasks) + (output, input_tasks, output_task, timeout_task) }; for task in input_tasks { task.abort(); } - for task in output_tasks { + if let Some(task) = output_task { + task.abort(); + let _ = task.await; + } + if let Some(task) = timeout_task { task.abort(); + let _ = task.await; } output.shutdown_drains().await; if self.shutdown_managed { @@ -451,6 +605,7 @@ mod release_tests { use crate::service::exec::process_instance::ProcessInstance; use nix::unistd::{pipe, Pid}; use std::os::fd::{AsRawFd, OwnedFd, RawFd}; + use std::sync::atomic::{AtomicUsize, Ordering}; fn fd_is_open(fd: RawFd) -> bool { (unsafe { nix::libc::fcntl(fd, nix::libc::F_GETFD) }) != -1 @@ -536,6 +691,15 @@ mod release_tests { ); } + #[tokio::test] + async fn release_aborts_the_timeout_task() { + let (state, _tracked_fds, _peers) = state_with_tracked_handle(); + let task = tokio::spawn(std::future::pending::<()>()); + state.set_timeout_task(task).await; + + assert!(state.release_resources().await); + } + #[tokio::test] async fn process_identity_refuses_a_changed_start_time() { use std::os::unix::process::ExitStatusExt; @@ -572,6 +736,83 @@ mod release_tests { ); } + #[tokio::test] + async fn aborting_an_unpublished_execution_kills_its_process() { + use std::os::unix::process::ExitStatusExt; + + let _test_guard = crate::reaper::reap_test_guard().await; + let mut child = std::process::Command::new("/bin/sleep") + .arg("30") + .spawn() + .expect("spawn sleep"); + let leader = Pid::from_raw(child.id() as i32); + let process = ProcessInstance::capture(leader).expect("read child identity"); + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let (exit, _exit_tx) = ExitSlot::pending_for_test(); + let state = ExecutionState::new( + ExecHandle::new(leader, stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"), + exit, + Some(process), + ); + + state.abort_unpublished().await; + + let status = tokio::task::spawn_blocking(move || { + let _fence = crate::reaper::reap_fence(); + child.wait().expect("aborted leader must exit") + }) + .await + .expect("wait task must not panic"); + assert_eq!( + status.signal(), + Some(nix::sys::signal::Signal::SIGKILL as i32) + ); + drop((stdin_peer, stdout_peer, stderr_peer)); + } + + #[tokio::test] + async fn direct_kill_does_not_signal_a_released_execution() { + use std::os::unix::process::ExitStatusExt; + + let _test_guard = crate::reaper::reap_test_guard().await; + let mut child = std::process::Command::new("/bin/sleep") + .arg("30") + .spawn() + .expect("spawn sleep"); + let leader = Pid::from_raw(child.id() as i32); + let process = ProcessInstance::capture(leader).expect("read child identity"); + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let (exit, _exit_tx) = ExitSlot::pending_for_test(); + let state = ExecutionState::new( + ExecHandle::new(leader, stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"), + exit, + Some(process), + ); + + assert!(state.release_resources().await); + assert!(!state.kill(nix::sys::signal::Signal::SIGTERM, false).await); + assert!(child.try_wait().expect("check child status").is_none()); + + child.kill().expect("kill test child"); + let status = tokio::task::spawn_blocking(move || { + let _fence = crate::reaper::reap_fence(); + child.wait().expect("wait for test child") + }) + .await + .expect("wait task must not panic"); + assert_eq!( + status.signal(), + Some(nix::sys::signal::Signal::SIGKILL as i32) + ); + drop((stdin_peer, stdout_peer, stderr_peer)); + } + #[tokio::test] async fn process_identity_signals_the_matching_process() { use std::os::unix::process::ExitStatusExt; @@ -643,4 +884,150 @@ mod release_tests { Some(nix::sys::signal::Signal::SIGKILL as i32) ); } + + /// Init is exempt from registry shutdown, not from an explicit Kill. + #[tokio::test] + async fn direct_kill_signals_an_init_session() { + use std::os::unix::process::ExitStatusExt; + + let _test_guard = crate::reaper::reap_test_guard().await; + let mut child = std::process::Command::new("/bin/sleep") + .arg("30") + .spawn() + .expect("spawn sleep"); + let leader = Pid::from_raw(child.id() as i32); + let process = ProcessInstance::capture(leader).expect("read child identity"); + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let (exit, _exit_tx) = ExitSlot::pending_for_test(); + let state = ExecutionState::new_init_session( + ExecHandle::new(leader, stdin, stdout, Some(stderr)) + .expect("test pipes must register with Tokio"), + exit, + Some(process), + ); + + assert!(!state + .signal_owned_process_if_current(nix::sys::signal::Signal::SIGTERM) + .await + .expect("shutdown must skip an init session")); + assert!(state.kill(nix::sys::signal::Signal::SIGTERM, false).await); + + let status = tokio::task::spawn_blocking(move || { + let _fence = crate::reaper::reap_fence(); + child.wait().expect("wait for test child") + }) + .await + .expect("wait task must not panic"); + assert_eq!( + status.signal(), + Some(nix::sys::signal::Signal::SIGTERM as i32) + ); + assert!(state.release_resources().await); + drop((stdin_peer, stdout_peer, stderr_peer)); + } + + #[tokio::test] + async fn state_waits_for_its_terminal_output_summary() { + let (state, _tracked_fds, peers) = state_with_tracked_handle(); + drop(peers); + + let summary = tokio::time::timeout( + std::time::Duration::from_secs(1), + state.wait_terminal_output_summary(), + ) + .await + .expect("closed output peers must finish the summary wait"); + assert!(summary.stdout.enabled); + assert!(summary.stderr.enabled); + } + + #[tokio::test] + async fn terminal_output_summary_seals_the_output_before_retention() { + let (state, _tracked_fds, peers) = state_with_tracked_handle(); + drop(peers); + + state.wait_terminal_output_summary().await; + + assert!(matches!( + state.attach("sealed-exec").await, + Err(ExecutionError::AlreadyAttached) + )); + } + + #[tokio::test] + async fn terminal_snapshot_caches_exit_with_completed_output() { + let (state, _tracked_fds, peers) = state_with_tracked_handle(); + drop(peers); + + let snapshot = state.wait_terminal_snapshot("test-exec").await; + assert_eq!(snapshot.exit.exit_code, 0); + assert!(snapshot.output.stdout.enabled); + assert!(snapshot.output.stderr.enabled); + } + + struct DeadInit { + diagnoses: Arc, + } + + impl InitHealthCheck for DeadInit { + fn is_running(&self) -> bool { + false + } + + fn diagnose_exit(&mut self) -> String { + self.diagnoses.fetch_add(1, Ordering::SeqCst); + "init exited".into() + } + } + + #[tokio::test] + async fn repeated_wait_caches_the_init_exit_diagnosis() { + let (stdin_peer, stdin) = pipe().unwrap(); + let (stdout, stdout_peer) = pipe().unwrap(); + let (stderr, stderr_peer) = pipe().unwrap(); + let handle = ExecHandle::new(Pid::from_raw(42_424), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); + let diagnoses = Arc::new(AtomicUsize::new(0)); + let health: Arc> = Arc::new(Mutex::new(DeadInit { + diagnoses: diagnoses.clone(), + })); + let state = ExecutionState::new_with_init_health( + handle, + health, + ExitSlot::settled_for_test(ExitStatus::Signal(nix::sys::signal::Signal::SIGKILL)), + None, + ); + + assert_eq!(state.wait_exit("exec").await.error_message, "init exited"); + assert_eq!(state.wait_exit("exec").await.error_message, "init exited"); + assert_eq!(diagnoses.load(Ordering::SeqCst), 1); + + drop((stdin_peer, stdout_peer, stderr_peer)); + } + + #[tokio::test] + async fn repeated_attach_does_not_retain_finished_forwarding_tasks() { + let (state, _tracked_fds, peers) = state_with_tracked_handle(); + + let first = state.attach("exec-1").await.unwrap(); + drop(first); + drop(peers); + + let second = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + match state.attach("exec-1").await { + Ok(receiver) => break receiver, + Err(ExecutionError::AlreadyAttached) => tokio::task::yield_now().await, + Err(error) => panic!("second attach must not fail: {error:?}"), + } + } + }) + .await + .expect("completed attach must release its output lease"); + drop(second); + + assert_eq!(state.output_task_count().await, 1); + } } diff --git a/src/guest/src/service/exec/timeout.rs b/src/guest/src/service/exec/timeout.rs index 186b5307b..e3d7deca8 100644 --- a/src/guest/src/service/exec/timeout.rs +++ b/src/guest/src/service/exec/timeout.rs @@ -44,7 +44,14 @@ impl TimeoutTarget { /// for the process to exit, then escalates to SIGKILL. SIGKILL is /// uncatchable, so a workload that installs `SIG_IGN`/handlers for /// SIGTERM (or SIGALRM, etc.) cannot outlive its deadline. -pub(super) fn start_timeout_watcher(target: TimeoutTarget, exec_id: String, timeout: Duration) { +/// +/// The handle is returned so a retained session can cancel the watcher instead +/// of leaving a task parked on a deadline its process already beat. +pub(super) fn start_timeout_watcher( + target: TimeoutTarget, + exec_id: String, + timeout: Duration, +) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { tokio::time::sleep(timeout).await; @@ -74,7 +81,7 @@ pub(super) fn start_timeout_watcher(target: TimeoutTarget, exec_id: String, time Ok(false) => info!(execution_id = %exec_id, "exited within grace after SIGTERM"), Err(error) => warn!(execution_id = %exec_id, %error, "timeout SIGKILL failed"), } - }); + }) } #[cfg(test)]