1 //! Example of instantiating a WASIp2 component with the use of resource
2 
3 /*
4 You can execute this example with:
5     cmake examples/
6     cargo run --example resource-component
7 */
8 
9 use std::collections::HashMap;
10 
11 use wasmtime::component::bindgen;
12 use wasmtime::component::{Component, Linker, ResourceTable};
13 use wasmtime::component::{HasSelf, Resource};
14 use wasmtime::{Engine, Result, Store};
15 use wasmtime_wasi::p2::add_to_linker_async;
16 use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView};
17 
18 pub struct ComponentRunStates {
19     // These two are required basically as a standard way to enable the impl of IoView and
20     // WasiView.
21     // impl of WasiView is required by [`wasmtime_wasi::p2::add_to_linker_sync`]
22     pub wasi_ctx: WasiCtx,
23     pub resource_table: ResourceTable,
24     // You can add other custom host states if needed
25 }
26 
27 impl WasiView for ComponentRunStates {
28     fn ctx(&mut self) -> WasiCtxView<'_> {
29         WasiCtxView {
30             ctx: &mut self.wasi_ctx,
31             table: &mut self.resource_table,
32         }
33     }
34 }
35 
36 impl ComponentRunStates {
37     pub fn new() -> Self {
38         // Create a WASI context and put it in a Store; all instances in the store
39         // share this context. `WasiCtx` provides a number of ways to
40         // configure what the target program will have access to.
41         ComponentRunStates {
42             wasi_ctx: WasiCtx::builder().build(),
43             resource_table: ResourceTable::new(),
44         }
45     }
46 }
47 
48 bindgen!({
49     path: "./examples/resource-component/kv-store.wit",
50     world: "kv-database",
51     // Interactions with `ResourceTable` can possibly trap so enable the ability
52     // to return traps from generated functions.
53     imports: { default: async | trappable },
54     exports: { default: async },
55     with: {
56         "example:kv-store/kvdb.connection": Connection
57     },
58 });
59 
60 pub struct Connection {
61     pub storage: HashMap<String, String>,
62 }
63 
64 impl KvDatabaseImports for ComponentRunStates {
65     async fn log(&mut self, msg: String) -> Result<(), wasmtime::Error> {
66         // provide host function to the component
67         println!("Log: {msg}");
68         Ok(())
69     }
70 }
71 
72 impl example::kv_store::kvdb::Host for ComponentRunStates {}
73 
74 impl example::kv_store::kvdb::HostConnection for ComponentRunStates {
75     async fn new(&mut self) -> Result<Resource<Connection>, wasmtime::Error> {
76         Ok(self.resource_table.push(Connection {
77             storage: HashMap::new(),
78         })?)
79     }
80 
81     async fn get(
82         &mut self,
83         resource: Resource<Connection>,
84         key: String,
85     ) -> Result<Option<String>, wasmtime::Error> {
86         let connection = self.resource_table.get(&resource)?;
87         Ok(connection.storage.get(&key).cloned())
88     }
89 
90     async fn set(
91         &mut self,
92         resource: Resource<Connection>,
93         key: String,
94         value: String,
95     ) -> Result<()> {
96         let connection = self.resource_table.get_mut(&resource)?;
97         connection.storage.insert(key, value);
98         Ok(())
99     }
100 
101     async fn remove(
102         &mut self,
103         resource: Resource<Connection>,
104         key: String,
105     ) -> Result<Option<String>> {
106         let connection = self.resource_table.get_mut(&resource)?;
107         Ok(connection.storage.remove(&key))
108     }
109 
110     async fn clear(&mut self, resource: Resource<Connection>) -> Result<(), wasmtime::Error> {
111         let large_string = self.resource_table.get_mut(&resource)?;
112         large_string.storage.clear();
113         Ok(())
114     }
115 
116     async fn drop(&mut self, resource: Resource<Connection>) -> Result<()> {
117         let _ = self.resource_table.delete(resource)?;
118         Ok(())
119     }
120 }
121 
122 #[tokio::main]
123 async fn main() -> Result<()> {
124     let engine = Engine::default();
125     let mut linker = Linker::new(&engine);
126     let state = ComponentRunStates::new();
127     let mut store = Store::new(&engine, state);
128 
129     KvDatabase::add_to_linker::<_, HasSelf<_>>(&mut linker, |s| s)?;
130     add_to_linker_async(&mut linker)?;
131 
132     // Instantiate our component with the imports we've created, and run its function
133     let component = Component::from_file(&engine, "target/wasm32-wasip2/debug/guest_kvdb.wasm")?;
134     let bindings = KvDatabase::instantiate_async(&mut store, &component, &linker).await?;
135     let result = bindings
136         .call_replace_value(&mut store, "hello", "world")
137         .await?;
138     assert_eq!(result, None);
139     let result = bindings
140         .call_replace_value(&mut store, "hello", "wasmtime")
141         .await?;
142     assert_eq!(result, Some("world".to_string()));
143     Ok(())
144 }
145