rb_sys/
tracking_allocator.rs

1//! Support for reporting Rust memory usage to the Ruby GC.
2
3use std::{
4    fmt::Formatter,
5    sync::{
6        atomic::{AtomicIsize, Ordering},
7        Arc,
8    },
9};
10
11#[cfg(ruby_engine = "mri")]
12mod mri {
13    use crate::{rb_gc_adjust_memory_usage, utils::is_ruby_vm_started};
14    use std::alloc::{GlobalAlloc, Layout, System};
15
16    /// A simple wrapper over [`System`] which reports memory usage to
17    /// the Ruby GC. This gives the GC a more accurate picture of the process'
18    /// memory usage so it can make better decisions about when to run.
19    #[derive(Debug)]
20    pub struct TrackingAllocator;
21
22    impl TrackingAllocator {
23        /// Create a new [`TrackingAllocator`].
24        #[allow(clippy::new_without_default)]
25        pub const fn new() -> Self {
26            Self
27        }
28
29        /// Create a new [`TrackingAllocator`] with default values.
30        pub const fn default() -> Self {
31            Self::new()
32        }
33
34        /// Adjust the memory usage reported to the Ruby GC by `delta`. Useful for
35        /// tracking allocations invisible to the Rust allocator, such as `mmap` or
36        /// direct `malloc` calls.
37        ///
38        /// # Example
39        /// ```
40        /// use rb_sys::TrackingAllocator;
41        ///
42        /// // Allocate 1024 bytes of memory using `mmap` or `malloc`...
43        /// TrackingAllocator::adjust_memory_usage(1024);
44        ///
45        /// // ...and then after the memory is freed, adjust the memory usage again.
46        /// TrackingAllocator::adjust_memory_usage(-1024);
47        /// ```
48        #[inline]
49        pub fn adjust_memory_usage(delta: isize) -> isize {
50            if delta == 0 {
51                return 0;
52            }
53
54            #[cfg(target_pointer_width = "32")]
55            let delta = delta as i32;
56
57            #[cfg(target_pointer_width = "64")]
58            let delta = delta as i64;
59
60            unsafe {
61                if is_ruby_vm_started() {
62                    rb_gc_adjust_memory_usage(delta);
63                    delta as isize
64                } else {
65                    0
66                }
67            }
68        }
69    }
70
71    unsafe impl GlobalAlloc for TrackingAllocator {
72        #[inline]
73        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
74            let ret = System.alloc(layout);
75            let delta = layout.size() as isize;
76
77            if !ret.is_null() && delta != 0 {
78                Self::adjust_memory_usage(delta);
79            }
80
81            ret
82        }
83
84        #[inline]
85        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
86            let ret = System.alloc_zeroed(layout);
87            let delta = layout.size() as isize;
88
89            if !ret.is_null() && delta != 0 {
90                Self::adjust_memory_usage(delta);
91            }
92
93            ret
94        }
95
96        #[inline]
97        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
98            System.dealloc(ptr, layout);
99            let delta = -(layout.size() as isize);
100
101            if delta != 0 {
102                Self::adjust_memory_usage(delta);
103            }
104        }
105
106        #[inline]
107        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
108            let ret = System.realloc(ptr, layout, new_size);
109            let delta = new_size as isize - layout.size() as isize;
110
111            if !ret.is_null() && delta != 0 {
112                Self::adjust_memory_usage(delta);
113            }
114
115            ret
116        }
117    }
118}
119
120#[cfg(not(ruby_engine = "mri"))]
121mod non_mri {
122    use std::alloc::{GlobalAlloc, Layout, System};
123
124    /// A simple wrapper over [`System`] as a fallback for non-MRI Ruby engines.
125    pub struct TrackingAllocator;
126
127    impl TrackingAllocator {
128        #[allow(clippy::new_without_default)]
129        pub const fn new() -> Self {
130            Self
131        }
132
133        pub const fn default() -> Self {
134            Self::new()
135        }
136
137        pub fn adjust_memory_usage(_delta: isize) -> isize {
138            0
139        }
140    }
141
142    unsafe impl GlobalAlloc for TrackingAllocator {
143        #[inline]
144        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
145            System.alloc(layout)
146        }
147
148        #[inline]
149        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
150            System.alloc_zeroed(layout)
151        }
152
153        #[inline]
154        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
155            System.dealloc(ptr, layout)
156        }
157
158        #[inline]
159        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
160            System.realloc(ptr, layout, new_size)
161        }
162    }
163}
164
165#[cfg(ruby_engine = "mri")]
166pub use mri::*;
167
168#[cfg(not(ruby_engine = "mri"))]
169pub use non_mri::*;
170
171/// Set the global allocator to [`TrackingAllocator`].
172///
173/// # Example
174/// ```
175/// // File: ext/my_gem/src/lib.rs
176/// use rb_sys::set_global_tracking_allocator;
177///
178/// set_global_tracking_allocator!();
179/// ```
180#[macro_export]
181macro_rules! set_global_tracking_allocator {
182    () => {
183        #[global_allocator]
184        static RUBY_GLOBAL_TRACKING_ALLOCATOR: $crate::tracking_allocator::TrackingAllocator =
185            $crate::tracking_allocator::TrackingAllocator;
186    };
187}
188
189#[derive(Debug)]
190#[repr(transparent)]
191struct MemsizeDelta(Arc<AtomicIsize>);
192
193impl MemsizeDelta {
194    fn new(delta: isize) -> Self {
195        let delta = TrackingAllocator::adjust_memory_usage(delta);
196        Self(Arc::new(AtomicIsize::new(delta)))
197    }
198
199    fn add(&self, delta: usize) {
200        if delta == 0 {
201            return;
202        }
203
204        let delta = TrackingAllocator::adjust_memory_usage(delta as _);
205        self.0.fetch_add(delta as _, Ordering::SeqCst);
206    }
207
208    fn sub(&self, delta: usize) {
209        if delta == 0 {
210            return;
211        }
212
213        let delta = TrackingAllocator::adjust_memory_usage(-(delta as isize));
214        self.0.fetch_add(delta, Ordering::SeqCst);
215    }
216
217    fn get(&self) -> isize {
218        self.0.load(Ordering::SeqCst)
219    }
220}
221
222impl Clone for MemsizeDelta {
223    fn clone(&self) -> Self {
224        Self(Arc::clone(&self.0))
225    }
226}
227
228impl Drop for MemsizeDelta {
229    fn drop(&mut self) {
230        let memsize = self.0.swap(0, Ordering::SeqCst);
231        TrackingAllocator::adjust_memory_usage(0 - memsize);
232    }
233}
234
235/// A guard which adjusts the memory usage reported to the Ruby GC by `delta`.
236/// This allows you to track resources which are invisible to the Rust
237/// allocator, such as items that are known to internally use `mmap` or direct
238/// `malloc` in their implementation.
239///
240/// Internally, it uses an [`Arc<AtomicIsize>`] to track the memory usage delta,
241/// and is safe to clone when `T` is [`Clone`].
242///
243/// # Example
244/// ```
245/// use rb_sys::tracking_allocator::ManuallyTracked;
246///
247/// type SomethingThatUsedMmap = ();
248///
249/// // Will tell the Ruby GC that 1024 bytes were allocated.
250/// let item = ManuallyTracked::new(SomethingThatUsedMmap, 1024);
251///
252/// // Will tell the Ruby GC that 1024 bytes were freed.
253/// std::mem::drop(item);
254/// ```
255pub struct ManuallyTracked<T> {
256    item: T,
257    memsize_delta: MemsizeDelta,
258}
259
260impl<T> ManuallyTracked<T> {
261    /// Create a new `ManuallyTracked<T>`, and immediately report that `memsize`
262    /// bytes were allocated.
263    pub fn wrap(item: T, memsize: usize) -> Self {
264        Self {
265            item,
266            memsize_delta: MemsizeDelta::new(memsize as _),
267        }
268    }
269
270    /// Increase the memory usage reported to the Ruby GC by `memsize` bytes.
271    pub fn increase_memory_usage(&self, memsize: usize) {
272        self.memsize_delta.add(memsize);
273    }
274
275    /// Decrease the memory usage reported to the Ruby GC by `memsize` bytes.
276    pub fn decrease_memory_usage(&self, memsize: usize) {
277        self.memsize_delta.sub(memsize);
278    }
279
280    /// Get the current memory usage delta.
281    pub fn memsize_delta(&self) -> isize {
282        self.memsize_delta.get()
283    }
284
285    /// Get a shared reference to the inner `T`.
286    pub fn get(&self) -> &T {
287        &self.item
288    }
289
290    /// Get a mutable reference to the inner `T`.
291    pub fn get_mut(&mut self) -> &mut T {
292        &mut self.item
293    }
294}
295
296impl ManuallyTracked<()> {
297    /// Create a new `ManuallyTracked<()>`, and immediately report that
298    /// `memsize` bytes were allocated.
299    pub fn new(memsize: usize) -> Self {
300        Self::wrap((), memsize)
301    }
302}
303
304impl Default for ManuallyTracked<()> {
305    fn default() -> Self {
306        Self::wrap((), 0)
307    }
308}
309
310impl<T: Clone> Clone for ManuallyTracked<T> {
311    fn clone(&self) -> Self {
312        Self {
313            item: self.item.clone(),
314            memsize_delta: self.memsize_delta.clone(),
315        }
316    }
317}
318
319impl<T: std::fmt::Debug> std::fmt::Debug for ManuallyTracked<T> {
320    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
321        f.debug_struct("ManuallyTracked")
322            .field("item", &self.item)
323            .field("memsize_delta", &self.memsize_delta)
324            .finish()
325    }
326}
rb_sys/tracking_allocator.rs

rb_sys/
tracking_allocator.rs
