Change CLI args

Add p core selector
Fix socket argument parsing
Better budget calculation

2 files changed, 37 insertions, 23 deletions

diff --git a/src/energy/budget.rs b/src/energy/budget.rs
index 6154187..3edce30 100644
--- a/src/energy/budget.rs
+++ b/src/energy/budget.rs
@@ -1,19 +1,22 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::atomic::AtomicU32};
 
 use crate::energy::EnergyService;
+use std::sync::Arc;
 
 type Pid = i32;
 
+const MAX_BUDGET: u64 = 100;
+
 pub trait BudgetPolicy: Send + 'static {
     fn calculate_budgets(&self, energy_service: &mut EnergyService) -> HashMap<Pid, u64>;
 }
 
 pub struct SimpleCappingPolicy {
-    power_cap: u64,
+    power_cap: Arc<AtomicU32>,
 }
 
 impl SimpleCappingPolicy {
-    pub fn new(power_cap: u64) -> Self {
+    pub fn new(power_cap: Arc<AtomicU32>) -> Self {
         Self { power_cap }
     }
 }
@@ -21,32 +24,43 @@ impl SimpleCappingPolicy {
 impl BudgetPolicy for SimpleCappingPolicy {
     fn calculate_budgets(&self, energy_service: &mut EnergyService) -> HashMap<Pid, u64> {
         let mut budgets = HashMap::new();
-        let process_energies = energy_service.all_process_energies();
+        let process_energies = energy_service.all_process_energy_deltas();
 
         // Total energy consumption across all processes
         //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 actual_energy > self.power_cap as f64 {
-            let ratio = self.power_cap as f64 / actual_energy;
+        let actual_energy = energy_service.last_energy_diff;
+        let energy_cap = self.power_cap.load(std::sync::atomic::Ordering::Relaxed) as f64
+            * energy_service.last_time_between_measurements.as_secs_f64();
+        println!("{actual_energy} {energy_cap}");
+        let base_energy_per_process =
+            energy_cap / process_energies.iter().filter(|(_, e)| **e > 0f64).count() as f64;
+        let ratio = energy_cap / actual_energy;
 
-            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 / actual_energy) * (1.0 - ratio));
-                let budget = (u64::MAX as f64 * scaling_factor) as u64;
-                budgets.insert(pid, budget);
-            }
-        } else {
-            // Under power cap, assign maximum budget to all
-            for pid in energy_service.active_processes() {
-                budgets.insert(*pid, u64::MAX);
-            }
+        for (pid, energy) in process_energies {
+            let budget = budgets.entry(pid).or_insert(0);
+            *budget = (*budget + (ratio * base_energy_per_process - energy as f64) as u64)
+                .min(MAX_BUDGET);
         }
 
+        // Simple proportional distribution if over cap
+        //if actual_energy > energy_cap {
+        // let ratio = energy_cap / actual_energy;
+        //
+        // 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 / actual_energy) * (1.0 - ratio));
+        //     let budget = (u64::MAX as f64 * scaling_factor) as u64;
+        //     budgets.insert(pid, budget);
+        // }
+        // } else {
+        //     // Under power cap, assign maximum budget to all
+        //     for pid in energy_service.active_processes() {
+        //         budgets.insert(*pid, u64::MAX);
+        //     }
+        // }
+
         budgets
     }
 }
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index 9262677..7223b12 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -161,7 +161,7 @@ impl Estimator for PerfEstimator {
 
         let time_running_ns = counts.time_running();
         if time_running_ns - counters.old_time == 0 || counts.iter().next().unwrap().1 == &0 {
-            println!("The counters are zero although the task has been scheduled!!");
+            //println!("The counters are zero although the task has been scheduled!!");
             return None;
         }
         let correction_factor = 10_000_000. / (time_running_ns - counters.old_time) as f64;