sql_cli/utils/
memory_tracker.rs

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

sql_cli/utils/
memory_tracker.rs
