Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add System::refresh_pids and System::refresh_pids_specifics methods #1170

Merged
merged 2 commits into from
Dec 17, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 59 additions & 1 deletion src/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
36 changes: 34 additions & 2 deletions src/unix/apple/system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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));
Expand All @@ -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,
Expand Down
33 changes: 30 additions & 3 deletions src/unix/freebsd/system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(),
Expand All @@ -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};
Expand All @@ -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,
Expand Down
22 changes: 22 additions & 0 deletions src/unix/linux/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -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(),
Expand Down
7 changes: 6 additions & 1 deletion src/unix/linux/system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
7 changes: 6 additions & 1 deletion src/unknown/system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
30 changes: 28 additions & 2 deletions src/windows/system.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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...
Expand Down
38 changes: 37 additions & 1 deletion tests/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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));
}
}
Loading