Restructure project

1 files changed, 174 insertions, 0 deletions

diff --git a/src/energy.rs b/src/energy.rs
new file mode 100644
index 0000000..c8da2cf
--- /dev/null
+++ b/src/energy.rs
@@ -0,0 +1,174 @@
+mod budget;
+mod trackers;
+
+use std::collections::{BTreeSet, HashMap};
+use std::sync::{mpsc, Arc};
+use std::thread;
+use std::time::Duration;
+
+use crate::Pid;
+use anyhow::Result;
+use dashmap::DashMap;
+
+pub use budget::BudgetPolicy;
+pub use trackers::{Estimator, KernelDriver, MockEstimator, PerfEstimator};
+
+pub enum Request {
+    NewTask(Pid),
+    RemoveTask(Pid),
+}
+
+pub struct ProcessInfo {
+    energy: u64,
+    last_update: std::time::Instant,
+    parent: Pid,
+}
+
+pub struct EnergyService {
+    estimator: Box<dyn Estimator>,
+    budget_policy: Box<dyn BudgetPolicy>,
+    active_processes: BTreeSet<Pid>,
+    process_info: HashMap<Pid, ProcessInfo>,
+    shared_budgets: Arc<DashMap<Pid, u64>>,
+    request_receiver: mpsc::Receiver<Request>,
+    update_interval: Duration,
+}
+
+impl EnergyService {
+    pub fn new(
+        estimator: Box<dyn Estimator>,
+        budget_policy: Box<dyn BudgetPolicy>,
+        shared_budgets: Arc<DashMap<Pid, u64>>,
+        request_receiver: mpsc::Receiver<Request>,
+        update_interval: Duration,
+    ) -> Self {
+        Self {
+            estimator,
+            budget_policy,
+            active_processes: BTreeSet::new(),
+            process_info: HashMap::new(),
+            shared_budgets,
+            request_receiver,
+            update_interval,
+        }
+    }
+
+    pub fn run(mut self) {
+        thread::spawn(move || {
+            loop {
+                // Process any incoming requests
+                self.handle_requests();
+
+                // Update energy measurements
+                self.update_measurements();
+
+                // Calculate and update budgets
+                self.update_budgets();
+
+                // Sleep for update interval
+                thread::sleep(self.update_interval);
+            }
+        });
+    }
+
+    fn handle_requests(&mut self) {
+        while let Ok(request) = self.request_receiver.try_recv() {
+            match request {
+                Request::NewTask(pid) => {
+                    self.estimator.start_trace(pid as u64);
+                    self.active_processes.insert(pid);
+                    let parent = (|| {
+                        let process = procfs::process::Process::new(pid)?;
+                        process.stat().map(|stat| stat.ppid)
+                    })()
+                    .unwrap_or_default();
+                    self.process_info.insert(
+                        pid,
+                        ProcessInfo {
+                            energy: 0,
+                            last_update: std::time::Instant::now(),
+                            parent,
+                        },
+                    );
+                    // Initialize with default budget
+                    self.shared_budgets.insert(pid, u64::MAX);
+                }
+                Request::RemoveTask(pid) => {
+                    self.estimator.stop_trace(pid as u64);
+                    self.active_processes.remove(&pid);
+                    self.process_info.remove(&pid);
+                    self.shared_budgets.remove(&pid);
+                }
+            }
+        }
+    }
+
+    fn update_measurements(&mut self) {
+        for &pid in &self.active_processes {
+            let energy = self.estimator.read_consumption(pid as u64);
+            if let Some(info) = self.process_info.get_mut(&pid) {
+                info.energy = energy;
+                info.last_update = std::time::Instant::now();
+            }
+        }
+    }
+
+    fn update_budgets(&mut self) {
+        let budgets = self.budget_policy.calculate_budgets(self);
+
+        // Update the shared budgets map
+        for (pid, budget) in budgets {
+            if let Some(mut entry) = self.shared_budgets.try_get_mut(&pid).try_unwrap() {
+                *entry = budget;
+            }
+        }
+    }
+
+    // Accessor methods for BudgetPolicy
+    pub fn active_processes(&self) -> &BTreeSet<Pid> {
+        &self.active_processes
+    }
+
+    pub fn process_energy(&self, pid: Pid) -> Option<u64> {
+        self.process_info.get(&pid).map(|info| info.energy)
+    }
+
+    pub fn all_process_energies(&self) -> HashMap<Pid, u64> {
+        self.process_info
+            .iter()
+            .map(|(&pid, info)| (pid, info.energy))
+            .collect()
+    }
+}
+
+pub fn start_energy_service(
+    use_mocking: bool,
+    power_cap: u64,
+    shared_budgets: Arc<DashMap<Pid, u64>>,
+) -> mpsc::SyncSender<Request> {
+    // Potentially convert back to bounded channel
+    let (request_sender, request_receiver) = mpsc::sync_channel(10000);
+
+    // Create the appropriate estimator based on configuration
+    let estimator: Box<dyn Estimator> = if use_mocking {
+        Box::new(PerfEstimator::default())
+    } else {
+        Box::new(KernelDriver::default())
+    };
+
+    // Create budget policy
+    let budget_policy = Box::new(budget::SimpleCappingPolicy::new(power_cap));
+
+    // Create and start the energy service
+    let service = EnergyService::new(
+        estimator,
+        budget_policy,
+        shared_budgets,
+        request_receiver,
+        Duration::from_millis(50), // 50ms update interval
+    );
+
+    service.run();
+
+    request_sender
+}