Skip to main content

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