Skip to content

Commit

Permalink
Correctly handle Vec::reserve
Browse files Browse the repository at this point in the history
  • Loading branch information
@GuillaumeGomez
GuillaumeGomez committed Nov 21, 2023
1 parent 18bad27 commit ea08e7f
Show file tree
Hide file tree
Showing 4 changed files with 114 additions and 96 deletions.
3 changes: 3 additions & 0 deletions src/unix/users.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ pub(crate) unsafe fn get_group_name(
{
// If there was not enough memory, we give it more.
if last_errno == libc::ERANGE as _ {
// Needs to be updated for `Vec::reserve` to actually add additional capacity.
// In here it's "fine" since we never read from `buffer`.
buffer.set_len(buffer.capacity());
buffer.reserve(2048);
continue;
}
Expand Down
23 changes: 20 additions & 3 deletions src/windows/cpu.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -501,19 +501,25 @@ pub(crate) fn get_frequencies(nb_cpus: usize) -> Vec<u64> {
}

pub(crate) fn get_physical_core_count() -> Option<usize> {
// we cannot use the number of cpus here to pre calculate the buf size
// GetLogicalCpuInformationEx with RelationProcessorCore passed to it not only returns
// the logical cores but also numa nodes
// We cannot use the number of cpus here to pre calculate the buf size.
// `GetLogicalCpuInformationEx` with `RelationProcessorCore` passed to it not only returns
// the logical cores but also numa nodes.
//
// GetLogicalCpuInformationEx: https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getlogicalprocessorinformationex

let mut needed_size = 0;
unsafe {
// This function call will always return an error as it only returns "success" when it
// has written at least one item in the buffer (which it cannot do here).
let _err = GetLogicalProcessorInformationEx(RelationAll, None, &mut needed_size);

let mut buf: Vec<u8> = Vec::with_capacity(needed_size as _);

loop {
// Needs to be updated for `Vec::reserve` to actually add additional capacity if
// `GetLogicalProcessorInformationEx` fails because the buffer isn't big enough.
buf.set_len(needed_size as _);

if GetLogicalProcessorInformationEx(
RelationAll,
Some(buf.as_mut_ptr().cast()),
Expand All @@ -540,6 +546,17 @@ pub(crate) fn get_physical_core_count() -> Option<usize> {
} else {
1
};
needed_size = match needed_size.checked_add(reserve as _) {
Some(new_size) => new_size,
None => {
sysinfo_debug!(
"get_physical_core_count: buffer size is too big ({} + {})",
needed_size,
reserve,
);
return None;
}
};
buf.reserve(reserve);
}

Expand Down
182 changes: 89 additions & 93 deletions src/windows/system.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,14 +201,14 @@ impl SystemInner {
let mut buffer_size = 512 * 1024;
let mut process_information: Vec<u8> = Vec::with_capacity(buffer_size);

loop {
let mut cb_needed = 0;
// reserve(n) ensures the Vec has capacity for n elements on top of len
// so we should reserve buffer_size - len. len will always be zero at this point
// this is a no-op on the first call as buffer_size == capacity
process_information.reserve(buffer_size);

unsafe {
unsafe {
loop {
let mut cb_needed = 0;
// reserve(n) ensures the Vec has capacity for n elements on top of len
// so we should reserve buffer_size - len. len will always be zero at this point
// this is a no-op on the first call as buffer_size == capacity
process_information.reserve(buffer_size);

let ntstatus = NtQuerySystemInformation(
SystemProcessInformation,
process_information.as_mut_ptr() as *mut _,
Expand Down Expand Up @@ -238,105 +238,101 @@ impl SystemInner {
}
}
}
}

// If we reach this point NtQuerySystemInformation succeeded
// and the buffer contents are initialized
unsafe {
// If we reach this point NtQuerySystemInformation succeeded
// and the buffer contents are initialized
process_information.set_len(buffer_size);
}

// Parse the data block to get process information
let mut process_ids = Vec::with_capacity(500);
let mut process_information_offset = 0;
loop {
let p = unsafe {
process_information
// Parse the data block to get process information
let mut process_ids = Vec::with_capacity(500);
let mut process_information_offset = 0;
loop {
let p = process_information
.as_ptr()
.offset(process_information_offset)
as *const SYSTEM_PROCESS_INFORMATION
};

process_ids.push(Wrap(p));
as *const SYSTEM_PROCESS_INFORMATION;

// read_unaligned is necessary to avoid
// misaligned pointer dereference: address must be a multiple of 0x8 but is 0x...
// under x86_64 wine (and possibly other systems)
let pi = unsafe { ptr::read_unaligned(p) };
process_ids.push(Wrap(p));

if pi.NextEntryOffset == 0 {
break;
}
// read_unaligned is necessary to avoid
// misaligned pointer dereference: address must be a multiple of 0x8 but is 0x...
// under x86_64 wine (and possibly other systems)
let pi = ptr::read_unaligned(p);

process_information_offset += pi.NextEntryOffset as isize;
}
let process_list = Wrap(UnsafeCell::new(&mut self.process_list));
let nb_cpus = if refresh_kind.cpu() {
self.cpus.len() as u64
} else {
0
};
if pi.NextEntryOffset == 0 {
break;
}

let now = get_now();
process_information_offset += pi.NextEntryOffset as isize;
}
let process_list = Wrap(UnsafeCell::new(&mut self.process_list));
let nb_cpus = if refresh_kind.cpu() {
self.cpus.len() as u64
} else {
0
};

#[cfg(feature = "multithread")]
use rayon::iter::ParallelIterator;

// TODO: instead of using parallel iterator only here, would be better to be
// able to run it over `process_information` directly!
let processes = into_iter(process_ids)
.filter_map(|pi| {
// as above, read_unaligned is necessary
let pi = unsafe { ptr::read_unaligned(pi.0) };
let pid = Pid(pi.UniqueProcessId as _);
if let Some(proc_) = unsafe { (*process_list.0.get()).get_mut(&pid) } {
let proc_ = &mut proc_.inner;
if proc_
.get_start_time()
.map(|start| start == proc_.start_time())
.unwrap_or(true)
{
if refresh_kind.memory() {
proc_.memory = pi.WorkingSetSize as _;
proc_.virtual_memory = pi.VirtualSize as _;
let now = get_now();

#[cfg(feature = "multithread")]
use rayon::iter::ParallelIterator;

// TODO: instead of using parallel iterator only here, would be better to be
// able to run it over `process_information` directly!
let processes = into_iter(process_ids)
.filter_map(|pi| {
// as above, read_unaligned is necessary
let pi = ptr::read_unaligned(pi.0);
let pid = Pid(pi.UniqueProcessId as _);
if let Some(proc_) = (*process_list.0.get()).get_mut(&pid) {
let proc_ = &mut proc_.inner;
if proc_
.get_start_time()
.map(|start| start == proc_.start_time())
.unwrap_or(true)
{
if refresh_kind.memory() {
proc_.memory = pi.WorkingSetSize as _;
proc_.virtual_memory = pi.VirtualSize as _;
}
proc_.update(refresh_kind, nb_cpus, now);
return None;
}
proc_.update(refresh_kind, nb_cpus, now);
return None;
// If the PID owner changed, we need to recompute the whole process.
sysinfo_debug!("owner changed for PID {}", proc_.pid());
}
// If the PID owner changed, we need to recompute the whole process.
sysinfo_debug!("owner changed for PID {}", proc_.pid());
}
let name = get_process_name(&pi, pid);
let (memory, virtual_memory) = if refresh_kind.memory() {
(pi.WorkingSetSize as _, pi.VirtualSize as _)
} else {
(0, 0)
};
let mut p = ProcessInner::new_full(
pid,
if pi.InheritedFromUniqueProcessId as usize != 0 {
Some(Pid(pi.InheritedFromUniqueProcessId as _))
let name = get_process_name(&pi, pid);
let (memory, virtual_memory) = if refresh_kind.memory() {
(pi.WorkingSetSize as _, pi.VirtualSize as _)
} else {
None
},
memory,
virtual_memory,
name,
now,
refresh_kind,
);
p.update(refresh_kind.without_memory(), nb_cpus, now);
Some(Process { inner: p })
})
.collect::<Vec<_>>();
for p in processes.into_iter() {
self.process_list.insert(p.pid(), p);
(0, 0)
};
let mut p = ProcessInner::new_full(
pid,
if pi.InheritedFromUniqueProcessId as usize != 0 {
Some(Pid(pi.InheritedFromUniqueProcessId as _))
} else {
None
},
memory,
virtual_memory,
name,
now,
refresh_kind,
);
p.update(refresh_kind.without_memory(), nb_cpus, now);
Some(Process { inner: p })
})
.collect::<Vec<_>>();
for p in processes.into_iter() {
self.process_list.insert(p.pid(), p);
}
self.process_list.retain(|_, v| {
let x = v.inner.updated;
v.inner.updated = false;
x
});
}
self.process_list.retain(|_, v| {
let x = v.inner.updated;
v.inner.updated = false;
x
});
}

pub(crate) fn processes(&self) -> &HashMap<Pid, Process> {
Expand Down
2 changes: 2 additions & 0 deletions src/windows/utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ pub(crate) fn get_reg_string_value(hkey: HKEY, path: &str, field_name: &str) ->
match new_key.get_value(&c_field_name, &mut buf, &mut buf_len) {
Ok(()) => break,
Err(err) if err.code() == Foundation::ERROR_MORE_DATA.to_hresult() => {
// Needs to be updated for `Vec::reserve` to actually add additional capacity.
buf.set_len(buf.capacity());
buf.reserve(buf_len as _);
}
_ => return None,
Expand Down

0 comments on commit ea08e7f

Please sign in to comment.