Fix compilation

2 files changed, 14 insertions, 11 deletions

diff --git a/src/energy.rs b/src/energy.rs
index 94e82f4..8c65f12 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -19,8 +19,9 @@ pub use budget::BudgetPolicy;
 pub use trackers::{KernelDriver, PerfEstimator};
 
 pub enum Request {
-    NewTask(Pid, TaskInfo),
+    NewTask(Pid, Arc<TaskInfo>),
     RemoveTask(Pid),
+    //TODO: the migrate message is probably redundant now
     MigrateTask {
         pid: Pid,
         new_cpu: i32,
@@ -41,7 +42,8 @@ impl TaskInfo {
         self.budget.load(std::sync::atomic::Ordering::Relaxed)
     }
     pub fn is_running_on_e_core(&self) -> bool {
-        self.running_on_e_core.load(std::sync::atomic::Ordering::Relaxed)
+        self.running_on_e_core
+            .load(std::sync::atomic::Ordering::Relaxed)
     }
     pub fn set_cpu(&self, cpu: i32) {
         self.cpu.store(cpu, std::sync::atomic::Ordering::Relaxed);
@@ -61,7 +63,7 @@ impl Default for TaskInfo {
         Self {
             cpu: Default::default(),
             budget: AtomicU64::new(u64::MAX),
-            running_on_e_core: Default::default()
+            running_on_e_core: Default::default(),
         }
     }
 }
@@ -151,12 +153,13 @@ impl EnergyService {
                 if self.process_info.read().unwrap().contains_key(&pid) {
                     self.handle_request(Request::MigrateTask {
                         pid,
-                        new_cpu: cpu,
-                        is_e_core,
+                        new_cpu: task_info.read_cpu(),
+                        is_e_core: task_info.is_running_on_e_core(),
                     });
                     return;
                 }
-                self.estimator.start_trace(pid as u64, is_e_core);
+                self.estimator
+                    .start_trace(pid as u64, task_info.is_running_on_e_core());
                 let parent = (|| {
                     let process = procfs::process::Process::new(pid)?;
                     process.stat().map(|stat| stat.ppid)
@@ -192,13 +195,13 @@ impl EnergyService {
                 is_e_core,
             } => {
                 if let Some(info) = self.process_info.write().unwrap().get_mut(&pid) {
-                    info.cpu = new_cpu;
-                    if info.running_on_e_core != is_e_core {
+                    info.task_info.set_cpu(new_cpu);
+                    if info.task_info.is_running_on_e_core() != is_e_core {
                         //TODO: maybe read consumption before removing?
                         self.estimator.stop_trace(pid as u64);
                         self.estimator.start_trace(pid as u64, is_e_core);
                     }
-                    info.running_on_e_core = is_e_core;
+                    info.task_info.set_running_on_e_core(is_e_core);
                 }
             }
         }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 7d2060b..300d98b 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -11,7 +11,7 @@ use std::collections::{HashMap, VecDeque};
 use std::mem::MaybeUninit;
 use std::ops::{Range, RangeInclusive};
 use std::process;
-use std::sync::atomic::{AtomicI32, AtomicU64};
+use std::sync::atomic::{AtomicBool, AtomicI32};
 use std::sync::mpsc::TrySendError;
 use std::sync::{mpsc, Arc, RwLock};
 use std::time::{Duration, Instant};
@@ -166,7 +166,7 @@ impl<'a> Scheduler<'a> {
                     // New task - register it with the energy service
                     let task_info = Arc::new(TaskInfo {
                         cpu: AtomicI32::new(task.cpu),
-                        running_on_e_core: is_e_core,
+                        running_on_e_core: AtomicBool::new(is_e_core),
                         ..Default::default()
                     });
                     e.insert(Task {