Merge pull request #343 from dezgeg/pgrep-threads

pgrep: Support -w option

Author: cakebaker

src/uu/pgrep/src/pgrep.rs

@@ -28,7 +28,8 @@ const USAGE: &str = help_usage!("pgrep.md");
 #[uucore::main]
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().try_get_matches_from(args)?;
-    let settings = process_matcher::get_match_settings(&matches)?;
+    let mut settings = process_matcher::get_match_settings(&matches)?;
+    settings.threads = matches.get_flag("lightweight");
 
     // Collect pids
     let pids = process_matcher::find_matching_pids(&settings);
@@ -82,7 +83,7 @@ pub fn uu_app() -> Command {
                 .hide_default_value(true),
             arg!(-l     --"list-name"           "list PID and process name"),
             arg!(-a     --"list-full"           "list PID and full command line"),
-            // arg!(-w     --lightweight           "list all TID"),
+            arg!(-w     --lightweight           "list all TID"),
         ])
         .args(process_matcher::clap_args(
             "Name of the program to find the PID of",

src/uu/pgrep/src/process.rs

@@ -3,7 +3,9 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
+use regex::Regex;
 use std::hash::Hash;
+use std::sync::LazyLock;
 use std::{
     collections::HashMap,
     fmt::{self, Display, Formatter},
@@ -460,6 +462,22 @@ pub fn walk_process() -> impl Iterator<Item = ProcessInformation> {
         .flat_map(ProcessInformation::try_from)
 }
 
+static THREAD_REGEX: LazyLock<Regex> = LazyLock::new(|| {
+    Regex::new(r"^/proc/[0-9]+$|^/proc/[0-9]+/task$|^/proc/[0-9]+/task/[0-9]+$").unwrap()
+});
+
+pub fn walk_threads() -> impl Iterator<Item = ProcessInformation> {
+    WalkDir::new("/proc/")
+        .min_depth(1)
+        .max_depth(3)
+        .follow_links(false)
+        .into_iter()
+        .filter_entry(|e| THREAD_REGEX.is_match(e.path().as_os_str().to_string_lossy().as_ref()))
+        .flatten()
+        .filter(|it| it.path().as_os_str().to_string_lossy().contains("/task/"))
+        .flat_map(ProcessInformation::try_from)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/uu/pgrep/src/process_matcher.rs

@@ -21,7 +21,7 @@ use uucore::{
 
 use uucore::error::{UResult, USimpleError};
 
-use crate::process::{walk_process, ProcessInformation, Teletype};
+use crate::process::{walk_process, walk_threads, ProcessInformation, Teletype};
 
 pub struct Settings {
     pub regex: Regex,
@@ -44,6 +44,7 @@ pub struct Settings {
     pub gid: Option<HashSet<u32>>,
     pub pgroup: Option<HashSet<u64>>,
     pub session: Option<HashSet<u64>>,
+    pub threads: bool,
 }
 
 pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
@@ -100,6 +101,7 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
             })
             .collect()
         }),
+        threads: false,
     };
 
     if !settings.newest
@@ -194,8 +196,14 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
     let filtered: Vec<ProcessInformation> = {
         let mut tmp_vec = Vec::new();
 
+        let pids = if settings.threads {
+            walk_threads().collect::<Vec<_>>()
+        } else {
+            walk_process().collect::<Vec<_>>()
+        };
         let our_pid = std::process::id() as usize;
-        for mut pid in walk_process().collect::<Vec<_>>() {
+
+        for mut pid in pids {
             if pid.pid == our_pid {
                 continue;
             }

tests/by-util/test_pgrep.rs

@@ -478,3 +478,21 @@ fn test_session() {
 fn test_nonexisting_session() {
     new_ucmd!().arg("--session=9999999999").fails();
 }
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_threads() {
+    std::thread::Builder::new()
+        .name("PgrepTest".to_owned())
+        .spawn(|| {
+            let thread_tid = unsafe { uucore::libc::gettid() };
+            new_ucmd!()
+                .arg("-w")
+                .arg("PgrepTest")
+                .succeeds()
+                .stdout_contains(thread_tid.to_string());
+        })
+        .unwrap()
+        .join()
+        .unwrap();
+}