Skip to content

pgrep, pkill: Unify process matching code #322

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cakebaker merged 4 commits into from
Feb 14, 2025
Merged

pgrep, pkill: Unify process matching code #322

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bd6e1f2
pgrep, pkill: Move regex to settings
dezgeg Feb 10, 2025
330d65b
pgrep, pkill: Move signal to settings
dezgeg Feb 10, 2025
021a2b7
pgrep, pkill: Extract function for finding all the processes
dezgeg Feb 10, 2025
ba00e9c
pgrep, pkill: Extract common process_matcher module
dezgeg Feb 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 5 additions & 231 deletions src/uu/pgrep/src/pgrep.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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") {
Expand Down Expand Up @@ -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())
Expand Down
Loading
Loading