re_memory/
peak_memory_stats.rs

1//! Monitor memory use periodically, and capture the stacktraces at the high water mark.
2
3use std::thread::JoinHandle;
4
5use crossbeam::channel::{Receiver, Sender};
6
7use crate::TrackingStatistics;
8use crate::accounting_allocator::{self, is_tracking_callstacks};
9
10struct Shutdown;
11
12/// Collects [`TrackingStatistics`] at the point of peak memory pressure.
13///
14/// Internally this runs a thread that periodically checks the current memory usage.
15/// At any high water mark, new [`TrackingStatistics`] is collected and stored.
16///
17/// This is primarily meant for developer investigation.
18pub struct PeakMemoryStats {
19    shutdown_tx: Sender<Shutdown>,
20    handle: JoinHandle<Option<TrackingStatistics>>,
21}
22
23impl PeakMemoryStats {
24    /// Start a background thread, tracking peak memory use.
25    pub fn start() -> Self {
26        if !is_tracking_callstacks() {
27            re_log::warn_once!(
28                "Callstack tracking is disabled - peak memory use will not be collected"
29            );
30        }
31
32        let (shutdown_tx, shutdown_rx) = crossbeam::channel::bounded(1);
33        let handle = std::thread::Builder::new()
34            .name("PeakMemoryStates".to_owned())
35            .spawn(move || collect_peak_memory_use(&shutdown_rx))
36            .expect("Failed to spawn PeakMemoryStates thread");
37        Self {
38            shutdown_tx,
39            handle,
40        }
41    }
42
43    pub fn finish(self) -> Option<TrackingStatistics> {
44        let Self {
45            shutdown_tx,
46            handle,
47        } = self;
48
49        shutdown_tx.send(Shutdown).ok();
50
51        match handle.join() {
52            Err(err) => {
53                re_log::warn_once!("Failed to collect PeakMemoryStates: {err:?}"); // NOLINT: err does not implement Display
54                None
55            }
56            Ok(result) => result,
57        }
58    }
59}
60
61fn current_memory_use() -> usize {
62    if let Some(stats) = crate::accounting_allocator::global_allocs() {
63        stats.size
64    } else {
65        re_log::error_once!("accounting_allocator is OFF. Can't collect peak memory use");
66        0
67    }
68}
69
70fn collect_peak_memory_use(shutdown_rx: &Receiver<Shutdown>) -> Option<TrackingStatistics> {
71    // How often we check RAM use.
72    let interval = std::time::Duration::from_secs(5);
73
74    let watermark_margin = 100 * 1024 * 1024;
75
76    let mut highest_memory_use = current_memory_use();
77    let mut highest_so_far = accounting_allocator::tracking_stats();
78
79    loop {
80        // Wait for either a shutdown signal or a 10-second timeout
81        crossbeam::select! {
82            recv(shutdown_rx) -> _ => {
83                // Shutdown signal received
84                return highest_so_far;
85            }
86            default(interval) => {
87                let current_memory_use = current_memory_use();
88
89                if highest_memory_use + watermark_margin < current_memory_use {
90                    highest_memory_use = current_memory_use;
91                    highest_so_far = accounting_allocator::tracking_stats();
92                }
93            }
94        }
95    }
96}
re_memory/peak_memory_stats.rs

re_memory/
peak_memory_stats.rs
