Merge pull request #322 from dezgeg/pgrep_pkill_unification2

pgrep, pkill: Unify process matching code

Author: cakebaker

src/uu/pgrep/src/pgrep.rs

@@ -5,36 +5,14 @@
 
 // Pid utils
 pub mod process;
+pub mod process_matcher;
 
-use clap::{arg, crate_version, Arg, ArgAction, ArgGroup, ArgMatches, Command};
-use process::{walk_process, ProcessInformation, Teletype};
-use regex::Regex;
-use std::{collections::HashSet, sync::OnceLock};
-#[cfg(unix)]
-use uucore::{display::Quotable, signals::signal_by_name_or_value};
-use uucore::{
-    error::{UResult, USimpleError},
-    format_usage, help_about, help_usage,
-};
+use clap::{arg, crate_version, Arg, ArgAction, ArgGroup, Command};
+use uucore::{error::UResult, format_usage, help_about, help_usage};
 
 const ABOUT: &str = help_about!("pgrep.md");
 const USAGE: &str = help_usage!("pgrep.md");
 
-static REGEX: OnceLock<Regex> = OnceLock::new();
-
-struct Settings {
-    exact: bool,
-    full: bool,
-    ignore_case: bool,
-    inverse: bool,
-    newest: bool,
-    oldest: bool,
-    older: Option<u64>,
-    parent: Option<Vec<u64>>,
-    runstates: Option<String>,
-    terminal: Option<HashSet<Teletype>>,
-}
-
 /// # Conceptual model of `pgrep`
 ///
 /// At first, `pgrep` command will check the patterns is legal.
@@ -50,67 +28,10 @@ struct Settings {
 #[uucore::main]
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().try_get_matches_from(args)?;
-
-    let pattern = try_get_pattern_from(&matches)?;
-    REGEX
-        .set(Regex::new(&pattern).map_err(|e| USimpleError::new(2, e.to_string()))?)
-        .unwrap();
-
-    let settings = Settings {
-        exact: matches.get_flag("exact"),
-        full: matches.get_flag("full"),
-        ignore_case: matches.get_flag("ignore-case"),
-        inverse: matches.get_flag("inverse"),
-        newest: matches.get_flag("newest"),
-        oldest: matches.get_flag("oldest"),
-        parent: matches
-            .get_many::<u64>("parent")
-            .map(|parents| parents.copied().collect()),
-        runstates: matches.get_one::<String>("runstates").cloned(),
-        older: matches.get_one::<u64>("older").copied(),
-        terminal: matches.get_many::<String>("terminal").map(|ttys| {
-            ttys.cloned()
-                .flat_map(Teletype::try_from)
-                .collect::<HashSet<_>>()
-        }),
-    };
-
-    if (!settings.newest
-        && !settings.oldest
-        && settings.runstates.is_none()
-        && settings.older.is_none()
-        && settings.parent.is_none()
-        && settings.terminal.is_none())
-        && pattern.is_empty()
-    {
-        return Err(USimpleError::new(
-            2,
-            "no matching criteria specified\nTry `pgrep --help' for more information.",
-        ));
-    }
-
-    // Parse signal
-    #[cfg(unix)]
-    let sig_num = parse_signal_value(matches.get_one::<String>("signal").unwrap())?;
+    let settings = process_matcher::get_match_settings(&matches)?;
 
     // Collect pids
-    let pids = {
-        let mut pids = collect_matched_pids(&settings);
-        #[cfg(unix)]
-        if matches.get_flag("require-handler") {
-            pids.retain(|pid| {
-                let mask =
-                    u64::from_str_radix(pid.clone().status().get("SigCgt").unwrap(), 16).unwrap();
-                mask & (1 << sig_num) != 0
-            });
-        }
-        if pids.is_empty() {
-            uucore::error::set_exit_code(1);
-            pids
-        } else {
-            process_flag_o_n(&settings, &mut pids)
-        }
-    };
+    let pids = process_matcher::find_matching_pids(&settings);
 
     // Processing output
     let output = if matches.get_flag("count") {
@@ -148,153 +69,6 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     Ok(())
 }
 
-/// Try to get the pattern from the command line arguments. Returns an empty string if no pattern
-/// is specified.
-fn try_get_pattern_from(matches: &ArgMatches) -> UResult<String> {
-    let pattern = match matches.get_many::<String>("pattern") {
-        Some(patterns) if patterns.len() > 1 => {
-            return Err(USimpleError::new(
-                2,
-                "only one pattern can be provided\nTry `pgrep --help' for more information.",
-            ))
-        }
-        Some(mut patterns) => patterns.next().unwrap(),
-        None => return Ok(String::new()),
-    };
-
-    let pattern = if matches.get_flag("ignore-case") {
-        &pattern.to_lowercase()
-    } else {
-        pattern
-    };
-
-    let pattern = if matches.get_flag("exact") {
-        &format!("^{}$", pattern)
-    } else {
-        pattern
-    };
-
-    Ok(pattern.to_string())
-}
-
-/// Collect pids with filter construct from command line arguments
-fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
-    // Filtration general parameters
-    let filtered: Vec<ProcessInformation> = {
-        let mut tmp_vec = Vec::new();
-
-        for mut pid in walk_process().collect::<Vec<_>>() {
-            let run_state_matched = match (&settings.runstates, pid.run_state()) {
-                (Some(arg_run_states), Ok(pid_state)) => {
-                    arg_run_states.contains(&pid_state.to_string())
-                }
-                (_, Err(_)) => false,
-                _ => true,
-            };
-
-            let binding = pid.status();
-            let name = binding.get("Name").unwrap();
-            let name = if settings.ignore_case {
-                name.to_lowercase()
-            } else {
-                name.into()
-            };
-            let pattern_matched = {
-                let want = if settings.exact {
-                    // Equals `Name` in /proc/<pid>/status
-                    // The `unwrap` operation must succeed
-                    // because the REGEX has been verified as correct in `uumain`.
-                    &name
-                } else if settings.full {
-                    // Equals `cmdline` in /proc/<pid>/cmdline
-                    &pid.cmdline
-                } else {
-                    // From manpage:
-                    // The process name used for matching is limited to the 15 characters present in the output of /proc/pid/stat.
-                    &pid.proc_stat()[..15]
-                };
-
-                REGEX.get().unwrap().is_match(want)
-            };
-
-            let tty_matched = match &settings.terminal {
-                Some(ttys) => ttys.contains(&pid.tty()),
-                None => true,
-            };
-
-            let arg_older = settings.older.unwrap_or(0);
-            let older_matched = pid.start_time().unwrap() >= arg_older;
-
-            // the PPID is the fourth field in /proc/<PID>/stat
-            // (https://www.kernel.org/doc/html/latest/filesystems/proc.html#id10)
-            let stat = pid.stat();
-            let ppid = stat.get(3);
-            let parent_matched = match (&settings.parent, ppid) {
-                (Some(parents), Some(ppid)) => parents.contains(&ppid.parse::<u64>().unwrap()),
-                _ => true,
-            };
-
-            if (run_state_matched
-                && pattern_matched
-                && tty_matched
-                && older_matched
-                && parent_matched)
-                ^ settings.inverse
-            {
-                tmp_vec.push(pid);
-            }
-        }
-        tmp_vec
-    };
-
-    filtered
-}
-
-/// Sorting pids for flag `-o` and `-n`.
-///
-/// This function can also be used as a filter to filter out process information.
-fn process_flag_o_n(
-    settings: &Settings,
-    pids: &mut [ProcessInformation],
-) -> Vec<ProcessInformation> {
-    if settings.oldest || settings.newest {
-        pids.sort_by(|a, b| {
-            b.clone()
-                .start_time()
-                .unwrap()
-                .cmp(&a.clone().start_time().unwrap())
-        });
-
-        let start_time = if settings.newest {
-            pids.first().cloned().unwrap().start_time().unwrap()
-        } else {
-            pids.last().cloned().unwrap().start_time().unwrap()
-        };
-
-        // There might be some process start at same time, so need to be filtered.
-        let mut filtered = pids
-            .iter()
-            .filter(|it| (*it).clone().start_time().unwrap() == start_time)
-            .collect::<Vec<_>>();
-
-        if settings.newest {
-            filtered.sort_by(|a, b| b.pid.cmp(&a.pid));
-        } else {
-            filtered.sort_by(|a, b| a.pid.cmp(&b.pid));
-        }
-
-        vec![filtered.first().cloned().unwrap().clone()]
-    } else {
-        pids.to_vec()
-    }
-}
-
-#[cfg(unix)]
-fn parse_signal_value(signal_name: &str) -> UResult<usize> {
-    signal_by_name_or_value(signal_name)
-        .ok_or_else(|| USimpleError::new(1, format!("Unknown signal {}", signal_name.quote())))
-}
-
 #[allow(clippy::cognitive_complexity)]
 pub fn uu_app() -> Command {
     Command::new(uucore::util_name())