limit_alloc/
lib.rs

1//! Allocator that allows to limit the available memory.
2//!
3//! This crate implements a few similar types, you can choose the best depending on your use case:
4//!
5//! * Use `ConstLimit` if you know the limit at compile time, because that makes the allocator
6//! zero-sized (as long as the inner allocator is also zero-sized).
7//! * Use `Limit` if you are not sure, or if you need more than one limit in the same application.
8//! This is needed because `ConstLimit` uses a static counter to store the allocated memory, so it
9//! is impossible to track the memory allocated by different instances of the allocator, we can
10//! only track the total allocated memory. The size of `Limit` is `1 * usize`.
11//! * Use `ArcLimit` if you need a `Limit` that implements `Clone`. Ideally you would have been
12//! able to use `Arc<Limit<A>>` instead, but `Arc<T>` cannot implement `GlobalAlloc`.
13//!
14//! Note on alignment: an allocation of 1 byte with alignment greater than 1, for example 2 bytes,
15//! will allocate 2 bytes because of padding. But this crate only counts 1 byte. So the limit may
16//! not be completely accurate.
17use std::alloc::{GlobalAlloc, Layout};
18use std::ptr;
19use std::sync::atomic::AtomicUsize;
20use std::sync::atomic::Ordering::SeqCst;
21use std::sync::Arc;
22
23pub struct Limit<A> {
24    remaining: AtomicUsize,
25    alloc: A,
26}
27
28impl<A: GlobalAlloc> Limit<A> {
29    pub const fn new(limit: usize, alloc: A) -> Self {
30        Self {
31            remaining: AtomicUsize::new(limit),
32            alloc,
33        }
34    }
35
36    /// Returns None if the memory limit would be exhausted after allocating.
37    ///
38    /// # Safety
39    ///
40    /// The same restrictions as `GlobalAlloc::alloc`.
41    pub unsafe fn try_alloc(&self, layout: Layout) -> Option<*mut u8> {
42        match self
43            .remaining
44            .fetch_update(SeqCst, SeqCst, |old| old.checked_sub(layout.size()))
45        {
46            Ok(_size) => {}
47            Err(_e) => return None,
48        }
49        let ret = self.alloc.alloc(layout);
50        if ret.is_null() {
51            // Nothing was actually allocated, so add back the size
52            self.remaining.fetch_add(layout.size(), SeqCst);
53        }
54
55        Some(ret)
56    }
57
58    /// Returns remaining memory in bytes. This value does not guarantee that an allocation of x
59    /// bytes will succeed.
60    pub fn remaining(&self) -> usize {
61        self.remaining.load(SeqCst)
62    }
63}
64
65unsafe impl<A: GlobalAlloc> GlobalAlloc for Limit<A> {
66    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
67        self.try_alloc(layout).unwrap_or(ptr::null_mut())
68    }
69
70    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
71        self.alloc.dealloc(ptr, layout);
72        self.remaining.fetch_add(layout.size(), SeqCst);
73    }
74}
75
76unsafe impl<'a, A: GlobalAlloc> GlobalAlloc for &'a Limit<A> {
77    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
78        Limit::alloc(self, layout)
79    }
80
81    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
82        Limit::dealloc(self, ptr, layout)
83    }
84}
85
86pub struct ArcLimit<A>(Arc<Limit<A>>);
87
88impl<A> Clone for ArcLimit<A> {
89    fn clone(&self) -> Self {
90        Self(Arc::clone(&self.0))
91    }
92}
93
94impl<A: GlobalAlloc> ArcLimit<A> {
95    pub fn new(l: Limit<A>) -> Self {
96        Self(Arc::new(l))
97    }
98}
99
100unsafe impl<A: GlobalAlloc> GlobalAlloc for ArcLimit<A> {
101    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
102        Limit::alloc(&self.0, layout)
103    }
104
105    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
106        Limit::dealloc(&self.0, ptr, layout)
107    }
108}
109
110/// Total memory allocated by `ConstLimit`, in bytes.
111static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
112
113#[derive(Clone)]
114pub struct ConstLimit<A, const L: usize> {
115    alloc: A,
116}
117
118impl<A: GlobalAlloc, const L: usize> ConstLimit<A, L> {
119    pub const fn new(alloc: A) -> Self {
120        Self { alloc }
121    }
122
123    /// Returns None if the memory limit would be exhausted after allocating.
124    ///
125    /// # Safety
126    ///
127    /// The same restrictions as `GlobalAlloc::alloc`.
128    pub unsafe fn try_alloc(&self, layout: Layout) -> Option<*mut u8> {
129        match ALLOCATED.fetch_update(SeqCst, SeqCst, |old| {
130            let new = old.checked_add(layout.size())?;
131            if new > L {
132                None
133            } else {
134                Some(new)
135            }
136        }) {
137            Ok(_size) => {}
138            Err(_e) => return None,
139        }
140        let ret = self.alloc.alloc(layout);
141        if ret.is_null() {
142            // Nothing was actually allocated, so subtract the size
143            ALLOCATED.fetch_sub(layout.size(), SeqCst);
144        }
145
146        Some(ret)
147    }
148
149    /// Returns remaining memory in bytes. This value does not guarantee that an allocation of x
150    /// bytes will succeed.
151    pub fn remaining(&self) -> usize {
152        L.checked_sub(ALLOCATED.load(SeqCst))
153            .expect("bug: allocated more memory than the limit")
154    }
155}
156
157unsafe impl<A: GlobalAlloc, const L: usize> GlobalAlloc for ConstLimit<A, L> {
158    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
159        self.try_alloc(layout).unwrap_or(ptr::null_mut())
160    }
161
162    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
163        self.alloc.dealloc(ptr, layout);
164        ALLOCATED.fetch_sub(layout.size(), SeqCst);
165    }
166}
limit_alloc/lib.rs

limit_alloc/
lib.rs
