Integrate wasi-runtime-config into wasmtime-cli (bytecodealliance#8970)

ahaoboy · Jul 19, 2024 · 4b842a0 · 4b842a0
1 parent 3447849
commit 4b842a0
Show file tree

Hide file tree

Showing 8 changed files with 169 additions and 1 deletion.
diff --git a/Cargo.toml b/Cargo.toml
@@ -350,6 +350,7 @@ default = [
   "wasi-nn",
   "wasi-threads",
   "wasi-http",
+  "wasi-runtime-config",
 
   # Most features of Wasmtime are enabled by default.
   "wat",
@@ -395,6 +396,7 @@ disable-logging = ["log/max_level_off", "tracing/max_level_off"]
 wasi-nn = ["dep:wasmtime-wasi-nn"]
 wasi-threads = ["dep:wasmtime-wasi-threads", "threads"]
 wasi-http = ["component-model", "dep:wasmtime-wasi-http", "dep:tokio", "dep:hyper"]
+wasi-runtime-config = ["dep:wasmtime-wasi-runtime-config"]
 pooling-allocator = ["wasmtime/pooling-allocator", "wasmtime-cli-flags/pooling-allocator"]
 component-model = [
   "wasmtime/component-model",

diff --git a/crates/cli-flags/src/lib.rs b/crates/cli-flags/src/lib.rs
@@ -284,6 +284,8 @@ wasmtime_option_group! {
         pub threads: Option<bool>,
         /// Enable support for WASI HTTP API (experimental)
         pub http: Option<bool>,
+        /// Enable support for WASI runtime config API (experimental)
+        pub runtime_config: Option<bool>,
         /// Inherit environment variables and file descriptors following the
         /// systemd listen fd specification (UNIX only)
         pub listenfd: Option<bool>,
@@ -321,6 +323,8 @@ wasmtime_option_group! {
         ///
         /// This option can be further overwritten with `--env` flags.
         pub inherit_env: Option<bool>,
+        /// Pass a wasi runtime config variable to the program.
+        pub runtime_config_var: Vec<WasiRuntimeConfigVariable>,
     }
 
     enum Wasi {
@@ -334,6 +338,12 @@ pub struct WasiNnGraph {
     pub dir: String,
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct WasiRuntimeConfigVariable {
+    pub key: String,
+    pub value: String,
+}
+
 /// Common options for commands that translate WebAssembly modules
 #[derive(Parser, Clone)]
 pub struct CommonOptions {

diff --git a/crates/cli-flags/src/opt.rs b/crates/cli-flags/src/opt.rs
@@ -5,7 +5,7 @@
 //! specifying options in a struct-like syntax where all other boilerplate about
 //! option parsing is contained exclusively within this module.
 
-use crate::WasiNnGraph;
+use crate::{WasiNnGraph, WasiRuntimeConfigVariable};
 use anyhow::{bail, Result};
 use clap::builder::{StringValueParser, TypedValueParser, ValueParserFactory};
 use clap::error::{Error, ErrorKind};
@@ -396,3 +396,18 @@ impl WasmtimeOptionValue for WasiNnGraph {
         })
     }
 }
+
+impl WasmtimeOptionValue for WasiRuntimeConfigVariable {
+    const VAL_HELP: &'static str = "=<name>=<val>";
+    fn parse(val: Option<&str>) -> Result<Self> {
+        let val = String::parse(val)?;
+        let mut parts = val.splitn(2, "=");
+        Ok(WasiRuntimeConfigVariable {
+            key: parts.next().unwrap().to_string(),
+            value: match parts.next() {
+                Some(part) => part.into(),
+                None => "".to_string(),
+            },
+        })
+    }
+}
diff --git a/crates/test-programs/src/bin/cli_serve_runtime_config.rs b/crates/test-programs/src/bin/cli_serve_runtime_config.rs
@@ -0,0 +1,29 @@
+use test_programs::config::wasi::config::runtime;
+use test_programs::proxy;
+use test_programs::wasi::http::types::{
+    Fields, IncomingRequest, OutgoingBody, OutgoingResponse, ResponseOutparam,
+};
+
+struct T;
+
+proxy::export!(T);
+
+impl proxy::exports::wasi::http::incoming_handler::Guest for T {
+    fn handle(_: IncomingRequest, outparam: ResponseOutparam) {
+        let fields = Fields::new();
+        let resp = OutgoingResponse::new(fields);
+        let body = resp.body().expect("outgoing response");
+
+        ResponseOutparam::set(outparam, Ok(resp));
+
+        let out = body.write().expect("outgoing stream");
+        let config = runtime::get("hello").unwrap().unwrap();
+        out.blocking_write_and_flush(config.as_bytes())
+            .expect("writing response");
+
+        drop(out);
+        OutgoingBody::finish(body, None).expect("outgoing-body.finish");
+    }
+}
+
+fn main() {}
diff --git a/crates/wasi-runtime-config/src/lib.rs b/crates/wasi-runtime-config/src/lib.rs
@@ -115,6 +115,13 @@ impl<'a> From<&'a WasiRuntimeConfigVariables> for WasiRuntimeConfig<'a> {
     }
 }
 
+impl<'a> WasiRuntimeConfig<'a> {
+    /// Create a new view into the `wasi-runtime-config` state.
+    pub fn new(vars: &'a WasiRuntimeConfigVariables) -> Self {
+        Self { vars }
+    }
+}
+
 impl generated::Host for WasiRuntimeConfig<'_> {
     fn get(&mut self, key: String) -> Result<Result<Option<String>, generated::ConfigError>> {
         Ok(Ok(self.vars.0.get(&key).map(|s| s.to_owned())))

diff --git a/src/commands/run.rs b/src/commands/run.rs
@@ -25,6 +25,8 @@ use wasmtime_wasi_threads::WasiThreadsCtx;
 
 #[cfg(feature = "wasi-http")]
 use wasmtime_wasi_http::WasiHttpCtx;
+#[cfg(feature = "wasi-runtime-config")]
+use wasmtime_wasi_runtime_config::{WasiRuntimeConfig, WasiRuntimeConfigVariables};
 
 fn parse_preloads(s: &str) -> Result<(String, PathBuf)> {
     let parts: Vec<&str> = s.splitn(2, '=').collect();
@@ -670,6 +672,38 @@ impl RunCommand {
             }
         }
 
+        if self.run.common.wasi.runtime_config == Some(true) {
+            #[cfg(not(feature = "wasi-runtime-config"))]
+            {
+                bail!("Cannot enable wasi-runtime-config when the binary is not compiled with this feature.");
+            }
+            #[cfg(all(feature = "wasi-runtime-config", feature = "component-model"))]
+            {
+                match linker {
+                    CliLinker::Core(_) => {
+                        bail!("Cannot enable wasi-runtime-config for core wasm modules");
+                    }
+                    CliLinker::Component(linker) => {
+                        let vars = WasiRuntimeConfigVariables::from_iter(
+                            self.run
+                                .common
+                                .wasi
+                                .runtime_config_var
+                                .iter()
+                                .map(|v| (v.key.clone(), v.value.clone())),
+                        );
+
+                        wasmtime_wasi_runtime_config::add_to_linker(linker, |h| {
+                            WasiRuntimeConfig::new(
+                                Arc::get_mut(h.wasi_runtime_config.as_mut().unwrap()).unwrap(),
+                            )
+                        })?;
+                        store.data_mut().wasi_runtime_config = Some(Arc::new(vars));
+                    }
+                }
+            }
+        }
+
         if self.run.common.wasi.threads == Some(true) {
             #[cfg(not(feature = "wasi-threads"))]
             {
@@ -814,6 +848,9 @@ struct Host {
     limits: StoreLimits,
     #[cfg(feature = "profiling")]
     guest_profiler: Option<Arc<wasmtime::GuestProfiler>>,
+
+    #[cfg(feature = "wasi-runtime-config")]
+    wasi_runtime_config: Option<Arc<WasiRuntimeConfigVariables>>,
 }
 
 impl Host {

diff --git a/src/commands/serve.rs b/src/commands/serve.rs
@@ -19,6 +19,8 @@ use wasmtime_wasi_http::{body::HyperOutgoingBody, WasiHttpCtx, WasiHttpView};
 
 #[cfg(feature = "wasi-nn")]
 use wasmtime_wasi_nn::wit::WasiNnCtx;
+#[cfg(feature = "wasi-runtime-config")]
+use wasmtime_wasi_runtime_config::{WasiRuntimeConfig, WasiRuntimeConfigVariables};
 
 struct Host {
     table: wasmtime::component::ResourceTable,
@@ -29,6 +31,9 @@ struct Host {
 
     #[cfg(feature = "wasi-nn")]
     nn: Option<WasiNnCtx>,
+
+    #[cfg(feature = "wasi-runtime-config")]
+    wasi_runtime_config: Option<WasiRuntimeConfigVariables>,
 }
 
 impl WasiView for Host {
@@ -147,6 +152,8 @@ impl ServeCommand {
 
             #[cfg(feature = "wasi-nn")]
             nn: None,
+            #[cfg(feature = "wasi-runtime-config")]
+            wasi_runtime_config: None,
         };
 
         if self.run.common.wasi.nn == Some(true) {
@@ -165,6 +172,21 @@ impl ServeCommand {
             }
         }
 
+        if self.run.common.wasi.runtime_config == Some(true) {
+            #[cfg(feature = "wasi-runtime-config")]
+            {
+                let vars = WasiRuntimeConfigVariables::from_iter(
+                    self.run
+                        .common
+                        .wasi
+                        .runtime_config_var
+                        .iter()
+                        .map(|v| (v.key.clone(), v.value.clone())),
+                );
+                host.wasi_runtime_config.replace(vars);
+            }
+        }
+
         let mut store = Store::new(engine, host);
 
         if self.run.common.wasm.timeout.is_some() {
@@ -228,6 +250,19 @@ impl ServeCommand {
             }
         }
 
+        if self.run.common.wasi.runtime_config == Some(true) {
+            #[cfg(not(feature = "wasi-runtime-config"))]
+            {
+                bail!("support for wasi-runtime-config was disabled at compile time");
+            }
+            #[cfg(feature = "wasi-runtime-config")]
+            {
+                wasmtime_wasi_runtime_config::add_to_linker(linker, |h| {
+                    WasiRuntimeConfig::from(h.wasi_runtime_config.as_ref().unwrap())
+                })?;
+            }
+        }
+
         if self.run.common.wasi.threads == Some(true) {
             bail!("support for wasi-threads is not available with components");
         }

diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs
@@ -1853,6 +1853,39 @@ stderr [1] :: after empty
         run_wasmtime(&["run", "--argv0=foo.wasm", CLI_ARGV0, "foo.wasm"])?;
         Ok(())
     }
+
+    #[tokio::test]
+    async fn cli_serve_runtime_config() -> Result<()> {
+        let server = WasmtimeServe::new(CLI_SERVE_RUNTIME_CONFIG_COMPONENT, |cmd| {
+            cmd.arg("-Scli");
+            cmd.arg("-Sruntime-config");
+            cmd.arg("-Sruntime-config-var=hello=world");
+        })?;
+
+        let resp = server
+            .send_request(
+                hyper::Request::builder()
+                    .uri("http://localhost/")
+                    .body(String::new())
+                    .context("failed to make request")?,
+            )
+            .await?;
+
+        assert!(resp.status().is_success());
+        assert_eq!(resp.body(), "world");
+        Ok(())
+    }
+
+    #[test]
+    fn cli_runtime_config() -> Result<()> {
+        run_wasmtime(&[
+            "run",
+            "-Sruntime-config",
+            "-Sruntime-config-var=hello=world",
+            RUNTIME_CONFIG_GET_COMPONENT,
+        ])?;
+        Ok(())
+    }
 }
 
 #[test]