Calculate budget deduction in scheduler

1 files changed, 19 insertions, 49 deletions

diff --git a/src/energy/budget.rs b/src/energy/budget.rs
index b049814..6356132 100644
--- a/src/energy/budget.rs
+++ b/src/energy/budget.rs
@@ -1,19 +1,17 @@
-use std::{collections::HashMap, sync::atomic::AtomicU32};
+use std::{sync::atomic::AtomicU32, time::Instant};
 
 use crate::energy::EnergyService;
 use std::sync::Arc;
 
-type Pid = i32;
+use super::rapl;
 
-const MAX_BUDGET: u64 = 30000;
-
-pub trait BudgetPolicy: Send + 'static {
-    fn calculate_budgets(&mut self, energy_service: &mut EnergyService) -> HashMap<Pid, u64>;
-}
+pub const MAX_BUDGET_MJ: i64 = 30_000;
 
 pub struct SimpleCappingPolicy {
     power_cap: Arc<AtomicU32>,
     last_ratio: f64,
+    last_energy: f64,
+    last_measurement: Instant,
 }
 
 impl SimpleCappingPolicy {
@@ -21,55 +19,27 @@ impl SimpleCappingPolicy {
         Self {
             power_cap,
             last_ratio: 1.,
+            last_energy: rapl::read_package_energy().unwrap(),
+            last_measurement: Instant::now(),
         }
     }
 }
 
-impl BudgetPolicy for SimpleCappingPolicy {
-    fn calculate_budgets(&mut self, energy_service: &mut EnergyService) -> HashMap<Pid, u64> {
-        let mut budgets = HashMap::new();
-        let process_energies = energy_service.all_process_energy_deltas();
-
-        // Total energy consumption across all processes
-        //let total_energy: u64 = process_energies.values().sum();
+impl SimpleCappingPolicy {
+    pub fn calculate_budgets(&mut self, energy_service: &mut EnergyService) -> i64 {
+        let active_processes = energy_service.active_processes.len();
+        let rapl = rapl::read_package_energy().unwrap();
+        let actual_energy = rapl - self.last_energy;
+        self.last_energy = rapl;
 
-        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 * self.last_ratio;
-        self.last_ratio = ratio.clamp(0.001, 100.0);
-
-        for (pid, energy) in process_energies {
-            let budget = budgets.entry(pid).or_insert(0);
-            *budget = (*budget + ((ratio * base_energy_per_process - energy) * 1000.) as u64)
-                //.min((ratio * base_energy_per_process * MAX_BUDGET_FACTOR * 1000.) as u64);
-                .min(MAX_BUDGET);
-            if energy != 0.0 {
-                //println!("budget: {budget} energy: {energy} ratio: {ratio} base: {base_energy_per_process}");
-            }
-        }
+            * self.last_measurement.elapsed().as_secs_f64();
+        self.last_measurement = Instant::now();
 
-        // 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);
-        //     }
-        // }
+        let base_energy_per_process = energy_cap / active_processes as f64;
+        let ratio = energy_cap / actual_energy * self.last_ratio;
+        self.last_ratio = ratio.clamp(0.001, 1000.0);
 
-        budgets
+        (((ratio * base_energy_per_process) * 1000.) as i64).min(MAX_BUDGET_MJ)
     }
 }