Merge pull request #172 from zidoshare/implementation-of-ps-commmand

Implementation of ps commmand

src/cgroups/common.rs

@@ -1,8 +1,8 @@
 use std::{
     env,
     fmt::{Debug, Display},
-    fs,
-    io::Write,
+    fs::{self, File},
+    io::{BufRead, BufReader, Write},
     path::{Path, PathBuf},
 };
 
@@ -36,6 +36,8 @@ pub trait CgroupManager {
     fn freeze(&self, state: FreezerState) -> Result<()>;
     /// Retrieve statistics for the cgroup
     fn stats(&self) -> Result<Stats>;
+    // Gets the PIDs inside the cgroup
+    fn get_all_pids(&self) -> Result<Vec<Pid>>;
 }
 
 #[derive(Debug)]
@@ -177,3 +179,37 @@ pub fn create_cgroup_manager<P: Into<PathBuf>>(
         _ => bail!("could not find cgroup filesystem"),
     }
 }
+
+pub fn get_all_pids(path: &Path) -> Result<Vec<Pid>> {
+    log::debug!("scan pids in folder: {:?}", path);
+    let mut result = vec![];
+    walk_dir(&path, &mut |p| {
+        let file_path = p.join(CGROUP_PROCS);
+        if file_path.exists() {
+            let file = File::open(file_path)?;
+            for line in BufReader::new(file).lines() {
+                if let Ok(line) = line {
+                    result.push(Pid::from_raw(line.parse::<i32>()?))
+                }
+            }
+        }
+        Ok(())
+    })?;
+    Ok(result)
+}
+
+fn walk_dir<F>(path: &Path, c: &mut F) -> Result<()>
+where
+    F: FnMut(&Path) -> Result<()>,
+{
+    c(&path)?;
+    for entry in fs::read_dir(path)? {
+        let entry = entry?;
+        let path = entry.path();
+
+        if path.is_dir() {
+            walk_dir(&path, c)?;
+        }
+    }
+    Ok(())
+}

src/cgroups/v1/manager.rs

@@ -16,7 +16,7 @@ use super::{
     perf_event::PerfEvent, pids::Pids, util, Controller,
 };
 
-use crate::cgroups::common::CGROUP_PROCS;
+use crate::cgroups::common::{self, CGROUP_PROCS};
 use crate::cgroups::stats::{Stats, StatsProvider};
 use crate::utils;
 use crate::{cgroups::common::CgroupManager, utils::PathBufExt};
@@ -101,6 +101,14 @@ impl Manager {
 }
 
 impl CgroupManager for Manager {
+    fn get_all_pids(&self) -> Result<Vec<Pid>> {
+        let devices = self.subsystems.get(&CtrlType::Devices);
+        if let Some(p) = devices {
+            common::get_all_pids(p)
+        } else {
+            bail!("subsystem does not exist")
+        }
+    }
     fn add_task(&self, pid: Pid) -> Result<()> {
         for subsys in &self.subsystems {
             match subsys.0 {

src/cgroups/v2/manager.rs

@@ -159,4 +159,8 @@ impl CgroupManager for Manager {
     fn stats(&self) -> Result<Stats> {
         Ok(Stats::default())
     }
+
+    fn get_all_pids(&self) -> Result<Vec<Pid>> {
+        common::get_all_pids(&self.full_path)
+    }
 }

src/cgroups/v2/systemd_manager.rs

@@ -259,6 +259,10 @@ impl CgroupManager for SystemDCGroupManager {
     fn stats(&self) -> Result<Stats> {
         Ok(Stats::default())
     }
+
+    fn get_all_pids(&self) -> Result<Vec<Pid>> {
+        common::get_all_pids(&self.full_path)
+    }
 }
 
 #[cfg(test)]

src/commands/mod.rs

@@ -6,6 +6,7 @@ pub mod info;
 pub mod kill;
 pub mod list;
 pub mod pause;
+pub mod ps;
 pub mod resume;
 pub mod run;
 pub mod spec_json;

src/commands/ps.rs

@@ -0,0 +1,86 @@
+use crate::{cgroups, container::Container, utils};
+use anyhow::{bail, Context, Result};
+use clap::{self, Clap};
+use std::{path::PathBuf, process::Command};
+
+/// display the processes inside a container
+#[derive(Clap, Debug)]
+pub struct Ps {
+    /// format to display processes: table or json (default: "table")
+    #[clap(short, long, default_value = "table")]
+    format: String,
+    pub container_id: String,
+    /// options will be passed to the ps utility
+    #[clap(setting = clap::ArgSettings::Last)]
+    ps_options: Vec<String>,
+}
+impl Ps {
+    pub fn exec(&self, root_path: PathBuf) -> Result<()> {
+        let container_root = root_path.join(&self.container_id);
+        if !container_root.exists() {
+            bail!("{} doesn't exist.", self.container_id)
+        }
+        let container = Container::load(container_root)?.refresh_status()?;
+        if container.root.exists() {
+            let config_absolute_path = container.root.join("config.json");
+            log::debug!("load spec from {:?}", config_absolute_path);
+            let spec = oci_spec::Spec::load(config_absolute_path)?;
+            log::debug!("spec: {:?}", spec);
+            let cgroups_path = utils::get_cgroup_path(
+                &spec.linux.context("no linux in spec")?.cgroups_path,
+                container.id(),
+            );
+            let systemd_cgroup = container
+                .systemd()
+                .context("could not determine cgroup manager")?;
+            let cmanager = cgroups::common::create_cgroup_manager(cgroups_path, systemd_cgroup)?;
+            let pids: Vec<i32> = cmanager
+                .get_all_pids()?
+                .iter()
+                .map(|pid| pid.as_raw())
+                .collect();
+
+            if self.format == "json" {
+                println!("{}", serde_json::to_string(&pids)?);
+            } else if self.format == "table" {
+                let default_ps_options = vec![String::from("-ef")];
+                let ps_options = if self.ps_options.is_empty() {
+                    &default_ps_options
+                } else {
+                    &self.ps_options
+                };
+                let output = Command::new("ps").args(ps_options).output()?;
+                if !output.status.success() {
+                    println!("{}", std::str::from_utf8(&output.stderr)?);
+                } else {
+                    let lines = std::str::from_utf8(&output.stdout)?;
+                    let lines: Vec<&str> = lines.split("\n").collect();
+                    let pid_index = get_pid_index(lines[0])?;
+                    println!("{}", &lines[0]);
+                    for line in &lines[1..] {
+                        if line.is_empty() {
+                            continue;
+                        }
+                        let fields: Vec<&str> = line.split_whitespace().collect();
+                        let pid: i32 = fields[pid_index].parse()?;
+                        if pids.contains(&pid) {
+                            println!("{}", line);
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+fn get_pid_index(title: &str) -> Result<usize> {
+    let titles = title.split_whitespace();
+
+    for (index, name) in titles.enumerate() {
+        if name == "PID" {
+            return Ok(index);
+        }
+    }
+    bail!("could't find PID field in ps output");
+}

src/main.rs

@@ -16,6 +16,7 @@ use youki::commands::info;
 use youki::commands::kill;
 use youki::commands::list;
 use youki::commands::pause;
+use youki::commands::ps;
 use youki::commands::resume;
 use youki::commands::run;
 use youki::commands::spec_json;
@@ -74,6 +75,8 @@ enum SubCommand {
     Resume(resume::Resume),
     #[clap(version = "0.0.1", author = "utam0k <k0ma@utam0k.jp>")]
     Events(events::Events),
+    #[clap(version = "0.0.1", author = "utam0k <k0ma@utam0k.jp>", setting=clap::AppSettings::AllowLeadingHyphen)]
+    Ps(ps::Ps),
 }
 
 /// This is the entry point in the container runtime. The binary is run by a high-level container runtime,
@@ -108,5 +111,6 @@ fn main() -> Result<()> {
         SubCommand::Pause(pause) => pause.exec(root_path, systemd_cgroup),
         SubCommand::Resume(resume) => resume.exec(root_path, systemd_cgroup),
         SubCommand::Events(events) => events.exec(root_path),
+        SubCommand::Ps(ps) => ps.exec(root_path),
     }
 }