1 //! Utilities for testing and fuzzing out-of-memory handling. 2 //! 3 //! Inspired by SpiderMonkey's `oomTest()` helper: 4 //! https://firefox-source-docs.mozilla.org/js/hacking_tips.html#how-to-debug-oomtest-failures 5 6 use backtrace::Backtrace; 7 use std::{alloc::GlobalAlloc, cell::Cell, mem, ptr, time}; 8 use wasmtime_core::error::{Error, OutOfMemory, Result, bail}; 9 10 /// An allocator for use with `OomTest`. 11 #[non_exhaustive] 12 pub struct OomTestAllocator; 13 14 impl OomTestAllocator { 15 /// Create a new OOM test allocator. 16 pub const fn new() -> Self { 17 OomTestAllocator 18 } 19 } 20 21 #[derive(Clone, Debug, Default, PartialEq, Eq)] 22 enum OomState { 23 /// We are in code that is not part of an OOM test. 24 #[default] 25 OutsideOomTest, 26 27 /// We are inside an OOM test and should inject an OOM when the counter 28 /// reaches zero. 29 OomOnAlloc(u32), 30 31 /// We are inside an OOM test and we already injected an OOM. 32 DidOom, 33 } 34 35 thread_local! { 36 static OOM_STATE: Cell<OomState> = const { Cell::new(OomState::OutsideOomTest) }; 37 } 38 39 /// Set the new OOM state, returning the old state. 40 fn set_oom_state(state: OomState) -> OomState { 41 OOM_STATE.with(|s| s.replace(state)) 42 } 43 44 /// RAII helper to set the OOM state within a block of code and reset it upon 45 /// exiting that block (even if exiting via panic unwinding). 46 struct ScopedOomState { 47 prev_state: OomState, 48 } 49 50 impl ScopedOomState { 51 fn new(state: OomState) -> Self { 52 ScopedOomState { 53 prev_state: set_oom_state(state), 54 } 55 } 56 57 /// Finish this OOM state scope early, resetting the OOM state to what it 58 /// was before this scope was created, and returning the previous state that 59 /// was just overwritten by the reset. 60 fn finish(&self) -> OomState { 61 set_oom_state(self.prev_state.clone()) 62 } 63 } 64 65 impl Drop for ScopedOomState { 66 fn drop(&mut self) { 67 set_oom_state(mem::take(&mut self.prev_state)); 68 } 69 } 70 71 unsafe impl GlobalAlloc for OomTestAllocator { 72 unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { 73 let old_state = set_oom_state(OomState::OutsideOomTest); 74 75 let new_state; 76 let ptr; 77 { 78 // NB: It's okay to log/backtrace/etc... in this block because the 79 // current state is `OutsideOomTest`, so any re-entrant allocations 80 // will be passed through to the system allocator. 81 82 match old_state { 83 OomState::OutsideOomTest => { 84 new_state = OomState::OutsideOomTest; 85 ptr = unsafe { std::alloc::System.alloc(layout) }; 86 } 87 88 OomState::OomOnAlloc(0) => { 89 log::trace!( 90 "injecting OOM for allocation: {layout:?}\nAllocation backtrace:\n{:?}", 91 Backtrace::new(), 92 ); 93 new_state = OomState::DidOom; 94 ptr = ptr::null_mut(); 95 } 96 97 OomState::OomOnAlloc(c) => { 98 new_state = OomState::OomOnAlloc(c - 1); 99 ptr = unsafe { std::alloc::System.alloc(layout) }; 100 } 101 102 // Don't panic on allocation-after-OOM attempts if we 103 // are already in the middle of panicking. That will 104 // cause an abort and we won't get as good of an error 105 // message for the original panic, which is most likely 106 // some kind of test failure. 107 OomState::DidOom if std::thread::panicking() => { 108 log::trace!( 109 "Ignoring attempt to allocate {layout:?} after OOM while panicking:\n{:?}", 110 Backtrace::new(), 111 ); 112 new_state = OomState::DidOom; 113 ptr = unsafe { std::alloc::System.alloc(layout) }; 114 } 115 116 OomState::DidOom => { 117 log::trace!( 118 "Attempt to allocate {layout:?} after OOM:\n{:?}", 119 Backtrace::new(), 120 ); 121 panic!("OOM test attempted to allocate after OOM: {layout:?}") 122 } 123 } 124 } 125 126 set_oom_state(new_state); 127 ptr 128 } 129 130 unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) { 131 unsafe { 132 std::alloc::System.dealloc(ptr, layout); 133 } 134 } 135 } 136 137 /// A test helper that checks that some code handles OOM correctly. 138 /// 139 /// `OomTest` will only work correctly when `OomTestAllocator` is configured as 140 /// the global allocator. 141 /// 142 /// `OomTest` does not support reentrancy, so you cannot run an `OomTest` within 143 /// an `OomTest`. 144 /// 145 /// # Example 146 /// 147 /// ```no_run 148 /// use std::time::Duration; 149 /// use wasmtime::Result; 150 /// use wasmtime_fuzzing::oom::{OomTest, OomTestAllocator}; 151 /// 152 /// #[global_allocator] 153 /// static GLOBAL_ALOCATOR: OomTestAllocator = OomTestAllocator::new(); 154 /// 155 /// #[test] 156 /// fn my_oom_test() -> Result<()> { 157 /// OomTest::new() 158 /// .max_iters(1_000_000) 159 /// .max_duration(Duration::from_secs(5)) 160 /// .test(|| { 161 /// todo!("insert code here that should handle OOM here...") 162 /// }) 163 /// } 164 /// ``` 165 pub struct OomTest { 166 max_iters: Option<u32>, 167 max_duration: Option<time::Duration>, 168 } 169 170 impl OomTest { 171 /// Create a new OOM test. 172 /// 173 /// By default there is no iteration or time limit, tests will be executed 174 /// until the pass (or fail). 175 pub fn new() -> Self { 176 let _ = env_logger::try_init(); 177 178 // NB: `std::backtrace::Backtrace` doesn't have ways to handle 179 // OOM. Ideally we would just disable the `"backtrace"` cargo feature, 180 // but workspace feature resolution doesn't play nice with that. 181 wasmtime_core::error::disable_backtrace(); 182 183 OomTest { 184 max_iters: None, 185 max_duration: None, 186 } 187 } 188 189 /// Configure the maximum number of times to run an OOM test. 190 pub fn max_iters(&mut self, max_iters: u32) -> &mut Self { 191 self.max_iters = Some(max_iters); 192 self 193 } 194 195 /// Configure the maximum duration of time to run an OOM text. 196 pub fn max_duration(&mut self, max_duration: time::Duration) -> &mut Self { 197 self.max_duration = Some(max_duration); 198 self 199 } 200 201 /// Repeatedly run the given test function, injecting OOMs at different 202 /// times and checking that it correctly handles them. 203 /// 204 /// The test function should not use threads, or else allocations may not be 205 /// tracked correctly and OOM injection may be incorrect. 206 /// 207 /// The test function should return an `Err(_)` if and only if it encounters 208 /// an OOM. 209 /// 210 /// Returns early once the test function returns `Ok(())` before an OOM has 211 /// been injected. 212 pub fn test(&self, test_func: impl Fn() -> Result<()>) -> Result<()> { 213 let start = time::Instant::now(); 214 215 for i in 0.. { 216 if self.max_iters.is_some_and(|n| i >= n) 217 || self.max_duration.is_some_and(|d| start.elapsed() >= d) 218 { 219 break; 220 } 221 222 log::trace!("=== Injecting OOM after {i} allocations ==="); 223 let (result, old_state) = { 224 let guard = ScopedOomState::new(OomState::OomOnAlloc(i)); 225 assert_eq!(guard.prev_state, OomState::OutsideOomTest); 226 227 let result = test_func(); 228 229 (result, guard.finish()) 230 }; 231 232 match (result, old_state) { 233 (_, OomState::OutsideOomTest) => unreachable!(), 234 235 // The test function completed successfully before we ran out of 236 // allocation fuel, so we're done. 237 (Ok(()), OomState::OomOnAlloc(_)) => break, 238 239 // We injected an OOM and the test function handled it 240 // correctly; continue to the next iteration. 241 (Err(e), OomState::DidOom) if self.is_oom_error(&e) => {} 242 243 // Missed OOMs. 244 (Ok(()), OomState::DidOom) => { 245 bail!("OOM test function missed an OOM: returned Ok(())"); 246 } 247 (Err(e), OomState::DidOom) => { 248 return Err( 249 e.context("OOM test function missed an OOM: returned non-OOM error") 250 ); 251 } 252 253 // Unexpected error. 254 (Err(e), OomState::OomOnAlloc(_)) => { 255 return Err( 256 e.context("OOM test function returned an error when there was no OOM") 257 ); 258 } 259 } 260 } 261 262 Ok(()) 263 } 264 265 fn is_oom_error(&self, e: &Error) -> bool { 266 e.is::<OutOfMemory>() 267 } 268 } 269