use crate::config::{Component, Target};
use crate::monitor::MonitorSample;
use std::fmt::Write as _;
use std::fs::{File, OpenOptions};
use std::io::{Seek, SeekFrom, Write};
use std::path::Path;
pub trait OutputSink {
fn send(&mut self, sample: &MonitorSample) -> crate::Result<()>;
}
#[derive(Clone, Copy, Debug)]
struct Column {
header: &'static str,
cell: fn(&MonitorSample) -> Cell,
}
#[derive(Clone, Copy, Debug, PartialEq)]
enum Cell {
Watts(f64),
Count(u64),
}
impl std::fmt::Display for Cell {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Cell::Watts(watts) => write!(f, "{watts:.2}"),
Cell::Count(count) => write!(f, "{count}"),
}
}
}
const TIMESTAMP: Column = Column {
header: "Timestamp",
cell: |sample| Cell::Count(sample.timestamp),
};
const TOTAL: Column = Column {
header: "Total Power (W)",
cell: |sample| Cell::Watts(sample.total_power()),
};
const CPU: Column = Column {
header: "CPU Power (W)",
cell: |sample| Cell::Watts(sample.cpu_power_or_zero()),
};
const GPU: Column = Column {
header: "GPU Power (W)",
cell: |sample| Cell::Watts(sample.gpu_power_or_zero()),
};
const USAGE: Column = Column {
header: "CPU Usage (%)",
cell: |sample| Cell::Watts(sample.cpu_usage),
};
const PROCESS: Column = Column {
header: "Process Power (W)",
cell: |sample| Cell::Watts(sample.target_power_or_zero()),
};
const APP: Column = Column {
header: "App Power (W)",
cell: |sample| Cell::Watts(sample.target_power_or_zero()),
};
const APP_PIDS: Column = Column {
header: "App PIDs",
cell: |sample| Cell::Count(sample.app_pid_count.unwrap_or(0) as u64),
};
#[derive(Clone, Copy, Debug)]
pub struct Schema {
columns: &'static [Column],
labelled: bool,
}
impl Schema {
#[must_use]
pub fn csv(component: Option<Component>, target: &Target) -> Self {
Self {
columns: match (component, target) {
(Some(Component::Cpu), _) => &[TIMESTAMP, CPU],
(Some(Component::Gpu), _) => &[TIMESTAMP, GPU],
(None, Target::System) => &[TIMESTAMP, TOTAL, CPU, GPU, USAGE],
(None, Target::Pid(_)) => &[TIMESTAMP, TOTAL, CPU, GPU, USAGE, PROCESS],
(None, Target::App(_)) => &[TIMESTAMP, TOTAL, CPU, GPU, USAGE, APP, APP_PIDS],
},
labelled: true,
}
}
#[must_use]
pub fn watts(component: Option<Component>) -> Self {
Self {
columns: match component {
Some(Component::Cpu) => &[CPU],
Some(Component::Gpu) => &[GPU],
None => &[TOTAL],
},
labelled: false,
}
}
#[must_use]
pub fn header(&self) -> String {
if !self.labelled {
return String::new();
}
self.columns
.iter()
.map(|column| column.header)
.collect::<Vec<_>>()
.join(",")
}
#[must_use]
pub fn row(&self, sample: &MonitorSample) -> String {
let mut out = String::new();
self.write_row(&mut out, sample);
out
}
fn write_row(&self, out: &mut String, sample: &MonitorSample) {
for (index, column) in self.columns.iter().enumerate() {
if index > 0 {
out.push(',');
}
let _ = write!(out, "{}", (column.cell)(sample));
}
}
}
pub struct FileWriter {
file: File,
schema: Schema,
overwrite: bool,
scratch: String,
}
impl FileWriter {
pub fn open(path: impl AsRef<Path>, schema: Schema, overwrite: bool) -> crate::Result<Self> {
let file = if overwrite {
OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)?
} else {
OpenOptions::new().create(true).append(true).open(path)?
};
Ok(Self {
file,
schema,
overwrite,
scratch: String::with_capacity(128),
})
}
pub fn csv(
path: impl AsRef<Path>,
config: &crate::config::MonitorConfig,
) -> crate::Result<Self> {
let schema = Schema::csv(config.component, &config.target);
let mut writer = Self::open(path, schema, false)?;
writer.write_header()?;
Ok(writer)
}
pub fn write_header(&mut self) -> crate::Result<()> {
let header = self.schema.header();
if header.is_empty() || self.overwrite {
return Ok(());
}
writeln!(self.file, "{header}")?;
Ok(())
}
}
impl OutputSink for FileWriter {
fn send(&mut self, sample: &MonitorSample) -> crate::Result<()> {
if self.overwrite {
self.file.set_len(0)?;
self.file.seek(SeekFrom::Start(0))?;
}
let row = &mut self.scratch;
row.clear();
self.schema.write_row(row, sample);
row.push('\n');
self.file.write_all(row.as_bytes())?;
Ok(())
}
}
#[derive(Default)]
pub struct OutputBundle {
sinks: Vec<Box<dyn OutputSink>>,
}
impl OutputBundle {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn push(&mut self, sink: impl OutputSink + 'static) {
self.sinks.push(Box::new(sink));
}
#[must_use]
pub fn with(mut self, sink: impl OutputSink + 'static) -> Self {
self.push(sink);
self
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.sinks.is_empty()
}
}
impl OutputSink for OutputBundle {
fn send(&mut self, sample: &MonitorSample) -> crate::Result<()> {
for sink in &mut self.sinks {
sink.send(sample)?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sample() -> MonitorSample {
MonitorSample {
timestamp: 1_700_000_000,
cpu_power: Some(12.345),
gpu_power: Some(6.0),
cpu_usage: 42.5,
target_power: None,
app_pid_count: None,
}
}
fn csv(sample: &MonitorSample, component: Option<Component>, target: &Target) -> String {
let schema = Schema::csv(component, target);
format!("{}\n{}\n", schema.header(), schema.row(sample))
}
#[test]
fn system_csv_has_header_and_row() {
assert_eq!(
csv(&sample(), None, &Target::System),
"Timestamp,Total Power (W),CPU Power (W),GPU Power (W),CPU Usage (%)\n\
1700000000,18.34,12.35,6.00,42.50\n"
);
}
#[test]
fn single_component_csv_reports_only_that_component() {
assert_eq!(
csv(&sample(), Some(Component::Cpu), &Target::System),
"Timestamp,CPU Power (W)\n1700000000,12.35\n"
);
assert_eq!(
csv(&sample(), Some(Component::Gpu), &Target::System),
"Timestamp,GPU Power (W)\n1700000000,6.00\n"
);
}
#[test]
fn process_and_app_targets_add_their_columns() {
let mut process = sample();
process.target_power = Some(3.5);
let rendered = csv(&process, None, &Target::Pid(42));
assert!(rendered.starts_with(
"Timestamp,Total Power (W),CPU Power (W),GPU Power (W),CPU Usage (%),Process Power (W)\n"
));
assert!(rendered.ends_with("1700000000,18.34,12.35,6.00,42.50,3.50\n"));
let mut app = sample();
app.target_power = Some(4.25);
app.app_pid_count = Some(7);
let rendered = csv(&app, None, &Target::App("firefox".into()));
assert!(rendered.contains("App Power (W),App PIDs\n"));
assert!(rendered.ends_with("1700000000,18.34,12.35,6.00,42.50,4.25,7\n"));
}
#[test]
fn unmeasured_components_are_written_as_zero() {
let mut sample = sample();
sample.cpu_power = None;
let rendered = csv(&sample, None, &Target::System);
assert!(rendered.ends_with("1700000000,6.00,0.00,6.00,42.50\n"));
}
#[test]
fn a_watts_schema_renders_one_bare_number_and_no_header() {
let total = Schema::watts(None);
assert_eq!(total.header(), "");
assert_eq!(total.row(&sample()), "18.34");
assert_eq!(Schema::watts(Some(Component::Cpu)).row(&sample()), "12.35");
assert_eq!(Schema::watts(Some(Component::Gpu)).row(&sample()), "6.00");
}
fn written(schema: Schema, overwrite: bool, samples: &[MonitorSample]) -> String {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("power.csv");
let mut writer = FileWriter::open(&path, schema, overwrite).unwrap();
writer.write_header().unwrap();
for sample in samples {
writer.send(sample).unwrap();
}
drop(writer);
std::fs::read_to_string(&path).unwrap()
}
#[test]
fn appending_writes_a_header_then_every_row() {
let schema = Schema::csv(None, &Target::System);
let mut first = sample();
first.timestamp = 1;
let mut second = sample();
second.timestamp = 2;
assert_eq!(
written(schema, false, &[first, second]),
"Timestamp,Total Power (W),CPU Power (W),GPU Power (W),CPU Usage (%)\n\
1,18.34,12.35,6.00,42.50\n\
2,18.34,12.35,6.00,42.50\n"
);
}
#[test]
fn overwrite_mode_keeps_only_the_latest_row_and_no_header() {
let schema = Schema::csv(None, &Target::System);
let mut first = sample();
first.timestamp = 1;
let mut second = sample();
second.timestamp = 2;
assert_eq!(
written(schema, true, &[first, second]),
"2,18.34,12.35,6.00,42.50\n"
);
}
#[test]
fn a_shorter_row_cannot_leave_a_tail_of_the_previous_one() {
let schema = Schema::watts(None);
let mut big = sample();
big.cpu_power = Some(1000.0);
let mut small = sample();
small.cpu_power = Some(1.0);
small.gpu_power = Some(0.0);
assert_eq!(written(schema, true, &[big, small]), "1.00\n");
}
#[test]
fn a_bundle_forwards_to_every_sink_it_holds() {
#[derive(Default)]
struct Counter(std::rc::Rc<std::cell::Cell<usize>>);
impl OutputSink for Counter {
fn send(&mut self, _sample: &MonitorSample) -> crate::Result<()> {
self.0.set(self.0.get() + 1);
Ok(())
}
}
let sends = std::rc::Rc::new(std::cell::Cell::new(0));
let mut outputs = OutputBundle::new();
assert!(outputs.is_empty());
outputs.send(&sample()).unwrap();
outputs.push(Counter(sends.clone()));
outputs.push(Counter(sends.clone()));
outputs.send(&sample()).unwrap();
assert_eq!(sends.get(), 2);
assert!(!outputs.is_empty());
}
#[test]
fn file_writer_csv_shortcut_and_bundle_with_builder() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("shortcut.csv");
let config = crate::config::MonitorConfig::default();
let writer = FileWriter::csv(&path, &config).unwrap();
let mut outputs = OutputBundle::new().with(writer);
outputs.send(&sample()).unwrap();
drop(outputs);
let content = std::fs::read_to_string(&path).unwrap();
assert!(
content.starts_with(
"Timestamp,Total Power (W),CPU Power (W),GPU Power (W),CPU Usage (%)\n"
)
);
assert!(content.contains("1700000000,18.34,12.35,6.00,42.50\n"));
}
}