Merge pull request #1473 from GuillaumeGomez/fix-tasks-update

Fix tasks update for linux processes

src/unix/linux/process.rs

@@ -508,7 +508,8 @@ fn update_existing_process(
     uptime: u64,
     info: &SystemInfo,
     refresh_kind: ProcessRefreshKind,
-) -> Result<(Option<Process>, Pid), ()> {
+    tasks: Option<HashSet<Pid>>,
+) -> Result<Option<Process>, ()> {
     let entry = &mut proc.inner;
     let data = if let Some(mut f) = entry.stat_file.take() {
         match get_all_data_from_file(&mut f, 1024) {
@@ -526,6 +527,8 @@ fn update_existing_process(
     } else {
         _get_stat_data(&entry.proc_path, &mut entry.stat_file)?
     };
+    entry.tasks = tasks;
+
     let parts = parse_stat_file(&data).ok_or(())?;
     let start_time_without_boot_time = compute_start_time_without_boot_time(&parts, info);
 
@@ -546,7 +549,7 @@ fn update_existing_process(
         );
 
         refresh_user_group_ids(entry, &mut proc_path, refresh_kind);
-        return Ok((None, entry.pid));
+        return Ok(None);
     }
     // If we're here, it means that the PID still exists but it's a different process.
     let p = retrieve_all_new_process_info(
@@ -560,9 +563,10 @@ fn update_existing_process(
     );
     *proc = p;
     // Since this PID is already in the HashMap, no need to add it again.
-    Ok((None, proc.inner.pid))
+    Ok(None)
 }
 
+#[allow(clippy::too_many_arguments)]
 pub(crate) fn _get_process_data(
     path: &Path,
     proc_list: &mut HashMap<Pid, Process>,
@@ -571,18 +575,20 @@ pub(crate) fn _get_process_data(
     uptime: u64,
     info: &SystemInfo,
     refresh_kind: ProcessRefreshKind,
-) -> Result<(Option<Process>, Pid), ()> {
+    tasks: Option<HashSet<Pid>>,
+) -> Result<Option<Process>, ()> {
     if let Some(ref mut entry) = proc_list.get_mut(&pid) {
-        return update_existing_process(entry, parent_pid, uptime, info, refresh_kind);
+        return update_existing_process(entry, parent_pid, uptime, info, refresh_kind, tasks);
     }
     let mut stat_file = None;
     let data = _get_stat_data(path, &mut stat_file)?;
     let parts = parse_stat_file(&data).ok_or(())?;
 
-    let mut p =
+    let mut new_process =
         retrieve_all_new_process_info(pid, parent_pid, &parts, path, info, refresh_kind, uptime);
-    p.inner.stat_file = stat_file;
-    Ok((Some(p), pid))
+    new_process.inner.stat_file = stat_file;
+    new_process.inner.tasks = tasks;
+    Ok(Some(new_process))
 }
 
 fn old_get_memory(entry: &mut ProcessInner, str_parts: &[&str], info: &SystemInfo) {
@@ -799,21 +805,20 @@ pub(crate) fn refresh_procs(
             .flatten()
             .filter(|e| filter_callback(e, filter))
             .filter_map(|e| {
-                let (mut p, _) = _get_process_data(
+                let proc_list = proc_list.get();
+                let new_process = _get_process_data(
                     e.path.as_path(),
-                    proc_list.get(),
+                    proc_list,
                     e.pid,
                     e.parent_pid,
                     uptime,
                     info,
                     refresh_kind,
+                    e.tasks,
                 )
                 .ok()?;
                 nb_updated.fetch_add(1, Ordering::Relaxed);
-                if let Some(ref mut p) = p {
-                    p.inner.tasks = e.tasks;
-                }
-                p
+                new_process
             })
             .collect::<Vec<_>>()
     };

tests/process.rs

@@ -116,7 +116,7 @@ fn test_environ() {
     s.refresh_processes_specifics(
         ProcessesToUpdate::Some(&[pid]),
         false,
-        sysinfo::ProcessRefreshKind::everything(),
+        ProcessRefreshKind::everything(),
     );
     p.kill().expect("Unable to kill process.");
 
@@ -374,7 +374,7 @@ fn test_refresh_process_doesnt_remove() {
 
     // Checks that the process is listed as it should.
     let mut s = System::new_with_specifics(
-        RefreshKind::nothing().with_processes(sysinfo::ProcessRefreshKind::nothing()),
+        RefreshKind::nothing().with_processes(ProcessRefreshKind::nothing()),
     );
     s.refresh_processes(ProcessesToUpdate::All, false);
 
@@ -649,7 +649,7 @@ fn test_process_iterator_lifetimes() {
     }
 
     let s = System::new_with_specifics(
-        sysinfo::RefreshKind::nothing().with_processes(sysinfo::ProcessRefreshKind::nothing()),
+        RefreshKind::nothing().with_processes(ProcessRefreshKind::nothing()),
     );
 
     let process: Option<&sysinfo::Process>;
@@ -731,8 +731,6 @@ fn test_process_creds() {
 // This test ensures that only the requested information is retrieved.
 #[test]
 fn test_process_specific_refresh() {
-    use sysinfo::{DiskUsage, ProcessRefreshKind};
-
     if !sysinfo::IS_SUPPORTED_SYSTEM || cfg!(feature = "apple-sandbox") {
         return;
     }
@@ -753,7 +751,7 @@ fn test_process_specific_refresh() {
         assert_eq!(p.memory(), 0);
         assert_eq!(p.virtual_memory(), 0);
         // These two won't be checked, too much lazyness in testing them...
-        assert_eq!(p.disk_usage(), DiskUsage::default());
+        assert_eq!(p.disk_usage(), sysinfo::DiskUsage::default());
         assert_eq!(p.cpu_usage(), 0.);
     }
 
@@ -1066,3 +1064,42 @@ fn test_exists() {
     s.refresh_processes_specifics(ProcessesToUpdate::Some(&[pid]), false, process_refresh_kind);
     assert!(!s.process(pid).unwrap().exists());
 }
+
+#[cfg(target_os = "linux")]
+#[test]
+fn test_tasks() {
+    use std::collections::HashSet;
+
+    fn get_tasks(system: &System) -> HashSet<Pid> {
+        let pid_to_process = system.processes();
+        let mut task_pids: HashSet<Pid> = HashSet::new();
+        for process in pid_to_process.values() {
+            if let Some(tasks) = process.tasks() {
+                task_pids.extend(tasks);
+            }
+        }
+        task_pids
+    }
+
+    let mut system = System::new_with_specifics(RefreshKind::nothing());
+    system.refresh_processes_specifics(ProcessesToUpdate::All, true, ProcessRefreshKind::nothing());
+
+    // Spawn a thread to increase the task count
+    let scheduler = std::thread::spawn(move || {
+        system.refresh_processes_specifics(
+            ProcessesToUpdate::All,
+            true,
+            ProcessRefreshKind::nothing(),
+        );
+
+        let mut system_new = System::new_with_specifics(RefreshKind::nothing());
+        system_new.refresh_processes_specifics(
+            ProcessesToUpdate::All,
+            true,
+            ProcessRefreshKind::nothing(),
+        );
+
+        assert_eq!(get_tasks(&system), get_tasks(&system_new));
+    });
+    scheduler.join().expect("Scheduler panicked");
+}