wasmtime_wizer/
wasmtime.rs

1use crate::{InstanceState, SnapshotVal, Wizer};
2use anyhow::Context;
3use wasmtime::{Extern, Instance, Module, Result, Store, Val};
4
5impl Wizer {
6    /// Initialize the given Wasm, snapshot it, and return the serialized
7    /// snapshot as a new, pre-initialized Wasm module.
8    pub async fn run<T: Send>(
9        &self,
10        store: &mut Store<T>,
11        wasm: &[u8],
12        instantiate: impl AsyncFnOnce(&mut Store<T>, &Module) -> Result<wasmtime::Instance>,
13    ) -> anyhow::Result<Vec<u8>> {
14        let (cx, instrumented_wasm) = self.instrument(wasm)?;
15
16        let engine = store.engine();
17        let module = wasmtime::Module::new(engine, &instrumented_wasm)
18            .context("failed to compile the Wasm module")?;
19        self.validate_init_func(&module)?;
20
21        let instance = instantiate(store, &module).await?;
22        self.initialize(store, &instance).await?;
23        self.snapshot(cx, &mut WasmtimeWizer { store, instance })
24            .await
25    }
26
27    /// Check that the module exports an initialization function, and that the
28    /// function has the correct type.
29    fn validate_init_func(&self, module: &wasmtime::Module) -> anyhow::Result<()> {
30        log::debug!("Validating the exported initialization function");
31        match module.get_export(self.get_init_func()) {
32            Some(wasmtime::ExternType::Func(func_ty)) => {
33                if func_ty.params().len() != 0 || func_ty.results().len() != 0 {
34                    anyhow::bail!(
35                        "the Wasm module's `{}` function export does not have type `[] -> []`",
36                        self.get_init_func()
37                    );
38                }
39            }
40            Some(_) => anyhow::bail!(
41                "the Wasm module's `{}` export is not a function",
42                self.get_init_func()
43            ),
44            None => anyhow::bail!(
45                "the Wasm module does not have a `{}` export",
46                self.get_init_func()
47            ),
48        }
49        Ok(())
50    }
51
52    /// Instantiate the module and call its initialization function.
53    async fn initialize<T: Send>(
54        &self,
55        store: &mut Store<T>,
56        instance: &wasmtime::Instance,
57    ) -> anyhow::Result<()> {
58        log::debug!("Calling the initialization function");
59
60        if let Some(export) = instance.get_export(&mut *store, "_initialize") {
61            if let Extern::Func(func) = export {
62                func.typed::<(), ()>(&store)?
63                    .call_async(&mut *store, ())
64                    .await
65                    .context("calling the Reactor initialization function")?;
66
67                if self.get_init_func() == "_initialize" {
68                    // Don't run `_initialize` twice if the it was explicitly
69                    // requested as the init function.
70                    return Ok(());
71                }
72            }
73        }
74
75        let init_func = instance
76            .get_typed_func::<(), ()>(&mut *store, self.get_init_func())
77            .expect("checked by `validate_init_func`");
78        init_func
79            .call_async(&mut *store, ())
80            .await
81            .with_context(|| format!("the `{}` function trapped", self.get_init_func()))?;
82
83        Ok(())
84    }
85}
86
87/// Impementation of [`InstanceState`] backed by Wasmtime.
88pub struct WasmtimeWizer<'a, T: 'static> {
89    /// The Wasmtime-based store that owns the `instance` field.
90    pub store: &'a mut Store<T>,
91    /// The instance that this will load state from.
92    pub instance: Instance,
93}
94
95impl<T: Send> InstanceState for WasmtimeWizer<'_, T> {
96    async fn global_get(&mut self, name: &str) -> SnapshotVal {
97        let global = self.instance.get_global(&mut *self.store, name).unwrap();
98        match global.get(&mut *self.store) {
99            Val::I32(x) => SnapshotVal::I32(x),
100            Val::I64(x) => SnapshotVal::I64(x),
101            Val::F32(x) => SnapshotVal::F32(x),
102            Val::F64(x) => SnapshotVal::F64(x),
103            Val::V128(x) => SnapshotVal::V128(x.as_u128()),
104            _ => panic!("unsupported global value type"),
105        }
106    }
107
108    async fn memory_contents(&mut self, name: &str, contents: impl FnOnce(&[u8]) + Send) {
109        let memory = self.instance.get_memory(&mut *self.store, name).unwrap();
110        contents(memory.data(&self.store))
111    }
112}