Merge pull request #413 from dezgeg/cgroup

pgrep/pkill/pidwait: Support --cgroup option

Author: cakebaker

src/uu/pgrep/src/process.rs

@@ -187,6 +187,37 @@ impl TryFrom<&String> for RunState {
     }
 }
 
+/// Represents an entry in `/proc/<pid>/cgroup`.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CgroupMembership {
+    pub hierarchy_id: u32,
+    pub controllers: Vec<String>,
+    pub cgroup_path: String,
+}
+
+impl TryFrom<&str> for CgroupMembership {
+    type Error = io::Error;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let parts: Vec<&str> = value.split(':').collect();
+        if parts.len() != 3 {
+            return Err(io::ErrorKind::InvalidData.into());
+        }
+
+        Ok(CgroupMembership {
+            hierarchy_id: parts[0]
+                .parse::<u32>()
+                .map_err(|_| io::ErrorKind::InvalidData)?,
+            controllers: if parts[1].is_empty() {
+                vec![]
+            } else {
+                parts[1].split(',').map(String::from).collect()
+            },
+            cgroup_path: parts[2].to_string(),
+        })
+    }
+}
+
 /// Process ID and its information
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct ProcessInformation {
@@ -383,6 +414,24 @@ impl ProcessInformation {
         read_link(format!("/proc/{}/root", self.pid))
     }
 
+    /// Returns cgroups (both v1 and v2) that the process belongs to.
+    pub fn cgroups(&mut self) -> Result<Vec<CgroupMembership>, io::Error> {
+        fs::read_to_string(format!("/proc/{}/cgroup", self.pid))?
+            .lines()
+            .map(CgroupMembership::try_from)
+            .collect()
+    }
+
+    /// Returns path to the v2 cgroup that the process belongs to.
+    pub fn cgroup_v2_path(&mut self) -> Result<String, io::Error> {
+        const V2_HIERARCHY_ID: u32 = 0;
+        self.cgroups()?
+            .iter()
+            .find(|cg| cg.hierarchy_id == V2_HIERARCHY_ID)
+            .map(|cg| cg.cgroup_path.clone())
+            .ok_or(io::ErrorKind::NotFound.into())
+    }
+
     /// Fetch run state from [ProcessInformation::cached_stat]
     ///
     /// - [The /proc Filesystem: Table 1-4](https://docs.kernel.org/filesystems/proc.html#id10)
@@ -625,4 +674,22 @@ mod tests {
         let mut pid_entry = ProcessInformation::current_process_info().unwrap();
         assert_eq!(pid_entry.root().unwrap(), PathBuf::from("/"));
     }
+
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn test_cgroups() {
+        let mut pid_entry = ProcessInformation::try_new("/proc/1".into()).unwrap();
+        if pid_entry.name().unwrap() == "systemd" {
+            let cgroups = pid_entry.cgroups().unwrap();
+            if let Some(membership) = cgroups.iter().find(|cg| cg.hierarchy_id == 0) {
+                let expected = CgroupMembership {
+                    hierarchy_id: 0,
+                    controllers: vec![],
+                    cgroup_path: "/init.scope".to_string(),
+                };
+                assert_eq!(membership, &expected);
+                assert_eq!(pid_entry.cgroup_v2_path().unwrap(), "/init.scope");
+            }
+        }
+    }
 }

src/uu/pgrep/src/process_matcher.rs

@@ -44,6 +44,7 @@ pub struct Settings {
     pub gid: Option<HashSet<u32>>,
     pub pgroup: Option<HashSet<u64>>,
     pub session: Option<HashSet<u64>>,
+    pub cgroup: Option<HashSet<String>>,
     pub threads: bool,
 }
 
@@ -101,6 +102,9 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
             })
             .collect()
         }),
+        cgroup: matches
+            .get_many::<String>("cgroup")
+            .map(|groups| groups.cloned().collect()),
         threads: false,
     };
 
@@ -115,6 +119,7 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
         && settings.gid.is_none()
         && settings.pgroup.is_none()
         && settings.session.is_none()
+        && settings.cgroup.is_none()
         && !settings.require_handler
         && pattern.is_empty()
     {
@@ -235,6 +240,10 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
             let parent_matched = any_matches(&settings.parent, pid.ppid().unwrap());
             let pgroup_matched = any_matches(&settings.pgroup, pid.pgid().unwrap());
             let session_matched = any_matches(&settings.session, pid.sid().unwrap());
+            let cgroup_matched = any_matches(
+                &settings.cgroup,
+                pid.cgroup_v2_path().unwrap_or("/".to_string()),
+            );
 
             let ids_matched = any_matches(&settings.uid, pid.uid().unwrap())
                 && any_matches(&settings.euid, pid.euid().unwrap())
@@ -265,6 +274,7 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
                 && parent_matched
                 && pgroup_matched
                 && session_matched
+                && cgroup_matched
                 && ids_matched
                 && handler_matched)
                 ^ settings.inverse
@@ -413,8 +423,7 @@ pub fn clap_args(pattern_help: &'static str, enable_v_flag: bool) -> Vec<Arg> {
         // arg!(-L --logpidfile           "fail if PID file is not locked"),
         arg!(-r --runstates <state>    "match runstates [D,S,Z,...]"),
         // arg!(-A --"ignore-ancestors"   "exclude our ancestors from results"),
-        // arg!(--cgroup <grp>            "match by cgroup v2 names")
-        //     .value_delimiter(','),
+        arg!(--cgroup <grp>            "match by cgroup v2 names").value_delimiter(','),
         // arg!(--ns <PID>                "match the processes that belong to the same namespace as <pid>"),
         // arg!(--nslist <ns>             "list which namespaces will be considered for the --ns option.")
         //     .value_delimiter(',')

tests/by-util/test_pgrep.rs

@@ -501,6 +501,29 @@ fn test_nonexisting_session() {
     new_ucmd!().arg("--session=9999999999").fails();
 }
 
+#[test]
+#[cfg(target_os = "linux")]
+fn test_nonexisting_cgroup() {
+    new_ucmd!()
+        .arg("--cgroup")
+        .arg("NONEXISTING")
+        .fails()
+        .stdout_is("");
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_cgroup() {
+    let init_is_systemd = new_ucmd!().arg("systemd").run().code() == 0;
+    if init_is_systemd {
+        new_ucmd!()
+            .arg("--cgroup")
+            .arg("/init.scope")
+            .succeeds()
+            .stdout_is("1\n");
+    }
+}
+
 #[test]
 #[cfg(target_os = "linux")]
 fn test_threads() {