Restructure project

6 files changed, 209 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
+    }
+}
diff --git a/src/energy/estimator.rs b/src/energy/estimator.rs
new file mode 100644
index 0000000..f7f1e5d
--- /dev/null
+++ b/src/energy/estimator.rs
@@ -0,0 +1,5 @@
+pub trait Estimator: Send + 'static {
+    fn start_trace(&mut self, pid: u64);
+    fn stop_trace(&mut self, pid: u64);
+    fn read_consumption(&mut self, pid: u64) -> u64;
+}
diff --git a/src/energy/trackers.rs b/src/energy/trackers.rs
new file mode 100644
index 0000000..592c926
--- /dev/null
+++ b/src/energy/trackers.rs
@@ -0,0 +1,13 @@
+mod kernel;
+mod mock;
+mod perf;
+
+pub use kernel::*;
+pub use mock::*;
+pub use perf::*;
+
+pub trait Estimator: Send + 'static {
+    fn start_trace(&mut self, pid: u64);
+    fn stop_trace(&mut self, pid: u64);
+    fn read_consumption(&mut self, pid: u64) -> u64;
+}
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
new file mode 100644
index 0000000..050805c
--- /dev/null
+++ b/src/energy/trackers/kernel.rs
@@ -0,0 +1,49 @@
+// energy/trackers/kernel.rs
+use iocuddle::*;
+use std::fs::File;
+use std::path::Path;
+
+use crate::energy::Estimator;
+
+#[derive(Debug)]
+pub struct KernelDriver {
+    file: File,
+}
+
+impl Default for KernelDriver {
+    fn default() -> Self {
+        KernelDriver::new("/dev/rust-pow")
+            .unwrap_or_else(|_| panic!("Failed to open kernel driver. Is the module loaded?"))
+    }
+}
+
+impl KernelDriver {
+    pub fn new(path: impl AsRef<Path>) -> Result<Self, std::io::Error> {
+        let file = File::open(path)?;
+        Ok(Self { file })
+    }
+}
+
+const PERF_MON: Group = Group::new(b'|');
+const START_TRACE: Ioctl<Write, &u64> = unsafe { PERF_MON.write(0x80) };
+const STOP_TRACE: Ioctl<Write, &u64> = unsafe { PERF_MON.write(0x81) };
+const READ_POWER: Ioctl<WriteRead, &u64> = unsafe { PERF_MON.write_read(0x82) };
+
+impl Estimator for KernelDriver {
+    fn start_trace(&mut self, pid: u64) {
+        let _ = START_TRACE.ioctl(&mut self.file, &pid);
+    }
+
+    fn stop_trace(&mut self, pid: u64) {
+        let _ = STOP_TRACE.ioctl(&mut self.file, &pid);
+    }
+
+    fn read_consumption(&mut self, pid: u64) -> u64 {
+        let mut arg = pid;
+        if let Ok(_) = READ_POWER.ioctl(&mut self.file, &mut arg) {
+            arg
+        } else {
+            0
+        }
+    }
+}
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
new file mode 100644
index 0000000..eb23421
--- /dev/null
+++ b/src/energy/trackers/mock.rs
@@ -0,0 +1,14 @@
+use super::Estimator;
+
+#[derive(Default)]
+pub struct MockEstimator;
+
+impl Estimator for MockEstimator {
+    fn start_trace(&mut self, _pid: u64) {}
+
+    fn stop_trace(&mut self, _pid: u64) {}
+
+    fn read_consumption(&mut self, _pid: u64) -> u64 {
+        14
+    }
+}
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
new file mode 100644
index 0000000..19ec0be
--- /dev/null
+++ b/src/energy/trackers/perf.rs
@@ -0,0 +1,79 @@
+// energy/trackers/perf.rs
+use std::collections::HashMap;
+
+use perf_event::{
+    events::{Event, Hardware},
+    Builder, Counter, Group,
+};
+
+use crate::energy::Estimator;
+
+#[derive(Default)]
+pub struct PerfEstimator {
+    registry: HashMap<u64, Counters>,
+}
+
+struct Counters {
+    group: Group,
+    counters: Vec<Counter>,
+}
+
+static EVENT_TYPES: &[(f32, Event)] = &[
+    (1.0, Event::Hardware(Hardware::CPU_CYCLES)),
+    (2.0, Event::Hardware(Hardware::INSTRUCTIONS)),
+];
+
+impl Estimator for PerfEstimator {
+    fn start_trace(&mut self, pid: u64) {
+        let Ok(mut group) = Group::new() else {
+            eprintln!("Failed to create performance counter group for PID {}", pid);
+            return;
+        };
+
+        let counters: Result<Vec<_>, _> = EVENT_TYPES
+            .iter()
+            .map(|(_, kind)| Builder::new().group(&mut group).kind(kind.clone()).build())
+            .collect();
+
+        let counters = match counters {
+            Ok(counters) => counters,
+            Err(e) => {
+                eprintln!(
+                    "Failed to create performance counter group for PID {}: {}",
+                    pid, e
+                );
+                return;
+            }
+        };
+
+        if let Err(e) = group.enable() {
+            eprintln!("Failed to enable performance counters: {}", e);
+            return;
+        }
+
+        let counters = Counters { counters, group };
+        self.registry.insert(pid, counters);
+    }
+
+    fn stop_trace(&mut self, pid: u64) {
+        self.registry.remove(&pid);
+    }
+
+    fn read_consumption(&mut self, pid: u64) -> u64 {
+        let Some(counters) = self.registry.get_mut(&pid) else {
+            return 0;
+        };
+
+        let counts = match counters.group.read() {
+            Ok(counts) => counts,
+            Err(_) => return 0,
+        };
+
+        let mut sum = 0;
+        for ((factor, _ty), count) in EVENT_TYPES.iter().zip(counts.iter()) {
+            sum += *factor as u64 * count.1;
+        }
+
+        sum
+    }
+}