Update to wasmtime 3.0 and enable wasi-nn calls (#209)

* Update to wasmtime 3.0 and enable wasi-nn calls

* Refactor handle_request

There are now two flavors of the handle request call. One is used for trap_test and the other is
used by the Viceroy framework. Also extended the call to bubble up the error so it could be easily included in the response body.

Co-authored-by: Matthew Tamayo <matthew@getmethodic.com>
Co-authored-by: Matthew Tamayo-Rios <100806329+mtr-fastly@users.noreply.github.com>

cli/src/main.rs

@@ -35,7 +35,7 @@ use {
 /// Create a new server, bind it to an address, and serve responses until an error occurs.
 pub async fn serve(opts: Opts) -> Result<(), Error> {
     // Load the wasm module into an execution context
-    let mut ctx = ExecuteCtx::new(opts.input(), opts.profiling_strategy())?
+    let mut ctx = ExecuteCtx::new(opts.input(), opts.profiling_strategy(), opts.wasi_modules())?
         .with_log_stderr(opts.log_stderr())
         .with_log_stdout(opts.log_stdout());
 

cli/src/opts.rs

@@ -1,13 +1,14 @@
 //! Command line arguments.
 
 use {
-    clap::Parser,
+    clap::{Parser, ValueEnum},
     std::net::{IpAddr, Ipv4Addr},
     std::{
+        collections::HashSet,
         net::SocketAddr,
         path::{Path, PathBuf},
     },
-    viceroy_lib::{Error, ProfilingStrategy},
+    viceroy_lib::{config::ExperimentalModule, Error, ProfilingStrategy},
 };
 
 // Command-line arguments for the Viceroy CLI.
@@ -43,6 +44,9 @@ pub struct Opts {
     // Whether to enable wasmtime's builtin profiler.
     #[arg(long = "profiler", value_parser = check_wasmtime_profiler_mode)]
     profiler: Option<ProfilingStrategy>,
+    /// Set of experimental WASI modules to link against.
+    #[arg(value_enum, long = "experimental_modules", required = false)]
+    experimental_modules: Vec<ExperimentalModuleArg>,
 }
 
 impl Opts {
@@ -83,6 +87,49 @@ impl Opts {
     pub fn profiling_strategy(&self) -> ProfilingStrategy {
         self.profiler.unwrap_or(ProfilingStrategy::None)
     }
+
+    // Set of experimental wasi modules to link against.
+    pub fn wasi_modules(&self) -> HashSet<ExperimentalModule> {
+        self.experimental_modules.iter().map(|x| x.into()).collect()
+    }
+}
+
+/// Enum of available (experimental) wasi modules
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Hash)]
+pub enum ExperimentalModuleArg {
+    WasiNn,
+}
+
+impl From<ExperimentalModuleArg> for ExperimentalModule {
+    fn from(arg: ExperimentalModuleArg) -> ExperimentalModule {
+        match arg {
+            ExperimentalModuleArg::WasiNn => ExperimentalModule::WasiNn,
+        }
+    }
+}
+
+impl From<&ExperimentalModuleArg> for ExperimentalModule {
+    fn from(arg: &ExperimentalModuleArg) -> ExperimentalModule {
+        match arg {
+            ExperimentalModuleArg::WasiNn => ExperimentalModule::WasiNn,
+        }
+    }
+}
+
+impl From<ExperimentalModule> for ExperimentalModuleArg {
+    fn from(module: ExperimentalModule) -> ExperimentalModuleArg {
+        match module {
+            ExperimentalModule::WasiNn => ExperimentalModuleArg::WasiNn,
+        }
+    }
+}
+
+impl From<&ExperimentalModule> for ExperimentalModuleArg {
+    fn from(module: &ExperimentalModule) -> ExperimentalModuleArg {
+        match module {
+            ExperimentalModule::WasiNn => ExperimentalModuleArg::WasiNn,
+        }
+    }
 }
 
 /// A parsing function used by [`Opts`][opts] to check that the input is a valid Wasm module in

cli/tests/integration/common.rs

@@ -3,7 +3,10 @@
 
 use futures::stream::StreamExt;
 use hyper::{service, Body as HyperBody, Request, Response, Server};
-use std::{convert::Infallible, future::Future, net::SocketAddr, path::PathBuf, sync::Arc};
+use std::{
+    collections::HashSet, convert::Infallible, future::Future, net::SocketAddr, path::PathBuf,
+    sync::Arc,
+};
 use tokio::sync::Mutex;
 use tracing_subscriber::filter::EnvFilter;
 use viceroy_lib::{
@@ -197,7 +200,7 @@ impl Test {
             .try_init()
             .ok();
 
-        let ctx = ExecuteCtx::new(&self.module_path, ProfilingStrategy::None)
+        let ctx = ExecuteCtx::new(&self.module_path, ProfilingStrategy::None, HashSet::new())
             .expect("failed to set up execution context")
             .with_backends(self.backends.clone())
             .with_dictionaries(self.dictionaries.clone())
@@ -250,6 +253,7 @@ impl Test {
                     .clone()
                     .handle_request(req.map(Into::into), addr.ip())
                     .await
+                    .map(|result| result.0)
                     .expect("failed to handle the request");
                 responses.push(resp);
             }