Repository: GuillaumeGomez/sysinfo
Pull request #1168 (2023-12-16 23:30:10)

Inline tasks at the same level as other processes
It allows to simplify the handling of processes in linux while also improving parallel computations.


Patch:
diff --git a/src/c_interface.rs b/src/c_interface.rs
index 1cf9ba9..0e4d48f 100644
--- a/src/c_interface.rs
+++ b/src/c_interface.rs
@@ -21,6 +21,9 @@ pub type CProcess = *const c_void;
 pub type RString = *const c_char;
 /// Callback used by [`processes`][crate::System#method.processes].
 pub type ProcessLoop = extern "C" fn(pid: PID, process: CProcess, data: *mut c_void) -> bool;
+#[cfg(target_os = "linux")]
+/// Callback used by [`tasks`][crate::Process#method.tasks].
+pub type ProcessPidLoop = extern "C" fn(pid: PID, data: *mut c_void) -> bool;
 /// Equivalent of [`Networks`][crate::Networks] struct.
 pub type CNetworks = *mut c_void;
 /// Equivalent of [`Disks`][crate::Disks] struct.
@@ -395,15 +398,15 @@ pub extern "C" fn sysinfo_process_by_pid(system: CSystem, pid: PID) -> CProcess
 #[no_mangle]
 pub extern "C" fn sysinfo_process_tasks(
     process: CProcess,
-    fn_pointer: Option<ProcessLoop>,
+    fn_pointer: Option<ProcessPidLoop>,
     data: *mut c_void,
 ) -> size_t {
     assert!(!process.is_null());
     if let Some(fn_pointer) = fn_pointer {
         unsafe {
             let process = process as *const Process;
-            for (pid, process) in (*process).tasks().iter() {
-                if !fn_pointer(pid.0, process as *const Process as CProcess, data) {
+            for pid in (*process).tasks().iter() {
+                if !fn_pointer(pid.0, data) {
                     break;
                 }
             }
diff --git a/src/common.rs b/src/common.rs
index 385cca5..9c9e95e 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -1212,7 +1212,7 @@ impl Process {
         any(target_os = "linux", target_os = "android"),
         not(feature = "unknown-ci")
     ))]
-    pub fn tasks(&self) -> &HashMap<Pid, Process> {
+    pub fn tasks(&self) -> &std::collections::HashSet<Pid> {
         &self.inner.tasks
     }
 }
diff --git a/src/sysinfo.h b/src/sysinfo.h
index f2d7955..e70b828 100644
--- a/src/sysinfo.h
+++ b/src/sysinfo.h
@@ -49,7 +49,7 @@ void        sysinfo_cpus_usage(CSystem system, unsigned int *length, float **cpu
 size_t      sysinfo_processes(CSystem system, bool (*fn_pointer)(PID, CProcess, void*),
                               void *data);
 #ifdef __linux__
-size_t      sysinfo_process_tasks(CProcess process, bool (*fn_pointer)(PID, CProcess, void*),
+size_t      sysinfo_process_tasks(CProcess process, bool (*fn_pointer)(PID, void*),
                                   void *data);
 #endif
 CProcess    sysinfo_process_by_pid(CSystem system, PID pid);
diff --git a/src/unix/apple/system.rs b/src/unix/apple/system.rs
index b036e7e..9dfce95 100644
--- a/src/unix/apple/system.rs
+++ b/src/unix/apple/system.rs
@@ -225,10 +225,8 @@ impl SystemInner {
             let port = self.port;
             time_interval = self.clock_info.as_mut().map(|c| c.get_time_interval(port));
         }
-        match {
-            let wrap = Wrap(UnsafeCell::new(&mut self.process_list));
-            update_process(&wrap, pid, time_interval, now, refresh_kind, true)
-        } {
+        let wrap = Wrap(UnsafeCell::new(&mut self.process_list));
+        match update_process(&wrap, pid, time_interval, now, refresh_kind, true) {
             Ok(Some(p)) => {
                 self.process_list.insert(p.pid(), p);
                 true
diff --git a/src/unix/linux/process.rs b/src/unix/linux/process.rs
index 1fa072a..5b9d43f 100644
--- a/src/unix/linux/process.rs
+++ b/src/unix/linux/process.rs
@@ -1,9 +1,10 @@
 // Take a look at the license at the top of the repository in the LICENSE file.
 
 use std::cell::UnsafeCell;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::ffi::OsStr;
 use std::fmt;
-use std::fs::{self, File};
+use std::fs::{self, DirEntry, File};
 use std::io::Read;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
@@ -14,7 +15,6 @@ use crate::sys::system::SystemInfo;
 use crate::sys::utils::{
     get_all_data, get_all_data_from_file, realpath, FileCounter, PathHandler, PathPush,
 };
-use crate::utils::into_iter;
 use crate::{DiskUsage, Gid, Pid, Process, ProcessRefreshKind, ProcessStatus, Signal, Uid};
 
 #[doc(hidden)]
@@ -111,7 +111,7 @@ pub(crate) struct ProcessInner {
     group_id: Option<Gid>,
     effective_group_id: Option<Gid>,
     pub(crate) status: ProcessStatus,
-    pub(crate) tasks: HashMap<Pid, Process>,
+    pub(crate) tasks: HashSet<Pid>,
     pub(crate) stat_file: Option<FileCounter>,
     old_read_bytes: u64,
     old_written_bytes: u64,
@@ -146,11 +146,7 @@ impl ProcessInner {
             group_id: None,
             effective_group_id: None,
             status: ProcessStatus::Unknown(0),
-            tasks: if pid.0 == 0 {
-                HashMap::with_capacity(1000)
-            } else {
-                HashMap::new()
-            },
+            tasks: HashSet::new(),
             stat_file: None,
             old_read_bytes: 0,
             old_written_bytes: 0,
@@ -286,17 +282,10 @@ pub(crate) fn compute_cpu_usage(p: &mut ProcessInner, total_time: f32, max_value
         / total_time
         * 100.)
         .min(max_value);
-
-    for task in p.tasks.values_mut() {
-        compute_cpu_usage(&mut task.inner, total_time, max_value);
-    }
 }
 
 pub(crate) fn unset_updated(p: &mut ProcessInner) {
     p.updated = false;
-    for task in p.tasks.values_mut() {
-        unset_updated(&mut task.inner);
-    }
 }
 
 pub(crate) fn set_time(p: &mut ProcessInner, utime: u64, stime: u64) {
@@ -433,7 +422,7 @@ fn update_proc_info(
 
 fn retrieve_all_new_process_info(
     pid: Pid,
-    proc_list: &ProcessInner,
+    parent_pid: Option<Pid>,
     parts: &[&str],
     path: &Path,
     info: &SystemInfo,
@@ -444,13 +433,12 @@ fn retrieve_all_new_process_info(
     let mut proc_path = PathHandler::new(path);
     let name = parts[ProcIndex::ShortExe as usize];
 
-    p.parent = if proc_list.pid.0 != 0 {
-        Some(proc_list.pid)
-    } else {
-        match Pid::from_str(parts[ProcIndex::ParentPid as usize]) {
+    p.parent = match parent_pid {
+        Some(parent_pid) if parent_pid.0 != 0 => Some(parent_pid),
+        _ => match Pid::from_str(parts[ProcIndex::ParentPid as usize]) {
             Ok(p) if p.0 != 0 => Some(p),
             _ => None,
-        }
+        },
     };
 
     p.start_time_without_boot_time = compute_start_time_without_boot_time(parts, info);
@@ -467,24 +455,15 @@ fn retrieve_all_new_process_info(
 
 pub(crate) fn _get_process_data(
     path: &Path,
-    proc_list: &mut ProcessInner,
-    parent_pid: Pid,
+    proc_list: &mut HashMap<Pid, Process>,
+    pid: Pid,
+    parent_pid: Option<Pid>,
     uptime: u64,
     info: &SystemInfo,
     refresh_kind: ProcessRefreshKind,
 ) -> Result<(Option<Process>, Pid), ()> {
-    let pid = match path.file_name().and_then(|x| x.to_str()).map(Pid::from_str) {
-        // If `pid` and `nb` are the same, it means the file is linking to itself so we skip it.
-        //
-        // It's because when reading `/proc/[PID]` folder, we then go through the folders inside it.
-        // Then, if we encounter a sub-folder with the same PID as the parent, then it's a link to
-        // the current folder we already did read so no need to do anything.
-        Some(Ok(nb)) if nb != parent_pid => nb,
-        _ => return Err(()),
-    };
-
     let data;
-    let parts = if let Some(ref mut entry) = proc_list.tasks.get_mut(&pid) {
+    let parts = if let Some(ref mut entry) = proc_list.get_mut(&pid) {
         let entry = &mut entry.inner;
         data = if let Some(mut f) = entry.stat_file.take() {
             match get_all_data_from_file(&mut f, 1024) {
@@ -522,15 +501,23 @@ pub(crate) fn _get_process_data(
         let data = _get_stat_data(path, &mut stat_file)?;
         let parts = parse_stat_file(&data).ok_or(())?;
 
-        let mut p =
-            retrieve_all_new_process_info(pid, proc_list, &parts, path, info, refresh_kind, uptime);
+        let mut p = retrieve_all_new_process_info(
+            pid,
+            parent_pid,
+            &parts,
+            path,
+            info,
+            refresh_kind,
+            uptime,
+        );
         p.inner.stat_file = stat_file;
         return Ok((Some(p), pid));
     };
 
     // If we're here, it means that the PID still exists but it's a different process.
-    let p = retrieve_all_new_process_info(pid, proc_list, &parts, path, info, refresh_kind, uptime);
-    match proc_list.tasks.get_mut(&pid) {
+    let p =
+        retrieve_all_new_process_info(pid, parent_pid, &parts, path, info, refresh_kind, uptime);
+    match proc_list.get_mut(&pid) {
         Some(ref mut entry) => **entry = p,
         // If it ever enters this case, it means that the process was removed from the HashMap
         // in-between with the usage of dark magic.
@@ -616,88 +603,118 @@ fn update_time_and_memory(
         );
         entry.run_time = uptime.saturating_sub(entry.start_time_without_boot_time);
     }
-    refresh_procs(
-        entry,
-        path.join("task"),
-        entry.pid,
-        uptime,
-        info,
-        refresh_kind,
-    );
+}
+
+struct ProcAndTasks {
+    pid: Pid,
+    path: PathBuf,
+    tasks: HashSet<Pid>,
+}
+
+fn get_all_pid_entries(
+    parent: Option<&OsStr>,
+    entry: DirEntry,
+    data: &mut Vec<ProcAndTasks>,
+) -> Option<Pid> {
+    let entry = entry.path();
+    let name = entry.file_name();
+
+    if name == parent {
+        // Needed because tasks have their own PID listed in the "task" folder.
+        return None;
+    }
+    let name = name?;
+    if !entry.is_dir() {
+        return None;
+    }
+    let pid = Pid::from(usize::from_str(&name.to_string_lossy()).ok()?);
+
+    let tasks_dir = Path::join(&entry, "task");
+    let mut tasks = HashSet::new();
+    if tasks_dir.is_dir() {
+        if let Ok(entries) = fs::read_dir(tasks_dir) {
+            for task in entries
+                .into_iter()
+                .filter_map(|entry| get_all_pid_entries(Some(name), entry.ok()?, data))
+            {
+                tasks.insert(task);
+            }
+        }
+    }
+    data.push(ProcAndTasks {
+        pid,
+        path: entry,
+        tasks,
+    });
+    Some(pid)
+}
+
+#[cfg(feature = "multithread")]
+#[inline]
+pub(crate) fn iter<T>(val: T) -> rayon::iter::IterBridge<T>
+where
+    T: rayon::iter::ParallelBridge,
+{
+    val.par_bridge()
+}
+
+#[cfg(not(feature = "multithread"))]
+#[inline]
+pub(crate) fn iter<T>(val: T) -> T
+where
+    T: Iterator,
+{
+    val
 }
 
 pub(crate) fn refresh_procs(
-    proc_list: &mut ProcessInner,
+    proc_list: &mut HashMap<Pid, Process>,
     path: &Path,
-    parent_pid: Pid,
     uptime: u64,
     info: &SystemInfo,
     refresh_kind: ProcessRefreshKind,
 ) -> bool {
-    let d = match fs::read_dir(path) {
-        Ok(d) => d,
-        Err(_) => return false,
-    };
-    let folders = d
-        .filter_map(|entry| {
-            let entry = entry.ok()?;
-            let entry = entry.path();
-
-            if entry.is_dir() {
-                Some(entry)
-            } else {
-                None
-            }
-        })
-        .collect::<Vec<_>>();
-    if parent_pid.0 == 0 {
+    #[cfg(feature = "multithread")]
+    use rayon::iter::ParallelIterator;
+
+    // FIXME: To prevent retrieving a task more than once (it can be listed in `/proc/[PID]/task`
+    // subfolder and directly in `/proc` at the same time), might be interesting to use a `HashSet`.
+    let procs = {
+        let d = match fs::read_dir(path) {
+            Ok(d) => d,
+            Err(_) => return false,
+        };
         let proc_list = Wrap(UnsafeCell::new(proc_list));
 
-        #[cfg(feature = "multithread")]
-        use rayon::iter::ParallelIterator;
-
-        into_iter(folders)
+        iter(d)
+            .map(|entry| {
+                let Ok(entry) = entry else { return Vec::new() };
+                let mut entries = Vec::new();
+                get_all_pid_entries(None, entry, &mut entries);
+                entries
+            })
+            .flatten()
             .filter_map(|e| {
-                let (p, _) = _get_process_data(
-                    e.as_path(),
+                let (mut p, _) = _get_process_data(
+                    e.path.as_path(),
                     proc_list.get(),
-                    parent_pid,
+                    e.pid,
+                    None,
                     uptime,
                     info,
                     refresh_kind,
                 )
                 .ok()?;
+                if let Some(ref mut p) = p {
+                    p.inner.tasks = e.tasks;
+                }
                 p
             })
             .collect::<Vec<_>>()
-    } else {
-        let mut updated_pids = Vec::with_capacity(folders.len());
-        let new_tasks = folders
-            .iter()
-            .filter_map(|e| {
-                let (p, pid) = _get_process_data(
-                    e.as_path(),
-                    proc_list,
-                    parent_pid,
-                    uptime,
-                    info,
-                    refresh_kind,
-                )
-                .ok()?;
-                updated_pids.push(pid);
-                p
-            })
-            .collect::<Vec<_>>();
-        // Sub-tasks are not cleaned up outside so we do it here directly.
-        proc_list
-            .tasks
-            .retain(|&pid, _| updated_pids.iter().any(|&x| x == pid));
-        new_tasks
-    }
-    .into_iter()
-    .for_each(|e| {
-        proc_list.tasks.insert(e.pid(), e);
-    });
+    };
+    for proc_ in procs {
+        proc_list.insert(proc_.pid(), proc_);
+    }
     true
 }
 
diff --git a/src/unix/linux/system.rs b/src/unix/linux/system.rs
index 9b0a3d9..ca87e05 100644
--- a/src/unix/linux/system.rs
+++ b/src/unix/linux/system.rs
@@ -3,7 +3,7 @@
 use crate::sys::cpu::{get_physical_core_count, CpusWrapper};
 use crate::sys::process::{_get_process_data, compute_cpu_usage, refresh_procs, unset_updated};
 use crate::sys::utils::{get_all_data, to_u64};
-use crate::{Cpu, CpuRefreshKind, LoadAvg, Pid, Process, ProcessInner, ProcessRefreshKind};
+use crate::{Cpu, CpuRefreshKind, LoadAvg, Pid, Process, ProcessRefreshKind};
 
 use libc::{self, c_char, sysconf, _SC_CLK_TCK, _SC_HOST_NAME_MAX, _SC_PAGESIZE};
 use std::cmp::min;
@@ -106,7 +106,7 @@ impl SystemInfo {
 }
 
 pub(crate) struct SystemInner {
-    process_list: Process,
+    process_list: HashMap<Pid, Process>,
     mem_total: u64,
     mem_free: u64,
     mem_available: u64,
@@ -151,7 +151,7 @@ impl SystemInner {
             (0., false, 0.)
         };
 
-        self.process_list.inner.tasks.retain(|_, proc_| {
+        self.process_list.retain(|_, proc_| {
             let proc_ = &mut proc_.inner;
             if !proc_.updated {
                 return false;
@@ -172,9 +172,7 @@ impl SystemInner {
 impl SystemInner {
     pub(crate) fn new() -> Self {
         Self {
-            process_list: Process {
-                inner: ProcessInner::new(Pid(0)),
-            },
+            process_list: HashMap::new(),
             mem_total: 0,
             mem_free: 0,
             mem_available: 0,
@@ -235,9 +233,8 @@ impl SystemInner {
     pub(crate) fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
         let uptime = self.uptime();
         refresh_procs(
-            &mut self.process_list.inner,
+            &mut self.process_list,
             Path::new("/proc"),
-            Pid(0),
             uptime,
             &self.info,
             refresh_kind,
@@ -254,14 +251,15 @@ impl SystemInner {
         let uptime = self.uptime();
         match _get_process_data(
             &Path::new("/proc/").join(pid.to_string()),
-            &mut self.process_list.inner,
-            Pid(0),
+            &mut self.process_list,
+            pid,
+            None,
             uptime,
             &self.info,
             refresh_kind,
         ) {
             Ok((Some(p), pid)) => {
-                self.process_list.inner.tasks.insert(pid, p);
+                self.process_list.insert(pid, p);
             }
             Ok(_) => {}
             Err(_e) => {
@@ -281,12 +279,12 @@ impl SystemInner {
             let total_time = total_time / self.cpus.len() as f32;
 
             let max_cpu_usage = self.get_max_process_cpu_usage();
-            if let Some(p) = self.process_list.inner.tasks.get_mut(&pid) {
+            if let Some(p) = self.process_list.get_mut(&pid) {
                 let p = &mut p.inner;
                 compute_cpu_usage(p, total_time, max_cpu_usage);
                 unset_updated(p);
             }
-        } else if let Some(p) = self.process_list.inner.tasks.get_mut(&pid) {
+        } else if let Some(p) = self.process_list.get_mut(&pid) {
             unset_updated(&mut p.inner);
         }
         true
@@ -297,11 +295,11 @@ impl SystemInner {
     // Need to be moved into a "common" file to avoid duplication.
 
     pub(crate) fn processes(&self) -> &HashMap<Pid, Process> {
-        &self.process_list.inner.tasks
+        &self.process_list
     }
 
     pub(crate) fn process(&self, pid: Pid) -> Option<&Process> {
-        self.process_list.inner.tasks.get(&pid)
+        self.process_list.get(&pid)
     }
 
     pub(crate) fn global_cpu_info(&self) -> &Cpu {
diff --git a/src/utils.rs b/src/utils.rs
index af61a82..41fd5dd 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -4,13 +4,7 @@
 /// Uses the `rayon::iter::IntoParallelIterator` trait.
 #[cfg(all(
     all(
-        any(
-            target_os = "linux",
-            target_os = "android",
-            target_os = "macos",
-            target_os = "windows",
-            target_os = "freebsd",
-        ),
+        any(target_os = "macos", target_os = "windows", target_os = "freebsd",),
         feature = "multithread"
     ),
     not(all(target_os = "macos", feature = "apple-sandbox")),
@@ -27,13 +21,7 @@ where
 /// Uses the `std::iter::IntoIterator` trait.
 #[cfg(all(
     all(
-        any(
-            target_os = "linux",
-            target_os = "android",
-            target_os = "macos",
-            target_os = "windows",
-            target_os = "freebsd",
-        ),
+        any(target_os = "macos", target_os = "windows", target_os = "freebsd",),
         not(feature = "multithread")
     ),
     not(feature = "unknown-ci"),


Test patch:
diff --git a/tests/process.rs b/tests/process.rs
index 76b516a..ad65733 100644
--- a/tests/process.rs
+++ b/tests/process.rs
@@ -393,24 +393,22 @@ fn test_refresh_tasks() {
     let mut s = System::new();
     s.refresh_processes();
 
-    assert!(s
-        .process(pid)
-        .unwrap()
-        .tasks()
-        .values()
-        .any(|t| t.name() == task_name));
+    assert!(s.process(pid).unwrap().tasks().iter().any(|task_pid| s
+        .process(*task_pid)
+        .map(|task| task.name() == task_name)
+        .unwrap_or(false)));
+    assert!(s.processes_by_exact_name(task_name).next().is_some());
 
     // Let's give some time to the system to clean up...
     std::thread::sleep(std::time::Duration::from_secs(2));
 
     s.refresh_processes();
 
-    assert!(!s
-        .process(pid)
-        .unwrap()
-        .tasks()
-        .values()
-        .any(|t| t.name() == task_name));
+    assert!(!s.process(pid).unwrap().tasks().iter().any(|task_pid| s
+        .process(*task_pid)
+        .map(|task| task.name() == task_name)
+        .unwrap_or(false)));
+    assert!(s.processes_by_exact_name(task_name).next().is_none());
 }
 
 // Checks that `refresh_process` is NOT removing dead processes.