wasmtime_wasi_keyvalue/
lib.rs

1//! # Wasmtime's [wasi-keyvalue] Implementation
2//!
3//! This crate provides a Wasmtime host implementation of the [wasi-keyvalue]
4//! API. With this crate, the runtime can run components that call APIs in
5//! [wasi-keyvalue] and provide components with access to key-value storages.
6//!
7//! Currently supported storage backends:
8//! * In-Memory (empty identifier)
9//!
10//! # Examples
11//!
12//! The usage of this crate is very similar to other WASI API implementations
13//! such as [wasi:cli] and [wasi:http].
14//!
15//! A common scenario is accessing KV store in a [wasi:cli] component.
16//! A standalone example of doing all this looks like:
17//!
18//! ```
19//! use wasmtime::{
20//!     component::{Linker, ResourceTable},
21//!     Config, Engine, Result, Store,
22//! };
23//! use wasmtime_wasi::{IoView, WasiCtx, WasiCtxBuilder, WasiView};
24//! use wasmtime_wasi_keyvalue::{WasiKeyValue, WasiKeyValueCtx, WasiKeyValueCtxBuilder};
25//!
26//! #[tokio::main]
27//! async fn main() -> Result<()> {
28//!     let mut config = Config::new();
29//!     config.async_support(true);
30//!     let engine = Engine::new(&config)?;
31//!
32//!     let mut store = Store::new(&engine, Ctx {
33//!         table: ResourceTable::new(),
34//!         wasi_ctx: WasiCtxBuilder::new().build(),
35//!         wasi_keyvalue_ctx: WasiKeyValueCtxBuilder::new().build(),
36//!     });
37//!
38//!     let mut linker = Linker::<Ctx>::new(&engine);
39//!     wasmtime_wasi::add_to_linker_async(&mut linker)?;
40//!     // add `wasi-keyvalue` world's interfaces to the linker
41//!     wasmtime_wasi_keyvalue::add_to_linker(&mut linker, |h: &mut Ctx| {
42//!         WasiKeyValue::new(&h.wasi_keyvalue_ctx, &mut h.table)
43//!     })?;
44//!
45//!     // ... use `linker` to instantiate within `store` ...
46//!
47//!     Ok(())
48//! }
49//!
50//! struct Ctx {
51//!     table: ResourceTable,
52//!     wasi_ctx: WasiCtx,
53//!     wasi_keyvalue_ctx: WasiKeyValueCtx,
54//! }
55//!
56//! impl IoView for Ctx {
57//!     fn table(&mut self) -> &mut ResourceTable { &mut self.table }
58//! }
59//! impl WasiView for Ctx {
60//!     fn ctx(&mut self) -> &mut WasiCtx { &mut self.wasi_ctx }
61//! }
62//! ```
63//!
64//! [wasi-keyvalue]: https://github.com/WebAssembly/wasi-keyvalue
65//! [wasi:cli]: https://docs.rs/wasmtime-wasi/latest
66//! [wasi:http]: https://docs.rs/wasmtime-wasi-http/latest
67
68#![deny(missing_docs)]
69
70mod generated {
71    wasmtime::component::bindgen!({
72        path: "wit",
73        world: "wasi:keyvalue/imports",
74        trappable_imports: true,
75        with: {
76            "wasi:keyvalue/store/bucket": crate::Bucket,
77        },
78        trappable_error_type: {
79            "wasi:keyvalue/store/error" => crate::Error,
80        },
81    });
82}
83
84use self::generated::wasi::keyvalue;
85use anyhow::Result;
86use std::collections::HashMap;
87use wasmtime::component::{Resource, ResourceTable, ResourceTableError};
88
89#[doc(hidden)]
90pub enum Error {
91    NoSuchStore,
92    AccessDenied,
93    Other(String),
94}
95
96impl From<ResourceTableError> for Error {
97    fn from(err: ResourceTableError) -> Self {
98        Self::Other(err.to_string())
99    }
100}
101
102#[doc(hidden)]
103pub struct Bucket {
104    in_memory_data: HashMap<String, Vec<u8>>,
105}
106
107/// Builder-style structure used to create a [`WasiKeyValueCtx`].
108#[derive(Default)]
109pub struct WasiKeyValueCtxBuilder {
110    in_memory_data: HashMap<String, Vec<u8>>,
111}
112
113impl WasiKeyValueCtxBuilder {
114    /// Creates a builder for a new context with default parameters set.
115    pub fn new() -> Self {
116        Default::default()
117    }
118
119    /// Preset data for the In-Memory provider.
120    pub fn in_memory_data<I, K, V>(mut self, data: I) -> Self
121    where
122        I: IntoIterator<Item = (K, V)>,
123        K: Into<String>,
124        V: Into<Vec<u8>>,
125    {
126        self.in_memory_data = data
127            .into_iter()
128            .map(|(k, v)| (k.into(), v.into()))
129            .collect();
130        self
131    }
132
133    /// Uses the configured context so far to construct the final [`WasiKeyValueCtx`].
134    pub fn build(self) -> WasiKeyValueCtx {
135        WasiKeyValueCtx {
136            in_memory_data: self.in_memory_data,
137        }
138    }
139}
140
141/// Capture the state necessary for use in the `wasi-keyvalue` API implementation.
142pub struct WasiKeyValueCtx {
143    in_memory_data: HashMap<String, Vec<u8>>,
144}
145
146impl WasiKeyValueCtx {
147    /// Convenience function for calling [`WasiKeyValueCtxBuilder::new`].
148    pub fn builder() -> WasiKeyValueCtxBuilder {
149        WasiKeyValueCtxBuilder::new()
150    }
151}
152
153/// A wrapper capturing the needed internal `wasi-keyvalue` state.
154pub struct WasiKeyValue<'a> {
155    ctx: &'a WasiKeyValueCtx,
156    table: &'a mut ResourceTable,
157}
158
159impl<'a> WasiKeyValue<'a> {
160    /// Create a new view into the `wasi-keyvalue` state.
161    pub fn new(ctx: &'a WasiKeyValueCtx, table: &'a mut ResourceTable) -> Self {
162        Self { ctx, table }
163    }
164}
165
166impl keyvalue::store::Host for WasiKeyValue<'_> {
167    fn open(&mut self, identifier: String) -> Result<Resource<Bucket>, Error> {
168        match identifier.as_str() {
169            "" => Ok(self.table.push(Bucket {
170                in_memory_data: self.ctx.in_memory_data.clone(),
171            })?),
172            _ => Err(Error::NoSuchStore),
173        }
174    }
175
176    fn convert_error(&mut self, err: Error) -> Result<keyvalue::store::Error> {
177        match err {
178            Error::NoSuchStore => Ok(keyvalue::store::Error::NoSuchStore),
179            Error::AccessDenied => Ok(keyvalue::store::Error::AccessDenied),
180            Error::Other(e) => Ok(keyvalue::store::Error::Other(e)),
181        }
182    }
183}
184
185impl keyvalue::store::HostBucket for WasiKeyValue<'_> {
186    fn get(&mut self, bucket: Resource<Bucket>, key: String) -> Result<Option<Vec<u8>>, Error> {
187        let bucket = self.table.get_mut(&bucket)?;
188        Ok(bucket.in_memory_data.get(&key).cloned())
189    }
190
191    fn set(&mut self, bucket: Resource<Bucket>, key: String, value: Vec<u8>) -> Result<(), Error> {
192        let bucket = self.table.get_mut(&bucket)?;
193        bucket.in_memory_data.insert(key, value);
194        Ok(())
195    }
196
197    fn delete(&mut self, bucket: Resource<Bucket>, key: String) -> Result<(), Error> {
198        let bucket = self.table.get_mut(&bucket)?;
199        bucket.in_memory_data.remove(&key);
200        Ok(())
201    }
202
203    fn exists(&mut self, bucket: Resource<Bucket>, key: String) -> Result<bool, Error> {
204        let bucket = self.table.get_mut(&bucket)?;
205        Ok(bucket.in_memory_data.contains_key(&key))
206    }
207
208    fn list_keys(
209        &mut self,
210        bucket: Resource<Bucket>,
211        cursor: Option<u64>,
212    ) -> Result<keyvalue::store::KeyResponse, Error> {
213        let bucket = self.table.get_mut(&bucket)?;
214        let keys: Vec<String> = bucket.in_memory_data.keys().cloned().collect();
215        let cursor = cursor.unwrap_or(0) as usize;
216        let keys_slice = &keys[cursor..];
217        Ok(keyvalue::store::KeyResponse {
218            keys: keys_slice.to_vec(),
219            cursor: None,
220        })
221    }
222
223    fn drop(&mut self, bucket: Resource<Bucket>) -> Result<()> {
224        self.table.delete(bucket)?;
225        Ok(())
226    }
227}
228
229impl keyvalue::atomics::Host for WasiKeyValue<'_> {
230    fn increment(
231        &mut self,
232        bucket: Resource<Bucket>,
233        key: String,
234        delta: u64,
235    ) -> Result<u64, Error> {
236        let bucket = self.table.get_mut(&bucket)?;
237        let value = bucket
238            .in_memory_data
239            .entry(key.clone())
240            .or_insert("0".to_string().into_bytes());
241        let current_value = String::from_utf8(value.clone())
242            .map_err(|e| Error::Other(e.to_string()))?
243            .parse::<u64>()
244            .map_err(|e| Error::Other(e.to_string()))?;
245        let new_value = current_value + delta;
246        *value = new_value.to_string().into_bytes();
247        Ok(new_value)
248    }
249}
250
251impl keyvalue::batch::Host for WasiKeyValue<'_> {
252    fn get_many(
253        &mut self,
254        bucket: Resource<Bucket>,
255        keys: Vec<String>,
256    ) -> Result<Vec<Option<(String, Vec<u8>)>>, Error> {
257        let bucket = self.table.get_mut(&bucket)?;
258        Ok(keys
259            .into_iter()
260            .map(|key| {
261                bucket
262                    .in_memory_data
263                    .get(&key)
264                    .map(|value| (key.clone(), value.clone()))
265            })
266            .collect())
267    }
268
269    fn set_many(
270        &mut self,
271        bucket: Resource<Bucket>,
272        key_values: Vec<(String, Vec<u8>)>,
273    ) -> Result<(), Error> {
274        let bucket = self.table.get_mut(&bucket)?;
275        for (key, value) in key_values {
276            bucket.in_memory_data.insert(key, value);
277        }
278        Ok(())
279    }
280
281    fn delete_many(&mut self, bucket: Resource<Bucket>, keys: Vec<String>) -> Result<(), Error> {
282        let bucket = self.table.get_mut(&bucket)?;
283        for key in keys {
284            bucket.in_memory_data.remove(&key);
285        }
286        Ok(())
287    }
288}
289
290/// Add all the `wasi-keyvalue` world's interfaces to a [`wasmtime::component::Linker`].
291pub fn add_to_linker<T: Send>(
292    l: &mut wasmtime::component::Linker<T>,
293    f: impl Fn(&mut T) -> WasiKeyValue<'_> + Send + Sync + Copy + 'static,
294) -> Result<()> {
295    keyvalue::store::add_to_linker_get_host(l, f)?;
296    keyvalue::atomics::add_to_linker_get_host(l, f)?;
297    keyvalue::batch::add_to_linker_get_host(l, f)?;
298    Ok(())
299}
wasmtime_wasi_keyvalue/lib.rs

wasmtime_wasi_keyvalue/
lib.rs
