use crate::{GpuInner, GpusInner};
use std::cmp::Ordering;
use std::fmt;
pub struct Gpus {
pub(crate) inner: GpusInner,
}
impl Gpus {
pub fn new() -> Result<Self, crate::Error> {
Ok(Self {
inner: GpusInner::new()?,
})
}
pub fn new_with_refreshed_list() -> Result<Self, crate::Error> {
let mut gpus = Self::new()?;
gpus.refresh(false);
Ok(gpus)
}
pub fn list(&self) -> &[Gpu] {
&self.inner.gpus
}
pub fn refresh(&mut self, remove_not_listed_gpus: bool) {
self.inner.refresh();
if remove_not_listed_gpus {
self.inner.gpus.retain_mut(|c| {
if !c.inner.updated {
return false;
}
c.inner.updated = false;
true
});
} else {
for gpu in &mut self.inner.gpus {
gpu.inner.updated = false;
}
}
}
}
impl std::ops::Deref for Gpus {
type Target = [Gpu];
fn deref(&self) -> &Self::Target {
self.list()
}
}
impl<'a> IntoIterator for &'a Gpus {
type Item = &'a Gpu;
type IntoIter = std::slice::Iter<'a, Gpu>;
fn into_iter(self) -> Self::IntoIter {
self.list().iter()
}
}
#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct PCI {
pub domain: u32,
pub bus: u32,
pub device: u32,
pub function: u32,
}
impl fmt::Display for PCI {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let Self {
domain,
bus,
device,
function,
} = self;
write!(f, "{domain:04x}:{bus:02x}:{device:02x}.{function}")
}
}
impl core::str::FromStr for PCI {
type Err = &'static str;
fn from_str(s: &str) -> Result<Self, Self::Err> {
fn get_next_u32<'a>(
iter: &mut impl Iterator<Item = &'a str>,
missing_msg: &'static str,
invalid_msg: &'static str,
) -> Result<u32, &'static str> {
let Some(value) = iter.next() else {
return Err(missing_msg);
};
value.parse::<u32>().map_err(|_| invalid_msg)
}
let mut iter = s.split(':');
let domain = get_next_u32(&mut iter, "missing domain", "invalid domain")?;
let bus = get_next_u32(&mut iter, "missing bus", "invalid bus")?;
let Some(last) = iter.next() else {
return Err("missing device");
};
if iter.next().is_some() {
return Err("unexpected `:` after bus");
};
let mut iter = last.split('.');
let device = get_next_u32(&mut iter, "missing device", "invalid device")?;
let function = get_next_u32(&mut iter, "missing function", "invalid function")?;
if iter.next().is_some() {
return Err("unexpected `:` after function");
};
Ok(Self {
domain,
bus,
device,
function,
})
}
}
pub struct Gpu {
pub(crate) inner: GpuInner,
}
impl PartialEq for Gpu {
fn eq(&self, other: &Self) -> bool {
self.pci() == other.pci()
}
}
impl Eq for Gpu {}
impl PartialOrd for Gpu {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Gpu {
fn cmp(&self, other: &Self) -> Ordering {
self.pci().cmp(other.pci())
}
}
impl Gpu {
pub fn pci(&self) -> &PCI {
self.inner.pci()
}
pub fn vendor(&self) -> Option<&str> {
self.inner.vendor()
}
pub fn model(&self) -> Option<&str> {
self.inner.model()
}
pub fn usage(&self) -> Option<f32> {
self.inner.usage()
}
pub fn total_memory(&self) -> Option<u64> {
self.inner.total_memory()
}
pub fn used_memory(&self) -> Option<u64> {
self.inner.used_memory()
}
}