Actually implement showing all tasks (threads)

keeperofdakeys · keeperofdakeys · commit 6e32299952b1 · 2016-03-06T03:37:15.000+10:30
diff --git a/src/procrs/error.rs b/src/procrs/error.rs
@@ -29,7 +29,11 @@ pub enum ProcFile {
     /// /proc/[pid]/stat file, contains various stats about the process.
     PidStat,
     /// /proc/[pid]/cmdline file, contains the cmdline given when starting the process.
-    PidCmdline
+    PidCmdline,
+
+    // TODO: Attach a pid to this directory
+    /// /proc/[pid]/task directory, contains threads of a process.
+    PidTaskDir,
 }
 
 impl Error for ProcFile {
@@ -45,7 +49,8 @@ impl Error for ProcFile {
             ProcFile::PidDir => "/proc/[pid] directory",
             ProcFile::PidStatus => "/proc/[pid]/status file",
             ProcFile::PidStat => "/proc/[pid]/stat file",
-            ProcFile::PidCmdline => "/proc/[pid]/cmdline file"
+            ProcFile::PidCmdline => "/proc/[pid]/cmdline file",
+            ProcFile::PidTaskDir => "/proc/[pid]/task",
         }
     }
 
diff --git a/src/procrs/pid/mod.rs b/src/procrs/pid/mod.rs
@@ -2,6 +2,7 @@ use std::io;
 use std::io::prelude::*;
 use std::fs::{self, File, ReadDir, DirEntry};
 use std::path::Path;
+use std::vec;
 use std::io::BufReader;
 use std::cmp::Ordering;
 use std::str::FromStr;
@@ -16,42 +17,54 @@ use self::status::PidStatus;
 use error::{ProcError, ProcFile, ProcOper};
 use TaskId;
 
-fn err_str<T: ToString>(err: T) -> String {
-    err.to_string()
-}
-
 /// A struct containing information about a process.
 ///
 /// This struct contains information from various files inside the
 /// /proc/[pid] directory (for the respective pid).
 #[derive(Debug)]
 pub struct Pid {
+    /// The tid of this process
+    pub pid: TaskId,
     /// The /proc/[pid]/stat file
     pub stat: Box<PidStat>,
     /// The /proc/[pid]/status file
     pub status: Box<PidStatus>,
     /// The /proc/[pid]/cmdline file
-    pub cmdline: Vec<String>
+    pub cmdline: Vec<String>,
+    /// If this is a thread, this is set to true.
+    /// Threads will never have tasks attached.
+    is_thread: bool,
+    /// Vec of threads under /proc/[pid]/tasks/[tid]
+    threads: Option<Vec<Pid>>,
 }
 
 impl Pid {
+    /// Create a new Pid struct for a process, given a pid.
     pub fn new(pid: TaskId) -> Result<Self, ProcError> {
-        let proc_dir = format!("/proc/{}", pid);
+        let pid_dir = Path::new("/proc");
+        Self::new_dir(pid_dir, pid)
+    }
+
+    fn new_dir(proc_dir: &Path, pid: TaskId) -> Result<Self, ProcError> {
+        let proc_dir = proc_dir.join(pid.to_string());
         let pid_stat = try!(PidStat::new(&proc_dir));
         let pid_status = try!(PidStatus::new(&proc_dir));
         let cmdline = try!(Self::read_cmdline(&proc_dir));
 
-        Ok(Pid{
+        Ok(Pid {
+            pid: pid,
             stat: Box::new(pid_stat),
             status: Box::new(pid_status),
-            cmdline: cmdline
+            cmdline: cmdline,
+            is_thread: false,
+            threads: None,
         })
     }
 
     /// Given a /proc/[pid] directory, read the respective /proc/[pid]/cmdline
     /// file and return them in a Vec.
-    fn read_cmdline(proc_dir: &str) -> Result<Vec<String>, ProcError> {
-        File::open(Path::new(proc_dir).join("cmdline"))
+    fn read_cmdline(proc_dir: &Path) -> Result<Vec<String>, ProcError> {
+        File::open(proc_dir.join("cmdline"))
             .map_err(|e| ProcError::new_err(ProcOper::Opening, ProcFile::PidCmdline, e))
             .and_then(|file| {
                 let mut contents = Vec::new();
@@ -86,6 +99,27 @@ impl Pid {
             PidQuery::NoneQuery => true
         }
     }
+
+    pub fn tasks(&mut self) -> Option<Vec<Pid>> {
+        self.tasks_query(PidQuery::NoneQuery)
+    }
+
+    // TODO: Work out if this really should return Option<_>
+    // or Option<Result<Vec<Pid>>>. Otherwise the error is uncaught.
+    pub fn tasks_query(&self, query: PidQuery) -> Option<Vec<Pid>> {
+        if self.is_thread {
+            return None;
+        }
+
+        PidIter::new_tid_query(self.pid, query.clone()).unwrap()
+            .filter(|p| {
+                let query = query.clone();
+                match *p {
+                    Ok(ref pid) => pid.query(&query),
+                    Err(_) => true
+                }
+            }).collect::<Result<Vec<_>, _>>().ok()
+    }
 }
 
 impl PartialEq for Pid {
@@ -113,21 +147,42 @@ impl Ord for Pid {
 /// non-trivial.
 pub struct PidIter {
     dir_iter: ReadDir,
-    query: PidQuery
+    query: PidQuery,
 }
 
 impl PidIter {
     /// Create a new iterator over all processes in /proc.
-    pub fn new() -> Result<Self, String> {
+    pub fn new() -> Result<Self, ProcError> {
         Self::new_query(PidQuery::NoneQuery)
     }
 
     /// Create a new iterator over all processes in /proc, but only yield
     /// processes that match the given query.
-    pub fn new_query(query: PidQuery) -> Result<Self, String> {
+    pub fn new_query(query: PidQuery) -> Result<Self, ProcError> {
         let proc_dir = Path::new("/proc");
-        let dir_iter = try!(fs::read_dir(proc_dir).map_err(err_str));
-        Ok(PidIter{
+        let dir_iter = try!(
+            fs::read_dir(proc_dir)
+                .map_err(|e|
+                    ProcError::new(ProcOper::Opening, ProcFile::ProcDir, Some(e), Some("PidIter"))
+                )
+        );
+        Ok(PidIter {
+            dir_iter: dir_iter,
+            query: query,
+        })
+    }
+
+    fn new_tid_query(pid: TaskId, query: PidQuery) -> Result<Self, ProcError> {
+        let dir_name = format!("/proc/{}/task", pid);
+        let task_dir = Path::new(&dir_name);
+        let dir_iter = try!(
+            fs::read_dir(task_dir)
+                .map_err(|e|
+                    ProcError::new(ProcOper::Opening, ProcFile::PidTaskDir,
+                        Some(e), Some("PidIter"))
+                )
+        );
+        Ok(PidIter {
             dir_iter: dir_iter,
             query: query
         })
@@ -136,19 +191,21 @@ impl PidIter {
     /// Given a DirEntry, try to create a Pid struct, and only return if
     /// it matches the query, and is complete.
     fn proc_dir_filter(entry_opt: Result<DirEntry, io::Error>, query: &PidQuery)
-        -> Option<Result<Pid, String>> {
-        // TODO: This sucks, find a better way
+        -> Option<Result<Pid, ProcError>> {
         let file = entry_opt
-            .map_err(err_str)
+            .map_err(|e|
+                ProcError::new(ProcOper::Reading, ProcFile::ProcDir, Some(e), Some("PidIter"))
+            )
             .and_then(|entry|
                 entry.file_name().into_string()
-                    .or(Err("Error parsing filename".to_owned()))
+                    .or(Err(ProcError::new_more(ProcOper::Parsing, ProcFile::ProcDir, Some("PidIter"))))
             );
 
-        if file.is_err() {
-            return None;
+        if let Err(e) = file{
+            return Some(Err(e));
         }
 
+        // Ensure filename is an integer (skip if not)
         match file.unwrap().parse() {
             Ok(pid) => {
                 // If an error is not hard (error opening or reading file),
@@ -158,7 +215,7 @@ impl PidIter {
                     Ok(prc) => prc,
                     Err(e) => {
                         if e.is_hard() {
-                            return Some(Err(e).map_err(err_str));
+                            return Some(Err(e));
                         } else {
                             return None;
                         }
@@ -175,7 +232,7 @@ impl PidIter {
 }
 
 impl Iterator for PidIter {
-    type Item = Result<Pid, String>;
+    type Item = Result<Pid, ProcError>;
 
     fn next(&mut self) -> Option<Self::Item> {
         for entry in self.dir_iter.by_ref() {
@@ -187,12 +244,69 @@ impl Iterator for PidIter {
         None
     }
 
-    // Size may be anywhere from 0 to number of dirs
+    /// Size may be anywhere from 0 to number of dirs.
     fn size_hint(&self) -> (usize, Option<usize>) {
         (0, self.dir_iter.size_hint().1)
     }
 }
 
+/// An Iterator over threads of processes in the system.
+///
+/// If a task disappears while scanning it, the partial Pid struct
+/// will not be yielded. An atomic view of processes on the system seems
+/// non-trivial.
+pub struct TidIter {
+    pid_iter: PidIter,
+    task_iter: Option<vec::IntoIter<Pid>>,
+    query: PidQuery,
+}
+
+impl TidIter {
+    /// Create a new iterator over all tasks in /proc.
+    pub fn new() -> Result<Self, ProcError> {
+            println!("{:?}", 3);
+        Self::new_query(PidQuery::NoneQuery)
+    }
+
+    /// Create a new iterator over all tasks in /proc, but only yield
+    /// those that match the given query.
+    pub fn new_query(query: PidQuery) -> Result<Self, ProcError> {
+        Ok(TidIter{
+            pid_iter: try!(PidIter::new_query(query.clone())),
+            task_iter: None,
+            query: query,
+        })
+    }
+}
+
+impl Iterator for TidIter {
+    type Item = Result<Pid, ProcError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if self.task_iter.is_none() {
+                let pid = match self.pid_iter.next() {
+                    Some(Ok(pid)) => pid,
+                    Some(Err(e)) => { return Some(Err(e)) },
+                    None => { return None; }
+                };
+                let tasks_vec = pid.tasks_query(self.query.clone());
+                if let Some(vec) = tasks_vec {
+                    self.task_iter = Some(vec.into_iter());
+                }
+                continue;
+            } else {
+                let next = self.task_iter.as_mut().unwrap().next();
+                match next {
+                    Some(pid) => { return Some(Ok(pid)); },
+                    None => { self.task_iter = None; },
+                };
+            }
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
 /// A list of query types for process querying.
 pub enum PidQuery {
     /// Query by pid
diff --git a/src/procrs/pid/stat.rs b/src/procrs/pid/stat.rs
@@ -142,9 +142,9 @@ macro_rules! stat_parse_opt_num {
 
 impl PidStat {
     /// Generate PidStat struct given a process directory.
-    pub fn new(pid_dir: &str) -> Result<Self, ProcError> {
+    pub fn new(pid_dir: &Path) -> Result<Self, ProcError> {
         let file = try!(
-            File::open(Path::new(pid_dir).join("stat"))
+            File::open(pid_dir.join("stat"))
                 .map_err(|e|
                     ProcError::new_err(ProcOper::Opening, ProcFile::PidStat, e)
                 )
diff --git a/src/procrs/pid/status.rs b/src/procrs/pid/status.rs
@@ -68,10 +68,10 @@ macro_rules! extract_line_opt {
 
 impl PidStatus {
     // Generate PidStatus struct given a process directory
-    pub fn new(pid_dir: &str) -> Result<Self, ProcError> {
+    pub fn new(pid_dir: &Path) -> Result<Self, ProcError> {
         // Try opening file
         let status_file = try!(
-            File::open(Path::new(pid_dir).join("status"))
+            File::open(pid_dir.join("status"))
                 .map_err(|e| ProcError::new_err(ProcOper::Opening, ProcFile::PidStatus, e))
         );
 
diff --git a/src/psq/main.rs b/src/psq/main.rs
@@ -7,6 +7,7 @@ use prettytable::row::Row;
 use prettytable::format::FormatBuilder;
 use std::collections::HashMap;
 use std::iter::repeat;
+use std::cmp::Ordering;
 use procrs::pid::*;
 use procrs::TaskId;
 use argparse::{ArgumentParser, StoreTrue, Store};
@@ -16,9 +17,18 @@ fn main() {
     let (query, long, perf, verbose, tree, threads) =
         (opts.query, opts.long, opts.perf, opts.verbose, opts.tree, opts.threads);
 
-    let pid_iter = PidIter::new_query(query);
-    let mut pids = pid_iter.unwrap()
-        .collect::<Result<Vec<_>, _>>().unwrap();
+    let mut pids: Vec<_> = match threads {
+        false => {
+            PidIter::new_query(query)
+                .unwrap()
+                .collect::<Result<_, _>>().unwrap()
+        },
+        true => {
+            TidIter::new_query(query)
+                .unwrap()
+                .collect::<Result<_, _>>().unwrap()
+        }
+    };
 
     let mut name_indent = HashMap::new();
 
@@ -32,7 +42,16 @@ fn main() {
     if opts.tree {
         pids = treeify_names(pids, &mut name_indent);
     } else {
-        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
+        pids.sort_by(|p1, p2| 
+            match threads {
+                false => p1.stat.pid.cmp(&p2.stat.pid),
+                true => {
+                    let cmp = p1.status.tgid.cmp(&p2.status.tgid);
+                    if let Ordering::Equal = cmp { return Ordering::Equal; }
+                    p1.stat.pid.cmp(&p2.stat.pid)
+                }
+            }
+        );
     };
     // Assume hertz is 100.
     // TODO: Look this up via syscall (no /proc value for it)
@@ -71,10 +90,10 @@ fn main() {
                 }
             };
             row.push(cell!(p.stat.ppid));
+            if long {
+            }
             match (long, perf) {
-                (_, false) => {
-                    row.push(cell!(name));
-                },
+                (_, false) => {},
                 (_, true) => {
                     let rss = p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned());
                     let raw_time = p.stat.utime + p.stat.stime;
@@ -101,11 +120,12 @@ fn main() {
     if threads {
         titles.push(cell!("Tid"));
     }
+    titles.push(cell!("Ppid"));
     // TODO: Possible remove Ppid from when long is false,
     // and have Cmd/Args as separate columns for long.
     match (long, perf) {
         (_, false) =>
-            titles.extend_from_slice(&[cell!("Ppid"), cell!("Cmd")]),
+            titles.extend_from_slice(&[cell!("Cmd")]),
         (_, true) =>
             titles.extend_from_slice(&[cell!("RSS"), cell!("Time"), cell!("Cmd")])
     };
