Fix budget reset

1 files changed, 2 insertions, 0 deletions

diff --git a/src/energy.rs b/src/energy.rs
index 439792d..35ead3b 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -145,7 +145,9 @@ impl EnergyService {
         match request {
             Request::NewTask(pid, task_info) => {
                 if let Some(info) = self.process_info.write().unwrap().get_mut(&pid) {
+                    let old_budget = task_info.read_budget();
                     info.task_info = task_info.clone();
+                    info.task_info.set_budget(old_budget);
                     return;
                 }
                 self.estimator.start_trace(