weak_alloc/
lib.rs

1use once_cell::sync::Lazy;
2use std::alloc::{GlobalAlloc, Layout, System};
3use std::fmt;
4use std::fmt::Write;
5use std::ops::Deref;
6use std::sync::Mutex;
7use weak_list::AllocHashSet;
8use weak_list::AllocMem;
9use weak_list::WeakList;
10use weak_list::WeakListHashSet;
11
12// TODO: interesting crates:
13// https://crates.io/crates/refbox
14// https://crates.io/crates/weak-table
15// https://crates.io/crates/provenant
16// https://crates.io/crates/rcgc
17// https://crates.io/crates/weak_list
18
19// Missing API:
20// * Manual or automatic garbage collection in case the Weak pointer is dropped
21// * Weak<T> will keep memory allocated for T even after dropping T, we must use Weak<Box<T>>
22
23/// A custom allocator that can be given ownership of data, returning a `WeakRef`.
24#[derive(Clone)]
25pub struct WeakAlloc<A> {
26    alloc: A,
27}
28
29/// List of values owned by the allocator
30static WEAK_LIST: Lazy<Mutex<WeakList<WeakListHashSet>>> =
31    Lazy::new(|| Mutex::new(WeakList::new()));
32// HashSet used as part of a manual realloc API, because the allocator cannot allocate memory while
33// the WEAK_LIST lock is held.
34static BIGGER_HASHSET: Lazy<Mutex<AllocHashSet>> =
35    Lazy::new(|| Mutex::new(AllocHashSet::with_capacity(1)));
36
37impl<A> WeakAlloc<A> {
38    pub const fn new(alloc: A) -> Self {
39        Self { alloc }
40    }
41}
42
43pub struct WeakRef<T: ?Sized, A: 'static + Clone + GlobalAlloc = System> {
44    weak: weak_list::WeakRef<T>,
45    alloc: WeakAlloc<A>,
46}
47
48impl<T, A: GlobalAlloc + Clone> Clone for WeakRef<T, A> {
49    fn clone(&self) -> Self {
50        Self {
51            weak: self.weak.clone(),
52            alloc: self.alloc.clone(),
53        }
54    }
55}
56
57impl<T: ?Sized + fmt::Debug, A: 'static + Clone + GlobalAlloc> fmt::Debug for WeakRef<T, A> {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        write!(f, "(Weak)")
60    }
61}
62
63impl<T: Send + Sync + 'static, A: 'static + Clone + GlobalAlloc> WeakRef<T, A> {
64    pub fn upgrade(&self) -> Option<ArcRef<T, A>> {
65        self.alloc.upgrade(self)
66    }
67}
68
69pub struct ArcRef<T: ?Sized, A: 'static + Clone + GlobalAlloc = System> {
70    arc: weak_list::ArcRef<T>,
71    alloc: WeakAlloc<A>,
72}
73
74impl<T, A: GlobalAlloc + Clone> Clone for ArcRef<T, A> {
75    fn clone(&self) -> Self {
76        Self {
77            arc: self.arc.clone(),
78            alloc: self.alloc.clone(),
79        }
80    }
81}
82
83impl<T: ?Sized, A: GlobalAlloc + Clone> AsRef<T> for ArcRef<T, A> {
84    fn as_ref(&self) -> &T {
85        &**self
86    }
87}
88
89impl<T: ?Sized, A: GlobalAlloc + Clone> Deref for ArcRef<T, A> {
90    type Target = T;
91
92    #[inline]
93    fn deref(&self) -> &Self::Target {
94        &self.arc
95    }
96}
97
98impl<T, A: GlobalAlloc + Clone> ArcRef<T, A> {
99    pub fn get_mut(this: &mut Self) -> Option<&mut T> {
100        weak_list::ArcRef::get_mut(&mut this.arc)
101    }
102
103    pub fn downgrade(this: &Self) -> WeakRef<T, A> {
104        WeakRef {
105            weak: weak_list::ArcRef::downgrade(&this.arc),
106            alloc: this.alloc.clone(),
107        }
108    }
109}
110
111impl<A> WeakAlloc<A>
112where
113    A: GlobalAlloc + Clone,
114{
115    /// Give ownership of a value to the allocator. The value may be deallocated at any time if no
116    /// other strong references to it exist.
117    ///
118    /// Returns a WeakRef that can be used to get back an Arc in the future using the upgrade
119    /// method, if the value still exists.
120    ///
121    /// Because of the way Arc is implemented, prefer giving T where the size of T is small.
122    /// Otherwise the backing allocation will not be deallocated when dropping the Arc, it will
123    /// only be deallocated after all the weak references go out of scope. So [u8; 1000] is
124    /// bad, but Box<[u8; 1000]> and Vec<[u8; 1000]> are good.
125    pub fn give<T: Send + Sync + 'static>(&self, element: T) -> WeakRef<T, A> {
126        let alloc_mem = AllocMem::default();
127        let mut big_hs = BIGGER_HASHSET.lock().unwrap();
128        let big_hs_cap = big_hs.capacity();
129        let big_hs_opt = Some(&mut *big_hs);
130        let mut lock = WEAK_LIST.lock().unwrap();
131        // Before pushing node to hashset, ensure the hashset has enough free memory
132        lock.realloc_hashset_if_needed_no_alloc(big_hs_opt);
133        let weak = lock.push_front_no_alloc(element, alloc_mem);
134        drop(lock);
135        let new_hs_cap = big_hs.capacity();
136
137        if new_hs_cap < big_hs_cap {
138            // Allocate more memory
139            big_hs.allocate_capacity(big_hs_cap * 2);
140        }
141
142        WeakRef {
143            weak,
144            alloc: self.clone(),
145        }
146    }
147
148    // TODO: instead of give and give_and_upgrade, make give always return an ArcRef and document
149    // give(x).downgrade() as the way to get a WeakRef?
150    /// Alternative to `A.give(x).upgrade().unwrap()` that never fails. There is a race condition
151    /// in that snippet if another thread fills the memory causing the allocator to call
152    /// `WEAK_LIST.pop_lru()` after the call to `give` but before the call to `upgrade`.
153    /// That race condition is avoided in this method by holding the lock slightly longer, so that
154    /// no other thread can modify the list before we upgrade to an ArcRef.
155    pub fn give_and_upgrade<T: Send + Sync + 'static>(&self, element: T) -> ArcRef<T, A> {
156        let alloc_mem = AllocMem::default();
157        let mut big_hs = BIGGER_HASHSET.lock().unwrap();
158        let big_hs_cap = big_hs.capacity();
159        let big_hs_opt = Some(&mut *big_hs);
160        let mut lock = WEAK_LIST.lock().unwrap();
161        // Before pushing node to hashset, ensure the hashset has enough free memory
162        lock.realloc_hashset_if_needed_no_alloc(big_hs_opt);
163        let weak = lock.push_front_no_alloc(element, alloc_mem);
164        // upgrade cannot fail because we hold a &mut WeakList, and we just pushed this item there.
165        // upgrade_quietly to avoid moving the element to the front of the list, because it already
166        // is at the front since the call to push_front_no_alloc.
167        let arc = weak.upgrade_quietly().unwrap();
168        drop(lock);
169        let new_hs_cap = big_hs.capacity();
170
171        if new_hs_cap < big_hs_cap {
172            // Allocate more memory
173            big_hs.allocate_capacity(big_hs_cap * 2);
174        }
175
176        ArcRef {
177            arc,
178            alloc: self.clone(),
179        }
180    }
181
182    pub fn upgrade<T: Send + Sync + 'static>(&self, w: &WeakRef<T, A>) -> Option<ArcRef<T, A>> {
183        let mut wl = WEAK_LIST.lock().unwrap();
184
185        w.weak.upgrade(&mut wl).map(|arc| ArcRef {
186            arc,
187            alloc: self.clone(),
188        })
189    }
190
191    /// Remove all the weak references from the WeakAlloc. This will deallocate all the WeakRefs
192    /// that do not have an active ArcRef.
193    pub fn clear(&self) {
194        let mut wl = WEAK_LIST.lock().unwrap();
195
196        wl.clear();
197    }
198
199    /// Try to allocate some memory without freeing any existing weak allocations. This can be used
200    /// to implement the equivalent to `try_give`: try to allocate a box with some contents but
201    /// only if there is enough memory. If there isn't enough memory this method will return a null
202    /// pointer and the code needed to initialize the box can be skipped (using `give` forces you
203    /// to initialize the value).
204    ///
205    /// # Safety
206    ///
207    /// The same restrictions as `GlobalAlloc::alloc`.
208    pub unsafe fn weak_alloc(&self, layout: Layout) -> *mut u8 {
209        self.alloc.alloc(layout)
210    }
211
212    /// Returns a reference to the inner allocator.
213    pub fn inner(&self) -> &A {
214        &self.alloc
215    }
216}
217
218unsafe impl<A> GlobalAlloc for WeakAlloc<A>
219where
220    A: GlobalAlloc,
221{
222    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
223        let mut ret = self.alloc.alloc(layout);
224        loop {
225            if !ret.is_null() {
226                break;
227            }
228
229            // Malloc returned null pointer!
230
231            // Mitigation in case the inner alloc returns a null pointers even though it has enough
232            // memory (bad align parameter?). Since this does not seem to be needed, it is
233            // commented out for now.
234            /*
235            // Before trying to free any weak allocations, try to allocate the same amount of
236            // memory using alignment of 1, to ensure that the null pointer is because of lack of
237            // memory and not because of an invalid alignment.
238            let align_1_layout = Layout::from_size_align(layout.size(), 1).unwrap();
239            let ptr_align_1 = self.alloc.alloc(align_1_layout);
240            if !ptr_align_1.is_null() {
241                // The returned pointer was not null. There are two possibilities:
242                // * The layout was indeed invalid
243                // * The system did not have enough memory at t0 but now it has enough memory at
244                // t1. In this case the expected behaviour would be to try to free a weak
245                // reference, because when the layout is valid we only want to return a null
246                // pointer after removing all the other weak allocations.
247                self.alloc.dealloc(ptr_align_1, align_1_layout);
248                // Return null pointer
249                return ret;
250            }
251            */
252
253            //log::error!("malloc returned null pointer");
254            // Free some weak allocation and try again
255            //WEAK_LIST.lock().unwrap().remove_all_unreachable();
256            let some_arc = WEAK_LIST.lock().unwrap().pop_lru();
257            if let Some(arc) = some_arc {
258                drop(arc);
259            } else {
260                // No more weak allocations to free, give up and return null pointer
261                // TODO: what if the inner alloc always returns a null pointer because of a wrong
262                // layout? We would delete all the weak allocations by mistake.
263                // Instrument layout to detect faulty layouts.
264                instrument::increase_null_ptr_layout_counter(layout);
265                return ret;
266            }
267
268            ret = self.alloc.alloc(layout);
269        }
270
271        ret
272    }
273
274    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
275        self.alloc.dealloc(ptr, layout);
276    }
277}
278
279/// Instrument failed allocations to detect applications that call malloc with an invalid argument.
280/// This is important because such calls to malloc will empty the WeakList (because the condition
281/// is to remove elements until malloc returns not null, and in that case malloc will always return
282/// null).
283pub mod instrument {
284    use super::*;
285
286    /// Histogram of `Layout`s that failed to allocate using the inner allocator.
287    struct SmallHist {
288        v: Vec<(Layout, u32)>,
289        other: u32,
290    }
291
292    const NUM_NULL_PTR_LAYOUT: usize = 4;
293    static NULL_PTR_LAYOUT_COUNTER: Lazy<Mutex<SmallHist>> = Lazy::new(|| {
294        Mutex::new(SmallHist {
295            v: Vec::with_capacity(NUM_NULL_PTR_LAYOUT),
296            other: 0,
297        })
298    });
299
300    pub fn increase_null_ptr_layout_counter(layout: Layout) {
301        let mut guard = NULL_PTR_LAYOUT_COUNTER.lock().unwrap();
302        let v = &mut guard.v;
303
304        let mut found = false;
305
306        // Increment layout counter if present.
307        for (l, ctr) in v.iter_mut() {
308            if *l == layout {
309                *ctr += 1;
310                found = true;
311                break;
312            }
313        }
314
315        // If the layout is not present yet, add it to the histogram with count=1.
316        // Unless the histogram is full, then add 1 to the "other" counter.
317        if !found {
318            if v.len() == NUM_NULL_PTR_LAYOUT {
319                let other = &mut guard.other;
320                *other += 1;
321            } else {
322                v.push((layout, 1));
323            }
324        }
325    }
326
327    pub fn dump_null_ptr_layout_counters() -> String {
328        let guard = NULL_PTR_LAYOUT_COUNTER.lock().unwrap();
329
330        if guard.v.is_empty() {
331            // Fast path: return empty string
332            return String::new();
333        }
334
335        // Need to drop guard to be able to allocate memory
336        drop(guard);
337
338        // Create a string with enough capacity to ensure that we do not allocate while the guard
339        // is held.
340        let mut buf = String::with_capacity(1024);
341        write!(
342            buf,
343            "Warn: the following layouts caused allocator to return null pointer: "
344        )
345        .unwrap();
346        let guard = NULL_PTR_LAYOUT_COUNTER.lock().unwrap();
347
348        writeln!(buf, "null_ptr layouts: {:?}", guard.v).unwrap();
349        if guard.other != 0 {
350            writeln!(buf, "number of other null_ptrs: {}", guard.other).unwrap();
351        }
352
353        drop(guard);
354
355        buf
356    }
357}