1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// These methods are described in:
// http://nadeausoftware.com/articles/2012/07/c_c_tip_how_get_process_resident_set_size_physical_memory_use

use libc;
use super::Result;

/// Get the current RSS memory of this process in KB
#[cfg(target_os = "linux")]
pub fn current_rss() -> Result<u64> {
    os::current_rss()
}

/// Get the current RSS memory of a process with given pid in KB
#[cfg(target_os = "linux")]
pub fn current_rss_of(pid: libc::pid_t) -> Result<u64> {
    os::current_rss_of(pid)
}

/// Get the max RSS memory of this process in KB
#[cfg(target_os = "linux")]
pub fn max_rss() -> u64 {
    os::max_rss()
}

#[cfg(target_os = "linux")]
mod os {
    use std::mem;
    use libc;
    use std::path::Path;
    use super::super::Result;
    use super::super::ProbeError;
    use super::super::file_to_string;

    #[inline]
    pub fn current_rss() -> Result<u64> {
        read_and_get_current_rss(&Path::new("/proc/self/statm"))
    }

    #[inline]
    pub fn current_rss_of(pid: libc::pid_t) -> Result<u64> {
        read_and_get_current_rss(&Path::new(&format!("/proc/{}/statm", pid)))
    }

    #[inline]
    pub fn read_and_get_current_rss(path: &Path) -> Result<u64> {
        let raw_data = try!(file_to_string(path));
        let segments: Vec<&str> = raw_data.split_whitespace().collect();

        if segments.len() < 2 {
            return Err(ProbeError::UnexpectedContent("Incorrect number of segments".to_owned()))
        }

        let pages: u64 = try!(segments[1].parse().map_err(|_| {
            ProbeError::UnexpectedContent("Could not parse segment".to_owned())
        }));

        // Value is in pages, needs to be multiplied by the page size to get a value in KB. We ask the OS
        // for this information using sysconf.
        let pagesize = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as u64 / 1024;

        Ok(pages * pagesize)
    }

    #[inline]
    pub fn max_rss() -> u64 {
        let mut rusage: libc::rusage = unsafe { mem::uninitialized() };
        unsafe { libc::getrusage(libc::RUSAGE_SELF, &mut rusage) };
        rusage.ru_maxrss as u64
    }
}

#[cfg(test)]
mod tests {
    use libc;
    use std::path::Path;
    use super::super::ProbeError;

    #[test]
    fn test_current_rss() {
        assert!(super::current_rss().is_ok());
        // See if it's a sort of sane value, between 1 and 250 mb
        assert!(super::current_rss().unwrap() > 1_000);
        assert!(super::current_rss().unwrap() < 250_000);
    }

    #[test]
    fn test_read_and_get_current_rss() {
        let path = Path::new("fixtures/linux/process_memory/proc_self_statm");
        let value = super::os::read_and_get_current_rss(&path).unwrap();
        assert_eq!(4552, value);
    }

    #[test]
    fn test_read_and_get_current_rss_wrong_path() {
        let path = Path::new("/nonsense");
        match super::os::read_and_get_current_rss(&path) {
            Err(ProbeError::IO(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_get_current_rss_incomplete() {
        let path = Path::new("fixtures/linux/process_memory/proc_self_statm_incomplete");
        match super::os::read_and_get_current_rss(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_read_and_get_current_rss_garbage() {
        let path = Path::new("fixtures/linux/process_memory/proc_self_statm_garbage");
        match super::os::read_and_get_current_rss(&path) {
            Err(ProbeError::UnexpectedContent(_)) => (),
            r => panic!("Unexpected result: {:?}", r)
        }
    }

    #[test]
    fn test_current_rss_of() {
        let pid = unsafe { libc::getpid() };
        assert!(super::current_rss_of(pid).is_ok());
        // See if it's a sort of sane value, between 1 and 250 mb
        assert!(super::current_rss_of(pid).unwrap() > 1_000);
        assert!(super::current_rss_of(pid).unwrap() < 250_000);
    }

    #[test]
    fn test_current_rss_of_invalid_pid() {
        assert!(super::current_rss_of(0).is_err());
    }

    #[test]
    fn test_max_rss() {
        // See if it's a sort of sane value, between 1 and 250 mb
        assert!(super::max_rss() > 1_000);
        assert!(super::max_rss() < 250_000);
    }
}