nice_assert_no_alloc/
lib.rs

1/* nice-assert-no-alloc -- A custom Rust allocator allowing to temporarily
2 * disable memory (de)allocations for a thread.
3 *
4 * Copyright (c) 2020 Florian Jung <flo@windfis.ch>
5 *
6 * Redistribution and use in source and binary forms, with or without
7 * modification, are permitted provided that the following conditions are met:
8 *
9 * Redistributions of source code must retain the above copyright notice, this
10 * list of conditions and the following disclaimer.
11 *
12 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR "AS IS" AND ANY EXPRESS OR IMPLIED
13 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
14 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
15 * EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
16 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
17 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
18 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
19 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
20 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
21 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
22 */
23
24#![doc = include_str!("../README.md")]
25
26use std::alloc::{GlobalAlloc, Layout, System};
27use std::cell::Cell;
28
29// check for mutually exclusive features.
30#[cfg(all(feature = "disable_release", feature = "warn_release"))]
31compile_error!("disable_release cannot be active at the same time with warn_release");
32
33#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
34thread_local! {
35    static ALLOC_FORBID_COUNT: Cell<u32> = Cell::new(0);
36    static ALLOC_PERMIT_COUNT: Cell<u32> = Cell::new(0);
37
38    #[cfg(any( all(feature="warn_debug", debug_assertions), all(feature="warn_release", not(debug_assertions)) ))]
39    static ALLOC_VIOLATION_COUNT: Cell<u32> = Cell::new(0);
40}
41
42#[cfg(all(feature = "disable_release", not(debug_assertions)))] // if disabled
43pub fn assert_no_alloc<T, F: FnOnce() -> T>(func: F) -> T {
44    // no-op
45    func()
46}
47
48#[cfg(all(feature = "disable_release", not(debug_assertions)))] // if disabled
49pub fn permit_alloc<T, F: FnOnce() -> T>(func: F) -> T {
50    // no-op
51    func()
52}
53
54#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
55/// Calls the `func` closure, but forbids any (de)allocations.
56///
57/// If a call to the allocator is made, the program will abort with an error,
58/// print a warning (depending on the `warn_debug` feature flag. Or ignore
59/// the situation, when compiled in `--release` mode with the `disable_release`
60///feature flag set (which is the default)).
61pub fn assert_no_alloc<T, F: FnOnce() -> T>(func: F) -> T {
62    // RAII guard for managing the forbid counter. This is to ensure correct behaviour
63    // when catch_unwind is used
64    struct Guard;
65    impl Guard {
66        fn new() -> Guard {
67            ALLOC_FORBID_COUNT.with(|c| c.set(c.get() + 1));
68            Guard
69        }
70    }
71    impl Drop for Guard {
72        fn drop(&mut self) {
73            ALLOC_FORBID_COUNT.with(|c| c.set(c.get() - 1));
74        }
75    }
76
77    #[cfg(any(
78        all(feature = "warn_debug", debug_assertions),
79        all(feature = "warn_release", not(debug_assertions))
80    ))] // if warn mode is selected
81    let old_violation_count = violation_count();
82
83    let guard = Guard::new(); // increment the forbid counter
84    let ret = func();
85    std::mem::drop(guard); // decrement the forbid counter
86
87    #[cfg(any(
88        all(feature = "warn_debug", debug_assertions),
89        all(feature = "warn_release", not(debug_assertions))
90    ))] // if warn mode is selected
91    if violation_count() > old_violation_count {
92        eprintln!("Tried to (de)allocate memory in a thread that forbids allocator calls!");
93    }
94
95    return ret;
96}
97
98#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
99/// Calls the `func` closure. Allocations are temporarily allowed, even if this
100/// code runs inside of assert_no_alloc.
101pub fn permit_alloc<T, F: FnOnce() -> T>(func: F) -> T {
102    // RAII guard for managing the permit counter
103    struct Guard;
104    impl Guard {
105        fn new() -> Guard {
106            ALLOC_PERMIT_COUNT.with(|c| c.set(c.get() + 1));
107            Guard
108        }
109    }
110    impl Drop for Guard {
111        fn drop(&mut self) {
112            ALLOC_PERMIT_COUNT.with(|c| c.set(c.get() - 1));
113        }
114    }
115
116    let guard = Guard::new(); // increment the forbid counter
117    let ret = func();
118    std::mem::drop(guard); // decrement the forbid counter
119
120    return ret;
121}
122
123#[cfg(any(
124    all(feature = "warn_debug", debug_assertions),
125    all(feature = "warn_release", not(debug_assertions))
126))] // if warn mode is selected
127/// Returns the count of allocation warnings emitted so far.
128///
129/// Only available when the `warn_debug` or `warn release` features are enabled.
130pub fn violation_count() -> u32 {
131    ALLOC_VIOLATION_COUNT.with(|c| c.get())
132}
133
134#[cfg(any(
135    all(feature = "warn_debug", debug_assertions),
136    all(feature = "warn_release", not(debug_assertions))
137))] // if warn mode is selected
138/// Resets the count of allocation warnings to zero.
139///
140/// Only available when the `warn_debug` or `warn release` features are enabled.
141pub fn reset_violation_count() {
142    ALLOC_VIOLATION_COUNT.with(|c| c.set(0));
143}
144
145#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
146/// The custom allocator that handles the checking.
147///
148/// To use this crate, you must add the following in your `main.rs`:
149/// ```rust
150/// use nice_assert_no_alloc::*;
151/// // ...
152/// #[cfg(debug_assertions)]
153/// #[global_allocator]
154/// static A: AllocDisabler = AllocDisabler;
155/// ```
156pub struct AllocDisabler;
157
158#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
159impl AllocDisabler {
160    fn check(&self, #[allow(unused)] layout: Layout) {
161        let forbid_count = ALLOC_FORBID_COUNT.with(|f| f.get());
162        let permit_count = ALLOC_PERMIT_COUNT.with(|p| p.get());
163        if forbid_count > permit_count {
164            #[cfg(any(
165                all(feature = "warn_debug", debug_assertions),
166                all(feature = "warn_release", not(debug_assertions))
167            ))] // if warn mode is selected
168            ALLOC_VIOLATION_COUNT.with(|c| c.set(c.get() + 1));
169
170            #[cfg(any(
171                all(not(feature = "warn_debug"), debug_assertions),
172                all(not(feature = "warn_release"), not(debug_assertions))
173            ))] // if abort mode is selected
174            {
175                #[cfg(all(feature = "log", feature = "backtrace"))]
176                permit_alloc(|| {
177                    log::error!(
178                        "Memory allocation of {} bytes failed from:\n{:?}",
179                        layout.size(),
180                        backtrace::Backtrace::new()
181                    )
182                });
183                #[cfg(all(feature = "log", not(feature = "backtrace")))]
184                permit_alloc(|| log::error!("Memory allocation of {} bytes failed", layout.size()));
185
186                #[cfg(all(not(feature = "log"), feature = "backtrace"))]
187                permit_alloc(|| {
188                    eprintln!(
189                        "Allocation failure from:\n{:?}",
190                        backtrace::Backtrace::new()
191                    )
192                });
193
194                // This handler can be overridden (although as of writing, the API to do so is still
195                // unstable) so we must always call this even when the log feature is enabled
196                std::alloc::handle_alloc_error(layout);
197            }
198        }
199    }
200}
201
202#[cfg(not(all(feature = "disable_release", not(debug_assertions))))] // if not disabled
203unsafe impl GlobalAlloc for AllocDisabler {
204    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
205        self.check(layout);
206        System.alloc(layout)
207    }
208
209    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
210        self.check(layout);
211        System.dealloc(ptr, layout)
212    }
213}
214
215/// Wrapper for objects whose Drop implementation shall be permitted
216/// to (de)allocate.
217///
218/// Typical usage:
219///
220/// ```rust
221/// # use nice_assert_no_alloc::*;
222/// let foo = PermitDrop::new(
223///     permit_alloc(||
224///         Box::new(42u32)
225///     )
226/// );
227/// ```
228///
229/// Here, creation of the Box is guarded by the explicit `permit_alloc` call,
230/// and destruction of the Box is guarded by PermitDrop. Neither creation nor
231/// destruction will cause an assertion failure from within `assert_no_alloc`.
232pub struct PermitDrop<T>(Option<T>);
233
234impl<T> PermitDrop<T> {
235    pub fn new(t: T) -> PermitDrop<T> {
236        permit_alloc(|| PermitDrop(Some(t)))
237    }
238}
239
240impl<T> std::ops::Deref for PermitDrop<T> {
241    type Target = T;
242    fn deref(&self) -> &T {
243        self.0.as_ref().unwrap()
244    }
245}
246
247impl<T> std::ops::DerefMut for PermitDrop<T> {
248    fn deref_mut(&mut self) -> &mut T {
249        self.0.as_mut().unwrap()
250    }
251}
252
253impl<I: Iterator> Iterator for PermitDrop<I> {
254    type Item = I::Item;
255    fn next(&mut self) -> Option<Self::Item> {
256        (**self).next()
257    }
258}
259
260impl<T> Drop for PermitDrop<T> {
261    fn drop(&mut self) {
262        let mut tmp = None;
263        std::mem::swap(&mut tmp, &mut self.0);
264        permit_alloc(|| {
265            std::mem::drop(tmp);
266        });
267    }
268}
nice_assert_no_alloc/lib.rs

nice_assert_no_alloc/
lib.rs
