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 wasi_common::sync::add_to_linker(&mut linker, |s: &mut WapcStore| &mut s.wasi_ctx).unwrap();
44
45 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 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 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#[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 #[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 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 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 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 #[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
292fn 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}