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

wasmtime_provider/
provider_async.rs
