diff --git a/host/src/process/cache.rs b/host/src/process/cache.rs index a5956ce..a997a45 100644 --- a/host/src/process/cache.rs +++ b/host/src/process/cache.rs @@ -130,7 +130,15 @@ impl Compiler { } } } -/// Compile programs in both command folders without blocking world startup. +fn priority(name: &str) -> usize { + [ + "sh", "git", "rg", "grep", "ls", "cat", "sed", "find", "cargo", + ] + .iter() + .position(|tool| *tool == name.strip_suffix(".wasm").unwrap_or(name)) + .unwrap_or(9) +} +/// Compile common tools first in both command folders without blocking world startup. pub(super) fn precompile( compiler: Arc, root: Arc, @@ -140,18 +148,20 @@ pub(super) fn precompile( // Leave CPU and file bandwidth for commands during startup. let slots = Arc::new(tokio::sync::Semaphore::new(2)); for folder in ["/bin", "/usr/local/bin"] { - let entries = match root.list(folder).await { + let mut entries = match root.list(folder).await { Ok(entries) => entries, Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue, Err(e) => return Err(e.into()), }; + entries.sort_by_key(|entry| priority(&entry.name)); for entry in entries.into_iter().filter(|entry| !entry.directory) { + // Reserve before spawning so task scheduling cannot reorder admission. + let permit = slots.clone().acquire_owned().await?; let root = root.clone(); let compiler = compiler.clone(); let path = format!("{folder}/{}", entry.name); - let slots = slots.clone(); tasks.spawn(async move { - let _permit = slots.acquire().await?; + let _permit = permit; let file = root .open( &path, @@ -186,6 +196,167 @@ mod tests { use crate::exec::{Command, Exec, Options}; use std::{future::Future, sync::atomic::Ordering}; + #[test] + fn common_tools_precede_other_programs() { + let mut names = [ + "z", "cat.wasm", "cargo", "sh.wasm", "git.wasm", "rg", "ls", "grep", "sed", "find", + "arch", + ]; + names.sort_by_key(|name| priority(name)); + assert_eq!( + names, + [ + "sh.wasm", "git.wasm", "rg", "grep", "ls", "cat.wasm", "sed", "find", "cargo", "z", + "arch" + ] + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn cold_compilation_does_not_spend_command_timeout() -> anyhow::Result<()> { + use std::time::Duration; + let temp = tempfile::tempdir()?; + let compiler = Compiler::new(&temp.path().join("cache"))?; + let world = crate::world::World::open(temp.path())?; + std::fs::create_dir_all(temp.path().join("world/bin"))?; + let program = temp.path().join("world/bin/program.wasm"); + std::fs::write( + &program, + wat::parse_str(r#"(module (memory (export "memory") 1) (func (export "_start")))"#)?, + )?; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (finish_tx, finish_rx) = std::sync::mpsc::channel(); + *compiler.before_compile.lock().unwrap() = Some(Box::new(move || { + let _ = started_tx.send(()); + let _ = finish_rx.recv(); + })); + let runtime = super::super::Runtime::with_compiler(world, compiler)?; + started_rx.await?; + let command_runtime = runtime.clone(); + let mut command = tokio::spawn(async move { + command_runtime + .exec( + Command::Argv(vec!["/bin/program".into()]), + Options { + timeout: Some(Duration::from_secs(2)), + ..Default::default() + }, + &mut |_, _| {}, + ) + .await + }); + let pending = tokio::time::timeout(Duration::from_millis(2100), &mut command).await; + // Always release the blocking compiler, even when the old deadline fails this assertion. + finish_tx.send(())?; + assert!( + pending.is_err(), + "Compilation consumed the execution deadline: {pending:?}" + ); + assert_eq!(command.await??, 0); + std::fs::write( + &program, + wat::parse_str( + r#"(module (memory (export "memory") 1) (func (export "_start") (loop br 0)))"#, + )?, + )?; + let result = runtime + .exec( + Command::Argv(vec!["/bin/program".into()]), + Options { + timeout: Some(Duration::from_millis(100)), + ..Default::default() + }, + &mut |_, _| {}, + ) + .await; + assert!( + result + .unwrap_err() + .to_string() + .contains("command timed out") + ); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn child_compilation_pauses_the_parent_deadline() -> anyhow::Result<()> { + use super::super::{GroupGuard, Request, Stdio, deadline::Deadline}; + use std::time::Duration; + let temp = tempfile::tempdir()?; + let world = crate::world::World::open(temp.path())?; + let compiler = Compiler::new(&temp.path().join("cache"))?; + let parent = wat::parse_str( + r#"(module (memory (export "memory") 1) (func (export "_start") (loop br 0)))"#, + )?; + std::fs::write(temp.path().join("world/parent.wasm"), &parent)?; + std::fs::write( + temp.path().join("world/child.wasm"), + wat::parse_str(r#"(module (memory (export "memory") 1) (func (export "_start")))"#)?, + )?; + compiler.get(parent).await?; + let runtime = super::super::Runtime::with_compiler(world, compiler.clone())?; + let deadline = Deadline::new(Duration::ZERO); + let request = |name: &str| Request { + program: name.into(), + argv: vec![name.into()], + cwd: "/".into(), + env: Vec::new(), + group: None, + stdin: Stdio::Null, + stdout: Stdio::Null, + stderr: Stdio::Null, + }; + let parent = runtime + .start( + request("/parent.wasm"), + 0, + 0, + 0, + Vec::new(), + None, + None, + Some(deadline.clone()), + ) + .await?; + let _guard = GroupGuard { + runtime: runtime.clone(), + group: parent.pid, + }; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (finish_tx, finish_rx) = std::sync::mpsc::channel(); + *compiler.before_compile.lock().unwrap() = Some(Box::new(move || { + let _ = started_tx.send(()); + let _ = finish_rx.recv(); + })); + let child_runtime = runtime.clone(); + let child = tokio::spawn(async move { + child_runtime + .start( + request("/child.wasm"), + parent.pid, + parent.pid, + parent.group, + Vec::new(), + None, + None, + None, + ) + .await + }); + let started = tokio::time::timeout(Duration::from_secs(5), started_rx).await; + let expired = deadline.expired(); + tokio::pin!(expired); + let paused = std::future::poll_fn(|cx| { + std::task::Poll::Ready(expired.as_mut().poll(cx).is_pending()) + }) + .await; + let _ = finish_tx.send(()); + started??; + assert!(paused, "Child compilation did not pause the parent budget"); + child.await??; + Ok(()) + } + #[test] fn compilation_survives_the_originating_runtime() { let temp = tempfile::tempdir().unwrap(); diff --git a/host/src/process/deadline.rs b/host/src/process/deadline.rs new file mode 100644 index 0000000..c9a1017 --- /dev/null +++ b/host/src/process/deadline.rs @@ -0,0 +1,87 @@ +//! Command-wide execution budgets exclude the union of component compilation waits, including child spawns. +use std::{sync::Arc, time::Duration}; +use tokio::{sync::watch, time::Instant}; + +#[derive(Clone, Copy)] +struct State { + deadline: Instant, + compiling: usize, + paused_at: Instant, +} + +pub(super) struct Deadline(watch::Sender); + +impl Deadline { + pub fn new(duration: Duration) -> Arc { + let now = Instant::now(); + Arc::new(Self( + watch::channel(State { + deadline: now + duration, + compiling: 0, + paused_at: now, + }) + .0, + )) + } + + pub fn pause(self: &Arc) -> Compilation { + self.0.send_modify(|state| { + if state.compiling == 0 { + state.paused_at = Instant::now(); + } + state.compiling += 1; + }); + Compilation(self.clone()) + } + + pub async fn expired(&self) { + let mut state = self.0.subscribe(); + loop { + let current = *state.borrow_and_update(); + if current.compiling == 0 { + tokio::select! { + biased; + changed = state.changed() => { changed.unwrap(); } + _ = tokio::time::sleep_until(current.deadline) => return, + } + } else { + state.changed().await.unwrap(); + } + } + } +} + +pub(super) struct Compilation(Arc); + +impl Drop for Compilation { + fn drop(&mut self) { + self.0.0.send_modify(|state| { + state.compiling -= 1; + if state.compiling == 0 { + state.deadline += state.paused_at.elapsed(); + } + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn overlapping_compilations_share_one_pause_and_drop_resumes() { + let deadline = Deadline::new(Duration::from_secs(2)); + let first = deadline.pause(); + let paused_at = deadline.0.borrow().paused_at; + let second = deadline.pause(); + assert_eq!(deadline.0.borrow().compiling, 2); + assert_eq!(deadline.0.borrow().paused_at, paused_at); + let original = deadline.0.borrow().deadline; + drop(first); + assert_eq!(deadline.0.borrow().compiling, 1); + assert_eq!(deadline.0.borrow().deadline, original); + drop(second); + assert_eq!(deadline.0.borrow().compiling, 0); + assert!(deadline.0.borrow().deadline >= original); + } +} diff --git a/host/src/process/mod.rs b/host/src/process/mod.rs index 8c1e850..921b08d 100644 --- a/host/src/process/mod.rs +++ b/host/src/process/mod.rs @@ -47,6 +47,7 @@ mod bindings { } use bindings::sausage::process::host::{Host, HostChild, Request, Stdio}; mod cache; +mod deadline; mod http; pub use http::Router; pub mod net; @@ -83,6 +84,7 @@ struct Entry { command: u32, group: u32, cancelled: Arc, + deadline: Option>, tag: sausage_vfs::journal::Tag, abort: AbortHandle, } @@ -329,16 +331,17 @@ impl Runtime { conversations: Vec, inherited_io: Option, tool_call: Option, + deadline: Option>, ) -> anyhow::Result { anyhow::ensure!(!request.argv.is_empty(), "empty argv"); - let cancelled = if parent == 0 { - Arc::new(AtomicBool::new(false)) + let (cancelled, deadline) = if parent == 0 { + (Arc::new(AtomicBool::new(false)), deadline) } else { self.entries .lock() .unwrap() .get(&parent) - .map(|entry| entry.cancelled.clone()) + .map(|entry| (entry.cancelled.clone(), entry.deadline.clone())) .ok_or_else(|| anyhow::anyhow!("parent exited"))? }; if !request.env.iter().any(|(k, _)| k == "PATH") { @@ -366,7 +369,9 @@ impl Runtime { let bytes = file .read(0, usize::try_from(file.stat().await?.size)?) .await?; + let compilation = deadline.as_ref().map(|deadline| deadline.pause()); let pre = self.compiler.get(bytes.clone()).await?; + drop(compilation); self.cache.lock().unwrap().insert(bytes); let pid = self.next.fetch_add(1, Ordering::Relaxed); let tag = if parent == 0 { @@ -528,6 +533,7 @@ impl Runtime { command, group, cancelled: cancelled.clone(), + deadline, tag, abort: task.abort_handle(), }, @@ -583,6 +589,7 @@ impl Host for State { self.conversations.clone(), Some(self.stdio.clone()), None, + None, ) .await { @@ -707,79 +714,77 @@ impl Exec for Arc { options: Options, output: &mut (dyn for<'chunk> FnMut(Stream, &'chunk [u8]) + Send), ) -> anyhow::Result { - let deadline = options - .timeout - .map(|duration| tokio::time::Instant::now() + duration); - let argv = match command { - Command::Shell(script) => vec![ - "/bin/sh".into(), - "--noprofile".into(), - "--norc".into(), - "-c".into(), - script, - ], - Command::Argv(argv) => argv, - }; - let env = options.env; - let start = self.start( - Request { - program: argv.first().cloned().unwrap_or_default(), - argv, - env, - cwd: options.cwd, - group: None, - stdin: Stdio::Pipe, - stdout: Stdio::Pipe, - stderr: Stdio::Pipe, - }, - 0, - 0, - 0, - options.conversations, - None, - options.tool_call, - ); - let mut child = if let Some(deadline) = deadline { - tokio::time::timeout_at(deadline, start) - .await - .map_err(|_| anyhow::anyhow!("command timed out"))?? - } else { - start.await? - }; - let _guard = GroupGuard { - runtime: self.clone(), - group: child.pid, - }; - drop(child.input); - let mut stdout = child.output.take().unwrap(); - let mut stderr = child.error.take().unwrap(); + let deadline = options.timeout.map(deadline::Deadline::new); let run = async { - let (mut out_open, mut err_open) = (true, true); - let (mut out, mut err) = ([0; 8192], [0; 8192]); - while out_open || err_open { - tokio::select! { - n = stdout.read(&mut out), if out_open => { let n = n?; out_open = n != 0; if n != 0 { output(Stream::Stdout, &out[..n]); } } - n = stderr.read(&mut err), if err_open => { let n = n?; err_open = n != 0; if n != 0 { output(Stream::Stderr, &err[..n]); } } + let argv = match command { + Command::Shell(script) => vec![ + "/bin/sh".into(), + "--noprofile".into(), + "--norc".into(), + "-c".into(), + script, + ], + Command::Argv(argv) => argv, + }; + let env = options.env; + let mut child = self + .start( + Request { + program: argv.first().cloned().unwrap_or_default(), + argv, + env, + cwd: options.cwd, + group: None, + stdin: Stdio::Pipe, + stdout: Stdio::Pipe, + stderr: Stdio::Pipe, + }, + 0, + 0, + 0, + options.conversations, + None, + options.tool_call, + deadline.clone(), + ) + .await?; + let _guard = GroupGuard { + runtime: self.clone(), + group: child.pid, + }; + drop(child.input); + let mut stdout = child.output.take().unwrap(); + let mut stderr = child.error.take().unwrap(); + let run = async { + let (mut out_open, mut err_open) = (true, true); + let (mut out, mut err) = ([0; 8192], [0; 8192]); + while out_open || err_open { + tokio::select! { + n = stdout.read(&mut out), if out_open => { let n = n?; out_open = n != 0; if n != 0 { output(Stream::Stdout, &out[..n]); } } + n = stderr.read(&mut err), if err_open => { let n = n?; err_open = n != 0; if n != 0 { output(Stream::Stderr, &err[..n]); } } + } } - } - loop { - if let Some(code) = *child.code.borrow_and_update() { - return Ok(code); + loop { + if let Some(code) = *child.code.borrow_and_update() { + return Ok(code); + } + child.code.changed().await?; } - child.code.changed().await?; + }; + let result = run.await; + if let Some(journal) = self.world.root.journal() { + journal.retain().await?; } + result }; - let result = if let Some(deadline) = deadline { - tokio::time::timeout_at(deadline, run) - .await - .map_err(|_| anyhow::anyhow!("command timed out"))? + if let Some(deadline) = &deadline { + tokio::select! { + result = run => result, + _ = deadline.expired() => anyhow::bail!("command timed out"), + } } else { run.await - }; - if let Some(journal) = self.world.root.journal() { - journal.retain().await?; } - result } } /// Construct the harness's shared executor over the same open world. diff --git a/host/src/process/services.rs b/host/src/process/services.rs index 9ddf59a..105ca59 100644 --- a/host/src/process/services.rs +++ b/host/src/process/services.rs @@ -34,6 +34,7 @@ pub fn supervise(runtime: Arc) -> Supervisor { Vec::new(), None, None, + None, ); let child = tokio::select! { _ = &mut stop => return, child = start => child }; if let Ok(mut child) = child {