Switch from wasi-common to wasmtime-wasi

Cargo.toml

@@ -38,7 +38,6 @@ futures = "0.3.24"
 url = "2.3.1"
 
 # Wasmtime dependencies
-wasi-common = "13.0.0"
 wasmtime = "20.0.0"
 wasmtime-wasi = "20.0.0"
 wasmtime-wasi-nn = "20.0.0"

cli/Cargo.toml

@@ -47,7 +47,7 @@ tracing-futures = { workspace = true }
 tracing-subscriber = { version = "^0.3.16", features = ["env-filter", "fmt"] }
 viceroy-lib = { path = "../lib", version = "^0.9.7" }
 wat = "^1.0.38"
-wasi-common = { workspace = true }
+wasmtime-wasi = { workspace = true }
 wasmtime = { workspace = true }
 libc = "^0.2.139"
 

cli/src/main.rs

@@ -16,7 +16,7 @@
 use std::path::PathBuf;
 use std::process::ExitCode;
 
-use wasi_common::I32Exit;
+use wasmtime_wasi::I32Exit;
 
 mod opts;
 

lib/Cargo.toml

@@ -52,7 +52,6 @@ toml = "^0.5.9"
 tracing = { workspace = true }
 tracing-futures = { workspace = true }
 url = { workspace = true }
-wasi-common = { workspace = true }
 wasmtime = { workspace = true }
 wasmtime-wasi = { workspace = true }
 wasmtime-wasi-nn = { workspace = true }

lib/src/execute.rs

@@ -31,8 +31,8 @@ use {
     },
     tokio::sync::oneshot::{self, Sender},
     tracing::{event, info, info_span, Instrument, Level},
-    wasi_common::I32Exit,
     wasmtime::{Engine, InstancePre, Linker, Module, ProfilingStrategy},
+    wasmtime_wasi::I32Exit,
 };
 
 pub const EPOCH_INTERRUPTION_PERIOD: Duration = Duration::from_micros(50);
@@ -364,7 +364,8 @@ impl ExecuteCtx {
         // due to wasmtime limitations, in particular the fact that `Instance` is not `Send`.
         // However, the fact that the module itself is created within `ExecuteCtx::new`
         // means that the heavy lifting happens only once.
-        let mut store = create_store(&self, session, profiler).map_err(ExecutionError::Context)?;
+        let mut store =
+            create_store(&self, session, profiler, |_| {}).map_err(ExecutionError::Context)?;
 
         let instance = self
             .instance_pre
@@ -446,11 +447,13 @@ impl ExecuteCtx {
                 vec![(program_name.to_string(), self.module.clone())],
             )
         });
-        let mut store = create_store(&self, session, profiler).map_err(ExecutionError::Context)?;
-        store.data_mut().wasi().push_arg(program_name)?;
-        for arg in args {
-            store.data_mut().wasi().push_arg(arg)?;
-        }
+        let mut store = create_store(&self, session, profiler, |builder| {
+            builder.arg(program_name);
+            for arg in args {
+                builder.arg(arg);
+            }
+        })
+        .map_err(ExecutionError::Context)?;
 
         let instance = self
             .instance_pre

lib/src/linking.rs

@@ -5,11 +5,9 @@ use {
         config::ExperimentalModule, execute::ExecuteCtx, logging::LogEndpoint, session::Session,
         wiggle_abi, Error,
     },
-    anyhow::Context,
     std::collections::HashSet,
-    wasi_common::{pipe::WritePipe, WasiCtx},
     wasmtime::{GuestProfiler, Linker, Store, UpdateDeadline},
-    wasmtime_wasi::WasiCtxBuilder,
+    wasmtime_wasi::{preview1::WasiP1Ctx, WasiCtxBuilder},
     wasmtime_wasi_nn::WasiNnCtx,
 };
 
@@ -44,15 +42,15 @@ impl wasmtime::ResourceLimiter for Limiter {
 }
 
 pub struct WasmCtx {
-    wasi: WasiCtx,
+    wasi: WasiP1Ctx,
     wasi_nn: WasiNnCtx,
     session: Session,
     guest_profiler: Option<Box<GuestProfiler>>,
     limiter: Limiter,
 }
 
 impl WasmCtx {
-    pub fn wasi(&mut self) -> &mut WasiCtx {
+    pub fn wasi(&mut self) -> &mut WasiP1Ctx {
         &mut self.wasi
     }
 
@@ -87,8 +85,13 @@ pub(crate) fn create_store(
     ctx: &ExecuteCtx,
     session: Session,
     guest_profiler: Option<GuestProfiler>,
+    extra_init: impl FnOnce(&mut WasiCtxBuilder),
 ) -> Result<Store<WasmCtx>, anyhow::Error> {
-    let wasi = make_wasi_ctx(ctx, &session).context("creating Wasi context")?;
+    let mut builder = make_wasi_ctx(ctx, &session);
+
+    extra_init(&mut builder);
+
+    let wasi = builder.build_p1();
     let (backends, registry) = wasmtime_wasi_nn::preload(&[])?;
     let wasi_nn = WasiNnCtx::new(backends, registry);
     let wasm_ctx = WasmCtx {
@@ -102,7 +105,7 @@ pub(crate) fn create_store(
     store.set_epoch_deadline(1);
     store.epoch_deadline_callback(|mut store| {
         if let Some(mut prof) = store.data_mut().guest_profiler.take() {
-            prof.sample(&store);
+            prof.sample(&store, std::time::Duration::ZERO);
             store.data_mut().guest_profiler = Some(prof);
         }
         Ok(UpdateDeadline::Yield(1))
@@ -111,38 +114,43 @@ pub(crate) fn create_store(
     Ok(store)
 }
 
-/// Constructs a fresh `WasiCtx` for _each_ incoming request.
-fn make_wasi_ctx(ctx: &ExecuteCtx, session: &Session) -> Result<WasiCtx, anyhow::Error> {
+/// Constructs a fresh `WasiP1Ctx` for _each_ incoming request.
+fn make_wasi_ctx(ctx: &ExecuteCtx, session: &Session) -> WasiCtxBuilder {
     let mut wasi_ctx = WasiCtxBuilder::new();
 
     // Viceroy provides the same `FASTLY_*` environment variables that the production
     // Compute platform provides:
 
     wasi_ctx
         // These variables are stubbed out for compatibility
-        .env("FASTLY_CACHE_GENERATION", "0")?
-        .env("FASTLY_CUSTOMER_ID", "0000000000000000000000")?
-        .env("FASTLY_POP", "XXX")?
-        .env("FASTLY_REGION", "Somewhere")?
-        .env("FASTLY_SERVICE_ID", "0000000000000000000000")?
-        .env("FASTLY_SERVICE_VERSION", "0")?
+        .env("FASTLY_CACHE_GENERATION", "0")
+        .env("FASTLY_CUSTOMER_ID", "0000000000000000000000")
+        .env("FASTLY_POP", "XXX")
+        .env("FASTLY_REGION", "Somewhere")
+        .env("FASTLY_SERVICE_ID", "0000000000000000000000")
+        .env("FASTLY_SERVICE_VERSION", "0")
         // signal that we're in a local testing environment
-        .env("FASTLY_HOSTNAME", "localhost")?
+        .env("FASTLY_HOSTNAME", "localhost")
         // request IDs start at 0 and increment, rather than being UUIDs, for ease of testing
-        .env("FASTLY_TRACE_ID", &format!("{:032x}", session.req_id()))?;
+        .env("FASTLY_TRACE_ID", &format!("{:032x}", session.req_id()));
 
     if ctx.log_stdout() {
-        wasi_ctx.stdout(Box::new(WritePipe::new(LogEndpoint::new(b"stdout"))));
+        wasi_ctx.stdout(wasmtime_wasi::AsyncStdoutStream::new(
+            wasmtime_wasi::pipe::AsyncWriteStream::new(usize::MAX, LogEndpoint::new(b"stdout")),
+        ));
     } else {
         wasi_ctx.inherit_stdout();
     }
 
     if ctx.log_stderr() {
-        wasi_ctx.stderr(Box::new(WritePipe::new(LogEndpoint::new(b"stderr"))));
+        wasi_ctx.stderr(wasmtime_wasi::AsyncStdoutStream::new(
+            wasmtime_wasi::pipe::AsyncWriteStream::new(usize::MAX, LogEndpoint::new(b"stderr")),
+        ));
     } else {
         wasi_ctx.inherit_stderr();
     }
-    Ok(wasi_ctx.build())
+
+    wasi_ctx
 }
 
 pub fn link_host_functions(
@@ -156,7 +164,7 @@ pub fn link_host_functions(
                 wasmtime_wasi_nn::witx::add_to_linker(linker, WasmCtx::wasi_nn)
             }
         })?;
-    wasmtime_wasi::add_to_linker(linker, WasmCtx::wasi)?;
+    wasmtime_wasi::preview1::add_to_linker_async(linker, WasmCtx::wasi)?;
     wiggle_abi::fastly_abi::add_to_linker(linker, WasmCtx::session)?;
     wiggle_abi::fastly_cache::add_to_linker(linker, WasmCtx::session)?;
     wiggle_abi::fastly_config_store::add_to_linker(linker, WasmCtx::session)?;

lib/src/logging.rs

@@ -4,6 +4,7 @@ use {
         io::{self, Write},
         sync::Mutex,
     },
+    tokio::io::AsyncWrite,
 };
 
 /// A logging endpoint, which for Viceroy is just a name.
@@ -71,3 +72,36 @@ impl Write for LogEndpoint {
         LOG_WRITER.lock().unwrap().flush()
     }
 }
+
+impl AsyncWrite for LogEndpoint {
+    fn poll_write(
+        mut self: std::pin::Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<Result<usize, io::Error>> {
+        // This could integrate better with Tokio, as the lock being unavailable would be a good
+        // place to return a pending result. However, that would require either holding the lock
+        // the entire time the buffer is being built in [`LogEndpoint::write_entry`], or building
+        // the buffer and possibly throwing it away if the lock is unavailable. Instead, this
+        // implementation does the simple thing, and allows some blocking for what should be a
+        // quick write to the log endpoint.
+        std::task::Poll::Ready(self.as_mut().write(buf))
+    }
+
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), io::Error>> {
+        match LOG_WRITER.try_lock() {
+            Ok(mut writer) => std::task::Poll::Ready(writer.flush()),
+            Err(_) => std::task::Poll::Pending,
+        }
+    }
+
+    fn poll_shutdown(
+        self: std::pin::Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), io::Error>> {
+        std::task::Poll::Ready(Ok(()))
+    }
+}