From be0b994fedfc5bdf36b75c6bdbb3f4385547d9e7 Mon Sep 17 00:00:00 2001 From: BatmanByte <300328404+BatmanByte@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:36:30 +0800 Subject: [PATCH] feat(exec): retain terminal sessions after exit An execution's registry entry disappeared as soon as its output stream ended, so a Wait or Attach arriving afterwards found nothing and the caller could not tell "never existed" from "already finished". The registry now holds three states per execution. Live behaves as before. Once an execution's exit and output summary are both in, it becomes Retained and its buffered output stays readable. After the retain grace it degrades to a Tombstone holding only the classified exit and a truncated diagnosis. Retention is bounded on grace, TTL, entry count and retained bytes, with LRU eviction, so a long-lived box cannot grow the registry without limit. Reservations close the window between spawn and register: an execution that cannot be published is SIGKILLed and torn down, and the reservation releases on drop so a request cancelled mid-spawn does not strand its id in a Reserved entry nothing prunes. Neither retirement path drops a retained session while a reader is still streaming it, because aborting that forwarder truncates the replay indistinguishably from a normal end of output: the caps skip it and the grace defers it until the reader detaches. Shutdown signals the retention pruner rather than aborting it, so an entry cannot be left tombstoned while still holding resources. The timeout watcher now returns its handle so a session that already exited cancels it rather than leaving a task parked on a deadline. --- src/guest/src/service/container.rs | 21 +- src/guest/src/service/exec/mod.rs | 466 ++++++++- src/guest/src/service/exec/output.rs | 266 ++++- src/guest/src/service/exec/registry.rs | 1255 +++++++++++++++++++++++- src/guest/src/service/exec/state.rs | 429 +++++++- src/guest/src/service/exec/timeout.rs | 11 +- 6 files changed, 2333 insertions(+), 115 deletions(-) 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)]