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