kumo_server_memory/
tracking.rs

1use crate::{NumBytes, Number};
2use backtrace::Backtrace;
3use kumo_prometheus::ShardedCountAndSize;
4use parking_lot::Mutex;
5use std::alloc::{GlobalAlloc, Layout};
6use std::cell::Cell;
7use std::collections::HashMap;
8use std::sync::atomic::Ordering::Relaxed;
9use std::sync::atomic::{AtomicBool, AtomicUsize};
10use std::sync::LazyLock;
11
12// Portions of this file are derived from the re_memory crate
13// which is Copyright (c) 2022 Rerun Technologies AB <opensource@rerun.io>
14// and used under the terms of its MIT License
15// <https://github.com/rerun-io/rerun/tree/main/crates/utils/re_memory>
16
17thread_local! {
18    static IN_TRACKER: Cell<bool> = const { Cell::new(false) };
19}
20
21#[derive(Default)]
22pub struct TrackingAllocator<A: GlobalAlloc> {
23    allocator: A,
24}
25
26impl<A: GlobalAlloc> TrackingAllocator<A> {
27    pub const fn new(allocator: A) -> Self {
28        Self { allocator }
29    }
30}
31
32static STATS: Stats = Stats::new();
33
34const SMALL_SIZE: usize = 128;
35const MEDIUM_SIZE: usize = 4 * 1024;
36
37const MEDIUM_RATE: u64 = 64;
38const BIG_RATE: u64 = 1;
39
40static BIG_TRACKER: LazyLock<Mutex<AllocationTracker>> =
41    LazyLock::new(|| Mutex::new(AllocationTracker::default()));
42static MEDIUM_TRACKER: LazyLock<Mutex<AllocationTracker>> =
43    LazyLock::new(|| Mutex::new(AllocationTracker::default()));
44
45// SAFETY: we're passing through the unsafe portions to an underlying
46// allocator, which we're relying on to uphold safety.
47// The additional logic we add here is safe and is merely tracking
48unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
49    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
50        let ptr = unsafe { self.allocator.alloc(layout) };
51        if !ptr.is_null() {
52            track_allocation(ptr, layout.size());
53        }
54        ptr
55    }
56
57    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
58        let ptr = unsafe { self.allocator.alloc_zeroed(layout) };
59        if !ptr.is_null() {
60            track_allocation(ptr, layout.size());
61        }
62        ptr
63    }
64
65    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
66        track_dealloc(ptr, layout.size());
67        unsafe { self.allocator.dealloc(ptr, layout) };
68    }
69
70    unsafe fn realloc(&self, old_ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
71        let track_callstacks = STATS.track_callstacks.load(Relaxed);
72        let mut saved_callstack = None;
73        track_dealloc_with(old_ptr, layout.size(), track_callstacks, |entry| {
74            saved_callstack = entry;
75        });
76
77        let new_ptr = unsafe { self.allocator.realloc(old_ptr, layout, new_size) };
78
79        if new_ptr.is_null() {
80            // Failed realloc leaves the original allocation live.
81            track_allocation_with(old_ptr, layout.size(), track_callstacks, || saved_callstack);
82        } else {
83            if let Some(entry) = saved_callstack {
84                // Dropping the entry may deallocate through this allocator;
85                // the guard suppresses re-entrant tracking of that drop.
86                IN_TRACKER.with(|in_track| {
87                    let previous = in_track.replace(true);
88                    drop(entry);
89                    in_track.set(previous);
90                });
91            }
92            track_allocation(new_ptr, new_size);
93        }
94        new_ptr
95    }
96}
97
98fn track_allocation(ptr: *mut u8, size: usize) {
99    track_allocation_with(ptr, size, STATS.track_callstacks.load(Relaxed), || {
100        Some(CallstackEntry {
101            size,
102            bt: Backtrace::new_unresolved(),
103        })
104    });
105}
106
107#[inline]
108fn track_allocation_with(
109    ptr: *mut u8,
110    size: usize,
111    track_callstacks: bool,
112    callstack: impl FnOnce() -> Option<CallstackEntry>,
113) {
114    STATS.live.add(ptr.addr(), size);
115
116    if !track_callstacks {
117        return;
118    }
119
120    if size < SMALL_SIZE {
121        STATS.small.add(size);
122        return;
123    }
124
125    IN_TRACKER.with(|in_track| {
126        if !in_track.get() {
127            in_track.set(true);
128
129            let hash = PtrHash::new(ptr);
130            let track = hash.should_sample_size(size);
131
132            if track {
133                let entry = callstack();
134                if size < MEDIUM_SIZE {
135                    STATS.medium.add(size);
136                    if let Some(entry) = entry {
137                        MEDIUM_TRACKER
138                            .lock()
139                            .track_allocation(hash, entry.size, entry.bt);
140                    }
141                } else {
142                    STATS.large.add(size);
143                    if let Some(entry) = entry {
144                        BIG_TRACKER
145                            .lock()
146                            .track_allocation(hash, entry.size, entry.bt);
147                    }
148                }
149            }
150
151            in_track.set(false);
152        }
153    });
154}
155
156fn track_dealloc(ptr: *mut u8, size: usize) {
157    track_dealloc_with(ptr, size, STATS.track_callstacks.load(Relaxed), drop);
158}
159
160#[inline]
161fn track_dealloc_with(
162    ptr: *mut u8,
163    size: usize,
164    track_callstacks: bool,
165    removed: impl FnOnce(Option<CallstackEntry>),
166) {
167    STATS.live.sub(ptr.addr(), size);
168
169    if !track_callstacks {
170        return;
171    }
172
173    if size < SMALL_SIZE {
174        STATS.small.sub(size);
175        return;
176    }
177
178    IN_TRACKER.with(|in_track| {
179        if !in_track.get() {
180            in_track.set(true);
181
182            let hash = PtrHash::new(ptr);
183            let track = hash.should_sample_size(size);
184
185            if track {
186                if size < MEDIUM_SIZE {
187                    let entry = MEDIUM_TRACKER.lock().track_dealloc(hash, size);
188                    STATS.medium.sub(size);
189                    removed(entry);
190                } else {
191                    STATS.large.sub(size);
192                    let entry = BIG_TRACKER.lock().track_dealloc(hash, size);
193                    removed(entry);
194                }
195            }
196
197            in_track.set(false);
198        }
199    });
200}
201
202/// Returns the stochastic sampling rate (really, an interval)
203/// that should be used for a given allocation size.
204fn stochastic_rate_by_size(size: usize) -> u64 {
205    if size < MEDIUM_SIZE {
206        MEDIUM_RATE
207    } else {
208        BIG_RATE
209    }
210}
211
212/// Given a pointer address, hash it into a 64-bit hash value.
213/// The hash re-distributes the bits which is important for
214/// the stochastic sampling approach used in this module.
215#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
216struct PtrHash(u64);
217
218impl PtrHash {
219    #[inline]
220    pub fn new(ptr: *mut u8) -> Self {
221        Self(ahash::RandomState::with_seeds(1, 2, 3, 4).hash_one(ptr))
222    }
223
224    /// Given an allocation size, returns true if we should sample
225    /// the associated allocation call stack based on the stochastic
226    /// rate configured for that allocation size.
227    pub fn should_sample_size(&self, size: usize) -> bool {
228        let rate = stochastic_rate_by_size(size);
229        self.should_sample_at_rate(rate)
230    }
231
232    /// Apply "stochastic sampling" at a specified "rate".
233    /// The rate is nominally a sampling interval.
234    /// The redistribution of the address bits by the hash
235    /// "randomizes" the bits and the rate/interval is used
236    /// as a mask
237    pub fn should_sample_at_rate(&self, rate: u64) -> bool {
238        self.0 & (rate - 1) == 0
239    }
240}
241
242struct CallstackEntry {
243    size: usize,
244    bt: Backtrace,
245}
246
247pub struct CallstackStats {
248    pub count: usize,
249    pub total_size: usize,
250    pub bt: Backtrace,
251    pub stochastic_rate: usize,
252}
253
254#[derive(Default)]
255struct AllocationTracker {
256    live_allocations: ahash::HashMap<PtrHash, CallstackEntry>,
257}
258
259impl AllocationTracker {
260    pub fn track_allocation(&mut self, ptr: PtrHash, size: usize, bt: Backtrace) {
261        self.live_allocations
262            .insert(ptr, CallstackEntry { size, bt });
263    }
264
265    pub fn track_dealloc(&mut self, ptr: PtrHash, _size: usize) -> Option<CallstackEntry> {
266        self.live_allocations.remove(&ptr)
267    }
268
269    pub fn top_callstacks(&self, max_stacks: usize) -> Vec<CallstackStats> {
270        let mut by_stack = HashMap::new();
271
272        for alloc in self.live_allocations.values() {
273            let key = alloc.bt.frames().iter().map(|f| f.ip()).collect::<Vec<_>>();
274            let entry = by_stack.entry(key).or_insert_with(|| CallstackStats {
275                count: 0,
276                total_size: 0,
277                bt: alloc.bt.clone(),
278                stochastic_rate: stochastic_rate_by_size(alloc.size) as usize,
279            });
280
281            entry.count += 1;
282            entry.total_size += alloc.size;
283        }
284
285        let mut stats = by_stack.into_values().collect::<Vec<_>>();
286        stats.sort_by(|a, b| b.total_size.cmp(&a.total_size));
287        stats.truncate(max_stacks);
288        stats.shrink_to_fit();
289        stats
290    }
291}
292
293struct AtomicCountAndSize {
294    /// Number of allocations.
295    pub count: AtomicUsize,
296
297    /// Number of bytes.
298    pub size: AtomicUsize,
299}
300
301impl AtomicCountAndSize {
302    pub const fn zero() -> Self {
303        Self {
304            count: AtomicUsize::new(0),
305            size: AtomicUsize::new(0),
306        }
307    }
308
309    /// Add an allocation.
310    fn add(&self, size: usize) {
311        self.count.fetch_add(1, Relaxed);
312        self.size.fetch_add(size, Relaxed);
313    }
314
315    /// Remove an allocation.
316    fn sub(&self, size: usize) {
317        self.count.fetch_sub(1, Relaxed);
318        self.size.fetch_sub(size, Relaxed);
319    }
320}
321
322#[derive(Debug, Clone, Copy)]
323pub struct CountAndSize {
324    pub count: Number,
325    pub size: NumBytes,
326}
327
328struct Stats {
329    live: ShardedCountAndSize,
330    track_callstacks: AtomicBool,
331    small: AtomicCountAndSize,
332    medium: AtomicCountAndSize,
333    large: AtomicCountAndSize,
334}
335
336impl Stats {
337    const fn new() -> Self {
338        Self {
339            live: ShardedCountAndSize::new(),
340            small: AtomicCountAndSize::zero(),
341            medium: AtomicCountAndSize::zero(),
342            large: AtomicCountAndSize::zero(),
343            track_callstacks: AtomicBool::new(false),
344        }
345    }
346}
347
348/// Returns the sampled number of bytes allocated via the global allocator. Not
349/// all of these may be resident. The RSS value will typically differ.
350/// Concurrent updates can make this overestimate or underestimate the
351/// instantaneous total.
352pub fn counted_usage() -> usize {
353    STATS.live.load_size()
354}
355
356pub fn set_tracking_callstacks(enable: bool) {
357    STATS.track_callstacks.store(enable, Relaxed);
358}
359
360pub struct TrackingStats {
361    pub small_threshold: NumBytes,
362    pub live: CountAndSize,
363    pub top_callstacks: Vec<CallstackStats>,
364}
365
366pub fn tracking_stats() -> TrackingStats {
367    const MAX_STACKS: usize = 128;
368
369    let mut top_callstacks = vec![];
370
371    IN_TRACKER.with(|in_track| {
372        if !in_track.get() {
373            in_track.set(true);
374            top_callstacks = BIG_TRACKER.lock().top_callstacks(MAX_STACKS);
375            top_callstacks.append(&mut MEDIUM_TRACKER.lock().top_callstacks(MAX_STACKS));
376
377            for stack in &mut top_callstacks {
378                stack.bt.resolve();
379            }
380
381            in_track.set(false);
382        }
383    });
384
385    let live = STATS.live.load();
386    TrackingStats {
387        small_threshold: SMALL_SIZE.into(),
388        live: CountAndSize {
389            count: live.count.into(),
390            size: live.size.into(),
391        },
392        top_callstacks,
393    }
394}
kumo_server_memory/tracking.rs

kumo_server_memory/
tracking.rs
