path: root/src/energy.rs

mod budget;
pub mod estimator;
pub mod rapl;
mod trackers;

use crate::energy::estimator::Estimator;
use std::collections::{BTreeSet, HashMap};
use std::ops::RangeInclusive;
use std::sync::atomic::{AtomicBool, AtomicI32, AtomicI64, AtomicU64};
use std::sync::{mpsc, Arc, RwLock};
use std::thread;
use std::time::{Duration, Instant};

use crate::freq::FrequencyKHZ;
use crate::socket;
use crate::Pid;

pub use budget::BudgetPolicy;
pub use trackers::{KernelDriver, PerfEstimator};

const IDLE_CONSUMPTION_W: f64 = 7.;
const UPDATE_INTERVAL: Duration = Duration::from_millis(10);

pub enum Request {
    NewTask(Pid, Arc<TaskInfo>),
    RemoveTask(Pid),
}
pub struct TaskInfo {
    pub cpu: AtomicI32,
    pub budget: AtomicU64,
    pub running_on_e_core: AtomicBool,
    pub last_scheduled: AtomicI64,
}

impl TaskInfo {
    pub fn read_cpu(&self) -> i32 {
        self.cpu.load(std::sync::atomic::Ordering::Relaxed)
    }
    pub fn read_budget(&self) -> u64 {
        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)
    }
    pub fn read_time_since_last_schedule(&self) -> Option<Duration> {
        let old_time = self
            .last_scheduled
            .load(std::sync::atomic::Ordering::Relaxed);
        if old_time == -1 {
            None
        } else {
            let now = chrono::Utc::now().timestamp_micros();
            Some(Duration::from_micros((now - old_time) as u64))
        }
    }
    pub fn read_time_since_last_schedule_raw(&self) -> i64 {
        self.last_scheduled
            .load(std::sync::atomic::Ordering::Relaxed)
    }
    pub fn set_cpu(&self, cpu: i32) {
        self.cpu.store(cpu, std::sync::atomic::Ordering::Relaxed);
    }
    pub fn set_last_scheduled_to_now(&self) {
        self.last_scheduled.store(
            chrono::Utc::now().timestamp_micros(),
            std::sync::atomic::Ordering::Relaxed,
        );
    }
    pub fn set_last_scheduled_raw(&self, last_scheduled: i64) {
        self.last_scheduled
            .store(last_scheduled, std::sync::atomic::Ordering::Relaxed);
    }
    pub fn set_budget(&self, budget: u64) {
        self.budget
            .store(budget, std::sync::atomic::Ordering::Relaxed);
    }
    pub fn set_running_on_e_core(&self, running_on_e_core: bool) {
        self.running_on_e_core
            .store(running_on_e_core, std::sync::atomic::Ordering::Relaxed);
    }
}

impl Default for TaskInfo {
    fn default() -> Self {
        Self {
            cpu: Default::default(),
            budget: AtomicU64::new(u64::MAX),
            running_on_e_core: Default::default(),
            last_scheduled: AtomicI64::new(-1),
        }
    }
}

#[derive(Clone)]
pub struct ProcessInfo {
    pub energy: f64,
    pub energy_delta: f64,
    pub tree_energy: f64,
    pub last_update: std::time::Instant,
    pub parent: Pid,
    pub task_info: Arc<TaskInfo>,
}

pub struct EnergyService {
    estimator: Box<dyn Estimator>,
    budget_policy: Option<Box<dyn BudgetPolicy>>,
    // contains the same data as the keys of process_info but having this reduces contention and
    // avoids unnecessary clone
    active_processes: BTreeSet<Pid>,
    process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
    request_receiver: mpsc::Receiver<Request>,
    update_interval: Duration,
    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
    rapl_offset: f64,
    last_energy_diff: f64,
    last_time_between_measurements: Duration,
    old_rapl: f64,
    system_energy: f64,
    bias: f64,
    graveyard: Vec<i32>,
    last_measurement: Instant,
}

impl EnergyService {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        estimator: Box<dyn Estimator>,
        budget_policy: Box<dyn BudgetPolicy>,
        process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
        request_receiver: mpsc::Receiver<Request>,
        update_interval: Duration,
        shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
        shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
        shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
    ) -> Self {
        Self {
            estimator,
            budget_policy: Some(budget_policy),
            active_processes: BTreeSet::new(),
            process_info,
            request_receiver,
            update_interval,
            shared_cpu_frequency_ranges,
            shared_policy_frequency_ranges,
            shared_cpu_current_frequencies,
            rapl_offset: rapl::read_package_energy().unwrap(),
            last_energy_diff: 0f64,
            last_time_between_measurements: Duration::new(0, 0),
            old_rapl: 0.,
            system_energy: 0.,
            bias: 1.,
            graveyard: Vec::with_capacity(100),
            last_measurement: Instant::now(),
        }
    }

    pub fn run(mut self) {
        thread::spawn(move || {
            let mut i = 0;
            loop {
                i += 1;
                // Process any incoming requests
                self.handle_requests();

                if i % 30 == 0 {
                    // Update energy measurements
                    self.update_measurements();

                    self.clear_graveyeard();

                    // 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() {
            self.handle_request(request);
        }
    }

    fn handle_request(&mut self, request: Request) {
        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();
                    let old_time = task_info.read_time_since_last_schedule_raw();
                    info.task_info = task_info.clone();
                    info.task_info.set_budget(old_budget);
                    info.task_info.set_last_scheduled_raw(old_time);
                    return;
                }
                if self
                    .estimator
                    .start_trace(
                        pid as u64,
                        task_info.read_cpu(),
                        task_info.is_running_on_e_core(),
                    )
                    .is_err()
                {
                    return;
                }
                let parent = (|| {
                    let process = procfs::process::Process::new(pid)?;
                    process.stat().map(|stat| stat.ppid)
                })()
                .unwrap_or_default();
                // We don't care whether the task has been scheduled before the counters are set up
                task_info.set_last_scheduled_raw(-1);
                self.process_info.write().unwrap().insert(
                    pid,
                    ProcessInfo {
                        energy: 0.,
                        energy_delta: 0.,
                        tree_energy: 0.,
                        last_update: std::time::Instant::now(),
                        parent,
                        task_info: task_info.clone(),
                    },
                );
                self.active_processes.insert(pid);
                if !self.process_info.read().unwrap().contains_key(&parent) && parent != 0 {
                    self.handle_request(Request::NewTask(parent, task_info));
                }
            }
            Request::RemoveTask(pid) => {
                if procfs::process::Process::new(pid).is_ok() {
                    return;
                }
                self.graveyard.push(pid);
            }
        }
    }

    fn update_measurements(&mut self) {
        let old_energy = self
            .process_info
            .read()
            .unwrap()
            .get(&1)
            .map(|info| info.tree_energy)
            .unwrap_or(0.);
        for pid in &self.active_processes {
            let mut process_info = self.process_info.write().unwrap();
            if let Some(info) = process_info.get_mut(pid) {
                if info
                    .task_info
                    .read_time_since_last_schedule()
                    .unwrap_or(UPDATE_INTERVAL)
                    >= UPDATE_INTERVAL
                {
                    continue;
                }
                if let Some(energy) = self.estimator.read_consumption(*pid as u64) {
                    info.energy_delta = energy * self.bias;
                    info.energy += energy * self.bias;
                    info.tree_energy += energy * self.bias;
                    self.estimator.update_information(
                        *pid as u64,
                        info.task_info.read_cpu(),
                        info.task_info.is_running_on_e_core(),
                    );
                    let mut parent = info.parent;
                    while let Some(info) = process_info.get_mut(&parent) {
                        info.tree_energy += energy * self.bias;
                        info.last_update = std::time::Instant::now();
                        parent = info.parent;
                    }
                }
            }
        }
        let elapsed = self.last_measurement.elapsed();
        self.last_time_between_measurements = elapsed;
        self.last_measurement = Instant::now();
        let rapl = rapl::read_package_energy().unwrap() - self.rapl_offset;
        let rapl_diff = rapl - self.old_rapl;
        self.last_energy_diff = rapl_diff;
        self.old_rapl = rapl;
        let power_comsumption_watt = rapl_diff / elapsed.as_secs_f64();
        let idle_consumption = elapsed.as_secs_f64() * IDLE_CONSUMPTION_W;
        if let Some(init) = self.process_info.write().unwrap().get_mut(&1) {
            let est_diff = init.tree_energy - old_energy + idle_consumption;
            // let offset_bias = (rapl / (init.tree_energy + idle_consumption)).clamp(0.1, 2.);
            let current_bias = if init.tree_energy - old_energy > idle_consumption * 0.5 {
                (rapl_diff / est_diff).clamp(0.1, 2.)
            } else {
                1.
            };
            // let current_bias = (offset_bias + diff_bias) * 0.5;
            let alpha: f64 = 10. * elapsed.as_secs_f64().recip();
            self.bias = (self.bias * (alpha.recip() * current_bias + ((alpha - 1.) / alpha)))
                .clamp(0.1, 5.);
            self.system_energy += est_diff;
            println!(
                "Energy estimation: {:.1} rapl: {:.1}, est diff: {:.1} rapl diff: {:.1}, bias: {:.1}, power consumption: {:.1}",
                self.system_energy, rapl, est_diff, rapl_diff, self.bias, power_comsumption_watt,
            );
        }
    }

    fn update_budgets(&mut 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 {
            if let Some(entry) = self.process_info.write().unwrap().get(&pid) {
                entry.task_info.set_budget(budget);
            }
        }
    }

    fn clear_graveyeard(&mut self) {
        for pid in self.graveyard.drain(..) {
            self.estimator.stop_trace(pid as u64);
            self.active_processes.remove(&pid);
            self.process_info.write().unwrap().remove(&pid);
        }
    }

    // Accessor methods for BudgetPolicy
    pub fn active_processes(&self) -> &BTreeSet<Pid> {
        &self.active_processes
    }

    fn package_energy(&mut self) -> f64 {
        rapl::read_package_energy().unwrap()
    }

    pub fn process_energy(&self, pid: Pid) -> Option<f64> {
        self.process_info
            .read()
            .unwrap()
            .get(&pid)
            .map(|info| info.energy)
    }

    pub fn all_process_energy_deltas(&self) -> HashMap<Pid, f64> {
        self.process_info
            .read()
            .unwrap()
            .iter()
            .map(|(&key, info)| (key, info.energy_delta))
            .collect()
    }
}

pub fn start_energy_service(
    use_mocking: bool,
    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
) -> std::io::Result<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::new(shared_cpu_current_frequencies.clone()))
    } else {
        Box::new(KernelDriver::default())
    };

    let process_info = Arc::new(RwLock::new(HashMap::new()));

    let power_cap = socket::start_logging_socket_service("/tmp/pm-sched", process_info.clone())?;

    // 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 {}

    // Create and start the energy service
    let service = EnergyService::new(
        estimator,
        budget_policy,
        process_info.clone(),
        request_receiver,
        UPDATE_INTERVAL,
        shared_cpu_frequency_ranges,
        shared_policy_frequency_ranges,
        shared_cpu_current_frequencies,
    );

    service.run();

    Ok(request_sender)
}