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#[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 wasmtime_wasi::p0::add_to_linker_sync(&mut linker, |s: &mut WapcStore| &mut s.wasi_ctx).unwrap();
47
48 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 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 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#[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 #[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 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 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 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 #[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
295fn 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}