Use rapl package energy to calculate new budgets

2 files changed, 19 insertions, 6 deletions

diff --git a/src/energy/budget.rs b/src/energy/budget.rs
index aada2eb..7a2c7e2 100644
--- a/src/energy/budget.rs
+++ b/src/energy/budget.rs
@@ -5,7 +5,7 @@ use crate::energy::EnergyService;
 type Pid = i32;
 
 pub trait BudgetPolicy: Send + 'static {
-    fn calculate_budgets(&self, energy_service: &EnergyService) -> HashMap<Pid, u64>;
+    fn calculate_budgets(&self, energy_service: &mut EnergyService) -> HashMap<Pid, u64>;
 }
 
 pub struct SimpleCappingPolicy {
@@ -19,21 +19,24 @@ impl SimpleCappingPolicy {
 }
 
 impl BudgetPolicy for SimpleCappingPolicy {
-    fn calculate_budgets(&self, energy_service: &EnergyService) -> HashMap<Pid, u64> {
+    fn calculate_budgets(&self, energy_service: &mut EnergyService) -> HashMap<Pid, u64> {
         let mut budgets = HashMap::new();
         let process_energies = energy_service.all_process_energies();
 
         // Total energy consumption across all processes
-        let total_energy: u64 = process_energies.values().sum();
+        //let total_energy: u64 = process_energies.values().sum();
+
+        let actual_energy =
+            energy_service.package_energy() / energy_service.update_interval.as_secs_f64();
 
         // Simple proportional distribution if over cap
-        if total_energy > self.power_cap {
-            let ratio = self.power_cap as f64 / total_energy as f64;
+        if actual_energy > self.power_cap as f64 {
+            let ratio = self.power_cap as f64 / actual_energy as f64;
 
             for (&pid, &energy) in &process_energies {
                 // Calculate a scaled budget based on the ratio
                 // Higher energy consumers get proportionally reduced budgets
-                let scaling_factor = 1.0 - ((energy as f64 / total_energy as f64) * (1.0 - ratio));
+                let scaling_factor = 1.0 - ((energy as f64 / actual_energy) * (1.0 - ratio));
                 let budget = (u64::MAX as f64 * scaling_factor) as u64;
                 budgets.insert(pid, budget);
             }
diff --git a/src/energy/rapl.rs b/src/energy/rapl.rs
index 293a0f7..34b128e 100644
--- a/src/energy/rapl.rs
+++ b/src/energy/rapl.rs
@@ -76,6 +76,16 @@ impl PackageEnergy {
         raw_value as f64 * self.scale
     }
 
+    /// Get the perf event type
+    pub fn get_type(&self) -> u32 {
+        self.event_type
+    }
+
+    /// Get the event config
+    pub fn get_event_config(&self) -> u64 {
+        self.event_config
+    }
+
     /// Check if perf_event_paranoid setting might cause permission issues
     pub fn check_paranoid() -> io::Result<i32> {
         let path = Path::new("/proc/sys/kernel/perf_event_paranoid");