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}