Add functionality to list connected GPU(s) (#140)

Cargo.toml

@@ -14,6 +14,7 @@ build = "build.rs"
 cfg-if = "1.0.0"
 libc = "0.2.131"
 home = "0.5.3"
+pciid-parser = "0.6.2"
 
 [build-dependencies.vergen]
 version = "7.3.2"

src/android/mod.rs

@@ -289,6 +289,10 @@ impl GeneralReadout for AndroidGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         Err(ReadoutError::NotImplemented)
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 impl MemoryReadout for AndroidMemoryReadout {

src/freebsd/mod.rs

@@ -297,6 +297,10 @@ impl GeneralReadout for FreeBSDGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         shared::disk_space(String::from("/"))
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 impl MemoryReadout for FreeBSDMemoryReadout {

src/linux/mod.rs

@@ -1,12 +1,15 @@
 #![allow(clippy::unnecessary_cast)]
+mod pci_devices;
 mod sysinfo_ffi;
 
+use self::pci_devices::get_pci_devices;
 use crate::extra;
 use crate::extra::get_entries;
 use crate::extra::path_extension;
 use crate::shared;
 use crate::traits::*;
 use itertools::Itertools;
+use pciid_parser::Database;
 use regex::Regex;
 use std::fs;
 use std::fs::read_dir;
@@ -543,6 +546,32 @@ impl GeneralReadout for LinuxGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         shared::disk_space(String::from("/"))
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        let db = match Database::read() {
+            Ok(db) => db,
+            _ => panic!("Could not read pci.ids file"),
+        };
+
+        let devices = get_pci_devices()?;
+        let mut gpus = vec![];
+
+        for device in devices {
+            if !device.is_gpu(&db) {
+                continue;
+            };
+
+            if let Some(sub_device_name) = device.get_sub_device_name(&db) {
+                gpus.push(sub_device_name);
+            };
+        }
+
+        if gpus.is_empty() {
+            Err(ReadoutError::MetricNotAvailable)
+        } else {
+            Ok(gpus)
+        }
+    }
 }
 
 impl MemoryReadout for LinuxMemoryReadout {

src/linux/pci_devices.rs

@@ -0,0 +1,109 @@
+use std::{
+    fs::{read_dir, read_to_string},
+    io,
+    path::PathBuf,
+};
+
+use pciid_parser::{schema::SubDeviceId, Database};
+
+use crate::extra::pop_newline;
+
+fn parse_device_hex(hex_str: &str) -> String {
+    pop_newline(hex_str).chars().skip(2).collect::<String>()
+}
+
+pub enum PciDeviceReadableValues {
+    Class,
+    Vendor,
+    Device,
+    SubVendor,
+    SubDevice,
+}
+
+impl PciDeviceReadableValues {
+    fn as_str(&self) -> &'static str {
+        match self {
+            PciDeviceReadableValues::Class => "class",
+            PciDeviceReadableValues::Vendor => "vendor",
+            PciDeviceReadableValues::Device => "device",
+            PciDeviceReadableValues::SubVendor => "subsystem_vendor",
+            PciDeviceReadableValues::SubDevice => "subsystem_device",
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct PciDevice {
+    base_path: PathBuf,
+}
+
+impl PciDevice {
+    fn new(base_path: PathBuf) -> PciDevice {
+        PciDevice { base_path }
+    }
+
+    fn _read_value(&self, readable_value: PciDeviceReadableValues) -> String {
+        let value_path = self.base_path.join(readable_value.as_str());
+
+        match read_to_string(&value_path) {
+            Ok(hex_string) => parse_device_hex(&hex_string),
+            _ => panic!("Could not find value: {:?}", value_path),
+        }
+    }
+
+    pub fn is_gpu(&self, db: &Database) -> bool {
+        let class_value = self._read_value(PciDeviceReadableValues::Class);
+        let first_pair = class_value.chars().take(2).collect::<String>();
+
+        match db.classes.get(&first_pair) {
+            Some(class) => class.name == "Display controller",
+            _ => false,
+        }
+    }
+
+    pub fn get_sub_device_name(&self, db: &Database) -> Option<String> {
+        let vendor_value = self._read_value(PciDeviceReadableValues::Vendor);
+        let sub_vendor_value = self._read_value(PciDeviceReadableValues::SubVendor);
+        let device_value = self._read_value(PciDeviceReadableValues::Device);
+        let sub_device_value = self._read_value(PciDeviceReadableValues::SubDevice);
+
+        let Some(vendor) = db.vendors.get(&vendor_value) else {
+            return None;
+        };
+
+        let Some(device) = vendor.devices.get(&device_value) else {
+            return None;
+        };
+
+        let sub_device_id = SubDeviceId {
+            subvendor: sub_vendor_value,
+            subdevice: sub_device_value,
+        };
+
+        if let Some(sub_device) = device.subdevices.get(&sub_device_id) {
+            let start = match sub_device.find('[') {
+                Some(i) => i + 1,
+                _ => panic!(
+                    "Could not find opening square bracket for sub device: {}",
+                    sub_device
+                ),
+            };
+            let end = sub_device.len() - 1;
+
+            Some(sub_device.chars().take(end).skip(start).collect::<String>())
+        } else {
+            None
+        }
+    }
+}
+
+pub fn get_pci_devices() -> Result<Vec<PciDevice>, io::Error> {
+    let devices_dir = read_dir("/sys/bus/pci/devices/")?;
+
+    let mut devices = vec![];
+    for device_entry in devices_dir.flatten() {
+        devices.push(PciDevice::new(device_entry.path()));
+    }
+
+    Ok(devices)
+}

src/macos/mod.rs

@@ -403,6 +403,10 @@ impl GeneralReadout for MacOSGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         shared::disk_space(String::from("/"))
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 impl MacOSGeneralReadout {

src/netbsd/mod.rs

@@ -333,6 +333,10 @@ impl GeneralReadout for NetBSDGeneralReadout {
             "Error while trying to get statfs structure.",
         )))
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 impl MemoryReadout for NetBSDMemoryReadout {

src/openwrt/mod.rs

@@ -214,6 +214,10 @@ impl GeneralReadout for OpenWrtGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         shared::disk_space(String::from("/"))
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 impl MemoryReadout for OpenWrtMemoryReadout {

src/traits.rs

@@ -483,6 +483,11 @@ impl GeneralReadout for MacOSGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         Ok((50000000,1000000000)) // Used / Total
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        // Get gpu(s) from list of connected pci devices
+        Ok(vec!(String::from("gpu1"), String::from("gpu2"))) // Return gpu sub-device names
+    }
 }
 
 ```
@@ -581,6 +586,9 @@ pub trait GeneralReadout {
     ///
     /// _e.g._ '1.2TB / 2TB'
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError>;
+
+    /// This function should return the sub device names of any _GPU(s)_ connected to the host machine.
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError>;
 }
 
 /// Holds the possible variants for battery status.

src/windows/mod.rs

@@ -341,6 +341,10 @@ impl GeneralReadout for WindowsGeneralReadout {
     fn disk_space(&self) -> Result<(u128, u128), ReadoutError> {
         Err(ReadoutError::NotImplemented)
     }
+
+    fn gpus(&self) -> Result<Vec<String>, ReadoutError> {
+        Err(ReadoutError::NotImplemented)
+    }
 }
 
 pub struct WindowsProductReadout {