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