ralloc 1.0.0

An efficient alternative platform-agnostic allocator.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

//! General error handling.

use prelude::*;

use core::sync::atomic::{self, AtomicPtr};
use core::mem;

use shim::config;

#[cfg(feature = "tls")]
use tls;

/// The global OOM handler.
static OOM_HANDLER: AtomicPtr<()> = AtomicPtr::new(config::default_oom_handler as *mut ());
#[cfg(feature = "tls")]
tls! {
    /// The thread-local OOM handler.
    static THREAD_OOM_HANDLER: MoveCell<Option<fn() -> !>> = MoveCell::new(None);
}

/// Call the OOM handler.
///
/// This is used one out-of-memory errors, and will never return. Usually, it simply consists
/// of aborting the process.
///
/// # An important note
///
/// This is for OOM-conditions, not malformed or too big allocations, but when the system is unable
/// to gather memory for the allocation (SBRK fails).
///
/// The rule of thumb is that this should be called, if and only if unwinding (which allocates)
/// will hit the same error.
pub fn oom() -> ! {
    // If TLS is enabled, we will use the thread-local OOM.
    #[cfg(feature = "tls")]
    {
        if let Some(handler) = THREAD_OOM_HANDLER.with(|x| x.replace(None)) {
            log!(DEBUG, "Calling the local OOM handler.");

            handler();
        }
    }

    log!(DEBUG, "Calling the global OOM handler.");

    unsafe {
        // LAST AUDIT: 2016-08-21 (Ticki).

        // Transmute the atomic pointer to a function pointer and call it.
        (mem::transmute::<_, fn() -> !>(OOM_HANDLER.load(atomic::Ordering::SeqCst)))()
    }
}

/// Set the OOM handler.
///
/// This is called when the process is out-of-memory.
#[inline]
pub fn set_oom_handler(handler: fn() -> !) {
    // Logging...
    log!(NOTE, "Setting the global OOM handler.");

    OOM_HANDLER.store(handler as *mut (), atomic::Ordering::SeqCst);
}

/// Override the OOM handler for the current thread.
///
/// # Panics
///
/// This might panic if a thread OOM handler already exists.
#[inline]
#[cfg(feature = "tls")]
pub fn set_thread_oom_handler(handler: fn() -> !) {
    // Logging...
    log!(NOTE, "Setting the thread OOM handler.");

    THREAD_OOM_HANDLER.with(|thread_oom| {
        // Replace it with the new handler.
        let res = thread_oom.replace(Some(handler));

        // Throw a warning if it overrides another handler.
        if res.is_some() {
            log!(WARNING, "An old thread OOM handler was overriden.");
        }
    });
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    #[should_panic]
    fn test_panic_oom() {
        fn panic() -> ! {
            panic!("cats are not cute.");
        }

        set_oom_handler(panic);
        oom();
    }

    #[test]
    #[should_panic]
    #[cfg(feature = "tls")]
    fn test_panic_thread_oom() {
        fn infinite() -> ! {
            #[allow(empty_loop)]
            loop {}
        }
        fn panic() -> ! {
            panic!("cats are not cute.");
        }

        set_oom_handler(infinite);
        set_thread_oom_handler(panic);
        oom();
    }
}