diff --git a/src/uu/pgrep/src/pgrep.rs b/src/uu/pgrep/src/pgrep.rs index c27cf3b..116674a 100644 --- a/src/uu/pgrep/src/pgrep.rs +++ b/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", diff --git a/src/uu/pgrep/src/process.rs b/src/uu/pgrep/src/process.rs index 85d7bc5..619a413 100644 --- a/src/uu/pgrep/src/process.rs +++ b/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 { .flat_map(ProcessInformation::try_from) } +static THREAD_REGEX: LazyLock = 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 { + 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::*; diff --git a/src/uu/pgrep/src/process_matcher.rs b/src/uu/pgrep/src/process_matcher.rs index 615a01a..9394dfb 100644 --- a/src/uu/pgrep/src/process_matcher.rs +++ b/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>, pub pgroup: Option>, pub session: Option>, + pub threads: bool, } pub fn get_match_settings(matches: &ArgMatches) -> UResult { @@ -100,6 +101,7 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult { }) .collect() }), + threads: false, }; if !settings.newest @@ -194,8 +196,14 @@ fn collect_matched_pids(settings: &Settings) -> Vec { let filtered: Vec = { let mut tmp_vec = Vec::new(); + let pids = if settings.threads { + walk_threads().collect::>() + } else { + walk_process().collect::>() + }; let our_pid = std::process::id() as usize; - for mut pid in walk_process().collect::>() { + + for mut pid in pids { if pid.pid == our_pid { continue; } diff --git a/tests/by-util/test_pgrep.rs b/tests/by-util/test_pgrep.rs index b03c6de..fca2b42 100644 --- a/tests/by-util/test_pgrep.rs +++ b/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(); +}