path: root/gatherer/src/thread.rs

use std::{
	fmt,
	fs::File,
	io::{BufRead, BufReader},
	sync::{Arc, atomic::Ordering},
	thread::JoinHandle,
	time::Duration,
};

use camino::{Utf8Path, Utf8PathBuf};
use time::OffsetDateTime;

use crate::{
	Database, State,
	db::DatabaseError,
	griph::{self, Style, TwoLineOrder},
};

#[derive(Clone, Debug)]
pub struct GathererConfig {
	pub network: String,
	pub graph_directory: Utf8PathBuf,
	pub graph_update: u8,
	pub graph_on_boot: bool,
}

impl GathererConfig {
	pub fn graph(&self, db: &Database, state: &State) {
		let root = &self.graph_directory;
		let mem = root.join("current_hostmeminfo.gif");
		let net = root.join("current_hostnetinfo.gif");
		let cpu = root.join("current_hostcpuinfo.gif");

		let meminfo = make_mem_graph(&db, &mem).unwrap();
		let net_upper = make_net_graph(&db, &net).unwrap();
		let cpu_upper = make_cpu_graph(&db, &cpu).unwrap();
		tracing::debug!("done generating graphs!");

		state
			.meminfo_upper_bound
			.store(meminfo.upper, Ordering::SeqCst);
		state.netinfo_upper_bound.store(net_upper, Ordering::SeqCst);
		state
			.meminfo_current
			.store(meminfo.current, Ordering::SeqCst);
		state.netinfo_upper_bound.store(net_upper, Ordering::SeqCst);
		state.cpuinfo_upper_bound.store(cpu_upper, Ordering::SeqCst);
		tracing::debug!("stored upper bounds in state");
	}
}

pub struct Gatherer {
	db: Arc<Database>,
	state: State,
	config: GathererConfig,
	hwnd: Option<JoinHandle<()>>,
}

impl Gatherer {
	pub fn new(db: Arc<Database>, state: State, config: GathererConfig) -> Self {
		Self {
			db,
			state,
			config,
			hwnd: None,
		}
	}

	pub fn start(&mut self) {
		let db = Arc::clone(&self.db);
		let state = self.state.clone();
		let config = self.config.clone();
		let hwnd = std::thread::spawn(|| task(db, state, config));
		tracing::info!("started gatherer thread!");

		self.hwnd = Some(hwnd);
	}
}

fn task(db: Arc<Database>, state: State, config: GathererConfig) {
	tracing::info!("starting gatherer thread");

	// If we collected a point less than a minute ago, like after
	// just being restarted, wait until it's been a minute
	let last_meminfo_time = db
		.get_last_host_meminfo()
		.unwrap()
		.map(|mem| mem.stamp)
		.unwrap_or_else(|| OffsetDateTime::UNIX_EPOCH);
	let since = OffsetDateTime::now_utc() - last_meminfo_time;
	if since < time::Duration::minutes(1) {
		let to_minute = time::Duration::minutes(1) - since;

		tracing::info!(
			"waiting for {}s to space a minute apart",
			to_minute.whole_seconds()
		);
		std::thread::sleep(to_minute.try_into().unwrap());
	}

	if config.graph_on_boot {
		tracing::debug!("generating graphs on boot");
		config.graph(&db, &state);
	}

	let mut last_netinfo: Option<Netinfo> = None;
	let mut last_cpuinfo: Option<Cpuinfo> = None;

	// this is a `let _` because otherwise the attribute was
	// making the compiler mad
	#[rustfmt::skip]
	let _ = loop {
		tracing::debug!("collecting stats");

		// Gather data
		let meminfo = Meminfo::current();
		let netinfo = Netinfo::current(&config.network);
		let cpuinfo = Cpuinfo::current();

		// Print traces, y'know, for tracing
		tracing::trace!("memory: {}MB used / {}MB total", meminfo.usage() / 1000, meminfo.total / 1000);
		tracing::trace!("net: rx {} / tx {}", data_human_fmt(netinfo.rx_bytes), data_human_fmt(netinfo.tx_bytes));
		tracing::trace!("cpu: user {} // nice {} // system {}", cpuinfo.user, cpuinfo.nice, cpuinfo.system);

		// Store stats in database
		if let Some(lni) = last_netinfo {
			let rx_delta = netinfo.rx_bytes - lni.rx_bytes;
			let tx_delta = netinfo.tx_bytes - lni.tx_bytes;

			db.insert_hostnet(60, rx_delta, tx_delta).unwrap();
		}
		last_netinfo = Some(netinfo);

		if let Some(lci) = last_cpuinfo {
			let user_delta = cpuinfo.user - lci.user;
			let nice_delta = cpuinfo.nice - lci.nice;
			let system_delta = cpuinfo.system - lci.system;

			db.insert_hostcpu(60, user_delta, nice_delta, system_delta).unwrap();
		}
		last_cpuinfo = Some(cpuinfo);

		db.insert_host_meminfo(meminfo.total, meminfo.avaialable).unwrap();

		let now = OffsetDateTime::now_utc();
		if now.minute() % config.graph_update == 0 {
			tracing::debug!("generating graphs on interval! interval = every {} minutes", config.graph_update);
			config.graph(&db, &state);
		}

		std::thread::sleep(Duration::from_secs(60));
	};
}

#[derive(Copy, Clone, Debug, Default)]
pub struct GraphInfo<T: Copy + fmt::Debug + Default> {
	pub upper: T,
	pub current: T,
}

impl<T: Copy + fmt::Debug + Default> GraphInfo<T> {
	pub fn new(upper: T, current: T) -> Self {
		Self { upper, current }
	}

	pub fn info_as<K: Copy + fmt::Debug + Default, F>(self, f: F) -> GraphInfo<K>
	where
		F: Fn(T) -> K,
	{
		GraphInfo {
			upper: f(self.upper),
			current: f(self.current),
		}
	}
}

pub fn make_mem_graph(db: &Database, path: &Utf8Path) -> Result<GraphInfo<u64>, DatabaseError> {
	tracing::debug!("generating meminfo graph");

	let infos = db.get_last_n_host_meminfo(256)?;
	let last_info = infos
		.get(0)
		.map(|db| GraphInfo::new(db.total_kb, db.total_kb.saturating_sub(db.available_kb)))
		.unwrap_or_default();

	let now = get_clean_now();
	let cleaned = clean_series(&infos, |mem| mem.stamp, now);
	let mut usages: Vec<Option<usize>> = cleaned
		.into_iter()
		.map(|mi| mi.map(|mi| mi.usage()))
		.collect();

	// Reversing here because we want latest valeus on on the
	// right side, so last in the array
	usages.reverse();

	let gif = griph::make_1line(0, last_info.upper, &usages, Style::Discontinuous);
	gif.save(path).unwrap();
	tracing::trace!("saved graph to {path}");

	Ok(last_info.info_as(|n| n as u64))
}

const NET_INCREMENTS: usize = 64;

pub fn make_net_graph(db: &Database, path: &Utf8Path) -> Result<u64, DatabaseError> {
	tracing::debug!("generating netinfo graph");

	let now = get_clean_now();
	let infos = db.get_last_n_hostnet(256)?;
	let cleaned = clean_series(&infos, |net| net.stamp, now);

	// 125 is (1000 / 8) so it converst Bytes to kiloBITS
	let mut rx_deltas = extract(&cleaned, |ni| ni.rx_bytes_per_sec() as usize / 125);
	let mut tx_deltas = extract(&cleaned, |ni| ni.tx_bytes_per_sec() as usize / 125);

	// Reversing to put latest values on the right side
	rx_deltas.reverse();
	tx_deltas.reverse();

	let rx_zeroed: Vec<usize> = cleaned
		.iter()
		.map(|m| m.map(|n| n.rx_bytes_per_sec() as usize / 125).unwrap_or(0))
		.collect();

	let tx_zeroed: Vec<usize> = cleaned
		.iter()
		.map(|m| m.map(|n| n.tx_bytes_per_sec() as usize / 125).unwrap_or(0))
		.collect();

	let rx_sum = rx_zeroed.iter().fold(0, |acc, x| acc + x);
	let tx_sum = tx_zeroed.iter().fold(0, |acc, x| acc + x);

	// Mixing the TX/RX delta so we can pick a range.
	let mut mixed = vec![0; 512];
	mixed[..256].copy_from_slice(&rx_zeroed);
	mixed[256..].copy_from_slice(&tx_zeroed);

	mixed.sort();
	let kinda_highest = mixed[511 - 32];
	// We add one here so we have less "surface area" covered by the lines
	// which makes the vertical divisions look btter
	let high_bound_whole = (kinda_highest as f32 / NET_INCREMENTS as f32)
		.ceil()
		.max(1.0) as usize;

	let high_bound = if high_bound_whole > 1 {
		high_bound_whole + 1
	} else {
		high_bound_whole
	} * NET_INCREMENTS;

	let order = if tx_sum > rx_sum {
		TwoLineOrder::SeriesAFirst
	} else {
		TwoLineOrder::SeriesBFirst
	};

	tracing::debug!("tx_sum = {tx_sum} // rx_sum = {rx_sum} // order = {order:?}");

	let gif = griph::make_2line(0, high_bound, &tx_deltas, &rx_deltas, order);
	gif.save(path).unwrap();
	tracing::trace!("saved graph to {path}");

	Ok(high_bound as u64)
}

pub fn make_cpu_graph(db: &Database, path: &Utf8Path) -> Result<u8, DatabaseError> {
	tracing::debug!("generating cpuinfo graph");

	let infos = db.get_last_n_hostcpu(256)?;

	let now = get_clean_now();
	let cleaned = clean_series(&infos, |cpu| cpu.stamp, now);

	// Usages are of the unit of hundreths of a second of CPU usage averaged over a minute.
	// A value of 1 means the cpu saw 1% CPU usage
	let usages = extract(&cleaned, |cpu| cpu.average_usage());

	let mut zeroed: Vec<f32> = usages.iter().map(|m| m.unwrap_or(0.0)).collect();
	zeroed.sort_by(|a, b| a.partial_cmp(&b).unwrap());
	let kinda_highest = zeroed[255 - 8];

	tracing::debug!(
		"kinda_highest = {kinda_highest} // highest = {}",
		zeroed[255]
	);

	// high_bound unit: hundreths of a second (1% / 1)
	// scaling by 1.33 here to drive the line down to decrease the apparent
	// "surface area" of the line
	let high_bound = (kinda_highest * 1.33).max(1.0);

	// can't scale the range, so we have to scale the inputs.
	// Finally, we multiply by ten to put it in a 0-1000 range
	let scale_factor = 100.0 / high_bound;
	let mut scaled: Vec<Option<usize>> = usages
		.iter()
		.map(|m| m.map(|v| (v * scale_factor) as usize * 10))
		.collect();

	// Reversing here because we want latest valeus on on the
	// right side, so last in the array
	scaled.reverse();

	let gif = griph::make_1line(0, 1000, &scaled, Style::Line);
	gif.save(path).unwrap();
	tracing::trace!("saved graph to {path}");

	Ok(high_bound as u8)
}

fn get_clean_now() -> OffsetDateTime {
	OffsetDateTime::now_utc().replace_second(0).unwrap()
}

fn clean_series<T, F>(series: &[T], time_extractor: F, end_time: OffsetDateTime) -> [Option<T>; 256]
where
	F: Fn(&T) -> OffsetDateTime,
	T: Clone,
{
	let mut res = [const { None }; 256];

	for value in series {
		let time = time_extractor(value);
		let delta = end_time - time;
		let mins = delta.whole_minutes();

		if mins >= 0 && mins < 256 {
			res[mins as usize] = Some(value.clone());
		}
	}

	res
}

fn extract<T, F, V>(series: &[Option<T>], extractor: F) -> [Option<V>; 256]
where
	F: Fn(&T) -> V,
{
	let mut res = [const { None }; 256];
	for (idx, maybe) in series.iter().enumerate() {
		if let Some(value) = maybe {
			res[idx] = Some(extractor(value));
		}
	}

	res
}

pub struct Meminfo {
	pub total: usize,
	pub free: usize,
	pub avaialable: usize,
}

impl Meminfo {
	pub fn current() -> Self {
		let procinfo = File::open("/proc/meminfo").unwrap();
		let bread = BufReader::new(procinfo);

		let mut meminfo = Meminfo {
			total: 0,
			free: 0,
			avaialable: 0,
		};

		for line in bread.lines() {
			let line = line.unwrap();

			if let Some((raw_key, raw_value_kb)) = line.split_once(':') {
				let value = if let Some(raw_value) = raw_value_kb.trim().strip_suffix(" kB") {
					if let Ok(parsed) = raw_value.parse() {
						parsed
					} else {
						continue;
					}
				} else {
					continue;
				};

				match raw_key.trim() {
					"MemTotal" => meminfo.total = value,
					"MemFree" => meminfo.free = value,
					"MemAvailable" => meminfo.avaialable = value,
					_ => (),
				}
			}
		}

		meminfo
	}

	pub fn usage(&self) -> usize {
		self.total - self.avaialable
	}
}

pub struct Netinfo {
	rx_bytes: usize,
	tx_bytes: usize,
}

impl Netinfo {
	pub fn current(network: &str) -> Self {
		let procinfo = File::open("/proc/net/dev").unwrap();
		let bread = BufReader::new(procinfo);

		let mut netinfo = Self {
			rx_bytes: 0,
			tx_bytes: 0,
		};

		let interface = format!("{network}:");
		for line in bread.lines() {
			let line = line.unwrap();
			let trim = line.trim();

			if let Some(data) = trim.strip_prefix(&interface) {
				let mut splits = data.split_whitespace();
				netinfo.rx_bytes = splits.next().unwrap().parse().unwrap();
				netinfo.tx_bytes = splits.skip(7).next().unwrap().parse().unwrap();
				break;
			}
		}

		netinfo
	}
}

/// Not strictly the correct name for this, but that's fine.
/// Structure for parsed-from /proc/stat
pub struct Cpuinfo {
	user: usize,
	nice: usize,
	system: usize,
}

impl Cpuinfo {
	pub fn current() -> Self {
		let procinfo = File::open("/proc/stat").unwrap();
		let bread = BufReader::new(procinfo);

		let mut cpuinfo = Cpuinfo {
			user: 0,
			nice: 0,
			system: 0,
		};

		for line in bread.lines() {
			let line = line.unwrap();
			let Some(data) = line.strip_prefix("cpu  ") else {
				continue;
			};

			let mut splits = data.split(' ');
			cpuinfo.user = splits.next().unwrap().parse().unwrap();
			cpuinfo.nice = splits.next().unwrap().parse().unwrap();
			cpuinfo.system = splits.next().unwrap().parse().unwrap();
		}

		cpuinfo
	}
}

fn data_human_fmt(bytes: usize) -> String {
	let (num, unit) = data_human(bytes);
	format!("{num}{unit}")
}

fn data_human(bytes: usize) -> (f32, &'static str) {
	const UNITS: &[&str] = &["B", "kB", "MB", "GB", "TB"];

	let mut wrk = bytes as f32;
	let mut unit_idx = 0;
	loop {
		if wrk < 1500.0 || unit_idx == UNITS.len() - 1 {
			return (wrk, UNITS[unit_idx]);
		}
		wrk /= 1000.0;
		unit_idx += 1;
	}
}