Use rapl package energy to calculate new budgets

1 files changed, 39 insertions, 5 deletions

diff --git a/src/energy.rs b/src/energy.rs
index c6cb84b..64ec267 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -11,7 +11,10 @@ use crate::Pid;
 use anyhow::Result;
 use dashmap::DashMap;
 
+use crate::energy::rapl::PackageEnergy;
 pub use budget::BudgetPolicy;
+use perf_event::events::Event;
+use perf_event::{Builder, Counter};
 pub use trackers::{Estimator, KernelDriver, MockEstimator, PerfEstimator};
 
 pub enum Request {
@@ -27,12 +30,14 @@ pub struct ProcessInfo {
 
 pub struct EnergyService {
     estimator: Box<dyn Estimator>,
-    budget_policy: Box<dyn BudgetPolicy>,
+    budget_policy: Option<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,
+    package_energy: PackageEnergy,
+    package_energy_counter: Counter,
 }
 
 impl EnergyService {
@@ -42,15 +47,19 @@ impl EnergyService {
         shared_budgets: Arc<DashMap<Pid, u64>>,
         request_receiver: mpsc::Receiver<Request>,
         update_interval: Duration,
+        package_energy: PackageEnergy,
+        package_energy_counter: Counter,
     ) -> Self {
         Self {
             estimator,
-            budget_policy,
+            budget_policy: Some(budget_policy),
             active_processes: BTreeSet::new(),
             process_info: HashMap::new(),
             shared_budgets,
             request_receiver,
             update_interval,
+            package_energy,
+            package_energy_counter,
         }
     }
 
@@ -115,7 +124,10 @@ impl EnergyService {
     }
 
     fn update_budgets(&mut self) {
-        let budgets = self.budget_policy.calculate_budgets(self);
+        // We can't call self.budget_policy.calculate_budgets(self) directly because the first self borrows immutable and the self second borrows mutable
+        let policy = self.budget_policy.take().unwrap();
+        let budgets = policy.calculate_budgets(self);
+        self.budget_policy = Some(policy);
 
         // Update the shared budgets map
         for (pid, budget) in budgets {
@@ -130,6 +142,11 @@ impl EnergyService {
         &self.active_processes
     }
 
+    fn package_energy(&mut self) -> f64 {
+        self.package_energy
+            .to_joules(self.package_energy_counter.read().unwrap_or(0))
+    }
+
     pub fn process_energy(&self, pid: Pid) -> Option<u64> {
         self.process_info.get(&pid).map(|info| info.energy)
     }
@@ -146,7 +163,7 @@ pub fn start_energy_service(
     use_mocking: bool,
     power_cap: u64,
     shared_budgets: Arc<DashMap<Pid, u64>>,
-) -> mpsc::SyncSender<Request> {
+) -> std::io::Result<mpsc::SyncSender<Request>> {
     // Potentially convert back to bounded channel
     let (request_sender, request_receiver) = mpsc::sync_channel(10000);
 
@@ -160,6 +177,21 @@ pub fn start_energy_service(
     // Create budget policy
     let budget_policy = Box::new(budget::SimpleCappingPolicy::new(power_cap));
 
+    // shouldn't be a problem because we are privileged
+    // if PackageEnergy::check_paranoid().unwrap_or(3) > 0 {}
+
+    let package_energy = PackageEnergy::new().unwrap();
+
+    // Any cpu should be fine because the package power should be the same for every cpu
+    let mut package_energy_counter = Builder::new()
+        .kind(Event::Raw(
+            package_energy.get_type(),
+            package_energy.get_event_config(),
+        ))
+        .any_cpu()
+        .build()?;
+    package_energy_counter.enable()?;
+
     // Create and start the energy service
     let service = EnergyService::new(
         estimator,
@@ -167,9 +199,11 @@ pub fn start_energy_service(
         shared_budgets,
         request_receiver,
         Duration::from_millis(50), // 50ms update interval
+        package_energy,
+        package_energy_counter,
     );
 
     service.run();
 
-    request_sender
+    Ok(request_sender)
 }