wasmtime_fuzzing/
oom.rs

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
6use backtrace::Backtrace;
7use rand::{RngExt, SeedableRng, rngs::SmallRng};
8use std::{
9    alloc::GlobalAlloc,
10    cell::Cell,
11    mem,
12    pin::Pin,
13    ptr,
14    task::{Context, Poll},
15    time,
16};
17use wasmtime_core::error::{Context as _, OutOfMemory, Result, bail, ensure};
18
19/// An allocator for use with `OomTest`.
20#[non_exhaustive]
21pub struct OomTestAllocator;
22
23impl OomTestAllocator {
24    /// Create a new OOM test allocator.
25    pub const fn new() -> Self {
26        OomTestAllocator
27    }
28}
29
30#[derive(Clone, Debug, Default, PartialEq, Eq)]
31enum OomState {
32    /// We are in code that is not part of an OOM test.
33    #[default]
34    OutsideOomTest,
35
36    /// We are inside an OOM test's dry run, counting total allocations.
37    Counting { count: u32 },
38
39    /// We are inside an OOM test and should inject an OOM when the counter
40    /// reaches zero.
41    OomOnAlloc {
42        counter: u32,
43        allow_alloc_after: bool,
44        alloc_succeeds_after: bool,
45    },
46
47    /// We are inside an OOM test and we already injected an OOM.
48    DidOom {
49        allow_alloc: bool,
50        alloc_succeeds_after: bool,
51    },
52}
53
54thread_local! {
55    static OOM_STATE: Cell<OomState> = const { Cell::new(OomState::OutsideOomTest) };
56}
57
58/// Set the new OOM state, returning the old state.
59fn set_oom_state(state: OomState) -> OomState {
60    OOM_STATE.with(|s| s.replace(state))
61}
62
63/// RAII helper to set the OOM state within a block of code and reset it upon
64/// exiting that block (even if exiting via panic unwinding).
65struct ScopedOomState<'a> {
66    state: &'a mut OomState,
67}
68
69impl<'a> ScopedOomState<'a> {
70    /// The OOM state will be initialized to `*state` on construction, and then
71    /// `state` will be updated to be the old OOM state that existed when the
72    /// original state is replaced on this type's drop.
73    fn new(state: &'a mut OomState) -> Self {
74        *state = set_oom_state(mem::take(state));
75        ScopedOomState { state }
76    }
77}
78
79impl Drop for ScopedOomState<'_> {
80    fn drop(&mut self) {
81        *self.state = set_oom_state(mem::take(self.state));
82    }
83}
84
85unsafe impl GlobalAlloc for OomTestAllocator {
86    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
87        let old_state = set_oom_state(OomState::OutsideOomTest);
88
89        let new_state;
90        let ptr;
91
92        'outside_oom_test: {
93            // NB: It's okay to log/backtrace/etc... in this block because the
94            // current state is `OutsideOomTest`, so any re-entrant allocations
95            // will be passed through to the system allocator.
96
97            // Don't panic on allocation-after-OOM attempts if we
98            // are already in the middle of panicking. That will
99            // cause an abort and we won't get as good of an error
100            // message for the original panic, which is most likely
101            // some kind of test failure.
102            if old_state == OomState::OutsideOomTest || std::thread::panicking() {
103                new_state = old_state;
104                ptr = unsafe { std::alloc::System.alloc(layout) };
105                break 'outside_oom_test;
106            }
107
108            let bt = Backtrace::new();
109            let bt = format!("{bt:?}");
110
111            // XXX: `env_logger` internally buffers writes in a `Vec` which
112            // means our OOM tests might sporadically fail when you enable
113            // logging to debug stuff, so simply let the allocation through if
114            // `env_logger` is on the stack.
115            if bt.contains("env_logger") {
116                new_state = old_state;
117                ptr = unsafe { std::alloc::System.alloc(layout) };
118                break 'outside_oom_test;
119            }
120
121            match old_state {
122                OomState::OutsideOomTest => unreachable!("handled above"),
123
124                OomState::Counting { count } => {
125                    new_state = OomState::Counting { count: count + 1 };
126                    ptr = unsafe { std::alloc::System.alloc(layout) };
127                }
128
129                OomState::OomOnAlloc {
130                    counter: 0,
131                    allow_alloc_after,
132                    alloc_succeeds_after,
133                } => {
134                    log::debug!(
135                        "injecting OOM for allocation: {layout:?}\nAllocation backtrace:\n{bt}"
136                    );
137                    new_state = OomState::DidOom {
138                        allow_alloc: allow_alloc_after,
139                        alloc_succeeds_after,
140                    };
141                    ptr = ptr::null_mut();
142                }
143
144                OomState::OomOnAlloc {
145                    counter: c,
146                    allow_alloc_after,
147                    alloc_succeeds_after,
148                } => {
149                    new_state = OomState::OomOnAlloc {
150                        counter: c - 1,
151                        allow_alloc_after,
152                        alloc_succeeds_after,
153                    };
154                    ptr = unsafe { std::alloc::System.alloc(layout) };
155                }
156
157                OomState::DidOom {
158                    allow_alloc,
159                    alloc_succeeds_after,
160                } => {
161                    log::trace!("Attempt to allocate {layout:?} after OOM:\n{bt}");
162                    if allow_alloc {
163                        new_state = OomState::DidOom {
164                            allow_alloc: true,
165                            alloc_succeeds_after,
166                        };
167                        ptr = if alloc_succeeds_after {
168                            unsafe { std::alloc::System.alloc(layout) }
169                        } else {
170                            ptr::null_mut()
171                        };
172                    } else {
173                        panic!(
174                            "OOM test attempted to allocate after OOM: {layout:?}\n\
175                             \n\
176                             Hint: if this is acceptable, configure the OOM test to allow allocation \
177                             after OOM with `OomTest::allow_alloc_after_oom`"
178                        )
179                    }
180                }
181            }
182        }
183
184        set_oom_state(new_state);
185        ptr
186    }
187
188    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
189        unsafe {
190            std::alloc::System.dealloc(ptr, layout);
191        }
192    }
193}
194
195/// A test helper that checks that some code handles OOM correctly.
196///
197/// `OomTest` will only work correctly when `OomTestAllocator` is configured as
198/// the global allocator.
199///
200/// `OomTest` does not support reentrancy, so you cannot run an `OomTest` within
201/// an `OomTest`.
202///
203/// # Example
204///
205/// ```no_run
206/// use std::time::Duration;
207/// use wasmtime::Result;
208/// use wasmtime_fuzzing::oom::{OomTest, OomTestAllocator};
209///
210/// #[global_allocator]
211/// static GLOBAL_ALLOCATOR: OomTestAllocator = OomTestAllocator::new();
212///
213/// #[test]
214/// fn my_oom_test() -> Result<()> {
215///     OomTest::new()
216///         .max_iters(1_000_000)
217///         .max_duration(Duration::from_secs(5))
218///         .allow_alloc_after_oom(true)
219///         .test(|| {
220///             todo!("insert code here that should handle OOM here...")
221///         })
222/// }
223/// ```
224///
225/// Use `test_with_reuse` when you also want to check that some context is still
226/// usable after the OOM, e.g. that nothing was left half-registered in state
227/// that is shared with an `Engine`:
228///
229/// ```no_run
230/// use wasmtime::Result;
231/// use wasmtime_fuzzing::oom::OomTest;
232///
233/// #[test]
234/// fn my_oom_reuse_test() -> Result<()> {
235///     OomTest::new().test_with_reuse(
236///         || todo!("create the context that is reused after the OOM..."),
237///         |context| todo!("insert code here that should handle OOM here..."),
238///         |context| todo!("check that `context` is still usable..."),
239///     )
240/// }
241/// ```
242pub struct OomTest {
243    max_iters: Option<u32>,
244    max_duration: Option<time::Duration>,
245    allow_alloc_after_oom: bool,
246    alloc_succeeds_after_oom: bool,
247    allow_missed_oom_errors: bool,
248    seed: u64,
249}
250
251impl OomTest {
252    /// Create a new OOM test.
253    ///
254    /// By default there is no iteration or time limit, tests will be executed
255    /// until the pass (or fail).
256    pub fn new() -> Self {
257        let _ = env_logger::try_init();
258
259        // NB: `std::backtrace::Backtrace` doesn't have ways to handle
260        // OOM. Ideally we would just disable the `"backtrace"` cargo feature,
261        // but workspace feature resolution doesn't play nice with that.
262        wasmtime_core::error::disable_backtrace();
263
264        OomTest {
265            max_iters: None,
266            max_duration: None,
267            allow_alloc_after_oom: false,
268            alloc_succeeds_after_oom: false,
269            allow_missed_oom_errors: false,
270            seed: 0,
271        }
272    }
273
274    /// Configure the maximum number of times to run an OOM test.
275    pub fn max_iters(&mut self, max_iters: u32) -> &mut Self {
276        self.max_iters = Some(max_iters);
277        self
278    }
279
280    /// Configure the maximum duration of time to run an OOM text.
281    pub fn max_duration(&mut self, max_duration: time::Duration) -> &mut Self {
282        self.max_duration = Some(max_duration);
283        self
284    }
285
286    /// Configure whether to allow allocation attempts after an OOM has already
287    /// been injected.
288    ///
289    /// The default is `false`.
290    pub fn allow_alloc_after_oom(&mut self, allow: bool) -> &mut Self {
291        self.allow_alloc_after_oom = allow;
292        self
293    }
294
295    /// Configure whether allocations after an OOM (assuming
296    /// `allow_alloc_after_oom(true)`) will succeed or will always return null.
297    ///
298    /// The default is `false`.
299    pub fn alloc_succeeds_after_oom(&mut self, succeeds: bool) -> &mut Self {
300        self.alloc_succeeds_after_oom = succeeds;
301        self
302    }
303
304    /// Configure whether to allow a test to pass if it does not return
305    /// `Err(OutOfMemory)` when a synthetic OOM was injected (perhaps because it
306    /// is exercising logic that is robust to OOM).
307    ///
308    /// The default is `false`.
309    pub fn allow_missed_oom_errors(&mut self, allow: bool) -> &mut Self {
310        self.allow_missed_oom_errors = allow;
311        self
312    }
313
314    /// Configure the seed for the RNG used in `fuzz` mode.
315    ///
316    /// The default is `0`.
317    pub fn seed(&mut self, seed: u64) -> &mut Self {
318        self.seed = seed;
319        self
320    }
321
322    /// Repeatedly run the given test function, injecting OOMs at different
323    /// times and checking that it correctly handles them.
324    ///
325    /// The test function should not use threads, or else allocations may not be
326    /// tracked correctly and OOM injection may be incorrect.
327    ///
328    /// The test function should return an `Err(_)` if and only if it encounters
329    /// an OOM.
330    ///
331    /// Returns early once the test function returns `Ok(())` before an OOM has
332    /// been injected.
333    pub fn test(&self, test_func: impl Fn() -> Result<()>) -> Result<()> {
334        let future = self.test_async(|| async { test_func() });
335        crate::block_on(future)
336    }
337
338    /// The same as `test` but `async`.
339    pub async fn test_async(&self, test_func: impl AsyncFn() -> Result<()>) -> Result<()> {
340        self.test_with_reuse_async(
341            async || Ok(()),
342            async |_: &mut ()| test_func().await,
343            async |_: ()| Ok(()),
344        )
345        .await
346    }
347
348    /// The same as `test`, but additionally exercising context that is reused
349    /// after the OOM is triggered.
350    ///
351    /// Each iteration runs three phases in sequence:
352    ///
353    /// 1. `make_context` creates the context that is shared between the two
354    ///    phases that follow. It is called afresh every iteration, so that each
355    ///    iteration begins from identical state.
356    ///
357    /// 2. `test_func` is the main OOM test, and is the only phase into which
358    ///    OOMs are injected. It is otherwise the same as the function given to
359    ///    `test`.
360    ///
361    /// 3. `post_oom_func` checks that we recovered from the OOM correctly and
362    ///    that we can still interact with the context afterwards, e.g. that
363    ///    nothing was left half-registered in shared state, rather than
364    ///    properly cleaned up, when the OOM happened.
365    ///
366    /// OOMs are never injected into `make_context` or `post_oom_func` and
367    /// neither's allocations are counted. An `Err(_)` returned from either is a
368    /// test failure.
369    pub fn test_with_reuse<T>(
370        &self,
371        make_context: impl Fn() -> Result<T>,
372        test_func: impl Fn(&mut T) -> Result<()>,
373        post_oom_func: impl Fn(T) -> Result<()>,
374    ) -> Result<()> {
375        let future = self.test_with_reuse_async(
376            async || make_context(),
377            async |context: &mut T| test_func(context),
378            async |context: T| post_oom_func(context),
379        );
380        crate::block_on(future)
381    }
382
383    /// The same as `test_with_reuse` but `async`.
384    pub async fn test_with_reuse_async<T>(
385        &self,
386        make_context: impl AsyncFn() -> Result<T>,
387        test_func: impl AsyncFn(&mut T) -> Result<()>,
388        post_oom_func: impl AsyncFn(T) -> Result<()>,
389    ) -> Result<()> {
390        let total_allocs =
391            Self::count_allocations(&make_context, &test_func, &post_oom_func).await?;
392
393        let start = time::Instant::now();
394
395        for i in 0..total_allocs {
396            if self.max_iters.is_some_and(|n| i >= n)
397                || self.max_duration.is_some_and(|d| start.elapsed() >= d)
398            {
399                break;
400            }
401
402            log::trace!("=== Injecting OOM after {i} allocations ===");
403            self.run_one_oom_injection(&make_context, &test_func, &post_oom_func, i)
404                .await?;
405        }
406
407        Ok(())
408    }
409
410    /// Similar to `test` but instead of exhaustively injecting OOMs on every
411    /// allocation, randomly chooses which allocation to OOM each iteration.
412    ///
413    /// Requires `max_iters` to be set.
414    pub fn fuzz(&self, test_func: impl Fn() -> Result<()>) -> Result<()> {
415        let future = self.fuzz_async(|| async { test_func() });
416        crate::block_on(future)
417    }
418
419    /// The same as `fuzz` but `async`.
420    pub async fn fuzz_async(&self, test_func: impl AsyncFn() -> Result<()>) -> Result<()> {
421        let max_iters = match self.max_iters {
422            Some(n) => n,
423            None => bail!("OomTest::fuzz requires max_iters to be set"),
424        };
425
426        let make_context = async || Ok(());
427        let test_func = async |_: &mut ()| test_func().await;
428        let post_oom_func = async |_: ()| Ok(());
429
430        let total_allocs =
431            Self::count_allocations(&make_context, &test_func, &post_oom_func).await?;
432
433        if total_allocs == 0 {
434            return Ok(());
435        }
436
437        let mut rng = SmallRng::seed_from_u64(self.seed);
438
439        for _ in 0..max_iters {
440            let i = rng.random_range(0..total_allocs);
441            log::trace!("=== Injecting OOM after {i} allocations (fuzz) ===");
442            self.run_one_oom_injection(&make_context, &test_func, &post_oom_func, i)
443                .await?;
444        }
445
446        Ok(())
447    }
448
449    /// Run the test function once, injecting OOM on the `i`th allocation, and
450    /// check that the result correctly reflects whether an OOM was hit.
451    async fn run_one_oom_injection<T>(
452        &self,
453        make_context: &impl AsyncFn() -> Result<T>,
454        test_func: &impl AsyncFn(&mut T) -> Result<()>,
455        post_oom_func: &impl AsyncFn(T) -> Result<()>,
456        i: u32,
457    ) -> Result<()> {
458        let mut context = make_context()
459            .await
460            .context("OOM test failed to create its reused context")?;
461
462        let (result, oom_state) = {
463            let future = std::pin::pin!(test_func(&mut context));
464            OomTestFuture::new(
465                future,
466                OomState::OomOnAlloc {
467                    counter: i,
468                    allow_alloc_after: self.allow_alloc_after_oom,
469                    alloc_succeeds_after: self.alloc_succeeds_after_oom,
470                },
471            )
472            .await
473        };
474
475        match (result, oom_state) {
476            (_, OomState::OutsideOomTest | OomState::Counting { .. }) => unreachable!(),
477
478            // The test function completed successfully before we reached
479            // this allocation point (allocation count may differ slightly
480            // between the counting run and OOM injection runs).
481            (Ok(()), OomState::OomOnAlloc { .. }) => {}
482
483            // We injected an OOM and the test function handled it correctly.
484            (Err(e), OomState::DidOom { .. }) if e.is::<OutOfMemory>() => {}
485
486            // Missed OOMs.
487            (Ok(()), OomState::DidOom { .. }) => {
488                ensure!(
489                    self.allow_missed_oom_errors,
490                    "OOM test function missed an OOM: returned Ok(())"
491                );
492            }
493            (Err(e), OomState::DidOom { .. }) => {
494                if !self.allow_missed_oom_errors {
495                    return Err(
496                        e.context("OOM test function missed an OOM: returned non-OOM error")
497                    );
498                }
499            }
500
501            // Unexpected error.
502            (Err(e), OomState::OomOnAlloc { .. }) => {
503                if !self.allow_missed_oom_errors {
504                    return Err(
505                        e.context("OOM test function returned an error when there was no OOM")
506                    );
507                }
508            }
509        }
510
511        post_oom_func(context)
512            .await
513            .context("OOM test's post-OOM function failed")
514    }
515
516    /// Run the test function without injecting OOMs, counting allocations.
517    ///
518    /// Repeats until the count stabilizes across two consecutive runs,
519    /// since the first run may change internal state (e.g. grow hash map
520    /// capacities) that subsequent runs reuse.
521    async fn count_allocations<T>(
522        make_context: &impl AsyncFn() -> Result<T>,
523        test_func: &impl AsyncFn(&mut T) -> Result<()>,
524        post_oom_func: &impl AsyncFn(T) -> Result<()>,
525    ) -> Result<u32> {
526        const MAX_DRY_RUN_ITERS: usize = 10;
527
528        let mut total_allocs = u32::MAX;
529        for i in 0..MAX_DRY_RUN_ITERS {
530            log::trace!("=== Counting allocations (dry run {i}) ===");
531
532            let mut context = make_context()
533                .await
534                .context("OOM test failed to create its reused context")?;
535
536            let (result, oom_state) = {
537                let future = std::pin::pin!(test_func(&mut context));
538                OomTestFuture::new(future, OomState::Counting { count: 0 }).await
539            };
540            result?;
541
542            post_oom_func(context)
543                .await
544                .context("OOM test's post-OOM function failed")?;
545
546            let count = match oom_state {
547                OomState::Counting { count } => count,
548                _ => unreachable!(),
549            };
550            log::trace!("=== Counted {count} allocations ===");
551            if count == total_allocs {
552                return Ok(total_allocs);
553            }
554            total_allocs = count;
555        }
556
557        bail!(
558            "allocation count did not stabilize after {MAX_DRY_RUN_ITERS} \
559             dry-run iterations (last count: {total_allocs})"
560        )
561    }
562}
563
564struct OomTestFuture<'a, F> {
565    inner: Pin<&'a mut F>,
566    state: OomState,
567}
568
569impl<'a, F> OomTestFuture<'a, F> {
570    fn new(inner: Pin<&'a mut F>, state: OomState) -> Self {
571        OomTestFuture { inner, state }
572    }
573}
574
575impl<F> Future for OomTestFuture<'_, F>
576where
577    F: Future<Output = Result<()>>,
578{
579    type Output = (Result<()>, OomState);
580
581    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
582        log::trace!("OomTestFuture::poll: Entered ---");
583
584        let result = {
585            let OomTestFuture { inner, state } = &mut *self;
586
587            let guard = ScopedOomState::new(state);
588            assert_eq!(*guard.state, OomState::OutsideOomTest);
589
590            inner.as_mut().poll(cx)
591        };
592
593        log::trace!("OomTestFuture::poll: inner result: {result:?}");
594        log::trace!("OomTestFuture::poll: old state: {:?}", self.state);
595
596        match result {
597            Poll::Pending => {
598                set_oom_state(OomState::OutsideOomTest);
599                Poll::Pending
600            }
601            Poll::Ready(result) => Poll::Ready((result, mem::take(&mut self.state))),
602        }
603    }
604}
wasmtime_fuzzing/oom.rs

wasmtime_fuzzing/
oom.rs
