Integrate wasi-keyvalue into wasmtime-cli (#9032)

* Integrate wasi-keyvalue into wasmtime-cli

* fix miri check

Author: iawia002

Cargo.toml

@@ -360,6 +360,7 @@ default = [
   "wasi-threads",
   "wasi-http",
   "wasi-runtime-config",
+  "wasi-keyvalue",
 
   # Most features of Wasmtime are enabled by default.
   "wat",
@@ -406,6 +407,7 @@ 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"]
+wasi-keyvalue = ["dep:wasmtime-wasi-keyvalue", "wasmtime-wasi-keyvalue/redis"]
 pooling-allocator = ["wasmtime/pooling-allocator", "wasmtime-cli-flags/pooling-allocator"]
 component-model = [
   "wasmtime/component-model",

crates/cli-flags/src/lib.rs

@@ -286,6 +286,8 @@ wasmtime_option_group! {
         pub http: Option<bool>,
         /// Enable support for WASI runtime config API (experimental)
         pub runtime_config: Option<bool>,
+        /// Enable support for WASI key-value API (experimental)
+        pub keyvalue: Option<bool>,
         /// Inherit environment variables and file descriptors following the
         /// systemd listen fd specification (UNIX only)
         pub listenfd: Option<bool>,
@@ -324,7 +326,18 @@ 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>,
+        pub runtime_config_var: Vec<KeyValuePair>,
+        /// Preset data for the In-Memory provider of WASI key-value API.
+        pub keyvalue_in_memory_data: Vec<KeyValuePair>,
+        /// Grant access to the given Redis host for the Redis provider of WASI
+        /// key-value API.
+        pub keyvalue_redis_host: Vec<String>,
+        /// Sets the connection timeout parameter for the Redis provider of WASI
+        /// key-value API.
+        pub keyvalue_redis_connection_timeout: Option<Duration>,
+        /// Sets the response timeout parameter for the Redis provider of WASI
+        /// key-value API.
+        pub keyvalue_redis_response_timeout: Option<Duration>,
     }
 
     enum Wasi {
@@ -339,7 +352,7 @@ pub struct WasiNnGraph {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub struct WasiRuntimeConfigVariable {
+pub struct KeyValuePair {
     pub key: String,
     pub value: String,
 }

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, WasiRuntimeConfigVariable};
+use crate::{KeyValuePair, WasiNnGraph};
 use anyhow::{bail, Result};
 use clap::builder::{StringValueParser, TypedValueParser, ValueParserFactory};
 use clap::error::{Error, ErrorKind};
@@ -397,12 +397,12 @@ impl WasmtimeOptionValue for WasiNnGraph {
     }
 }
 
-impl WasmtimeOptionValue for WasiRuntimeConfigVariable {
+impl WasmtimeOptionValue for KeyValuePair {
     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 {
+        Ok(KeyValuePair {
             key: parts.next().unwrap().to_string(),
             value: match parts.next() {
                 Some(part) => part.into(),

crates/test-programs/src/bin/cli_serve_keyvalue.rs

@@ -0,0 +1,30 @@
+use test_programs::keyvalue::wasi::keyvalue;
+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 bucket = keyvalue::store::open("").unwrap();
+        let data = bucket.get("hello").unwrap().unwrap();
+        out.blocking_write_and_flush(&data)
+            .expect("writing response");
+
+        drop(out);
+        OutgoingBody::finish(body, None).expect("outgoing-body.finish");
+    }
+}
+
+fn main() {}

crates/test-programs/src/bin/keyvalue_main.rs

@@ -1,10 +1,7 @@
 use test_programs::keyvalue::wasi::keyvalue::{atomics, batch, store};
 
 fn main() {
-    let identifier = std::env::var_os("IDENTIFIER")
-        .unwrap()
-        .into_string()
-        .unwrap();
+    let identifier = std::env::var("IDENTIFIER").unwrap_or("".to_string());
     let bucket = store::open(&identifier).unwrap();
 
     if identifier != "" {

crates/wasi-keyvalue/Cargo.toml

@@ -13,13 +13,13 @@ workspace = true
 [dependencies]
 anyhow = { workspace = true }
 wasmtime = { workspace = true, features = ["runtime", "async", "component-model"] }
+wasmtime-wasi = { workspace = true }
 async-trait = { workspace = true }
 url = { workspace = true }
 redis = { workspace = true, optional = true, features = ["tokio-comp"] }
 
 [dev-dependencies]
 test-programs-artifacts = { workspace = true }
-wasmtime-wasi = { workspace = true }
 tokio = { workspace = true, features = ["macros"] }
 
 [features]

crates/wasi-keyvalue/src/bindings.rs

@@ -0,0 +1,26 @@
+wasmtime::component::bindgen!({
+    path: "wit",
+    world: "wasi:keyvalue/imports",
+    trappable_imports: true,
+    async: true,
+    with: {
+        "wasi:keyvalue/store/bucket": crate::Bucket,
+    },
+    trappable_error_type: {
+        "wasi:keyvalue/store/error" => crate::Error,
+    },
+});
+
+pub(crate) mod sync {
+    wasmtime::component::bindgen!({
+        path: "wit",
+        world: "wasi:keyvalue/imports",
+        trappable_imports: true,
+        with: {
+            "wasi:keyvalue/store/bucket": crate::Bucket,
+        },
+        trappable_error_type: {
+            "wasi:keyvalue/store/error" => crate::Error,
+        },
+    });
+}

crates/wasi-keyvalue/src/lib.rs

@@ -41,7 +41,7 @@
 //!     let mut linker = Linker::<Ctx>::new(&engine);
 //!     wasmtime_wasi::add_to_linker_async(&mut linker)?;
 //!     // add `wasi-runtime-config` world's interfaces to the linker
-//!     wasmtime_wasi_keyvalue::add_to_linker(&mut linker, |h: &mut Ctx| {
+//!     wasmtime_wasi_keyvalue::add_to_linker_async(&mut linker, |h: &mut Ctx| {
 //!         WasiKeyValue::new(&h.wasi_keyvalue_ctx, &mut h.table)
 //!     })?;
 //!
@@ -68,29 +68,17 @@
 
 #![deny(missing_docs)]
 
+mod bindings;
 mod provider;
-mod generated {
-    wasmtime::component::bindgen!({
-        path: "wit",
-        world: "wasi:keyvalue/imports",
-        trappable_imports: true,
-        async: true,
-        with: {
-            "wasi:keyvalue/store/bucket": crate::Bucket,
-        },
-        trappable_error_type: {
-            "wasi:keyvalue/store/error" => crate::Error,
-        },
-    });
-}
 
-use self::generated::wasi::keyvalue;
+use self::bindings::{sync::wasi::keyvalue as keyvalue_sync, wasi::keyvalue};
 use anyhow::Result;
 use async_trait::async_trait;
 use std::collections::HashMap;
 use std::fmt::Display;
 use url::Url;
 use wasmtime::component::{Resource, ResourceTable, ResourceTableError};
+use wasmtime_wasi::runtime::in_tokio;
 
 #[doc(hidden)]
 pub enum Error {
@@ -411,7 +399,12 @@ impl keyvalue::batch::Host for WasiKeyValue<'_> {
 }
 
 /// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
-pub fn add_to_linker<T: Send>(
+///
+/// This function will add the `async` variant of all interfaces into the
+/// `Linker` provided. By `async` this means that this function is only
+/// compatible with [`Config::async_support(true)`][wasmtime::Config::async_support].
+/// For embeddings with async support disabled see [`add_to_linker_sync`] instead.
+pub fn add_to_linker_async<T: Send>(
     l: &mut wasmtime::component::Linker<T>,
     f: impl Fn(&mut T) -> WasiKeyValue<'_> + Send + Sync + Copy + 'static,
 ) -> Result<()> {
@@ -421,6 +414,104 @@ pub fn add_to_linker<T: Send>(
     Ok(())
 }
 
+impl keyvalue_sync::store::Host for WasiKeyValue<'_> {
+    fn open(&mut self, identifier: String) -> Result<Resource<Bucket>, Error> {
+        in_tokio(async { keyvalue::store::Host::open(self, identifier).await })
+    }
+
+    fn convert_error(&mut self, err: Error) -> Result<keyvalue_sync::store::Error> {
+        match err {
+            Error::NoSuchStore => Ok(keyvalue_sync::store::Error::NoSuchStore),
+            Error::AccessDenied => Ok(keyvalue_sync::store::Error::AccessDenied),
+            Error::Other(e) => Ok(keyvalue_sync::store::Error::Other(e)),
+        }
+    }
+}
+
+impl keyvalue_sync::store::HostBucket for WasiKeyValue<'_> {
+    fn get(&mut self, bucket: Resource<Bucket>, key: String) -> Result<Option<Vec<u8>>, Error> {
+        in_tokio(async { keyvalue::store::HostBucket::get(self, bucket, key).await })
+    }
+
+    fn set(&mut self, bucket: Resource<Bucket>, key: String, value: Vec<u8>) -> Result<(), Error> {
+        in_tokio(async { keyvalue::store::HostBucket::set(self, bucket, key, value).await })
+    }
+
+    fn delete(&mut self, bucket: Resource<Bucket>, key: String) -> Result<(), Error> {
+        in_tokio(async { keyvalue::store::HostBucket::delete(self, bucket, key).await })
+    }
+
+    fn exists(&mut self, bucket: Resource<Bucket>, key: String) -> Result<bool, Error> {
+        in_tokio(async { keyvalue::store::HostBucket::exists(self, bucket, key).await })
+    }
+
+    fn list_keys(
+        &mut self,
+        bucket: Resource<Bucket>,
+        cursor: Option<u64>,
+    ) -> Result<keyvalue_sync::store::KeyResponse, Error> {
+        in_tokio(async {
+            let resp = keyvalue::store::HostBucket::list_keys(self, bucket, cursor).await?;
+            Ok(keyvalue_sync::store::KeyResponse {
+                keys: resp.keys,
+                cursor: resp.cursor,
+            })
+        })
+    }
+
+    fn drop(&mut self, bucket: Resource<Bucket>) -> Result<()> {
+        keyvalue::store::HostBucket::drop(self, bucket)
+    }
+}
+
+impl keyvalue_sync::atomics::Host for WasiKeyValue<'_> {
+    fn increment(
+        &mut self,
+        bucket: Resource<Bucket>,
+        key: String,
+        delta: u64,
+    ) -> Result<u64, Error> {
+        in_tokio(async { keyvalue::atomics::Host::increment(self, bucket, key, delta).await })
+    }
+}
+
+impl keyvalue_sync::batch::Host for WasiKeyValue<'_> {
+    fn get_many(
+        &mut self,
+        bucket: Resource<Bucket>,
+        keys: Vec<String>,
+    ) -> Result<Vec<Option<(String, Vec<u8>)>>, Error> {
+        in_tokio(async { keyvalue::batch::Host::get_many(self, bucket, keys).await })
+    }
+
+    fn set_many(
+        &mut self,
+        bucket: Resource<Bucket>,
+        key_values: Vec<(String, Vec<u8>)>,
+    ) -> Result<(), Error> {
+        in_tokio(async { keyvalue::batch::Host::set_many(self, bucket, key_values).await })
+    }
+
+    fn delete_many(&mut self, bucket: Resource<Bucket>, keys: Vec<String>) -> Result<(), Error> {
+        in_tokio(async { keyvalue::batch::Host::delete_many(self, bucket, keys).await })
+    }
+}
+
+/// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
+///
+/// This function will add the `sync` variant of all interfaces into the
+/// `Linker` provided. For embeddings with async support see
+/// [`add_to_linker_async`] instead.
+pub fn add_to_linker_sync<T>(
+    l: &mut wasmtime::component::Linker<T>,
+    f: impl Fn(&mut T) -> WasiKeyValue<'_> + Send + Sync + Copy + 'static,
+) -> Result<()> {
+    keyvalue_sync::store::add_to_linker_get_host(l, f)?;
+    keyvalue_sync::atomics::add_to_linker_get_host(l, f)?;
+    keyvalue_sync::batch::add_to_linker_get_host(l, f)?;
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     #[test]

crates/wasi-keyvalue/src/provider/inmemory.rs

@@ -1,4 +1,4 @@
-use crate::{generated::wasi::keyvalue::store::KeyResponse, to_other_error, Error, Host};
+use crate::{bindings::wasi::keyvalue::store::KeyResponse, to_other_error, Error, Host};
 use async_trait::async_trait;
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};

crates/wasi-keyvalue/src/provider/redis.rs

@@ -1,4 +1,4 @@
-use crate::{generated::wasi::keyvalue::store::KeyResponse, Error, Host};
+use crate::{bindings::wasi::keyvalue::store::KeyResponse, Error, Host};
 use anyhow::Result;
 use async_trait::async_trait;
 use redis::{aio::MultiplexedConnection, AsyncCommands, RedisError};
@@ -91,7 +91,7 @@ impl Host for Redis {
         for (key, value) in key_values {
             pipe.set(key, value).ignore();
         }
-        pipe.query_async(&mut self.conn).await?;
+        let _: () = pipe.query_async(&mut self.conn).await?;
         Ok(())
     }
 
@@ -100,7 +100,7 @@ impl Host for Redis {
         for key in keys {
             pipe.del(key).ignore();
         }
-        pipe.query_async(&mut self.conn).await?;
+        let _: () = pipe.query_async(&mut self.conn).await?;
         Ok(())
     }
 }