use std::{
collections::HashMap,
time::{Duration, Instant},
};
use crate::{
hardware::iokit::{IOKit, IOKitImpl},
Result,
};
#[derive(Debug, Clone, PartialEq)]
pub enum SensorLocation {
Cpu,
Gpu,
Memory,
Storage,
Battery,
Heatsink,
Ambient,
Other(String),
}
#[derive(Debug, Clone)]
pub struct Fan {
pub name: String,
pub speed_rpm: u32,
pub min_speed: u32,
pub max_speed: u32,
pub percentage: f64,
}
#[derive(Debug, Clone)]
pub struct TemperatureConfig {
pub poll_interval_ms: u64,
pub throttling_threshold: f64,
pub auto_refresh: bool,
}
impl Default for TemperatureConfig {
fn default() -> Self {
Self {
poll_interval_ms: 1000, throttling_threshold: 80.0, auto_refresh: true,
}
}
}
#[derive(Debug)]
pub struct Temperature {
pub sensors: HashMap<String, f64>,
pub fans: Vec<Fan>,
pub is_throttling: bool,
pub cpu_power: Option<f64>,
pub config: TemperatureConfig,
io_kit: IOKitImpl,
last_refresh: Instant,
}
impl Temperature {
pub fn new() -> Self {
Self {
sensors: HashMap::new(),
fans: Vec::new(),
is_throttling: false,
cpu_power: None,
config: TemperatureConfig::default(),
io_kit: IOKitImpl,
last_refresh: Instant::now() - Duration::from_secs(60), }
}
pub fn with_config(config: TemperatureConfig) -> Self {
Self {
sensors: HashMap::new(),
fans: Vec::new(),
is_throttling: false,
cpu_power: None,
config,
io_kit: IOKitImpl,
last_refresh: Instant::now() - Duration::from_secs(60), }
}
fn should_refresh(&self) -> bool {
self.last_refresh.elapsed().as_millis() as u64 > self.config.poll_interval_ms
}
pub fn refresh(&mut self) -> Result<()> {
if let Ok(thermal_info) = self.io_kit.get_thermal_info() {
self.sensors.insert("CPU".to_string(), thermal_info.cpu_temp);
self.sensors.insert("GPU".to_string(), thermal_info.gpu_temp);
if let Some(temp) = thermal_info.heatsink_temp {
self.sensors.insert("Heatsink".to_string(), temp);
}
if let Some(temp) = thermal_info.ambient_temp {
self.sensors.insert("Ambient".to_string(), temp);
}
if let Some(temp) = thermal_info.battery_temp {
self.sensors.insert("Battery".to_string(), temp);
}
self.is_throttling = thermal_info.is_throttling;
self.cpu_power = thermal_info.cpu_power;
}
if let Ok(fan_infos) = self.io_kit.get_all_fans() {
self.fans.clear();
for (i, fan_info) in fan_infos.iter().enumerate() {
self.fans.push(Fan {
name: format!("Fan {}", i),
speed_rpm: fan_info.speed_rpm,
min_speed: fan_info.min_speed,
max_speed: fan_info.max_speed,
percentage: fan_info.percentage,
});
}
}
self.last_refresh = Instant::now();
Ok(())
}
pub fn cpu_temperature(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors.get("CPU").cloned().ok_or_else(|| {
crate::Error::Temperature("CPU temperature sensor not available".to_string())
})
}
pub fn gpu_temperature(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors.get("GPU").cloned().ok_or_else(|| {
crate::Error::Temperature("GPU temperature sensor not available".to_string())
})
}
pub fn heatsink_temperature(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors.get("Heatsink").cloned().ok_or_else(|| {
crate::Error::Temperature("Heatsink temperature sensor not available".to_string())
})
}
pub fn ambient_temperature(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors.get("Ambient").cloned().ok_or_else(|| {
crate::Error::Temperature("Ambient temperature sensor not available".to_string())
})
}
pub fn battery_temperature(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors.get("Battery").cloned().ok_or_else(|| {
crate::Error::Temperature("Battery temperature sensor not available".to_string())
})
}
pub fn list_sensors(&mut self) -> Result<Vec<(String, SensorLocation)>> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
let mut result = Vec::new();
for name in self.sensors.keys() {
let location = match name.as_str() {
"CPU" => SensorLocation::Cpu,
"GPU" => SensorLocation::Gpu,
"Heatsink" => SensorLocation::Heatsink,
"Ambient" => SensorLocation::Ambient,
"Battery" => SensorLocation::Battery,
"Memory" => SensorLocation::Memory,
"Storage" => SensorLocation::Storage,
_ => SensorLocation::Other(name.to_string()),
};
result.push((name.clone(), location));
}
Ok(result)
}
pub fn get_sensor_temperature(&mut self, name: &str) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.sensors
.get(name)
.cloned()
.ok_or_else(|| crate::Error::Temperature(format!("Sensor {} not found", name)))
}
pub fn fan_count(&mut self) -> Result<usize> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
Ok(self.fans.len())
}
pub fn get_fans(&mut self) -> Result<&Vec<Fan>> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
Ok(&self.fans)
}
pub fn get_fan(&mut self, index: usize) -> Result<&Fan> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.fans
.get(index)
.ok_or_else(|| crate::Error::Temperature(format!("Fan with index {} not found", index)))
}
pub fn cpu_power(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
self.cpu_power.ok_or_else(|| {
crate::Error::Temperature("CPU power information not available".to_string())
})
}
pub fn is_throttling(&mut self) -> Result<bool> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh()?;
}
if let Ok(throttling) = self.io_kit.check_thermal_throttling() {
return Ok(throttling);
}
let cpu_temp = self.cpu_temperature()?;
Ok(cpu_temp > self.config.throttling_threshold)
}
pub fn get_thermal_metrics(&mut self) -> Result<ThermalMetrics> {
self.refresh()?;
Ok(ThermalMetrics {
cpu_temperature: self.sensors.get("CPU").cloned(),
gpu_temperature: self.sensors.get("GPU").cloned(),
heatsink_temperature: self.sensors.get("Heatsink").cloned(),
ambient_temperature: self.sensors.get("Ambient").cloned(),
battery_temperature: self.sensors.get("Battery").cloned(),
is_throttling: self.is_throttling,
cpu_power: self.cpu_power,
fans: self.fans.clone(),
})
}
pub async fn cpu_temperature_async(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
self.sensors.get("CPU").cloned().ok_or_else(|| {
crate::Error::Temperature("CPU temperature sensor not available".to_string())
})
}
pub async fn gpu_temperature_async(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
self.sensors.get("GPU").cloned().ok_or_else(|| {
crate::Error::Temperature("GPU temperature sensor not available".to_string())
})
}
pub async fn heatsink_temperature_async(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
self.sensors.get("Heatsink").cloned().ok_or_else(|| {
crate::Error::Temperature("Heatsink temperature sensor not available".to_string())
})
}
pub async fn ambient_temperature_async(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
self.sensors.get("Ambient").cloned().ok_or_else(|| {
crate::Error::Temperature("Ambient temperature sensor not available".to_string())
})
}
pub async fn battery_temperature_async(&mut self) -> Result<f64> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
self.sensors.get("Battery").cloned().ok_or_else(|| {
crate::Error::Temperature("Battery temperature sensor not available".to_string())
})
}
pub async fn refresh_async(&mut self) -> Result<()> {
let io_kit = self.io_kit.clone();
let thermal_info =
tokio::task::spawn_blocking(move || io_kit.get_thermal_info())
.await
.map_err(|e| crate::Error::Temperature(format!("Task join error: {}", e)))??;
self.sensors.insert("CPU".to_string(), thermal_info.cpu_temp);
self.sensors.insert("GPU".to_string(), thermal_info.gpu_temp);
if let Some(temp) = thermal_info.heatsink_temp {
self.sensors.insert("Heatsink".to_string(), temp);
}
if let Some(temp) = thermal_info.ambient_temp {
self.sensors.insert("Ambient".to_string(), temp);
}
if let Some(temp) = thermal_info.battery_temp {
self.sensors.insert("Battery".to_string(), temp);
}
self.is_throttling = thermal_info.is_throttling;
self.cpu_power = thermal_info.cpu_power;
let io_kit = self.io_kit.clone();
let fan_infos = tokio::task::spawn_blocking(move || io_kit.get_all_fans())
.await
.map_err(|e| crate::Error::Temperature(format!("Task join error: {}", e)))??;
self.fans.clear();
for (i, fan_info) in fan_infos.iter().enumerate() {
self.fans.push(Fan {
name: format!("Fan {}", i),
speed_rpm: fan_info.speed_rpm,
min_speed: fan_info.min_speed,
max_speed: fan_info.max_speed,
percentage: fan_info.percentage,
});
}
self.last_refresh = Instant::now();
Ok(())
}
pub async fn get_thermal_metrics_async(&mut self) -> Result<ThermalMetrics> {
self.refresh_async().await?;
Ok(ThermalMetrics {
cpu_temperature: self.sensors.get("CPU").cloned(),
gpu_temperature: self.sensors.get("GPU").cloned(),
heatsink_temperature: self.sensors.get("Heatsink").cloned(),
ambient_temperature: self.sensors.get("Ambient").cloned(),
battery_temperature: self.sensors.get("Battery").cloned(),
is_throttling: self.is_throttling,
cpu_power: self.cpu_power,
fans: self.fans.clone(),
})
}
pub async fn is_throttling_async(&mut self) -> Result<bool> {
if self.config.auto_refresh && self.should_refresh() {
self.refresh_async().await?;
}
let io_kit = self.io_kit.clone();
match tokio::task::spawn_blocking(move || io_kit.check_thermal_throttling())
.await
.map_err(|e| crate::Error::Temperature(format!("Task join error: {}", e)))?
{
Ok(throttling) => Ok(throttling),
Err(_) => {
let cpu_temp = self.cpu_temperature_async().await?;
Ok(cpu_temp > self.config.throttling_threshold)
},
}
}
}
#[derive(Debug, Clone)]
pub struct ThermalMetrics {
pub cpu_temperature: Option<f64>,
pub gpu_temperature: Option<f64>,
pub heatsink_temperature: Option<f64>,
pub ambient_temperature: Option<f64>,
pub battery_temperature: Option<f64>,
pub is_throttling: bool,
pub cpu_power: Option<f64>,
pub fans: Vec<Fan>,
}
impl Default for Temperature {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use std::{thread, time::Duration};
use super::*;
#[test]
fn test_temperature_new() {
let temp = Temperature::new();
assert!(temp.sensors.is_empty());
assert!(temp.fans.is_empty());
assert!(!temp.is_throttling);
assert_eq!(temp.cpu_power, None);
}
#[test]
fn test_with_config() {
let config = TemperatureConfig {
poll_interval_ms: 5000,
throttling_threshold: 90.0,
auto_refresh: false,
};
let temp = Temperature::with_config(config);
assert_eq!(temp.config.poll_interval_ms, 5000);
assert_eq!(temp.config.throttling_threshold, 90.0);
assert!(!temp.config.auto_refresh);
}
#[test]
fn test_should_refresh() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 10,
throttling_threshold: 80.0,
auto_refresh: true,
});
assert!(temp.should_refresh());
temp.last_refresh = Instant::now();
assert!(!temp.should_refresh());
thread::sleep(Duration::from_millis(20));
assert!(temp.should_refresh());
}
#[test]
fn test_cpu_temperature() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.sensors.insert("CPU".to_string(), 42.5);
let result = temp.cpu_temperature().unwrap();
assert_eq!(result, 42.5);
}
#[test]
fn test_gpu_temperature() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.sensors.insert("GPU".to_string(), 55.0);
let result = temp.gpu_temperature().unwrap();
assert_eq!(result, 55.0);
}
#[test]
fn test_additional_sensors() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.sensors.insert("Heatsink".to_string(), 45.0);
temp.sensors.insert("Ambient".to_string(), 32.0);
temp.sensors.insert("Battery".to_string(), 38.0);
assert_eq!(temp.heatsink_temperature().unwrap(), 45.0);
assert_eq!(temp.ambient_temperature().unwrap(), 32.0);
assert_eq!(temp.battery_temperature().unwrap(), 38.0);
}
#[test]
fn test_list_sensors() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.sensors.insert("CPU".to_string(), 42.5);
temp.sensors.insert("GPU".to_string(), 55.0);
temp.sensors.insert("Heatsink".to_string(), 45.0);
temp.sensors.insert("Ambient".to_string(), 32.0);
temp.sensors.insert("Custom".to_string(), 27.0);
let sensors = temp.list_sensors().unwrap();
assert_eq!(sensors.len(), 5);
let has_cpu = sensors
.iter()
.any(|(name, location)| name == "CPU" && matches!(location, SensorLocation::Cpu));
let has_gpu = sensors
.iter()
.any(|(name, location)| name == "GPU" && matches!(location, SensorLocation::Gpu));
let has_heatsink = sensors.iter().any(|(name, location)| {
name == "Heatsink" && matches!(location, SensorLocation::Heatsink)
});
let has_ambient = sensors.iter().any(|(name, location)| {
name == "Ambient" && matches!(location, SensorLocation::Ambient)
});
let has_custom = sensors.iter().any(|(name, location)| {
name == "Custom"
&& if let SensorLocation::Other(custom_name) = location {
custom_name == "Custom"
} else {
false
}
});
assert!(has_cpu);
assert!(has_gpu);
assert!(has_heatsink);
assert!(has_ambient);
assert!(has_custom);
}
#[test]
fn test_fan_functions() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.fans.push(Fan {
name: "Fan 0".to_string(),
speed_rpm: 2000,
min_speed: 1000,
max_speed: 4000,
percentage: 33.3,
});
temp.fans.push(Fan {
name: "Fan 1".to_string(),
speed_rpm: 3000,
min_speed: 1200,
max_speed: 5000,
percentage: 47.4,
});
assert_eq!(temp.fan_count().unwrap(), 2);
let fans = temp.get_fans().unwrap();
assert_eq!(fans.len(), 2);
assert_eq!(fans[0].speed_rpm, 2000);
assert_eq!(fans[1].speed_rpm, 3000);
let fan0 = temp.get_fan(0).unwrap();
assert_eq!(fan0.name, "Fan 0");
assert_eq!(fan0.speed_rpm, 2000);
assert_eq!(fan0.min_speed, 1000);
assert_eq!(fan0.max_speed, 4000);
assert!(fan0.percentage > 33.0 && fan0.percentage < 34.0);
let fan1 = temp.get_fan(1).unwrap();
assert_eq!(fan1.name, "Fan 1");
assert_eq!(fan1.speed_rpm, 3000);
assert!(temp.get_fan(2).is_err());
}
#[test]
fn test_is_throttling_property() {
let temp = Temperature::new();
assert!(!temp.is_throttling);
let mut temp = Temperature::new();
temp.is_throttling = true;
assert!(temp.is_throttling);
}
#[test]
fn test_is_throttling_temperature_heuristic() {
let threshold = 80.0;
let cpu_temp = 75.0;
assert!(cpu_temp < threshold);
let cpu_temp = 85.0;
assert!(cpu_temp > threshold);
}
#[test]
fn test_get_thermal_metrics() {
let mut temp = Temperature::with_config(TemperatureConfig {
poll_interval_ms: 1000,
throttling_threshold: 80.0,
auto_refresh: false,
});
temp.sensors.insert("CPU".to_string(), 42.5);
temp.sensors.insert("GPU".to_string(), 55.0);
temp.sensors.insert("Heatsink".to_string(), 45.0);
temp.is_throttling = false;
temp.cpu_power = Some(28.5);
temp.fans.push(Fan {
name: "Fan 0".to_string(),
speed_rpm: 2000,
min_speed: 1000,
max_speed: 4000,
percentage: 33.3,
});
let metrics = temp.get_thermal_metrics().unwrap();
let heatsink_temp = metrics.heatsink_temperature;
assert!(heatsink_temp.is_some(), "Heatsink temperature should be present");
assert!(metrics.cpu_temperature.is_some(), "CPU temperature should be present");
assert!(metrics.gpu_temperature.is_some(), "GPU temperature should be present");
assert!(!metrics.is_throttling, "System should not be throttling");
assert!(metrics.cpu_power.is_some(), "CPU power should be present");
assert!(!metrics.fans.is_empty(), "There should be at least one fan");
assert!(metrics.fans[0].speed_rpm > 0, "Fan speed should be greater than 0");
}
}