sql_cli/utils/
memory_tracker.rs

1/// Get current process memory usage in MB
2#[must_use]
3pub fn get_memory_mb() -> usize {
4    // Use /proc/self/status on Linux
5    #[cfg(target_os = "linux")]
6    {
7        if let Ok(status) = std::fs::read_to_string("/proc/self/status") {
8            for line in status.lines() {
9                if line.starts_with("VmRSS:") {
10                    if let Some(kb_str) = line.split_whitespace().nth(1) {
11                        if let Ok(kb) = kb_str.parse::<usize>() {
12                            return kb / 1024;
13                        }
14                    }
15                }
16            }
17        }
18    }
19
20    // Fallback or other platforms
21    0
22}
23
24/// Get current process memory usage in KB (cross-platform)
25#[must_use]
26pub fn get_process_memory_kb() -> Option<usize> {
27    #[cfg(target_os = "linux")]
28    {
29        std::fs::read_to_string("/proc/self/status")
30            .ok()?
31            .lines()
32            .find(|line| line.starts_with("VmRSS:"))
33            .and_then(|line| {
34                line.split_whitespace()
35                    .nth(1)
36                    .and_then(|s| s.parse::<usize>().ok())
37            })
38    }
39
40    #[cfg(target_os = "macos")]
41    {
42        use std::process::Command;
43        if let Ok(output) = Command::new("ps")
44            .args(&["-o", "rss=", "-p", &std::process::id().to_string()])
45            .output()
46        {
47            if let Ok(s) = String::from_utf8(output.stdout) {
48                if let Ok(kb) = s.trim().parse::<usize>() {
49                    return Some(kb);
50                }
51            }
52        }
53        None
54    }
55
56    #[cfg(target_os = "windows")]
57    {
58        use std::mem;
59        use winapi::um::processthreadsapi::GetCurrentProcess;
60        use winapi::um::psapi::{GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS};
61
62        unsafe {
63            let h_process = GetCurrentProcess();
64            let mut pmc: PROCESS_MEMORY_COUNTERS = mem::zeroed();
65            pmc.cb = mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32;
66
67            if GetProcessMemoryInfo(h_process, &mut pmc as *mut _ as *mut _, pmc.cb) != 0 {
68                // WorkingSetSize is in bytes, convert to KB
69                Some((pmc.WorkingSetSize / 1024) as usize)
70            } else {
71                None
72            }
73        }
74    }
75
76    #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
77    {
78        None
79    }
80}
81
82// Use thread-local storage instead of global static
83thread_local! {
84    static MEMORY_LOG: std::cell::RefCell<Vec<(String, usize)>> = const { std::cell::RefCell::new(Vec::new()) };
85}
86
87/// Track memory at a specific point
88#[must_use]
89pub fn track_memory(label: &str) -> usize {
90    let mb = get_memory_mb();
91
92    MEMORY_LOG.with(|log| {
93        let mut log = log.borrow_mut();
94
95        // Calculate delta from last entry
96        let delta = if let Some((_, last_mb)) = log.last() {
97            let diff = (mb as i32) - (*last_mb as i32);
98            if diff != 0 {
99                format!(" ({diff:+} MB)")
100            } else {
101                String::new()
102            }
103        } else {
104            String::new()
105        };
106
107        log.push((label.to_string(), mb));
108
109        // Keep last 30 entries
110        if log.len() > 30 {
111            log.remove(0);
112        }
113
114        tracing::info!("MEMORY[{}]: {} MB{}", label, mb, delta);
115    });
116
117    mb
118}
119
120/// Get memory history for display
121#[must_use]
122pub fn get_memory_history() -> Vec<(String, usize)> {
123    MEMORY_LOG.with(|log| log.borrow().clone())
124}
125
126/// Format memory history as a string
127#[must_use]
128pub fn format_memory_history() -> String {
129    MEMORY_LOG.with(|log| {
130        let log = log.borrow();
131        let mut output = String::from("Memory History:\n");
132
133        for (i, (label, mb)) in log.iter().enumerate() {
134            let delta = if i > 0 {
135                let prev_mb = log[i - 1].1;
136                let diff = (*mb as i32) - (prev_mb as i32);
137                if diff != 0 {
138                    format!(" ({diff:+} MB)")
139                } else {
140                    String::new()
141                }
142            } else {
143                String::new()
144            };
145
146            output.push_str(&format!("  {label}: {mb} MB{delta}\n"));
147        }
148
149        output
150    })
151}