submilli_engine/runtime/
watchdog.rs1use std::sync::mpsc::{self, RecvTimeoutError, Sender};
4use std::thread::{self, JoinHandle};
5use std::time::Duration;
6
7use wasmtime::Engine;
8
9pub struct Watchdog {
12 cancel: Sender<()>,
13 handle: Option<JoinHandle<()>>,
14}
15
16impl Drop for Watchdog {
17 fn drop(&mut self) {
18 let _ = self.cancel.send(());
20 let _ = self.handle.take();
21 }
22}
23
24pub fn arm(engine: &Engine, timeout: Duration) -> std::io::Result<Watchdog> {
26 #[cfg(test)]
27 if FAIL_SPAWN.get() {
28 return Err(std::io::Error::from_raw_os_error(11));
29 }
30 let (tx, rx) = mpsc::channel::<()>();
31 let engine = engine.clone();
32 let handle = thread::Builder::new().spawn(move || match rx.recv_timeout(timeout) {
33 Err(RecvTimeoutError::Timeout) => engine.increment_epoch(),
34 Ok(()) | Err(RecvTimeoutError::Disconnected) => {}
35 })?;
36 Ok(Watchdog {
37 cancel: tx,
38 handle: Some(handle),
39 })
40}
41
42#[cfg(test)]
43thread_local! {
44 static FAIL_SPAWN: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
45}
46
47#[cfg(test)]
48mod tests {
49 use super::*;
50 use crate::{FileId, RuntimeConfig, compile_script};
51
52 struct SpawnFailure(bool);
53
54 impl SpawnFailure {
55 fn inject() -> Self {
56 Self(FAIL_SPAWN.replace(true))
57 }
58 }
59
60 impl Drop for SpawnFailure {
61 fn drop(&mut self) {
62 FAIL_SPAWN.set(self.0);
63 }
64 }
65
66 #[test]
67 fn spawn_failure_preserves_io_cause_and_healthy_follow_up() {
68 let config = RuntimeConfig {
69 timeout: Some(Duration::from_secs(60)),
70 ..Default::default()
71 };
72 let compiled = compile_script(
73 "function main(): number { return 42; }",
74 "watchdog.ts",
75 FileId(0),
76 &[],
77 &[],
78 )
79 .unwrap();
80 {
81 let _failure = SpawnFailure::inject();
82 let error = pollster::block_on(config.run_compiled(&compiled)).unwrap_err();
83 assert_eq!(error.to_string(), "starting execution timeout watchdog");
84 assert_eq!(
85 error
86 .downcast_ref::<std::io::Error>()
87 .unwrap()
88 .raw_os_error(),
89 Some(11)
90 );
91 }
92 let result = pollster::block_on(config.run_compiled(&compiled)).unwrap();
93 assert_eq!(result.value.as_deref(), Some("42"));
94 }
95
96 #[test]
97 fn spawn_failure_precedes_guest_top_level_execution() {
98 let config = RuntimeConfig {
99 timeout: Some(Duration::from_secs(60)),
100 ..Default::default()
101 };
102 let compiled = compile_script(
103 "throw new Error(\"guest top level ran\"); function main(): void {}",
104 "watchdog.ts",
105 FileId(0),
106 &[],
107 &[],
108 )
109 .unwrap();
110 let _failure = SpawnFailure::inject();
111 let error = pollster::block_on(config.run(&compiled.wasm)).unwrap_err();
112 assert!(error.is::<std::io::Error>());
113 assert!(!format!("{error:#}").contains("guest top level ran"));
114 }
115
116 #[test]
117 fn disabled_timeout_does_not_spawn() {
118 let config = RuntimeConfig::default();
119 let engine = config.engine().unwrap();
120 let _failure = SpawnFailure::inject();
121 assert!(config.arm_timeout(&engine).unwrap().is_none());
122 }
123
124 #[test]
125 fn expired_watchdog_interrupts_guest() {
126 let config = RuntimeConfig {
127 timeout: Some(Duration::ZERO),
128 ..Default::default()
129 };
130 let engine = config.engine().unwrap();
131 let mut store = config.store(&engine, ()).unwrap();
132 let mut watchdog = arm(&engine, Duration::ZERO).unwrap();
133 watchdog.handle.take().unwrap().join().unwrap();
134 let error = call_guest(&mut store).unwrap_err();
135 assert!(matches!(
136 error.downcast_ref::<wasmtime::Trap>(),
137 Some(wasmtime::Trap::Interrupt)
138 ));
139 }
140
141 #[test]
142 fn dropped_watchdog_disarms_and_releases_thread() {
143 let config = RuntimeConfig {
144 timeout: Some(Duration::from_secs(3600)),
145 ..Default::default()
146 };
147 let engine = config.engine().unwrap();
148 let mut store = config.store(&engine, ()).unwrap();
149 let mut watchdog = arm(&engine, config.timeout.unwrap()).unwrap();
150 let handle = watchdog.handle.take().unwrap();
151 drop(watchdog);
152 handle.join().unwrap();
153 call_guest(&mut store).unwrap();
154 }
155
156 fn call_guest(store: &mut wasmtime::Store<()>) -> wasmtime::Result<()> {
157 use wasm_encoder::*;
158 let mut module = wasm_encoder::Module::new();
159 let mut types = TypeSection::new();
160 types.ty().function([], []);
161 module.section(&types);
162 let mut functions = FunctionSection::new();
163 functions.function(0);
164 module.section(&functions);
165 let mut exports = ExportSection::new();
166 exports.export("run", ExportKind::Func, 0);
167 module.section(&exports);
168 let mut code = CodeSection::new();
169 let mut function = Function::new([(1, ValType::I32)]);
171 function.instruction(&Instruction::I32Const(1000));
172 function.instruction(&Instruction::LocalSet(0));
173 function.instruction(&Instruction::Loop(BlockType::Empty));
174 function.instruction(&Instruction::LocalGet(0));
175 function.instruction(&Instruction::I32Const(1));
176 function.instruction(&Instruction::I32Sub);
177 function.instruction(&Instruction::LocalTee(0));
178 function.instruction(&Instruction::BrIf(0));
179 function.instruction(&Instruction::End);
180 function.instruction(&Instruction::End);
181 code.function(&function);
182 module.section(&code);
183 let module = wasmtime::Module::new(store.engine(), module.finish())?;
184 let instance = wasmtime::Instance::new(&mut *store, &module, &[])?;
185 instance
186 .get_typed_func::<(), ()>(&mut *store, "run")?
187 .call(store, ())
188 }
189}