wasmtime_provider/
provider_async.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use log::{error, info};
5use parking_lot::RwLock;
6use tracing::trace;
7#[cfg(feature = "wasi")]
8use wapc::WasiParams;
9use wapc::{wapc_functions, ModuleStateAsync, WebAssemblyEngineProviderAsync};
10use wasmtime::{AsContextMut, Engine, Instance, InstancePre, Linker, Module, Store, TypedFunc};
11
12use crate::errors::{Error, Result};
13use crate::store_async::WapcStoreAsync;
14use crate::{callbacks_async, EpochDeadlines};
15
16struct EngineInner {
17  instance: Arc<RwLock<Instance>>,
18  guest_call_fn: TypedFunc<(i32, i32), i32>,
19  host: Arc<ModuleStateAsync>,
20}
21
22/// A pre initialized [`WasmtimeEngineProviderAsync`]
23///
24/// Can be used to quickly create a new instance of [`WasmtimeEngineProviderAsync`]
25///
26/// Refer to [`WasmtimeEngineProviderBuilder::build_async_pre`](crate::WasmtimeEngineProviderBuilder::build_async_pre) to create an instance of this struct.
27#[allow(missing_debug_implementations)]
28#[derive(Clone)]
29pub struct WasmtimeEngineProviderAsyncPre {
30  module: Module,
31  #[cfg(feature = "wasi")]
32  wasi_params: WasiParams,
33  engine: Engine,
34  linker: Linker<WapcStoreAsync>,
35  instance_pre: InstancePre<WapcStoreAsync>,
36  epoch_deadlines: Option<EpochDeadlines>,
37}
38
39impl WasmtimeEngineProviderAsyncPre {
40  #[cfg(feature = "wasi")]
41  pub(crate) fn new(
42    engine: Engine,
43    module: Module,
44    wasi: Option<WasiParams>,
45    epoch_deadlines: Option<EpochDeadlines>,
46  ) -> Result<Self> {
47    let mut linker: Linker<WapcStoreAsync> = Linker::new(&engine);
48
49    let wasi_params = wasi.unwrap_or_default();
50    wasmtime_wasi::p1::add_to_linker_async(&mut linker, |s: &mut WapcStoreAsync| &mut s.wasi_ctx).unwrap();
51    // Also register the older `wasi_unstable` (Preview 0) ABI, used by some
52    // older toolchains (e.g. older TinyGo versions).
53    wasmtime_wasi::p0::add_to_linker_async(&mut linker, |s: &mut WapcStoreAsync| &mut s.wasi_ctx).unwrap();
54
55    // register all the waPC host functions
56    callbacks_async::add_to_linker(&mut linker)?;
57
58    let instance_pre = linker.instantiate_pre(&module)?;
59
60    Ok(Self {
61      module,
62      wasi_params,
63      engine,
64      linker,
65      instance_pre,
66      epoch_deadlines,
67    })
68  }
69
70  #[cfg(not(feature = "wasi"))]
71  pub(crate) fn new(engine: Engine, module: Module, epoch_deadlines: Option<EpochDeadlines>) -> Result<Self> {
72    let mut linker: Linker<WapcStoreAsync> = Linker::new(&engine);
73
74    // register all the waPC host functions
75    callbacks_async::add_to_linker(&mut linker)?;
76
77    let instance_pre = linker.instantiate_pre(&module)?;
78
79    Ok(Self {
80      module,
81      engine,
82      linker,
83      instance_pre,
84      epoch_deadlines,
85    })
86  }
87
88  /// Create an instance of [`WasmtimeEngineProviderAsync`] ready to be consumed
89  ///
90  /// Note: from micro-benchmarking, this method is 10 microseconds faster than
91  /// `WasmtimeEngineProviderAsync::clone`.
92  pub fn rehydrate(&self) -> Result<WasmtimeEngineProviderAsync> {
93    let engine = self.engine.clone();
94
95    #[cfg(feature = "wasi")]
96    let wapc_store = WapcStoreAsync::new(&self.wasi_params, None)?;
97    #[cfg(not(feature = "wasi"))]
98    let wapc_store = WapcStoreAsync::new(None);
99
100    let store = Store::new(&engine, wapc_store);
101
102    Ok(WasmtimeEngineProviderAsync {
103      module: self.module.clone(),
104      inner: None,
105      engine,
106      epoch_deadlines: self.epoch_deadlines,
107      linker: self.linker.clone(),
108      instance_pre: self.instance_pre.clone(),
109      store,
110      #[cfg(feature = "wasi")]
111      wasi_params: self.wasi_params.clone(),
112    })
113  }
114}
115
116/// A waPC engine provider that encapsulates the Wasmtime WebAssembly runtime.
117/// This can be used inside of async contexts.
118///
119/// Refer to
120/// [`WasmtimeEngineProviderBuilder::build_async`](crate::WasmtimeEngineProviderBuilder::build_async) to create an instance of this struct.
121///
122/// ## Example
123///
124/// ```rust
125/// use wasmtime_provider::WasmtimeEngineProviderBuilder;
126/// use wapc::WapcHostAsync;
127/// use std::error::Error;
128///
129/// // Sample host callback that prints the operation a WASM module requested.
130/// async fn host_callback(
131///   id: u64,
132///   bd: String,
133///   ns: String,
134///   op: String,
135///   payload: Vec<u8>,
136/// ) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
137///   println!(
138///     "Guest {} invoked '{}->{}:{}' on the host with a payload of '{}'",
139///     id,
140///     bd,
141///     ns,
142///     op,
143///     ::std::str::from_utf8(&payload).unwrap()
144///   );
145///   Ok(vec![])
146/// }
147///
148/// #[tokio::main]
149/// pub async fn main() -> Result<(), Box<dyn Error>> {
150///   let callback: Box<wapc::HostCallbackAsync> = Box::new(move |id, bd, ns, op, payload| {
151///     let fut = host_callback(id, bd, ns, op, payload);
152///     Box::pin(fut)
153///   });
154///
155///   let file = "../../wasm/crates/wasm-basic/build/wasm_basic.wasm";
156///   let module_bytes = std::fs::read(file)?;
157///
158///   let engine = WasmtimeEngineProviderBuilder::new()
159///     .module_bytes(&module_bytes)
160///     .build_async()?;
161///   let host = WapcHostAsync::new(Box::new(engine), Some(callback)).await?;
162///
163///   let res = host.call("ping", b"payload bytes").await?;
164///   assert_eq!(res, b"payload bytes");
165///
166///   Ok(())
167/// }
168/// ```
169#[allow(missing_debug_implementations)]
170pub struct WasmtimeEngineProviderAsync {
171  module: Module,
172  #[cfg(feature = "wasi")]
173  wasi_params: WasiParams,
174  inner: Option<EngineInner>,
175  engine: Engine,
176  linker: Linker<WapcStoreAsync>,
177  store: Store<WapcStoreAsync>,
178  instance_pre: InstancePre<WapcStoreAsync>,
179  epoch_deadlines: Option<EpochDeadlines>,
180}
181
182impl Clone for WasmtimeEngineProviderAsync {
183  fn clone(&self) -> Self {
184    let engine = self.engine.clone();
185
186    #[cfg(feature = "wasi")]
187    let wapc_store = WapcStoreAsync::new(&self.wasi_params, None).unwrap();
188    #[cfg(not(feature = "wasi"))]
189    let wapc_store = WapcStoreAsync::new(None);
190
191    let store = Store::new(&engine, wapc_store);
192
193    match &self.inner {
194      Some(state) => {
195        let mut new = Self {
196          module: self.module.clone(),
197          inner: None,
198          engine,
199          epoch_deadlines: self.epoch_deadlines,
200          linker: self.linker.clone(),
201          instance_pre: self.instance_pre.clone(),
202          store,
203          #[cfg(feature = "wasi")]
204          wasi_params: self.wasi_params.clone(),
205        };
206
207        tokio::runtime::Handle::current().block_on(async {
208          new.init(state.host.clone()).await.unwrap();
209        });
210
211        new
212      }
213      None => Self {
214        module: self.module.clone(),
215        inner: None,
216        engine,
217        epoch_deadlines: self.epoch_deadlines,
218        linker: self.linker.clone(),
219        instance_pre: self.instance_pre.clone(),
220        store,
221        #[cfg(feature = "wasi")]
222        wasi_params: self.wasi_params.clone(),
223      },
224    }
225  }
226}
227
228#[async_trait]
229impl WebAssemblyEngineProviderAsync for WasmtimeEngineProviderAsync {
230  async fn init(
231    &mut self,
232    host: Arc<ModuleStateAsync>,
233  ) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
234    // create the proper store, now we have a value for `host`
235    #[cfg(feature = "wasi")]
236    let wapc_store = WapcStoreAsync::new(&self.wasi_params, Some(host.clone()))?;
237    #[cfg(not(feature = "wasi"))]
238    let wapc_store = WapcStoreAsync::new(Some(host.clone()));
239
240    self.store = Store::new(&self.engine, wapc_store);
241
242    let instance = self.instance_pre.instantiate_async(&mut self.store).await?;
243
244    let instance_ref = Arc::new(RwLock::new(instance));
245    let gc = guest_call_fn(&mut self.store, &instance_ref)?;
246    self.inner = Some(EngineInner {
247      instance: instance_ref,
248      guest_call_fn: gc,
249      host,
250    });
251    self.initialize().await?;
252    Ok(())
253  }
254
255  async fn call(
256    &mut self,
257    op_length: i32,
258    msg_length: i32,
259  ) -> std::result::Result<i32, Box<dyn std::error::Error + Send + Sync>> {
260    if let Some(deadlines) = &self.epoch_deadlines {
261      // the deadline counter must be set before invoking the wasm function
262      self.store.set_epoch_deadline(deadlines.wapc_func);
263    }
264
265    let engine_inner = self.inner.as_ref().unwrap();
266    let call = engine_inner
267      .guest_call_fn
268      .call_async(&mut self.store, (op_length, msg_length))
269      .await;
270
271    match call {
272      Ok(result) => Ok(result),
273      Err(err) => {
274        error!("Failure invoking guest module handler: {err:?}");
275        let mut guest_error = err.to_string();
276        if let Some(trap) = err.downcast_ref::<wasmtime::Trap>() {
277          if matches!(trap, wasmtime::Trap::Interrupt) {
278            "guest code interrupted, execution deadline exceeded".clone_into(&mut guest_error);
279          }
280        }
281        engine_inner.host.set_guest_error(guest_error).await;
282        Ok(0)
283      }
284    }
285  }
286
287  async fn replace(&mut self, module: &[u8]) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>> {
288    info!(
289      "HOT SWAP - Replacing existing WebAssembly module with new buffer, {} bytes",
290      module.len()
291    );
292
293    let module = Module::new(&self.engine, module)?;
294    self.module = module;
295    self.instance_pre = self.linker.instantiate_pre(&self.module)?;
296    let new_instance = self.instance_pre.instantiate_async(&mut self.store).await?;
297    if let Some(inner) = self.inner.as_mut() {
298      *inner.instance.write() = new_instance;
299      let gc = guest_call_fn(&mut self.store, &inner.instance)?;
300      inner.guest_call_fn = gc;
301    }
302
303    Ok(self.initialize().await?)
304  }
305}
306
307impl WasmtimeEngineProviderAsync {
308  async fn initialize(&mut self) -> Result<()> {
309    for starter in wapc_functions::REQUIRED_STARTS.iter() {
310      if let Some(deadlines) = &self.epoch_deadlines {
311        // the deadline counter must be set before invoking the wasm function
312        self.store.set_epoch_deadline(deadlines.wapc_init);
313      }
314
315      let engine_inner = self.inner.as_ref().unwrap();
316      if engine_inner
317        .instance
318        .read()
319        .get_export(&mut self.store, starter)
320        .is_some()
321      {
322        // Need to get a `wasmtime::TypedFunc` because its `call` method
323        // can return a Trap error. Non-typed functions instead return a
324        // generic `anyhow::Error` that doesn't allow nice handling of
325        // errors
326        let starter_func: TypedFunc<(), ()> = engine_inner.instance.read().get_typed_func(&mut self.store, starter)?;
327
328        if let Err(err) = starter_func.call_async(&mut self.store, ()).await {
329          trace!(function = starter, ?err, "handling error returned by init function");
330          if let Some(trap) = err.downcast_ref::<wasmtime::Trap>() {
331            if matches!(trap, wasmtime::Trap::Interrupt) {
332              return Err(Error::InitializationFailedTimeout((*starter).to_owned()));
333            }
334            return Err(Error::InitializationFailed(err.to_string()));
335          }
336
337          // WASI programs built by tinygo have to be written with a `main` function, even if it's empty.
338          // Starting from tinygo >= 0.35.0, the `main` function calls the WASI process exit function,
339          // which is handled by wasmtime as an Error.
340          //
341          // We must check if this error can be converted into a WASI exit
342          // error and, if the exit code is 0, we can ignore it. Otherwise the waPC initialization
343          // will fail.
344          #[cfg(feature = "wasi")]
345          if let Some(exit_err) = err.downcast_ref::<wasmtime_wasi::I32Exit>() {
346            if exit_err.0 != 0 {
347              return Err(Error::InitializationFailed(err.to_string()));
348            }
349            trace!("ignoring successful exit trap generated by WASI");
350            continue;
351          }
352
353          return Err(Error::InitializationFailed(err.to_string()));
354        };
355      }
356    }
357    Ok(())
358  }
359}
360
361// Called once, then the result is cached. This returns a `Func` that corresponds
362// to the `__guest_call` export
363fn guest_call_fn(store: impl AsContextMut, instance: &Arc<RwLock<Instance>>) -> Result<TypedFunc<(i32, i32), i32>> {
364  instance
365    .read()
366    .get_typed_func::<(i32, i32), i32>(store, wapc_functions::GUEST_CALL)
367    .map_err(|_| Error::GuestCallNotFound)
368}
wasmtime_provider/provider_async.rs

wasmtime_provider/
provider_async.rs
