wasmtime_provider/
provider.rs

1use std::sync::Arc;
2
3use log::{error, info};
4use parking_lot::RwLock;
5use tracing::trace;
6#[cfg(feature = "wasi")]
7use wapc::WasiParams;
8use wapc::{wapc_functions, ModuleState, WebAssemblyEngineProvider};
9use wasmtime::{AsContextMut, Engine, Instance, InstancePre, Linker, Module, Store, TypedFunc};
10
11use crate::callbacks;
12use crate::errors::{Error, Result};
13use crate::store::WapcStore;
14use crate::EpochDeadlines;
15
16struct EngineInner {
17  instance: Arc<RwLock<Instance>>,
18  guest_call_fn: TypedFunc<(i32, i32), i32>,
19  host: Arc<ModuleState>,
20}
21
22/// A pre initialized WasmtimeEngineProvider
23///
24/// Can be used to quickly create a new instance of WasmtimeEngineProvider
25///
26/// Refer to [`WasmtimeEngineProviderBuilder::build_pre`](crate::WasmtimeEngineProviderBuilder::build_pre) to create an instance of this struct.
27#[allow(missing_debug_implementations)]
28#[derive(Clone)]
29pub struct WasmtimeEngineProviderPre {
30  module: Module,
31  #[cfg(feature = "wasi")]
32  wasi_params: WasiParams,
33  engine: Engine,
34  linker: Linker<WapcStore>,
35  instance_pre: InstancePre<WapcStore>,
36  epoch_deadlines: Option<EpochDeadlines>,
37}
38
39impl WasmtimeEngineProviderPre {
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<WapcStore> = Linker::new(&engine);
48
49    let wasi_params = wasi.unwrap_or_default();
50    wasi_common::sync::add_to_linker(&mut linker, |s: &mut WapcStore| &mut s.wasi_ctx).unwrap();
51
52    // register all the waPC host functions
53    callbacks::add_to_linker(&mut linker)?;
54
55    let instance_pre = linker.instantiate_pre(&module)?;
56
57    Ok(Self {
58      module,
59      wasi_params,
60      engine,
61      linker,
62      instance_pre,
63      epoch_deadlines,
64    })
65  }
66
67  #[cfg(not(feature = "wasi"))]
68  pub(crate) fn new(engine: Engine, module: Module, epoch_deadlines: Option<EpochDeadlines>) -> Result<Self> {
69    let mut linker: Linker<WapcStore> = Linker::new(&engine);
70
71    // register all the waPC host functions
72    callbacks::add_to_linker(&mut linker)?;
73
74    let instance_pre = linker.instantiate_pre(&module)?;
75
76    Ok(Self {
77      module,
78      engine,
79      linker,
80      instance_pre,
81      epoch_deadlines,
82    })
83  }
84
85  /// Create an instance of [`WasmtimeEngineProvider`] ready to be consumed
86  ///
87  /// Note: from micro-benchmarking, this method is 10 microseconds faster than
88  /// `WasmtimeEngineProvider::clone`.
89  pub fn rehydrate(&self) -> Result<WasmtimeEngineProvider> {
90    let engine = self.engine.clone();
91
92    #[cfg(feature = "wasi")]
93    let wapc_store = WapcStore::new(&self.wasi_params, None)?;
94    #[cfg(not(feature = "wasi"))]
95    let wapc_store = WapcStore::new(None);
96
97    let store = Store::new(&engine, wapc_store);
98
99    Ok(WasmtimeEngineProvider {
100      module: self.module.clone(),
101      inner: None,
102      engine,
103      epoch_deadlines: self.epoch_deadlines,
104      linker: self.linker.clone(),
105      instance_pre: self.instance_pre.clone(),
106      store,
107      #[cfg(feature = "wasi")]
108      wasi_params: self.wasi_params.clone(),
109    })
110  }
111}
112
113/// A waPC engine provider that encapsulates the Wasmtime WebAssembly runtime
114#[allow(missing_debug_implementations)]
115pub struct WasmtimeEngineProvider {
116  module: Module,
117  #[cfg(feature = "wasi")]
118  wasi_params: WasiParams,
119  inner: Option<EngineInner>,
120  engine: Engine,
121  linker: Linker<WapcStore>,
122  store: Store<WapcStore>,
123  instance_pre: InstancePre<WapcStore>,
124  epoch_deadlines: Option<EpochDeadlines>,
125}
126
127impl Clone for WasmtimeEngineProvider {
128  fn clone(&self) -> Self {
129    let engine = self.engine.clone();
130
131    #[cfg(feature = "wasi")]
132    let wapc_store = WapcStore::new(&self.wasi_params, None).unwrap();
133    #[cfg(not(feature = "wasi"))]
134    let wapc_store = WapcStore::new(None);
135
136    let store = Store::new(&engine, wapc_store);
137
138    match &self.inner {
139      Some(state) => {
140        let mut new = Self {
141          module: self.module.clone(),
142          inner: None,
143          engine,
144          epoch_deadlines: self.epoch_deadlines,
145          linker: self.linker.clone(),
146          instance_pre: self.instance_pre.clone(),
147          store,
148          #[cfg(feature = "wasi")]
149          wasi_params: self.wasi_params.clone(),
150        };
151        new.init(state.host.clone()).unwrap();
152        new
153      }
154      None => Self {
155        module: self.module.clone(),
156        inner: None,
157        engine,
158        epoch_deadlines: self.epoch_deadlines,
159        linker: self.linker.clone(),
160        instance_pre: self.instance_pre.clone(),
161        store,
162        #[cfg(feature = "wasi")]
163        wasi_params: self.wasi_params.clone(),
164      },
165    }
166  }
167}
168
169impl WebAssemblyEngineProvider for WasmtimeEngineProvider {
170  fn init(
171    &mut self,
172    host: Arc<ModuleState>,
173  ) -> std::result::Result<(), Box<(dyn std::error::Error + Send + Sync + 'static)>> {
174    // create the proper store, now we have a value for `host`
175    #[cfg(feature = "wasi")]
176    let wapc_store = WapcStore::new(&self.wasi_params, Some(host.clone()))?;
177    #[cfg(not(feature = "wasi"))]
178    let wapc_store = WapcStore::new(Some(host.clone()));
179
180    self.store = Store::new(&self.engine, wapc_store);
181
182    let instance = self.instance_pre.instantiate(&mut self.store)?;
183
184    let instance_ref = Arc::new(RwLock::new(instance));
185    let gc = guest_call_fn(&mut self.store, &instance_ref)?;
186    self.inner = Some(EngineInner {
187      instance: instance_ref,
188      guest_call_fn: gc,
189      host,
190    });
191    self.initialize()?;
192    Ok(())
193  }
194
195  fn call(
196    &mut self,
197    op_length: i32,
198    msg_length: i32,
199  ) -> std::result::Result<i32, Box<(dyn std::error::Error + Send + Sync + 'static)>> {
200    if let Some(deadlines) = &self.epoch_deadlines {
201      // the deadline counter must be set before invoking the wasm function
202      self.store.set_epoch_deadline(deadlines.wapc_func);
203    }
204
205    let engine_inner = self.inner.as_ref().unwrap();
206    let call = engine_inner
207      .guest_call_fn
208      .call(&mut self.store, (op_length, msg_length));
209
210    match call {
211      Ok(result) => Ok(result),
212      Err(err) => {
213        error!("Failure invoking guest module handler: {:?}", err);
214        let mut guest_error = err.to_string();
215        if let Some(trap) = err.downcast_ref::<wasmtime::Trap>() {
216          if matches!(trap, wasmtime::Trap::Interrupt) {
217            "guest code interrupted, execution deadline exceeded".clone_into(&mut guest_error);
218          }
219        }
220        engine_inner.host.set_guest_error(guest_error);
221        Ok(0)
222      }
223    }
224  }
225
226  fn replace(
227    &mut self,
228    module: &[u8],
229  ) -> std::result::Result<(), Box<(dyn std::error::Error + Send + Sync + 'static)>> {
230    info!(
231      "HOT SWAP - Replacing existing WebAssembly module with new buffer, {} bytes",
232      module.len()
233    );
234
235    let module = Module::new(&self.engine, module)?;
236    self.module = module;
237    self.instance_pre = self.linker.instantiate_pre(&self.module)?;
238    let new_instance = self.instance_pre.instantiate(&mut self.store)?;
239    if let Some(inner) = self.inner.as_mut() {
240      *inner.instance.write() = new_instance;
241      let gc = guest_call_fn(&mut self.store, &inner.instance)?;
242      inner.guest_call_fn = gc;
243    }
244
245    Ok(self.initialize()?)
246  }
247}
248
249impl WasmtimeEngineProvider {
250  fn initialize(&mut self) -> Result<()> {
251    for starter in wapc_functions::REQUIRED_STARTS.iter() {
252      trace!(function = starter, "calling init function");
253      if let Some(deadlines) = &self.epoch_deadlines {
254        // the deadline counter must be set before invoking the wasm function
255        self.store.set_epoch_deadline(deadlines.wapc_init);
256      }
257
258      let engine_inner = self.inner.as_ref().unwrap();
259      if engine_inner
260        .instance
261        .read()
262        .get_export(&mut self.store, starter)
263        .is_some()
264      {
265        // Need to get a `wasmtime::TypedFunc` because its `call` method
266        // can return a Trap error. Non-typed functions instead return a
267        // generic `anyhow::Error` that doesn't allow nice handling of
268        // errors
269        let starter_func: TypedFunc<(), ()> = engine_inner.instance.read().get_typed_func(&mut self.store, starter)?;
270
271        if let Err(err) = starter_func.call(&mut self.store, ()) {
272          trace!(function = starter, ?err, "handling error returned by init function");
273          if let Some(trap) = err.downcast_ref::<wasmtime::Trap>() {
274            if matches!(trap, wasmtime::Trap::Interrupt) {
275              return Err(Error::InitializationFailedTimeout((*starter).to_owned()));
276            }
277            return Err(Error::InitializationFailed(err.to_string()));
278          }
279
280          // WASI programs built by tinygo have to be written with a `main` function, even if it's empty.
281          // Starting from tinygo >= 0.35.0, the `main` function calls the WASI process exit function,
282          // which is handled by wasmtime as an Error.
283          //
284          // We must check if this error can be converted into a WASI exit
285          // error and, if the exit code is 0, we can ignore it. Otherwise the waPC initialization
286          // will fail.
287          #[cfg(feature = "wasi")]
288          if let Some(exit_err) = err.downcast_ref::<wasi_common::I32Exit>() {
289            if exit_err.0 != 0 {
290              return Err(Error::InitializationFailed(err.to_string()));
291            }
292            trace!("ignoring successful exit trap generated by WASI");
293            continue;
294          }
295
296          return Err(Error::InitializationFailed(err.to_string()));
297        };
298      }
299    }
300    Ok(())
301  }
302}
303
304// Called once, then the result is cached. This returns a `Func` that corresponds
305// to the `__guest_call` export
306fn guest_call_fn(store: impl AsContextMut, instance: &Arc<RwLock<Instance>>) -> Result<TypedFunc<(i32, i32), i32>> {
307  instance
308    .read()
309    .get_typed_func::<(i32, i32), i32>(store, wapc_functions::GUEST_CALL)
310    .map_err(|_| Error::GuestCallNotFound)
311}