236 lines
9.2 KiB
Rust
236 lines
9.2 KiB
Rust
//! Shares compiled programs and single-flight compilation across worlds in one host process.
|
|
use super::{Compiled, State, adapter, bindings};
|
|
use std::{
|
|
collections::HashMap,
|
|
path::{Path, PathBuf},
|
|
sync::{Arc, Mutex, OnceLock},
|
|
};
|
|
use tokio::sync::watch;
|
|
type Compilation = watch::Receiver<Option<Result<Compiled, String>>>;
|
|
use wasmtime::{
|
|
Engine,
|
|
component::{Component, Linker},
|
|
};
|
|
use wasmtime_wasi::p2::bindings::CommandPre;
|
|
|
|
static DIRECTORY: OnceLock<PathBuf> = OnceLock::new();
|
|
static SHARED: Mutex<Option<Arc<Compiler>>> = Mutex::new(None);
|
|
|
|
/// Select the host user's cache before opening a world.
|
|
pub fn configure(directory: PathBuf) -> anyhow::Result<()> {
|
|
DIRECTORY
|
|
.set(directory)
|
|
.map_err(|_| anyhow::anyhow!("cache directory already configured"))
|
|
}
|
|
fn directory() -> anyhow::Result<PathBuf> {
|
|
if let Some(dir) = DIRECTORY.get() {
|
|
return Ok(dir.clone());
|
|
}
|
|
if let Some(dir) = std::env::var_os("SAUSAGE_CACHE") {
|
|
return Ok(dir.into());
|
|
}
|
|
if let Some(dir) = std::env::var_os("XDG_CACHE_HOME") {
|
|
return Ok(PathBuf::from(dir).join("sausage"));
|
|
}
|
|
Ok(
|
|
PathBuf::from(std::env::var_os("HOME").ok_or_else(|| anyhow::anyhow!("HOME is not set"))?)
|
|
.join(".cache/sausage"),
|
|
)
|
|
}
|
|
pub(super) fn shared() -> anyhow::Result<Arc<Compiler>> {
|
|
let mut shared = SHARED.lock().unwrap();
|
|
if shared.is_none() {
|
|
*shared = Some(Compiler::new(&directory()?)?);
|
|
}
|
|
Ok(shared.as_ref().unwrap().clone())
|
|
}
|
|
pub(super) struct Compiler {
|
|
pub engine: Engine,
|
|
programs: Mutex<HashMap<Vec<u8>, Compilation>>,
|
|
#[cfg(test)]
|
|
before_compile: Mutex<Option<Box<dyn FnOnce() + Send>>>,
|
|
#[cfg(test)]
|
|
compilations: std::sync::atomic::AtomicUsize,
|
|
}
|
|
impl Compiler {
|
|
pub fn new(directory: &Path) -> anyhow::Result<Arc<Self>> {
|
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
|
let mut config = wasmtime::Config::new();
|
|
config.epoch_interruption(true);
|
|
config.cranelift_opt_level(wasmtime::OptLevel::None);
|
|
let mut cache = wasmtime::CacheConfig::new();
|
|
cache.with_directory(directory);
|
|
config.cache(Some(wasmtime::Cache::new(cache)?));
|
|
Ok(Arc::new(Self {
|
|
engine: Engine::new(&config)?,
|
|
programs: Mutex::new(HashMap::new()),
|
|
#[cfg(test)]
|
|
before_compile: Mutex::new(None),
|
|
#[cfg(test)]
|
|
compilations: std::sync::atomic::AtomicUsize::new(0),
|
|
}))
|
|
}
|
|
pub async fn get(self: &Arc<Self>, bytes: Vec<u8>) -> anyhow::Result<Compiled> {
|
|
let mut result = self
|
|
.programs
|
|
.lock()
|
|
.unwrap()
|
|
.entry(bytes.clone())
|
|
.or_insert_with(|| {
|
|
let (tx, rx) = watch::channel(None);
|
|
let compiler = self.clone();
|
|
// The job owns completion even if a command or a precompile waiter is cancelled.
|
|
tokio::spawn(async move {
|
|
let result =
|
|
tokio::task::spawn_blocking(move || compiler.compile(&bytes)).await;
|
|
tx.send_replace(Some(
|
|
result
|
|
.map_err(anyhow::Error::from)
|
|
.and_then(|r| r)
|
|
.map_err(|e| format!("{e:#}")),
|
|
));
|
|
});
|
|
rx
|
|
})
|
|
.clone();
|
|
loop {
|
|
if let Some(result) = result.borrow().clone() {
|
|
return result.map_err(|e| anyhow::anyhow!(e));
|
|
}
|
|
result.changed().await?;
|
|
}
|
|
}
|
|
fn compile(&self, bytes: &[u8]) -> anyhow::Result<Compiled> {
|
|
#[cfg(test)]
|
|
{
|
|
self.compilations
|
|
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
|
if let Some(hook) = self.before_compile.lock().unwrap().take() {
|
|
hook();
|
|
}
|
|
}
|
|
if bytes.get(..8) == Some(b"\0asm\x01\0\0\0") {
|
|
let module = wasmtime::Module::new(&self.engine, bytes)?;
|
|
let mut linker = wasmtime::Linker::new(&self.engine);
|
|
sausage_vfs::preview1::add_to_linker(&mut linker, |s: &mut State| {
|
|
s.preview1.as_mut().unwrap()
|
|
})?;
|
|
Ok(Compiled::Module(linker.instantiate_pre(&module)?))
|
|
} else {
|
|
let component = Component::new(&self.engine, bytes)?;
|
|
let mut linker = Linker::new(&self.engine);
|
|
wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
|
|
wasmtime_wasi_http::p2::add_only_http_to_linker_async(&mut linker)?;
|
|
adapter::add_to_linker(&mut linker, |s: &mut State| adapter::View {
|
|
adapter: &mut s.fs,
|
|
table: &mut s.table,
|
|
})?;
|
|
bindings::sausage::process::host::add_to_linker::<_, wasmtime::component::HasSelf<_>>(
|
|
&mut linker,
|
|
|s| s,
|
|
)?;
|
|
Ok(Compiled::Component(CommandPre::new(
|
|
linker.instantiate_pre(&component)?,
|
|
)?))
|
|
}
|
|
}
|
|
}
|
|
/// Compile programs in both command folders without blocking world startup.
|
|
pub(super) fn precompile(
|
|
compiler: Arc<Compiler>,
|
|
root: Arc<dyn sausage_vfs::Layer>,
|
|
) -> tokio::task::JoinHandle<anyhow::Result<usize>> {
|
|
tokio::spawn(async move {
|
|
let mut tasks = tokio::task::JoinSet::new();
|
|
// 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 {
|
|
Ok(entries) => entries,
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue,
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
for entry in entries.into_iter().filter(|entry| !entry.directory) {
|
|
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 file = root
|
|
.open(
|
|
&path,
|
|
sausage_vfs::Open {
|
|
follow: true,
|
|
..Default::default()
|
|
},
|
|
)
|
|
.await?;
|
|
let bytes = file.read(0, file.stat().await?.size as usize).await?;
|
|
if bytes.get(..4) != Some(b"\0asm") {
|
|
return Ok::<_, anyhow::Error>(false);
|
|
}
|
|
compiler.get(bytes).await?;
|
|
Ok(true)
|
|
});
|
|
}
|
|
}
|
|
let mut count = 0;
|
|
while let Some(result) = tasks.join_next().await {
|
|
if result?? {
|
|
count += 1;
|
|
}
|
|
}
|
|
Ok(count)
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::exec::{Command, Exec, Options};
|
|
use std::{future::Future, sync::atomic::Ordering};
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn command_joins_precompile_even_after_waiter_cancellation() {
|
|
let temp = tempfile::tempdir().unwrap();
|
|
let compiler = Compiler::new(&temp.path().join("cache")).unwrap();
|
|
let world = crate::world::World::open(&temp.path().join("world")).unwrap();
|
|
let bytes =
|
|
wat::parse_str(r#"(module (memory (export "memory") 1) (func (export "_start")))"#)
|
|
.unwrap();
|
|
std::fs::create_dir_all(temp.path().join("world/world/bin")).unwrap();
|
|
std::fs::write(temp.path().join("world/world/bin/program.wasm"), &bytes).unwrap();
|
|
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 || {
|
|
started_tx.send(()).unwrap();
|
|
finish_rx.recv().unwrap();
|
|
}));
|
|
// Startup returns while the compiler remains blocked behind the test gate.
|
|
let runtime = super::super::Runtime::with_compiler(world, compiler.clone()).unwrap();
|
|
started_rx.await.unwrap();
|
|
let background = compiler.clone();
|
|
let precompile = tokio::spawn(async move { background.get(bytes).await });
|
|
let mut output = |_: crate::exec::Stream, _: &[u8]| {};
|
|
let command = runtime.exec(
|
|
Command::Argv(vec!["/bin/program".into()]),
|
|
Options::default(),
|
|
&mut output,
|
|
);
|
|
tokio::pin!(command);
|
|
let pending = std::future::poll_fn(|cx| {
|
|
std::task::Poll::Ready(command.as_mut().poll(cx).is_pending())
|
|
})
|
|
.await;
|
|
assert!(pending);
|
|
assert_eq!(compiler.compilations.load(Ordering::Relaxed), 1);
|
|
precompile.abort();
|
|
let _ = precompile.await;
|
|
finish_tx.send(()).unwrap();
|
|
assert_eq!(command.await.unwrap(), 0);
|
|
assert_eq!(compiler.compilations.load(Ordering::Relaxed), 1);
|
|
assert!(!temp.path().join("world/host/compile-cache").exists());
|
|
}
|
|
}
|