1 files changed, 49 insertions, 0 deletions
diff --git a/src/energy/budget.rs b/src/energy/budget.rs
new file mode 100644
index 0000000..aada2eb
--- /dev/null
+++ b/src/energy/budget.rs
@@ -0,0 +1,49 @@
+use std::collections::HashMap;
+
+use crate::energy::EnergyService;
+
+type Pid = i32;
+
+pub trait BudgetPolicy: Send + 'static {
+    fn calculate_budgets(&self, energy_service: &EnergyService) -> HashMap<Pid, u64>;
+}
+
+pub struct SimpleCappingPolicy {
+    power_cap: u64,
+}
+
+impl SimpleCappingPolicy {
+    pub fn new(power_cap: u64) -> Self {
+        Self { power_cap }
+    }
+}
+
+impl BudgetPolicy for SimpleCappingPolicy {
+    fn calculate_budgets(&self, energy_service: &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();
+
+        // Simple proportional distribution if over cap
+        if total_energy > self.power_cap {
+            let ratio = self.power_cap as f64 / total_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 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
+    }
+}