pidof: Implement --with-workers

src/uu/pidof/src/pidof.rs

@@ -45,20 +45,22 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     Ok(())
 }
 
-fn get_executable_name(process: &mut ProcessInformation) -> String {
+fn match_process_name(
+    process: &mut ProcessInformation,
+    name_to_match: &str,
+    with_workers: bool,
+) -> bool {
     let binding = process.cmdline.split(' ').collect::<Vec<_>>();
     let mut path = binding.first().unwrap().to_string();
 
     if path.is_empty() {
+        if !with_workers {
+            return false;
+        }
         path.clone_from(&process.status()["Name"]);
     };
 
-    PathBuf::from(path)
-        .file_name()
-        .unwrap()
-        .to_str()
-        .unwrap()
-        .to_string()
+    PathBuf::from(path).file_name().unwrap().to_str().unwrap() == name_to_match
 }
 
 fn collect_matched_pids(matches: &ArgMatches) -> Vec<usize> {
@@ -67,6 +69,7 @@ fn collect_matched_pids(matches: &ArgMatches) -> Vec<usize> {
         .unwrap()
         .cloned()
         .collect();
+    let with_workers = matches.get_flag("with-workers");
 
     let collected = walk_process().collect::<Vec<_>>();
     let arg_omit_pid = matches
@@ -80,7 +83,7 @@ fn collect_matched_pids(matches: &ArgMatches) -> Vec<usize> {
         .flat_map(|program| {
             let mut processed = Vec::new();
             for mut process in collected.clone() {
-                let contains = program == get_executable_name(&mut process);
+                let contains = match_process_name(&mut process, &program, with_workers);
                 let should_omit = arg_omit_pid.contains(&process.pid);
 
                 if contains && !should_omit {
@@ -174,6 +177,13 @@ pub fn uu_app() -> Command {
                 .help("Show thread ids instead of process ids")
                 .action(ArgAction::SetTrue),
         )
+        .arg(
+            Arg::new("with-workers")
+                .short('w')
+                .long("with-workers")
+                .help("Show kernel worker threads as well")
+                .action(ArgAction::SetTrue),
+        )
     // .arg(
     //     Arg::new("x")
     //         .short('x')

tests/by-util/test_pidof.rs

@@ -24,8 +24,9 @@ fn test_find_init() {
 
 #[test]
 #[cfg(target_os = "linux")]
-fn test_find_kthreadd() {
-    new_ucmd!().arg("kthreadd").succeeds();
+fn test_find_kthreadd_only_with_w_flag() {
+    new_ucmd!().arg("kthreadd").fails();
+    new_ucmd!().arg("-w").arg("kthreadd").succeeds();
 }
 
 #[test]
@@ -37,15 +38,20 @@ fn test_no_pid_found() {
 #[test]
 #[cfg(target_os = "linux")]
 fn test_quiet() {
-    new_ucmd!().arg("kthreadd").arg("-q").succeeds().no_output();
+    new_ucmd!()
+        .arg("kthreadd")
+        .arg("-q")
+        .arg("-w")
+        .succeeds()
+        .no_output();
 }
 
 #[test]
 #[cfg(target_os = "linux")]
 fn test_single_shot() {
     for arg in ["-s", "--single-shot"] {
         let binding = new_ucmd!()
-            .args(&[arg, "kthreadd", "kthreadd", "kthreadd"])
+            .args(&[arg, "-w", "kthreadd", "kthreadd", "kthreadd"])
             .succeeds();
         let output = binding.stdout_str().trim_end();
 
@@ -63,7 +69,7 @@ fn test_single_shot() {
 #[cfg(target_os = "linux")]
 fn test_omit_pid() {
     for arg in ["-o=1000", "--omit-pid=1000"] {
-        new_ucmd!().arg(arg).arg("kthreadd").succeeds();
+        new_ucmd!().arg(arg).arg("-w").arg("kthreadd").succeeds();
     }
 }
 
@@ -76,7 +82,7 @@ fn test_separator() {
 
     for arg in ["-S", "-d", "--separator"] {
         new_ucmd!()
-            .args(&[arg, "separator", "kthreadd", "kthreadd"])
+            .args(&[arg, "separator", "-w", "kthreadd", "kthreadd"])
             .succeeds()
             .stdout_matches(re);
     }