diff --git a/src/backends/wslc/common/src/wsl_container_runner.rs b/src/backends/wslc/common/src/wsl_container_runner.rs index e296c9bcf..ab4e14e95 100644 --- a/src/backends/wslc/common/src/wsl_container_runner.rs +++ b/src/backends/wslc/common/src/wsl_container_runner.rs @@ -29,6 +29,7 @@ use wxc_common::sandbox_process::StdioMode; use wxc_common::script_runner::ScriptRunner; use wxc_common::string_util::{to_wide, CoTaskMemPWSTR}; +use crate::container_steps::sdk_error; use crate::policy_mapping; use crate::stream_buffer::{stream_pair, StreamReader, StreamWriter}; use crate::wslc_bindings::*; @@ -593,16 +594,6 @@ enum TarFormat { Unknown, } -/// Create a ScriptResponse error from an HRESULT failure with optional SDK error message. -fn sdk_error(context: &str, hr: HRESULT, sdk_msg: &str) -> ScriptResponse { - let msg = if sdk_msg.is_empty() { - format!("{}: HRESULT 0x{:08X}", context, hr as u32) - } else { - format!("{}: {} (HRESULT 0x{:08X})", context, sdk_msg, hr as u32) - }; - ScriptResponse::error(&msg) -} - /// Builds a user-facing prerequisite error for the components `WslcGetMissingComponents` /// reports as missing. `missing` may combine multiple bits, and the guidance is branched /// per-component so a user missing only `VirtualMachinePlatform` isn't told to update WSL diff --git a/src/backends/wslc/daemon/src/control_server.rs b/src/backends/wslc/daemon/src/control_server.rs index 577b73eec..f28df9518 100644 --- a/src/backends/wslc/daemon/src/control_server.rs +++ b/src/backends/wslc/daemon/src/control_server.rs @@ -24,9 +24,9 @@ use std::time::Duration; use anyhow::{bail, Context, Result}; use serde::de::DeserializeOwned; use serde::Serialize; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::windows::named_pipe::{NamedPipeServer, ServerOptions}; -use tokio::sync::{Notify, Semaphore}; +use tokio::sync::{oneshot, Semaphore}; use tokio::task::JoinSet; use tokio::time::timeout; use windows::core::{PCWSTR, PWSTR}; @@ -41,10 +41,10 @@ use windows::Win32::Security::{ use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; use wslc_common::daemon_protocol::{ - encode_frame, DaemonRequest, DaemonResponse, ErrKind, StreamFrame, MAX_FRAME_SIZE, + encode_frame, DaemonRequest, DaemonResponse, StreamFrame, MAX_FRAME_SIZE, }; -use crate::session_manager::SessionHandle; +use crate::session_manager::{SessionHandle, WorkerError}; /// Upper bound on concurrently-serviced client connections. At capacity the /// accept loop applies backpressure (a new connection waits for a slot) instead @@ -85,15 +85,22 @@ pub fn bind(pipe_name: &str) -> Result<(OwnerOnlySecurity, NamedPipeServer)> { /// `active_clients` tracks in-flight requests so the idle watchdog does not tear /// the daemon down mid-request. Concurrency is bounded by a semaphore, and on /// shutdown all in-flight handlers are drained before returning so the caller -/// can release the WSLc session without racing a live handler. +/// can release the WSLc session without racing a live handler. `activity` is a +/// monotonic connection counter the watchdog compares across polls to catch +/// bursts that start and finish between two of its samples. pub async fn run( session: SessionHandle, pipe_name: String, security: OwnerOnlySecurity, first_instance: NamedPipeServer, - active_clients: Arc, - shutdown: Arc, + signals: crate::DaemonSignals, ) -> Result<()> { + let crate::DaemonSignals { + active_clients, + activity, + shutdown, + draining, + } = signals; let mut server = first_instance; let limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_CLIENTS)); let mut clients: JoinSet<()> = JoinSet::new(); @@ -106,6 +113,11 @@ pub async fn run( loop { tokio::select! { + // Bias toward shutdown: once the watchdog signals, prefer tearing + // down over accepting a connection that raced the final idle sample. + biased; + + _ = shutdown.notified() => break, connect = server.connect() => { if let Err(e) = connect { eprintln!("[wslc-daemon] pipe connect error: {e}"); @@ -118,7 +130,19 @@ pub async fn run( } continue; } + // The watchdog may have entered the draining state after its + // final sample but before this connection arrived. Refuse it + // rather than provision into a session that is about to be + // released, handing the client an ID that teardown invalidates. + // The client re-spawns a fresh daemon once our record is gone. + if draining.load(Ordering::SeqCst) { + break; + } let connected = server; + // Record the connection so an idle streak that spans this + // request is invalidated even if it completes between polls. + activity.fetch_add(1, Ordering::SeqCst); + // Recreate the next listening instance before servicing this one // so the next client is not refused. Transient failures are // retried with bounded backoff rather than tearing the whole @@ -147,7 +171,6 @@ pub async fn run( } // Reap finished handlers so the JoinSet does not accumulate. Some(_) = clients.join_next() => {} - _ = shutdown.notified() => break, } } @@ -361,7 +384,7 @@ async fn handle_client(mut pipe: NamedPipeServer, session: SessionHandle) -> Res DaemonRequest::Provision(config) => { let resp = match session.provision(config).await { Ok(sandbox_id) => DaemonResponse::Provisioned { sandbox_id }, - Err(e) => err_response(ErrKind::Backend, e), + Err(e) => worker_err_response(e), }; write_frame(&mut pipe, &resp).await?; } @@ -384,7 +407,15 @@ async fn handle_client(mut pipe: NamedPipeServer, session: SessionHandle) -> Res Ok(()) } -/// Exec: admit with `Ok`, then stream the run's outcome as [`StreamFrame`]s. +/// Exec: validate-then-admit, then stream the run's outcome as [`StreamFrame`]s. +/// +/// The sandbox is validated (exists + started) *before* the `Ok` admission is +/// written, and — critically — admission is **atomic** with the start of the +/// run on the worker thread (see [`SessionHandle::exec`]): the worker validates +/// and begins running within one command handler, so no `Stop`/`Deprovision` +/// can invalidate the checked state between the admission and the run. An +/// unknown/not-started sandbox therefore comes back as a pre-admission typed +/// [`DaemonResponse::Err`] rather than a post-admission stream `Error` frame. /// /// TODO(fill-in): bidirectional live stdio (client `Stdin` frames -> process, /// process stdout/stderr -> `Stdout`/`Stderr` frames). The skeleton runs the @@ -394,34 +425,60 @@ async fn handle_exec( session: SessionHandle, config: wslc_common::daemon_protocol::ExecConfig, ) -> Result<()> { - write_frame(&mut pipe, &DaemonResponse::Ok).await?; - let terminal = match session.exec(config).await { - Ok(code) => StreamFrame::Exit { code }, - Err(e) => StreamFrame::Error { - message: format!("{e:#}"), + // Await the worker's admission decision before writing anything: a rejected + // exec is a pre-admission typed error, never a post-admission stream frame. + write_exec_result(&mut pipe, session.exec(config).await).await +} + +/// Turn an exec **admission** outcome into the client's frame sequence, generic +/// over the transport so the protocol can be exercised over an in-memory duplex +/// in tests. On rejection it writes a single typed [`DaemonResponse::Err`]; on +/// admission it writes `Ok` then awaits completion and writes exactly one +/// terminal [`StreamFrame`] — `Exit` on success, `Error` on a run failure or a +/// dropped completion channel. +async fn write_exec_result( + pipe: &mut S, + admission: Result>, WorkerError>, +) -> Result<()> { + let done = match admission { + Ok(done) => done, + Err(e) => { + write_frame(pipe, &worker_err_response(e)).await?; + return Ok(()); + } + }; + write_frame(pipe, &DaemonResponse::Ok).await?; + let terminal = match done.await { + Ok(Ok(code)) => StreamFrame::Exit { code }, + Ok(Err(e)) => StreamFrame::Error { + message: e.to_string(), + }, + Err(_) => StreamFrame::Error { + message: "WSLc worker dropped the exec reply channel".to_string(), }, }; - write_frame(&mut pipe, &terminal).await?; + write_frame(pipe, &terminal).await?; Ok(()) } -/// Map a `Result<()>` to `Ok` / `Err` response. -fn ok_or_err(result: Result<()>) -> DaemonResponse { +/// Map a worker `Result<()>` to an `Ok` / typed `Err` response. +fn ok_or_err(result: Result<(), WorkerError>) -> DaemonResponse { match result { Ok(()) => DaemonResponse::Ok, - Err(e) => err_response(ErrKind::Backend, e), + Err(e) => worker_err_response(e), } } -fn err_response(kind: ErrKind, e: anyhow::Error) -> DaemonResponse { +/// Build a [`DaemonResponse::Err`] carrying the worker error's protocol `kind`. +fn worker_err_response(e: WorkerError) -> DaemonResponse { DaemonResponse::Err { - kind, - message: format!("{e:#}"), + kind: e.kind(), + message: e.to_string(), } } /// Read one length-prefixed frame and deserialise it. -async fn read_frame(pipe: &mut NamedPipeServer) -> Result { +async fn read_frame(pipe: &mut S) -> Result { let mut len_buf = [0u8; 4]; pipe.read_exact(&mut len_buf).await?; let len = u32::from_le_bytes(len_buf) as usize; @@ -434,9 +491,104 @@ async fn read_frame(pipe: &mut NamedPipeServer) -> Result(pipe: &mut NamedPipeServer, msg: &T) -> Result<()> { +async fn write_frame(pipe: &mut S, msg: &T) -> Result<()> { let frame = encode_frame(msg)?; pipe.write_all(&frame).await?; pipe.flush().await?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::duplex; + use wslc_common::daemon_protocol::ErrKind; + + /// A rejected admission (unknown sandbox) round-trips as a single typed + /// `DaemonResponse::Err { NotProvisioned }` through frame encode → transport + /// → decode, with no terminal frame following it. + #[tokio::test] + async fn exec_rejection_round_trips_as_typed_error() { + let (mut server, mut client) = duplex(64 * 1024); + let admission = Err(WorkerError::NotProvisioned("wslc:nope".to_string())); + + write_exec_result(&mut server, admission).await.unwrap(); + drop(server); + + let resp: DaemonResponse = read_frame(&mut client).await.unwrap(); + match resp { + DaemonResponse::Err { kind, message } => { + assert_eq!(kind, ErrKind::NotProvisioned); + assert!(message.contains("wslc:nope"), "message was {message:?}"); + } + other => panic!("expected a typed Err response, got {other:?}"), + } + // A rejected exec is a single frame: nothing else follows. + assert!(read_frame::<_, StreamFrame>(&mut client).await.is_err()); + } + + /// The `NotStarted` admission contract has an SDK-free regression here: a + /// provisioned-but-not-started sandbox surfaces as a typed pre-admission + /// error, exercised without constructing a real container handle. + #[tokio::test] + async fn exec_not_started_round_trips_as_typed_error() { + let (mut server, mut client) = duplex(64 * 1024); + let admission = Err(WorkerError::NotStarted("wslc:cold".to_string())); + + write_exec_result(&mut server, admission).await.unwrap(); + drop(server); + + let resp: DaemonResponse = read_frame(&mut client).await.unwrap(); + assert_eq!( + resp, + DaemonResponse::Err { + kind: ErrKind::NotStarted, + message: "sandbox wslc:cold is not started".to_string(), + } + ); + } + + /// Dropping the completion sender after admission must produce exactly one + /// terminal `StreamFrame::Error` — never a hang or a malformed stream. + #[tokio::test] + async fn dropped_completion_channel_yields_single_error_terminal() { + let (mut server, mut client) = duplex(64 * 1024); + let (done_tx, done_rx) = oneshot::channel::>(); + drop(done_tx); + + write_exec_result(&mut server, Ok(done_rx)).await.unwrap(); + drop(server); + + let admit: DaemonResponse = read_frame(&mut client).await.unwrap(); + assert_eq!(admit, DaemonResponse::Ok); + let terminal: StreamFrame = read_frame(&mut client).await.unwrap(); + match terminal { + StreamFrame::Error { message } => { + assert!( + message.contains("dropped the exec reply channel"), + "message was {message:?}" + ); + } + other => panic!("expected a terminal Error frame, got {other:?}"), + } + // Exactly one terminal frame is emitted. + assert!(read_frame::<_, StreamFrame>(&mut client).await.is_err()); + } + + /// A successful run writes admission `Ok` then a single `Exit` terminal + /// carrying the process exit code. + #[tokio::test] + async fn successful_exec_writes_ok_then_exit() { + let (mut server, mut client) = duplex(64 * 1024); + let (done_tx, done_rx) = oneshot::channel::>(); + done_tx.send(Ok(7)).unwrap(); + + write_exec_result(&mut server, Ok(done_rx)).await.unwrap(); + drop(server); + + let admit: DaemonResponse = read_frame(&mut client).await.unwrap(); + assert_eq!(admit, DaemonResponse::Ok); + let terminal: StreamFrame = read_frame(&mut client).await.unwrap(); + assert_eq!(terminal, StreamFrame::Exit { code: 7 }); + } +} diff --git a/src/backends/wslc/daemon/src/main.rs b/src/backends/wslc/daemon/src/main.rs index 8928168d2..b2d3a3773 100644 --- a/src/backends/wslc/daemon/src/main.rs +++ b/src/backends/wslc/daemon/src/main.rs @@ -26,7 +26,7 @@ mod control_server; mod session_manager; #[cfg(windows)] -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; #[cfg(windows)] use std::sync::Arc; #[cfg(windows)] @@ -107,23 +107,27 @@ async fn run() -> Result<()> { }; write_daemon_record(&record).context("publish daemon record")?; - let shutdown = Arc::new(Notify::new()); - let active_clients = Arc::new(AtomicUsize::new(0)); + // Shared control-plane signals: `shutdown` (idle teardown), `active_clients` + // (in-flight request count), `activity` (monotonic connection counter the + // watchdog compares across polls to catch bursts between samples), and + // `draining` (set by the watchdog before it signals shutdown so the accept + // loop refuses a client that raced the final idle sample). + let signals = DaemonSignals { + active_clients: Arc::new(AtomicUsize::new(0)), + activity: Arc::new(AtomicU64::new(0)), + shutdown: Arc::new(Notify::new()), + draining: Arc::new(AtomicBool::new(false)), + }; let server = tokio::spawn(control_server::run( session.clone(), pipe_name.clone(), security, first_instance, - active_clients.clone(), - shutdown.clone(), + signals.clone(), )); - let watchdog = tokio::spawn(idle_watchdog( - session.clone(), - active_clients, - shutdown.clone(), - )); + let watchdog = tokio::spawn(idle_watchdog(session.clone(), signals)); // Wait for the control server to finish: it stops accepting when `shutdown` // fires and then drains its in-flight handlers, so once this returns no @@ -153,41 +157,87 @@ async fn run() -> Result<()> { .context("control server failed") } +/// Shared control-plane signals wired between the accept loop and the idle +/// watchdog. Bundled so both tasks take one handle instead of a long argument +/// list; every field is an `Arc`, so `Clone` is cheap. +#[cfg(windows)] +#[derive(Clone)] +struct DaemonSignals { + /// Count of client requests currently being serviced. + active_clients: Arc, + /// Monotonic connection counter the watchdog compares across polls. + activity: Arc, + /// Idle-teardown signal (a retained permit via `notify_one`). + shutdown: Arc, + /// Set by the watchdog before it signals shutdown so the accept loop refuses + /// a client that connected after the final idle sample. + draining: Arc, +} + /// Poll the live-container count; once it has been zero for `IDLE_TIMEOUT` with -/// no in-flight requests, notify shutdown. +/// no in-flight requests and no new connections since the previous poll, notify +/// shutdown. +/// +/// Idle is only declared when three signals agree: the container count is zero, +/// no client request is in flight (`active_clients`), and the monotonic +/// `activity` counter has not advanced since the last poll. The last guard +/// closes the window where a client connects and completes entirely between two +/// polls — the count and `active_clients` would both read zero again, but the +/// bumped `activity` generation still resets the idle streak. /// /// The signal is delivered with [`Notify::notify_one`], not `notify_waiters`: /// the control server only awaits `shutdown.notified()` inside its `select!`, so /// a wakeup raised while it is executing another branch (accepting or spawning a /// handler) would be dropped by `notify_waiters` (it retains no permit when no /// task is parked). `notify_one` stores a permit, so the server's next -/// `notified()` completes regardless of when the signal was raised — the -/// watchdog can then return without leaving the daemon alive forever. +/// `notified()` completes regardless of when the signal was raised. #[cfg(windows)] -async fn idle_watchdog( - session: session_manager::SessionHandle, - active_clients: Arc, - shutdown: Arc, -) { +async fn idle_watchdog(session: session_manager::SessionHandle, signals: DaemonSignals) { + let DaemonSignals { + active_clients, + activity, + shutdown, + draining, + } = signals; let mut idle_for = Duration::ZERO; + let mut last_activity = activity.load(Ordering::SeqCst); loop { tokio::time::sleep(IDLE_POLL).await; - match session.container_count().await { - // A provision holds the count at zero while it runs, so also require - // no in-flight client requests before counting as idle. - Ok(0) if active_clients.load(Ordering::SeqCst) == 0 => { - idle_for += IDLE_POLL; - if idle_for >= IDLE_TIMEOUT { - shutdown.notify_one(); - return; - } - } - Ok(_) => idle_for = Duration::ZERO, + // Query the count first: it is serialized on the worker, so it cannot + // observe zero while a provision that will make it non-zero is in flight. + let count = match session.container_count().await { + Ok(count) => count, Err(_) => { // Worker gone: nothing left to serve. + draining.store(true, Ordering::SeqCst); shutdown.notify_one(); return; } + }; + // Read the monotonic generation last — after count and active_clients — + // so a client that connects and completes entirely within this poll + // (bumping `activity`, then dropping `active_clients` back to zero) is + // still observed here as `generation != last_activity`, instead of + // slipping through with both zero-reads while the generation bump goes + // unsampled. + let active = active_clients.load(Ordering::SeqCst); + let generation = activity.load(Ordering::SeqCst); + let idle = count == 0 && active == 0 && generation == last_activity; + last_activity = generation; + + if idle { + idle_for += IDLE_POLL; + if idle_for >= IDLE_TIMEOUT { + // Enter the draining state *before* signalling shutdown so the + // accept loop, once it wakes, refuses any client that connected + // after this final sample instead of provisioning it into a + // session that is about to be released. + draining.store(true, Ordering::SeqCst); + shutdown.notify_one(); + return; + } + } else { + idle_for = Duration::ZERO; } } } @@ -215,4 +265,39 @@ mod tests { .await .expect("a permit raised before the waiter parked must not be lost"); } + + /// Regression for the post-sample shutdown race: on an idle session the + /// watchdog must enter the draining state *before* it notifies shutdown, so + /// the accept loop refuses any client that connects after the final sample + /// instead of provisioning it into a session that is about to be released. + #[tokio::test(start_paused = true)] + async fn idle_shutdown_enters_draining_before_notifying() { + use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; + use std::sync::Arc; + + let session = crate::session_manager::spawn().unwrap(); + let signals = super::DaemonSignals { + active_clients: Arc::new(AtomicUsize::new(0)), + activity: Arc::new(AtomicU64::new(0)), + shutdown: Arc::new(Notify::new()), + draining: Arc::new(AtomicBool::new(false)), + }; + let shutdown = signals.shutdown.clone(); + let draining = signals.draining.clone(); + + // No containers and no clients: the watchdog runs to the idle timeout, + // sets `draining`, then notifies shutdown, and returns. + super::idle_watchdog(session.clone(), signals).await; + + assert!( + draining.load(Ordering::SeqCst), + "watchdog must enter the draining state on idle shutdown" + ); + // The shutdown permit is retained for the accept loop's next `notified()`. + tokio::time::timeout(Duration::from_secs(5), shutdown.notified()) + .await + .expect("idle shutdown must leave a retained permit"); + + session.shutdown().await.unwrap(); + } } diff --git a/src/backends/wslc/daemon/src/session_manager.rs b/src/backends/wslc/daemon/src/session_manager.rs index 799dfceb9..460db98f1 100644 --- a/src/backends/wslc/daemon/src/session_manager.rs +++ b/src/backends/wslc/daemon/src/session_manager.rs @@ -31,11 +31,11 @@ use tokio::sync::{mpsc, oneshot}; use wslc_common::container_steps::{self, ProcessSettings}; use wslc_common::daemon_protocol::{ - DeprovisionConfig, ExecConfig, NetworkMode, ProvisionConfig, StartConfig, StopConfig, + DeprovisionConfig, ErrKind, ExecConfig, NetworkMode, ProvisionConfig, StartConfig, StopConfig, }; use wslc_common::policy_mapping; use wslc_common::wslc_bindings::{ - WslcContainerGuard, WslcContainerNetworkingMode, WslcSdk, WslcSessionGuard, + WslcContainer, WslcContainerGuard, WslcContainerNetworkingMode, WslcSdk, WslcSessionGuard, }; use wxc_common::logger::{Logger, Mode}; use wxc_common::models::ScriptResponse; @@ -57,31 +57,78 @@ fn sr_err(resp: ScriptResponse) -> anyhow::Error { anyhow::anyhow!(resp.error_message) } +/// A typed worker failure. The control server maps [`WorkerError::kind`] onto +/// the protocol's [`ErrKind`] so clients can react (e.g. distinguish an unknown +/// sandbox from a backend fault) without string-matching the message. +#[derive(Debug)] +pub enum WorkerError { + /// The referenced sandbox id is unknown to the daemon. + NotProvisioned(String), + /// The sandbox exists but has not been started. + NotStarted(String), + /// A backend/SDK-level failure, or an internal worker/channel fault. + Backend(anyhow::Error), +} + +impl WorkerError { + /// The protocol classification the control server returns for this error. + pub fn kind(&self) -> ErrKind { + match self { + WorkerError::NotProvisioned(_) => ErrKind::NotProvisioned, + WorkerError::NotStarted(_) => ErrKind::NotStarted, + WorkerError::Backend(_) => ErrKind::Backend, + } + } +} + +impl std::fmt::Display for WorkerError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + WorkerError::NotProvisioned(id) => write!(f, "unknown sandbox {id}"), + WorkerError::NotStarted(id) => write!(f, "sandbox {id} is not started"), + WorkerError::Backend(e) => write!(f, "{e:#}"), + } + } +} + +impl std::error::Error for WorkerError {} + +impl From for WorkerError { + fn from(e: anyhow::Error) -> Self { + WorkerError::Backend(e) + } +} + /// A unit of work dispatched from an async pipe handler to the WSLc worker /// thread. Each variant carries a `oneshot` reply channel the worker fulfils. pub enum WorkerCommand { Provision { config: ProvisionConfig, - reply: oneshot::Sender>, + reply: oneshot::Sender>, }, Start { config: StartConfig, - reply: oneshot::Sender>, + reply: oneshot::Sender>, }, - /// Run a command to completion, returning its exit code. Live stdio - /// streaming is layered on in the fill-in phase; for now this is - /// request/response only. + /// Validate the sandbox (exists + started) and, if admitted, run the + /// command to completion. The two replies make admission **atomic** with the + /// run: because the worker services this whole command on its single thread + /// without yielding, no `Stop`/`Deprovision` can interleave between the + /// validation and the run. `admit` carries the pre-run decision (so an + /// unknown/not-started sandbox is a pre-admission typed error, never a + /// post-admission stream `Error`); `done` carries the run's exit code. Exec { config: ExecConfig, - reply: oneshot::Sender>, + admit: oneshot::Sender>, + done: oneshot::Sender>, }, Stop { config: StopConfig, - reply: oneshot::Sender>, + reply: oneshot::Sender>, }, Deprovision { config: DeprovisionConfig, - reply: oneshot::Sender>, + reply: oneshot::Sender>, }, /// Report the current live-container count (drives the idle watchdog). ContainerCount { reply: oneshot::Sender }, @@ -97,63 +144,78 @@ pub struct SessionHandle { impl SessionHandle { /// Provision a container, returning its minted `sandbox_id`. - pub async fn provision(&self, config: ProvisionConfig) -> Result { + pub async fn provision(&self, config: ProvisionConfig) -> Result { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::Provision { config, reply })?; rx.await.map_err(worker_gone)? } /// Start a provisioned container. - pub async fn start(&self, config: StartConfig) -> Result<()> { + pub async fn start(&self, config: StartConfig) -> Result<(), WorkerError> { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::Start { config, reply })?; rx.await.map_err(worker_gone)? } - /// Run a command in a started container to completion. - pub async fn exec(&self, config: ExecConfig) -> Result { - let (reply, rx) = oneshot::channel(); - self.send(WorkerCommand::Exec { config, reply })?; - rx.await.map_err(worker_gone)? + /// Admit and run a command in a started container. Awaits the worker's + /// **admission** decision first: on rejection (unknown/not-started sandbox) + /// this returns the typed error *before* the caller writes any admission to + /// the client. On admission it returns the completion receiver, which + /// resolves to the run's exit code. Admission and the start of the run are + /// atomic on the worker thread, so no lifecycle command can invalidate the + /// checked state between the two. + pub async fn exec( + &self, + config: ExecConfig, + ) -> Result>, WorkerError> { + let (admit, admit_rx) = oneshot::channel(); + let (done, done_rx) = oneshot::channel(); + self.send(WorkerCommand::Exec { + config, + admit, + done, + })?; + admit_rx.await.map_err(worker_gone)??; + Ok(done_rx) } /// Stop a running container. - pub async fn stop(&self, config: StopConfig) -> Result<()> { + pub async fn stop(&self, config: StopConfig) -> Result<(), WorkerError> { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::Stop { config, reply })?; rx.await.map_err(worker_gone)? } /// Deprovision (delete) a container. - pub async fn deprovision(&self, config: DeprovisionConfig) -> Result<()> { + pub async fn deprovision(&self, config: DeprovisionConfig) -> Result<(), WorkerError> { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::Deprovision { config, reply })?; rx.await.map_err(worker_gone)? } /// Current number of live containers (0 means the daemon is idle). - pub async fn container_count(&self) -> Result { + pub async fn container_count(&self) -> Result { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::ContainerCount { reply })?; rx.await.map_err(worker_gone) } /// Ask the worker to release everything and stop. Awaits confirmation. - pub async fn shutdown(&self) -> Result<()> { + pub async fn shutdown(&self) -> Result<(), WorkerError> { let (reply, rx) = oneshot::channel(); self.send(WorkerCommand::Shutdown { reply })?; rx.await.map_err(worker_gone) } - fn send(&self, cmd: WorkerCommand) -> Result<()> { + fn send(&self, cmd: WorkerCommand) -> Result<(), WorkerError> { self.tx .send(cmd) - .map_err(|_| anyhow::anyhow!("WSLc worker thread is gone")) + .map_err(|_| WorkerError::Backend(anyhow::anyhow!("WSLc worker thread is gone"))) } } -fn worker_gone(_e: oneshot::error::RecvError) -> anyhow::Error { - anyhow::anyhow!("WSLc worker dropped the reply channel") +fn worker_gone(_e: oneshot::error::RecvError) -> WorkerError { + WorkerError::Backend(anyhow::anyhow!("WSLc worker dropped the reply channel")) } /// Per-container bookkeeping held by the worker: whether the container is @@ -216,7 +278,7 @@ impl Worker { Ok(()) } - fn provision(&mut self, config: ProvisionConfig) -> Result { + fn provision(&mut self, config: ProvisionConfig) -> Result { self.ensure_session()?; let sdk = self.sdk.as_ref().expect("session ensured"); @@ -278,11 +340,11 @@ impl Worker { Ok(sandbox_id) } - fn start(&mut self, config: StartConfig) -> Result<()> { + fn start(&mut self, config: StartConfig) -> Result<(), WorkerError> { // Existence check first, so an unknown sandbox errors without needing the // SDK (keeps the no-WSL unit tests self-contained). if !self.containers.contains_key(&config.sandbox_id) { - anyhow::bail!("unknown sandbox {}", config.sandbox_id); + return Err(WorkerError::NotProvisioned(config.sandbox_id)); } let sdk = self .sdk @@ -301,14 +363,21 @@ impl Worker { Ok(()) } - fn exec(&mut self, config: ExecConfig) -> Result { - let (container, started) = match self.containers.get(&config.sandbox_id) { - Some(e) => (e.container.as_raw(), e.started), - None => anyhow::bail!("unknown sandbox {}", config.sandbox_id), - }; - if !started { - anyhow::bail!("sandbox {} is not started", config.sandbox_id); + /// Validate that a sandbox exists and is started, returning the live handle + /// needed to run. Sole owner of the exists+started invariant: [`exec`] trusts + /// the handle it is given and never re-checks, because the worker services + /// admission and the run on one thread without yielding between them. + fn validate_exec(&self, sandbox_id: &str) -> Result { + match self.containers.get(sandbox_id) { + None => Err(WorkerError::NotProvisioned(sandbox_id.to_string())), + Some(entry) if !entry.started => Err(WorkerError::NotStarted(sandbox_id.to_string())), + Some(entry) => Ok(entry.container.as_raw()), } + } + + /// Run a command in a sandbox whose existence/started state was already + /// confirmed by [`validate_exec`]; `container` is that validated handle. + fn exec(&mut self, config: ExecConfig, container: WslcContainer) -> Result { let sdk = self .sdk .as_ref() @@ -347,15 +416,18 @@ impl Worker { }; } self.containers.remove(&config.sandbox_id); - anyhow::bail!( + return Err(WorkerError::Backend(anyhow::anyhow!( "exec on sandbox {} could not be confirmed terminated; the container was \ quarantined", config.sandbox_id - ); + ))); } if outcome.timed_out { - anyhow::bail!("exec timed out after {}ms", config.timeout_ms); + return Err(WorkerError::Backend(anyhow::anyhow!( + "exec timed out after {}ms", + config.timeout_ms + ))); } // NOTE: outcome.stdout/stderr are captured but not yet forwarded — live // stdio streaming over the control pipe is a later fill-in; the PR1 @@ -363,9 +435,9 @@ impl Worker { Ok(outcome.exit_code) } - fn stop(&mut self, config: StopConfig) -> Result<()> { + fn stop(&mut self, config: StopConfig) -> Result<(), WorkerError> { if !self.containers.contains_key(&config.sandbox_id) { - anyhow::bail!("unknown sandbox {}", config.sandbox_id); + return Err(WorkerError::NotProvisioned(config.sandbox_id)); } let sdk = self .sdk @@ -384,10 +456,10 @@ impl Worker { Ok(()) } - fn deprovision(&mut self, config: DeprovisionConfig) -> Result<()> { + fn deprovision(&mut self, config: DeprovisionConfig) -> Result<(), WorkerError> { let container_raw = match self.containers.get(&config.sandbox_id) { Some(e) => e.container.as_raw(), - None => anyhow::bail!("unknown sandbox {}", config.sandbox_id), + None => return Err(WorkerError::NotProvisioned(config.sandbox_id)), }; if let Some(sdk) = self.sdk.as_ref() { @@ -467,8 +539,38 @@ pub fn spawn() -> Result { WorkerCommand::Start { config, reply } => { let _ = reply.send(worker.start(config)); } - WorkerCommand::Exec { config, reply } => { - let _ = reply.send(worker.exec(config)); + WorkerCommand::Exec { + config, + admit, + done, + } => { + // Validate and run in one handler so admission is atomic + // with the start of the run: the worker never yields + // between the two, so no Stop/Deprovision can interleave. + match worker.validate_exec(&config.sandbox_id) { + Err(e) => { + let _ = admit.send(Err(e)); + } + // Only run if the admission receiver is still there: + // if the client handler was dropped before it read + // admission, the blocking exec would otherwise starve + // every other lifecycle command for its full timeout. + Ok(container) if admit.send(Ok(())).is_ok() => { + let sandbox_id = config.sandbox_id.clone(); + let outcome = worker.exec(config, container); + if let Err(orphaned) = done.send(outcome) { + // The client handler is gone (e.g. its + // post-admission Ok write failed) but the run + // already happened. Record the result so a + // completed exec is never silently lost. + worker.logger.warning_line(&format!( + "exec on {sandbox_id} completed after the client \ + disconnected; orphaned result: {orphaned:?}" + )); + } + } + Ok(_) => {} + } } WorkerCommand::Stop { config, reply } => { let _ = reply.send(worker.stop(config)); @@ -589,6 +691,36 @@ mod tests { handle.shutdown().await.unwrap(); } + #[tokio::test] + async fn unknown_sandbox_maps_to_not_provisioned_kind() { + let handle = spawn().unwrap(); + let err = handle + .start(StartConfig { + sandbox_id: "wslc:does-not-exist".to_string(), + }) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrKind::NotProvisioned); + handle.shutdown().await.unwrap(); + } + + #[tokio::test] + async fn exec_unknown_sandbox_admission_is_not_provisioned() { + let handle = spawn().unwrap(); + let err = handle + .exec(ExecConfig { + sandbox_id: "wslc:does-not-exist".to_string(), + script_code: "echo hi".to_string(), + working_directory: String::new(), + env: Vec::new(), + timeout_ms: 0, + }) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrKind::NotProvisioned); + handle.shutdown().await.unwrap(); + } + // ---- Full lifecycle integration test (WSL2 host only) ---- // // Exercises the real SDK path end to end: provision (boot VM + create @@ -630,6 +762,9 @@ mod tests { timeout_ms: 30_000, }) .await + .unwrap() + .await + .unwrap() .unwrap(); assert_eq!(code, 0);