Merge pull request #1170 from GuillaumeGomez/refresh-pids

Add `System::refresh_pids` and `System::refresh_pids_specifics` methods

src/common.rs

@@ -265,7 +265,65 @@ impl System {
     /// s.refresh_processes_specifics(ProcessRefreshKind::new());
     /// ```
     pub fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
-        self.inner.refresh_processes_specifics(refresh_kind)
+        self.inner.refresh_processes_specifics(None, refresh_kind)
+    }
+
+    /// Gets specified processes and updates their information.
+    ///
+    /// It does the same as:
+    ///
+    /// ```no_run
+    /// # use sysinfo::{Pid, ProcessRefreshKind, System, UpdateKind};
+    /// # let mut system = System::new();
+    /// system.refresh_pids_specifics(
+    ///     &[Pid::from(1), Pid::from(2)],
+    ///     ProcessRefreshKind::new()
+    ///         .with_memory()
+    ///         .with_cpu()
+    ///         .with_disk_usage()
+    ///         .with_exe(UpdateKind::OnlyIfNotSet),
+    /// );
+    /// ```
+    ///
+    /// ⚠️ On Linux, `sysinfo` keeps the `stat` files open by default. You can change this behaviour
+    /// by using [`set_open_files_limit`][crate::set_open_files_limit].
+    ///
+    /// Example:
+    ///
+    /// ```no_run
+    /// use sysinfo::System;
+    ///
+    /// let mut s = System::new_all();
+    /// s.refresh_processes();
+    /// ```
+    pub fn refresh_pids(&mut self, pids: &[Pid]) {
+        self.refresh_pids_specifics(
+            pids,
+            ProcessRefreshKind::new()
+                .with_memory()
+                .with_cpu()
+                .with_disk_usage()
+                .with_exe(UpdateKind::OnlyIfNotSet),
+        );
+    }
+
+    /// Gets specified processes and updates the specified information.
+    ///
+    /// ⚠️ On Linux, `sysinfo` keeps the `stat` files open by default. You can change this behaviour
+    /// by using [`set_open_files_limit`][crate::set_open_files_limit].
+    ///
+    /// ```no_run
+    /// use sysinfo::{Pid, ProcessRefreshKind, System};
+    ///
+    /// let mut s = System::new_all();
+    /// s.refresh_pids_specifics(&[Pid::from(1), Pid::from(2)], ProcessRefreshKind::new());
+    /// ```
+    pub fn refresh_pids_specifics(&mut self, pids: &[Pid], refresh_kind: ProcessRefreshKind) {
+        if pids.is_empty() {
+            return;
+        }
+        self.inner
+            .refresh_processes_specifics(Some(pids), refresh_kind)
     }
 
     /// Refreshes *only* the process corresponding to `pid`. Returns `false` if the process doesn't

src/unix/apple/system.rs

@@ -164,10 +164,19 @@ impl SystemInner {
     }
 
     #[cfg(any(target_os = "ios", feature = "apple-sandbox"))]
-    pub(crate) fn refresh_processes_specifics(&mut self, _refresh_kind: ProcessRefreshKind) {}
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        _filter: Option<&[Pid]>,
+        _refresh_kind: ProcessRefreshKind,
+    ) {
+    }
 
     #[cfg(all(target_os = "macos", not(feature = "apple-sandbox")))]
-    pub(crate) fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        filter: Option<&[Pid]>,
+        refresh_kind: ProcessRefreshKind,
+    ) {
         use crate::utils::into_iter;
 
         unsafe {
@@ -177,6 +186,26 @@ impl SystemInner {
             }
         }
         if let Some(pids) = get_proc_list() {
+            #[inline(always)]
+            fn real_filter(e: Pid, filter: &[Pid]) -> bool {
+                filter.contains(&e)
+            }
+
+            #[inline(always)]
+            fn empty_filter(_e: Pid, _filter: &[Pid]) -> bool {
+                true
+            }
+
+            #[allow(clippy::type_complexity)]
+            let (filter, filter_callback): (
+                &[Pid],
+                &(dyn Fn(Pid, &[Pid]) -> bool + Sync + Send),
+            ) = if let Some(filter) = filter {
+                (filter, &real_filter)
+            } else {
+                (&[], &empty_filter)
+            };
+
             let now = get_now();
             let port = self.port;
             let time_interval = self.clock_info.as_mut().map(|c| c.get_time_interval(port));
@@ -188,6 +217,9 @@ impl SystemInner {
 
                 into_iter(pids)
                     .flat_map(|pid| {
+                        if !filter_callback(pid, filter) {
+                            return None;
+                        }
                         match update_process(wrap, pid, time_interval, now, refresh_kind, false) {
                             Ok(x) => x,
                             _ => None,

src/unix/freebsd/system.rs

@@ -64,8 +64,12 @@ impl SystemInner {
         self.cpus.refresh(refresh_kind)
     }
 
-    pub(crate) fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
-        unsafe { self.refresh_procs(refresh_kind) }
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        filter: Option<&[Pid]>,
+        refresh_kind: ProcessRefreshKind,
+    ) {
+        unsafe { self.refresh_procs(filter, refresh_kind) }
     }
 
     pub(crate) fn refresh_process_specifics(
@@ -250,7 +254,7 @@ impl SystemInner {
 }
 
 impl SystemInner {
-    unsafe fn refresh_procs(&mut self, refresh_kind: ProcessRefreshKind) {
+    unsafe fn refresh_procs(&mut self, filter: Option<&[Pid]>, refresh_kind: ProcessRefreshKind) {
         let mut count = 0;
         let kvm_procs = libc::kvm_getprocs(
             self.system_info.kd.as_ptr(),
@@ -263,6 +267,26 @@ impl SystemInner {
             return;
         }
 
+        #[inline(always)]
+        fn real_filter(e: &libc::kinfo_proc, filter: &[Pid]) -> bool {
+            filter.contains(&Pid(e.ki_pid))
+        }
+
+        #[inline(always)]
+        fn empty_filter(_e: &libc::kinfo_proc, _filter: &[Pid]) -> bool {
+            true
+        }
+
+        #[allow(clippy::type_complexity)]
+        let (filter, filter_callback): (
+            &[Pid],
+            &(dyn Fn(&libc::kinfo_proc, &[Pid]) -> bool + Sync + Send),
+        ) = if let Some(filter) = filter {
+            (filter, &real_filter)
+        } else {
+            (&[], &empty_filter)
+        };
+
         let new_processes = {
             #[cfg(feature = "multithread")]
             use rayon::iter::{ParallelIterator, ParallelIterator as IterTrait};
@@ -278,6 +302,9 @@ impl SystemInner {
             let proc_list = utils::WrapMap(UnsafeCell::new(&mut self.process_list));
 
             IterTrait::filter_map(crate::utils::into_iter(kvm_procs), |kproc| {
+                if !filter_callback(kproc, filter) {
+                    return None;
+                }
                 super::process::get_process_data(
                     kproc,
                     &proc_list,

src/unix/linux/process.rs

@@ -675,11 +675,32 @@ pub(crate) fn refresh_procs(
     path: &Path,
     uptime: u64,
     info: &SystemInfo,
+    filter: Option<&[Pid]>,
     refresh_kind: ProcessRefreshKind,
 ) -> bool {
     #[cfg(feature = "multithread")]
     use rayon::iter::ParallelIterator;
 
+    #[inline(always)]
+    fn real_filter(e: &ProcAndTasks, filter: &[Pid]) -> bool {
+        filter.contains(&e.pid)
+    }
+
+    #[inline(always)]
+    fn empty_filter(_e: &ProcAndTasks, _filter: &[Pid]) -> bool {
+        true
+    }
+
+    #[allow(clippy::type_complexity)]
+    let (filter, filter_callback): (
+        &[Pid],
+        &(dyn Fn(&ProcAndTasks, &[Pid]) -> bool + Sync + Send),
+    ) = if let Some(filter) = filter {
+        (filter, &real_filter)
+    } else {
+        (&[], &empty_filter)
+    };
+
     // 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 = {
@@ -697,6 +718,7 @@ pub(crate) fn refresh_procs(
                 entries
             })
             .flatten()
+            .filter(|e| filter_callback(e, filter))
             .filter_map(|e| {
                 let (mut p, _) = _get_process_data(
                     e.path.as_path(),

src/unix/linux/system.rs

@@ -230,13 +230,18 @@ impl SystemInner {
         self.refresh_cpus(false, refresh_kind);
     }
 
-    pub(crate) fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        filter: Option<&[Pid]>,
+        refresh_kind: ProcessRefreshKind,
+    ) {
         let uptime = self.uptime();
         refresh_procs(
             &mut self.process_list,
             Path::new("/proc"),
             uptime,
             &self.info,
+            filter,
             refresh_kind,
         );
         self.clear_procs(refresh_kind);

src/unknown/system.rs

@@ -27,7 +27,12 @@ impl SystemInner {
 
     pub(crate) fn refresh_cpu_specifics(&mut self, _refresh_kind: CpuRefreshKind) {}
 
-    pub(crate) fn refresh_processes_specifics(&mut self, _refresh_kind: ProcessRefreshKind) {}
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        _filter: Option<&[Pid]>,
+        _refresh_kind: ProcessRefreshKind,
+    ) {
+    }
 
     pub(crate) fn refresh_process_specifics(
         &mut self,

src/windows/system.rs

@@ -196,7 +196,11 @@ impl SystemInner {
     }
 
     #[allow(clippy::cast_ptr_alignment)]
-    pub(crate) fn refresh_processes_specifics(&mut self, refresh_kind: ProcessRefreshKind) {
+    pub(crate) fn refresh_processes_specifics(
+        &mut self,
+        filter: Option<&[Pid]>,
+        refresh_kind: ProcessRefreshKind,
+    ) {
         // Windows 10 notebook requires at least 512KiB of memory to make it in one go
         let mut buffer_size = 512 * 1024;
         let mut process_information: Vec<u8> = Vec::with_capacity(buffer_size);
@@ -239,6 +243,26 @@ impl SystemInner {
                 }
             }
 
+            #[inline(always)]
+            fn real_filter(e: Pid, filter: &[Pid]) -> bool {
+                filter.contains(&e)
+            }
+
+            #[inline(always)]
+            fn empty_filter(_e: Pid, _filter: &[Pid]) -> bool {
+                true
+            }
+
+            #[allow(clippy::type_complexity)]
+            let (filter, filter_callback): (
+                &[Pid],
+                &(dyn Fn(Pid, &[Pid]) -> bool + Sync + Send),
+            ) = if let Some(filter) = filter {
+                (filter, &real_filter)
+            } else {
+                (&[], &empty_filter)
+            };
+
             // If we reach this point NtQuerySystemInformation succeeded
             // and the buffer contents are initialized
             process_information.set_len(buffer_size);
@@ -252,7 +276,9 @@ impl SystemInner {
                     .offset(process_information_offset)
                     as *const SYSTEM_PROCESS_INFORMATION;
 
-                process_ids.push(Wrap(p));
+                if filter_callback(Pid((*p).UniqueProcessId as _), filter) {
+                    process_ids.push(Wrap(p));
+                }
 
                 // read_unaligned is necessary to avoid
                 // misaligned pointer dereference: address must be a multiple of 0x8 but is 0x...

tests/process.rs

@@ -474,7 +474,7 @@ fn test_wait_child() {
         std::process::Command::new("waitfor")
             .arg("/t")
             .arg("300")
-            .arg("RefreshProcess")
+            .arg("WaitChild")
             .stdout(std::process::Stdio::null())
             .spawn()
             .unwrap()
@@ -736,3 +736,39 @@ fn test_process_specific_refresh() {
     update_specific_and_check!(exe, with_exe, , None);
     update_specific_and_check!(cwd, with_cwd, , None);
 }
+
+#[test]
+fn test_refresh_pids() {
+    if !sysinfo::IS_SUPPORTED || cfg!(feature = "apple-sandbox") {
+        return;
+    }
+    let self_pid = sysinfo::get_current_pid().expect("failed to get current pid");
+    let mut s = System::new();
+
+    let mut p = if cfg!(target_os = "windows") {
+        std::process::Command::new("waitfor")
+            .arg("/t")
+            .arg("3")
+            .arg("RefreshPids")
+            .stdout(std::process::Stdio::null())
+            .spawn()
+            .unwrap()
+    } else {
+        std::process::Command::new("sleep")
+            .arg("3")
+            .stdout(std::process::Stdio::null())
+            .spawn()
+            .unwrap()
+    };
+
+    let child_pid = Pid::from_u32(p.id() as _);
+    let pids = &[child_pid, self_pid];
+    std::thread::sleep(std::time::Duration::from_millis(500));
+    s.refresh_pids(pids);
+    p.kill().expect("Unable to kill process.");
+
+    assert_eq!(s.processes().len(), 2);
+    for pid in s.processes().keys() {
+        assert!(pids.contains(pid));
+    }
+}