ntex_bytes/
tracker.rs

1//! Memory attribution for buffer allocations.
2use std::sync::atomic::{AtomicUsize, Ordering::Acquire, Ordering::Relaxed, Ordering::Release};
3use std::{fmt, ptr::NonNull, sync::atomic::fence};
4
5use crate::stvec::abort;
6
7const MAX_REFS: usize = isize::MAX as usize;
8
9/// Shared counter of the memory held by buffer allocations.
10///
11/// An allocation attributed to a tracker charges its size, including the
12/// allocation header, until the allocation is freed or returned to the
13/// page cache. The tracker state lives while any handle or attributed
14/// allocation exists, so a handle can outlive the buffers and the
15/// buffers can outlive all handles.
16// `BytesMut::tracker` borrows the allocation header field as a handle
17#[cfg_attr(docsrs, doc(cfg(feature = "tracker")))]
18#[repr(transparent)]
19pub struct BytesTracker(NonNull<TrackerInner>);
20
21const _: () = assert!(size_of::<BytesTracker>() == size_of::<*const TrackerInner>());
22
23/// Point-in-time values of a [`BytesTracker`].
24///
25/// The counters are read one by one, the snapshot is exact only when
26/// no allocation is attributed or released concurrently.
27#[cfg_attr(docsrs, doc(cfg(feature = "tracker")))]
28#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
29#[non_exhaustive]
30pub struct BytesTrackerSnapshot {
31    /// Memory held by attributed allocations, in bytes.
32    pub live_bytes: usize,
33    /// Number of attributed allocations.
34    pub live_allocations: usize,
35}
36
37pub(crate) struct TrackerInner {
38    // owner handles + attributed allocations
39    refs: AtomicUsize,
40    // owner handles
41    handles: AtomicUsize,
42    live_bytes: AtomicUsize,
43}
44
45// SAFETY: the shared state consists of atomics only
46unsafe impl Send for BytesTracker {}
47// SAFETY: the shared state consists of atomics only
48unsafe impl Sync for BytesTracker {}
49
50impl BytesTracker {
51    /// Creates a new tracker.
52    pub fn new() -> Self {
53        let inner = Box::new(TrackerInner {
54            refs: AtomicUsize::new(1),
55            handles: AtomicUsize::new(1),
56            live_bytes: AtomicUsize::new(0),
57        });
58        BytesTracker(NonNull::from(Box::leak(inner)))
59    }
60
61    /// Returns the current values of the tracker.
62    pub fn snapshot(&self) -> BytesTrackerSnapshot {
63        let inner = self.inner();
64        let handles = inner.handles.load(Relaxed);
65        let refs = inner.refs.load(Relaxed);
66        BytesTrackerSnapshot {
67            live_bytes: inner.live_bytes.load(Relaxed),
68            // fields are read one by one, a racing clone can make it negative
69            live_allocations: refs.saturating_sub(handles),
70        }
71    }
72
73    /// Returns `true` if both handles refer to the same tracker.
74    pub fn ptr_eq(&self, other: &BytesTracker) -> bool {
75        self.0 == other.0
76    }
77
78    pub(crate) fn as_inner(&self) -> *const TrackerInner {
79        self.0.as_ptr()
80    }
81
82    fn inner(&self) -> &TrackerInner {
83        // SAFETY: the handle holds a reference
84        unsafe { self.0.as_ref() }
85    }
86}
87
88impl Default for BytesTracker {
89    fn default() -> Self {
90        Self::new()
91    }
92}
93
94impl Clone for BytesTracker {
95    fn clone(&self) -> Self {
96        let inner = self.inner();
97        if inner.refs.fetch_add(1, Relaxed) > MAX_REFS {
98            abort();
99        }
100        inner.handles.fetch_add(1, Relaxed);
101        BytesTracker(self.0)
102    }
103}
104
105impl Drop for BytesTracker {
106    fn drop(&mut self) {
107        self.inner().handles.fetch_sub(1, Relaxed);
108        // SAFETY: the handle holds a reference
109        unsafe { TrackerInner::release(self.0.as_ptr()) }
110    }
111}
112
113impl fmt::Debug for BytesTracker {
114    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115        let snapshot = self.snapshot();
116        f.debug_struct("BytesTracker")
117            .field("live_bytes", &snapshot.live_bytes)
118            .field("live_allocations", &snapshot.live_allocations)
119            .finish()
120    }
121}
122
123impl TrackerInner {
124    /// Attributes an allocation of `size` bytes, it holds a reference
125    /// until [`TrackerInner::uncharge`].
126    pub(crate) fn charge(&self, size: usize) {
127        if self.refs.fetch_add(1, Relaxed) > MAX_REFS {
128            abort();
129        }
130        self.live_bytes.fetch_add(size, Relaxed);
131    }
132
133    /// Adjusts the size of an attributed allocation.
134    pub(crate) fn resize(&self, old: usize, new: usize) {
135        if new > old {
136            self.live_bytes.fetch_add(new - old, Relaxed);
137        } else {
138            self.live_bytes.fetch_sub(old - new, Relaxed);
139        }
140    }
141
142    /// Releases an allocation of `size` bytes.
143    ///
144    /// # Safety
145    ///
146    /// `this` must point to a live tracker and the caller must own
147    /// a reference taken by [`TrackerInner::charge`].
148    pub(crate) unsafe fn uncharge(this: *const TrackerInner, size: usize) {
149        // SAFETY: the caller owns a reference
150        unsafe { (*this).live_bytes.fetch_sub(size, Relaxed) };
151        // SAFETY: the caller owns a reference
152        unsafe { TrackerInner::release(this) }
153    }
154
155    /// # Safety
156    ///
157    /// The caller must own a reference, it is consumed.
158    unsafe fn release(this: *const TrackerInner) {
159        // SAFETY: the caller owns a reference
160        if unsafe { (*this).refs.fetch_sub(1, Release) } != 1 {
161            return;
162        }
163        fence(Acquire);
164        // SAFETY: last reference, the state was created by Box in BytesTracker::new
165        drop(unsafe { Box::from_raw(this.cast_mut()) });
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    fn snap(bytes: usize, allocs: usize) -> BytesTrackerSnapshot {
174        BytesTrackerSnapshot {
175            live_bytes: bytes,
176            live_allocations: allocs,
177        }
178    }
179
180    #[test]
181    fn charge_resize_uncharge() {
182        let tracker = BytesTracker::new();
183        assert_eq!(tracker.snapshot(), snap(0, 0));
184
185        let inner = tracker.as_inner();
186        // SAFETY: the tracker is alive
187        unsafe { (*inner).charge(100) };
188        unsafe { (*inner).charge(50) };
189        assert_eq!(tracker.snapshot(), snap(150, 2));
190
191        unsafe { (*inner).resize(100, 300) };
192        assert_eq!(tracker.snapshot(), snap(350, 2));
193        unsafe { (*inner).resize(300, 200) };
194        unsafe { (*inner).resize(200, 200) };
195        assert_eq!(tracker.snapshot(), snap(250, 2));
196
197        unsafe { TrackerInner::uncharge(inner, 200) };
198        assert_eq!(tracker.snapshot(), snap(50, 1));
199        unsafe { TrackerInner::uncharge(inner, 50) };
200        assert_eq!(tracker.snapshot(), snap(0, 0));
201        assert_eq!(
202            format!("{tracker:?}"),
203            "BytesTracker { live_bytes: 0, live_allocations: 0 }"
204        );
205    }
206
207    #[test]
208    fn handles() {
209        let tracker = BytesTracker::default();
210        let tracker2 = tracker.clone();
211        assert!(tracker.ptr_eq(&tracker2));
212        assert!(!tracker.ptr_eq(&BytesTracker::new()));
213
214        // handles are not allocations
215        assert_eq!(tracker.snapshot(), snap(0, 0));
216        unsafe { (*tracker.as_inner()).charge(10) };
217        assert_eq!(tracker2.snapshot(), snap(10, 1));
218        drop(tracker2);
219        assert_eq!(tracker.snapshot(), snap(10, 1));
220        unsafe { TrackerInner::uncharge(tracker.as_inner(), 10) };
221        assert_eq!(tracker.snapshot(), snap(0, 0));
222    }
223
224    #[test]
225    fn allocation_outlives_handles() {
226        let tracker = BytesTracker::new();
227        let inner = tracker.as_inner();
228        unsafe { (*inner).charge(10) };
229        unsafe { (*inner).charge(20) };
230        drop(tracker);
231        // the allocations keep the state alive, the last one frees it
232        unsafe { TrackerInner::uncharge(inner, 10) };
233        unsafe { TrackerInner::uncharge(inner, 20) };
234    }
235
236    #[test]
237    fn concurrent() {
238        const THREADS: usize = 4;
239        const ITERS: usize = if cfg!(miri) { 20 } else { 10_000 };
240
241        let tracker = BytesTracker::new();
242        let threads: Vec<_> = (0..THREADS)
243            .map(|_| {
244                let tracker = tracker.clone();
245                std::thread::spawn(move || {
246                    let inner = tracker.as_inner();
247                    for i in 0..ITERS {
248                        let handle = tracker.clone();
249                        unsafe { (*inner).charge(i) };
250                        unsafe { (*inner).resize(i, i + 1) };
251                        drop(handle);
252                        unsafe { TrackerInner::uncharge(inner, i + 1) };
253                    }
254                    // left for the main thread to release
255                    unsafe { (*inner).charge(1) };
256                })
257            })
258            .collect();
259        for t in threads {
260            t.join().unwrap();
261        }
262        assert_eq!(tracker.snapshot(), snap(THREADS, THREADS));
263
264        let inner = tracker.as_inner();
265        drop(tracker);
266        for _ in 0..THREADS {
267            unsafe { TrackerInner::uncharge(inner, 1) };
268        }
269    }
270}
ntex_bytes/tracker.rs

ntex_bytes/
tracker.rs
