From 967d288d8b26ad90c2f9196ef73d534d1400d8fc Mon Sep 17 00:00:00 2001
From: Tuomas Tynkkynen <tuomas.tynkkynen@iki.fi>
Date: Thu, 6 Feb 2025 20:29:42 +0200
Subject: [PATCH] pidof: Support '-t' flag

This flag makes pidof print thread ids of threads belonging to the
matching processes.
---
 src/uu/pgrep/src/process.rs | 47 +++++++++++++++++++++++++++++++++++++
 src/uu/pidof/src/pidof.rs   | 21 +++++++++++++----
 tests/by-util/test_pidof.rs | 24 +++++++++++++++++++
 3 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/src/uu/pgrep/src/process.rs b/src/uu/pgrep/src/process.rs
index 46fa407f..c207cd3b 100644
--- a/src/uu/pgrep/src/process.rs
+++ b/src/uu/pgrep/src/process.rs
@@ -186,6 +186,8 @@ pub struct ProcessInformation {
     cached_stat: Option<Rc<Vec<String>>>,
 
     cached_start_time: Option<u64>,
+
+    cached_thread_ids: Option<Rc<Vec<usize>>>,
 }
 
 impl ProcessInformation {
@@ -328,6 +330,31 @@ impl ProcessInformation {
 
         Teletype::Unknown
     }
+
+    pub fn thread_ids(&mut self) -> Rc<Vec<usize>> {
+        if let Some(c) = &self.cached_thread_ids {
+            return Rc::clone(c);
+        }
+
+        let tids_dir = format!("/proc/{}/task", self.pid);
+        let result = Rc::new(
+            WalkDir::new(tids_dir)
+                .min_depth(1)
+                .max_depth(1)
+                .follow_links(false)
+                .into_iter()
+                .flatten()
+                .flat_map(|it| {
+                    it.path()
+                        .file_name()
+                        .map(|it| it.to_str().unwrap().parse::<usize>().unwrap())
+                })
+                .collect::<Vec<_>>(),
+        );
+
+        self.cached_thread_ids = Some(Rc::clone(&result));
+        Rc::clone(&result)
+    }
 }
 impl TryFrom<DirEntry> for ProcessInformation {
     type Error = io::Error;
@@ -446,6 +473,26 @@ mod tests {
         assert!(result.contains(&pid_entry.tty()));
     }
 
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn test_thread_ids() {
+        let main_tid = unsafe { uucore::libc::gettid() };
+        std::thread::spawn(move || {
+            let mut pid_entry = ProcessInformation::try_new(
+                PathBuf::from_str(&format!("/proc/{}", current_pid())).unwrap(),
+            )
+            .unwrap();
+            let thread_ids = pid_entry.thread_ids();
+
+            assert!(thread_ids.contains(&(main_tid as usize)));
+
+            let new_thread_tid = unsafe { uucore::libc::gettid() };
+            assert!(thread_ids.contains(&(new_thread_tid as usize)));
+        })
+        .join()
+        .unwrap();
+    }
+
     #[test]
     fn test_stat_split() {
         let case = "32 (idle_inject/3) S 2 0 0 0 -1 69238848 0 0 0 0 0 0 0 0 -51 0 1 0 34 0 0 18446744073709551615 0 0 0 0 0 0 0 2147483647 0 0 0 0 17 3 50 1 0 0 0 0 0 0 0 0 0 0 0";
diff --git a/src/uu/pidof/src/pidof.rs b/src/uu/pidof/src/pidof.rs
index d11c686a..12dc8f7a 100644
--- a/src/uu/pidof/src/pidof.rs
+++ b/src/uu/pidof/src/pidof.rs
@@ -33,7 +33,7 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
 
     let output = collected
         .into_iter()
-        .map(|it| it.pid.to_string())
+        .map(|it| it.to_string())
         .collect::<Vec<_>>()
         .join(arg_separator);
 
@@ -61,7 +61,7 @@ fn get_executable_name(process: &mut ProcessInformation) -> String {
         .to_string()
 }
 
-fn collect_matched_pids(matches: &ArgMatches) -> Vec<ProcessInformation> {
+fn collect_matched_pids(matches: &ArgMatches) -> Vec<usize> {
     let program_names: Vec<_> = matches
         .get_many::<String>("program-name")
         .unwrap()
@@ -84,16 +84,20 @@ fn collect_matched_pids(matches: &ArgMatches) -> Vec<ProcessInformation> {
                 let should_omit = arg_omit_pid.contains(&process.pid);
 
                 if contains && !should_omit {
-                    processed.push(process);
+                    if matches.get_flag("t") {
+                        processed.extend_from_slice(&process.thread_ids());
+                    } else {
+                        processed.push(process.pid);
+                    }
                 }
             }
 
-            processed.sort_by(|a, b| b.pid.cmp(&a.pid));
+            processed.sort_by(|a, b| b.cmp(a));
 
             let flag_s = matches.get_flag("s");
             if flag_s {
                 match processed.first() {
-                    Some(first) => vec![first.clone()],
+                    Some(first) => vec![*first],
                     None => Vec::new(),
                 }
             } else {
@@ -164,6 +168,13 @@ pub fn uu_app() -> Command {
                 .help("Only return one PID")
                 .action(ArgAction::SetTrue),
         )
+        .arg(
+            Arg::new("t")
+                .short('t')
+                .long("lightweight")
+                .help("Show thread ids instead of process ids")
+                .action(ArgAction::SetTrue),
+        )
     // .arg(
     //     Arg::new("x")
     //         .short('x')
diff --git a/tests/by-util/test_pidof.rs b/tests/by-util/test_pidof.rs
index 53580699..31ebb753 100644
--- a/tests/by-util/test_pidof.rs
+++ b/tests/by-util/test_pidof.rs
@@ -81,3 +81,27 @@ fn test_separator() {
             .stdout_matches(re);
     }
 }
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_threads() {
+    let main_tid = unsafe { uucore::libc::gettid() };
+    std::thread::spawn(move || {
+        let argv0 = std::env::args().next().unwrap();
+        let our_name = std::path::Path::new(argv0.as_str())
+            .file_name()
+            .unwrap()
+            .to_str()
+            .unwrap();
+
+        let new_thread_tid = unsafe { uucore::libc::gettid() };
+        new_ucmd!()
+            .arg("-t")
+            .arg(our_name)
+            .succeeds()
+            .stdout_contains(main_tid.to_string())
+            .stdout_contains(new_thread_tid.to_string());
+    })
+    .join()
+    .unwrap();
+}