For #94 - virtualize procfs

Author: Lars T Hansen

src/main.rs

@@ -9,6 +9,7 @@ mod jobs;
 mod nvidia;
 mod process;
 mod procfs;
+mod procfsapi;
 mod ps;
 mod slurm;
 mod util;

src/process.rs

@@ -4,7 +4,7 @@ use crate::command::{self, CmdError};
 use crate::TIMEOUT_SECONDS;
 use crate::util;
 
-#[derive(PartialEq)]
+#[derive(PartialEq, Debug)]
 pub struct Process {
     pub pid: usize,
     pub uid: usize,

src/procfs.rs

@@ -1,17 +1,11 @@
 /// Collect CPU process information without GPU information, from files in /proc.
-
-extern crate libc;
-extern crate page_size;
-extern crate users;
+///
+/// The underlying system is virtualized through ProcfsAPI.
 
 use crate::process;
+use crate::procfsapi::ProcfsAPI;
 
 use std::collections::HashMap;
-use std::fs;
-use std::path;
-use std::os::linux::fs::MetadataExt;
-use std::time::{SystemTime, UNIX_EPOCH};
-use users::{uid_t, get_user_by_uid};
 
 /// Obtain process information via /proc and return a vector of structures with all the information
 /// we need.  In the returned vector, pids uniquely tag the records.
@@ -22,50 +16,44 @@ use users::{uid_t, get_user_by_uid};
 /// This function uniformly uses /proc, even though in some cases there are system calls that
 /// provide the same information.
 
-pub fn get_process_information() -> Result<Vec<process::Process>, String> {
+pub fn get_process_information(fs: &dyn ProcfsAPI) -> Result<Vec<process::Process>, String> {
 
     // The boot time is the `btime` field of /proc/stat.  It is measured in seconds since epoch.  We
     // need this to compute the process's real time, which we need to compute ps-compatible cpu
     // utilization.
 
     let mut boot_time = 0;
-    if let Ok(s) = fs::read_to_string(path::Path::new("/proc/stat")) {
-        for l in s.split('\n') {
-            if l.starts_with("btime ") {
-                let fields = l.split_ascii_whitespace().collect::<Vec<&str>>();
-                boot_time = parse_usize_field(&fields, 1, &l, "stat", 0, "btime")? as u64;
-                break;
-            }
-        }
-        if boot_time == 0 {
-            return Err(format!("Could not find btime in /proc/stat: {s}"));
+    let stat_s = fs.read_to_string("stat")?;
+    for l in stat_s.split('\n') {
+        if l.starts_with("btime ") {
+            let fields = l.split_ascii_whitespace().collect::<Vec<&str>>();
+            boot_time = parse_usize_field(&fields, 1, &l, "stat", 0, "btime")? as u64;
+            break;
         }
-    } else {
-        return Err(format!("Could not open or read /proc/stat"));
-    };
+    }
+    if boot_time == 0 {
+        return Err(format!("Could not find btime in /proc/stat: {stat_s}"));
+    }
 
     // The total RAM installed is in the `MemTotal` field of /proc/meminfo.  We need this to compute
     // ps-compatible relative memory use.
 
     let mut memtotal_kib = 0;
-    if let Ok(s) = fs::read_to_string(path::Path::new("/proc/meminfo")) {
-        for l in s.split('\n') {
-            if l.starts_with("MemTotal: ") {
-                // We expect "MemTotal:\s+(\d+)\s+kB", roughly
-                let fields = l.split_ascii_whitespace().collect::<Vec<&str>>();
-                if fields.len() != 3 || fields[2] != "kB" {
-                    return Err(format!("Unexpected MemTotal in /proc/meminfo: {l}"));
-                }
-                memtotal_kib = parse_usize_field(&fields, 1, &l, "meminfo", 0, "MemTotal")?;
-                break;
+    let meminfo_s = fs.read_to_string("meminfo")?;
+    for l in meminfo_s.split('\n') {
+        if l.starts_with("MemTotal: ") {
+            // We expect "MemTotal:\s+(\d+)\s+kB", roughly
+            let fields = l.split_ascii_whitespace().collect::<Vec<&str>>();
+            if fields.len() != 3 || fields[2] != "kB" {
+                return Err(format!("Unexpected MemTotal in /proc/meminfo: {l}"));
             }
+            memtotal_kib = parse_usize_field(&fields, 1, &l, "meminfo", 0, "MemTotal")?;
+            break;
         }
-        if memtotal_kib == 0 {
-            return Err(format!("Could not find MemTotal in /proc/meminfo: {s}"));
-        }
-    } else {
-        return Err(format!("Could not open or read /proc/meminfo"));
-    };
+    }
+    if memtotal_kib == 0 {
+        return Err(format!("Could not find MemTotal in /proc/meminfo: {meminfo_s}"));
+    }
 
     // Enumerate all pids, and collect the uids while we're here.
     //
@@ -76,30 +64,14 @@ pub fn get_process_information() -> Result<Vec<process::Process>, String> {
     // Note that a pid may disappear between the time we see it here and the time we get around to
     // reading it, later, and that new pids may appear meanwhile.  We should ignore both issues.
 
-    let mut pids = vec![];
-    if let Ok(dir) = fs::read_dir("/proc") {
-        for dirent in dir {
-            if let Ok(dirent) = dirent {
-                if let Ok(meta) = dirent.metadata() {
-                    let uid = meta.st_uid();
-                    if let Some(name) = dirent.path().file_name() {
-                        if let Ok(pid) = name.to_string_lossy().parse::<usize>() {
-                            pids.push((pid, uid));
-                        }
-                    }
-                }
-            }
-        }
-    } else {
-        return Err(format!("Could not open /proc"));
-    };
+    let pids = fs.read_proc_pids()?;
 
     // Collect remaining system data from /proc/{pid}/stat for the enumerated pids.
 
-    let kib_per_page = page_size::get() / 1024;
+    let kib_per_page = fs.page_size();
     let mut result = vec![];
     let mut user_table = UserTable::new();
-    let clock_ticks_per_sec: usize = unsafe { libc::sysconf(libc::_SC_CLK_TCK) as usize };
+    let clock_ticks_per_sec = fs.clock_ticks_per_sec();
     for (pid, uid) in pids {
 
         // Basic system variables.
@@ -111,7 +83,7 @@ pub fn get_process_information() -> Result<Vec<process::Process>, String> {
         let comm;
         let utime;
         let stime;
-        if let Ok(line) = fs::read_to_string(path::Path::new(&format!("/proc/{pid}/stat"))) {
+        if let Ok(line) = fs.read_to_string(&format!("{pid}/stat")) {
             // The comm field is a little tricky, it must be extracted first as the contents between
             // the first '(' and the last ')' in the line.
             let commstart = line.find('(');
@@ -132,12 +104,13 @@ pub fn get_process_information() -> Result<Vec<process::Process>, String> {
             let cstime = parse_usize_field(&fields, 14, &line, "stat", pid, "cstime")? / clock_ticks_per_sec;
             bsdtime = utime + stime + cutime + cstime;
             let start_time = (parse_usize_field(&fields, 19, &line, "stat", pid, "starttime")? / clock_ticks_per_sec) as u64;
-            let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
+            let now = fs.now();
             realtime = now - (boot_time + start_time);
         } else {
 	    // This is *usually* benign - the process may have gone away since we enumerated the
 	    // /proc directory.  It is *possibly* indicative of a permission problem, but that
 	    // problem would be so pervasive that diagnosing it here is not right.
+            println!("E");
 	    continue;
         }
 
@@ -152,12 +125,13 @@ pub fn get_process_information() -> Result<Vec<process::Process>, String> {
 
         let size_kib;
         let rss_kib;
-        if let Ok(s) = fs::read_to_string(path::Path::new(&format!("/proc/{pid}/statm"))) {
+        if let Ok(s) = fs.read_to_string(&format!("{pid}/statm")) {
             let fields = s.split_ascii_whitespace().collect::<Vec<&str>>();
             rss_kib = parse_usize_field(&fields, 1, &s, "statm", pid, "resident set size")? * kib_per_page;
             size_kib = parse_usize_field(&fields, 5, &s, "statm", pid, "data size")? * kib_per_page;
         } else {
 	    // This is *usually* benign - see above.
+            println!("F");
 	    continue;
         }
 
@@ -172,7 +146,7 @@ pub fn get_process_information() -> Result<Vec<process::Process>, String> {
 
         let pcpu = (((utime + stime) as f64 * 1000.0 / realtime as f64)).ceil() / 10.0;
         let pmem = f64::min(((rss_kib as f64) * 1000.0 / (memtotal_kib as f64)).ceil() / 10.0, 99.9);
-        let user = user_table.lookup(uid);
+        let user = user_table.lookup(fs, uid);
 
         result.push(process::Process {
             pid,
@@ -212,7 +186,7 @@ fn parse_usize_field(fields: &[&str], ix: usize, line: &str, file: &str, pid: us
 // The UserTable optimizes uid -> name lookup.
 
 struct UserTable {
-    ht: HashMap<uid_t, String>,
+    ht: HashMap<u32, String>,
 }
 
 impl UserTable {
@@ -222,11 +196,10 @@ impl UserTable {
         }
     }
 
-    fn lookup(&mut self, uid: uid_t) -> String {
+    fn lookup(&mut self, fs: &dyn ProcfsAPI, uid: u32) -> String {
         if let Some(name) = self.ht.get(&uid) {
             name.clone()
-        } else if let Some(u) = get_user_by_uid(uid) {
-            let name = u.name().to_string_lossy().to_string();
+        } else if let Some(name) = fs.user_by_uid(uid) {
             self.ht.insert(uid, name.clone());
             name
         } else {

src/procfsapi.rs

@@ -0,0 +1,161 @@
+extern crate libc;
+extern crate page_size;
+extern crate users;
+
+use users::get_user_by_uid;
+
+use std::collections::HashMap;
+use std::fs;
+use std::path;
+use std::os::linux::fs::MetadataExt;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+pub trait ProcfsAPI {
+    // Open /proc/<path> (which can have multiple path elements, eg, {PID}/filename), read it, and
+    // return its entire contents as a string.  Return a sensible error message if the file can't
+    // be opened or read.
+    fn read_to_string(&self, path: &str) -> Result<String, String>;
+
+    // Return (pid,uid) for every file /proc/{PID}.  Return a sensible error message in case
+    // something goes really, really wrong, but otherwise try to make the best of it.
+    fn read_proc_pids(&self) -> Result<Vec<(usize,u32)>, String>;
+
+    // Try to figure out the user's name from system tables, this may be an expensive operation.
+    fn user_by_uid(&self, uid: u32) -> Option<String>;
+
+    fn clock_ticks_per_sec(&self) -> usize;
+
+    fn page_size(&self) -> usize;
+
+    fn now(&self) -> u64;
+}
+
+// RealFS is used to actually access /proc and other system tables
+
+pub struct RealFS {
+}
+
+impl RealFS {
+    pub fn new() -> RealFS {
+        RealFS{}
+    }
+}
+
+impl ProcfsAPI for RealFS {
+    fn read_to_string(&self, path: &str) -> Result<String, String> {
+        match fs::read_to_string(path::Path::new("/proc/{path}")) {
+            Ok(s) => Ok(s),
+            Err(_) => Err(format!("Unable to read /proc/{path}")),
+        }
+    }
+
+    fn read_proc_pids(&self) -> Result<Vec<(usize,u32)>, String> {
+        let mut pids = vec![];
+        if let Ok(dir) = fs::read_dir("/proc") {
+            for dirent in dir {
+                if let Ok(dirent) = dirent {
+                    if let Ok(meta) = dirent.metadata() {
+                        let uid = meta.st_uid();
+                        if let Some(name) = dirent.path().file_name() {
+                            if let Ok(pid) = name.to_string_lossy().parse::<usize>() {
+                                pids.push((pid, uid));
+                            }
+                        }
+                    }
+                }
+            }
+        } else {
+            return Err(format!("Could not open /proc"));
+        };
+        Ok(pids)
+    }
+
+    fn user_by_uid(&self, uid: u32) -> Option<String> {
+        if let Some(u) = get_user_by_uid(uid) {
+            Some(u.name().to_string_lossy().to_string())
+        } else {
+            None
+        }
+    }
+
+    fn clock_ticks_per_sec(&self) -> usize {
+        unsafe { libc::sysconf(libc::_SC_CLK_TCK) as usize }
+    }
+
+    fn page_size(&self) -> usize {
+        page_size::get() / 1024
+    }
+
+    fn now(&self) -> u64 {
+        unix_now()
+    }
+}
+
+pub fn unix_now() -> u64 {
+    SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs()
+}
+
+// MockFS is used for testing
+
+#[cfg(test)]
+pub struct MockFS {
+    files: HashMap<String,String>,
+    pids: Vec<(usize,u32)>,
+    users: HashMap<u32,String>,
+    ticks_per_sec: usize,
+    pagesz: usize,
+    now: u64,
+}
+
+#[cfg(test)]
+impl MockFS {
+    pub fn new(
+        files: HashMap<String,String>,
+        pids: Vec<(usize,u32)>,
+        users: HashMap<u32,String>,
+        now: u64,
+    ) -> MockFS {
+        MockFS {
+            files,
+            pids,
+            users,
+            ticks_per_sec: 100,
+            pagesz: 4,
+            now,
+        }
+    }
+}
+
+#[cfg(test)]
+impl ProcfsAPI for MockFS {
+    fn read_to_string(&self, path: &str) -> Result<String, String> {
+        match self.files.get(path) {
+            Some(s) => Ok(s.clone()),
+            None    => Err(format!("Unable to read /proc/{path}")),
+        }
+    }
+
+    fn read_proc_pids(&self) -> Result<Vec<(usize,u32)>, String> {
+        Ok(self.pids.clone())
+    }
+
+    fn user_by_uid(&self, uid: u32) -> Option<String> {
+        match self.users.get(&uid) {
+            Some(s) => Some(s.clone()),
+            None => None
+        }
+    }
+
+    fn clock_ticks_per_sec(&self) -> usize {
+        self.ticks_per_sec
+    }
+
+    fn page_size(&self) -> usize {
+        self.pagesz
+    }
+
+    fn now(&self) -> u64 {
+        self.now
+    }
+}
+