Restructure project

3 files changed, 142 insertions, 0 deletions

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
+    }
+}