Low-budget stats

2 files changed, 168 insertions, 7 deletions

diff --git a/src/db.rs b/src/db.rs
new file mode 100644
index 0000000..d4038cb
--- /dev/null
+++ b/src/db.rs
@@ -0,0 +1,64 @@
+use std::sync::Mutex;
+
+use camino::Utf8PathBuf;
+use rusqlite::{params, Connection, OptionalExtension};
+use time::OffsetDateTime;
+
+use crate::Meminfo;
+
+
+pub struct Database{ 
+    db_path: Utf8PathBuf,
+    conn: Mutex<Connection>
+}
+
+impl Database {
+    pub fn new(db_path: Utf8PathBuf) -> Self {
+        Self {
+            conn:Mutex::new(Connection::open(&db_path).unwrap()),
+            db_path,
+        }
+    }
+
+    pub fn create_tables(&self) {
+        let conn = self.conn.lock().unwrap();
+        conn.execute(CREATE_TABLE_HOSTMEM, params![]).unwrap();
+    }
+
+    pub fn insert_host_meminfo(&self, meminfo: Meminfo) {
+        let conn = self.conn.lock().unwrap();
+
+        conn.execute("INSERT INTO stats_hostmem(total_kb, available_kb) VALUES (?1, ?2)", params![meminfo.total, meminfo.avaialable]).unwrap();
+    }
+
+    pub fn get_last_host_meminfo(&self) -> DbMeminfo {
+        let conn = self.conn.lock().unwrap();
+
+        conn.query_row("SELECT * FROM stats_hostmem ORDER BY stamp DESC LIMIT 1", [], |row| {
+            let (stamp, total_kb, available_kb) = row.try_into().unwrap();
+
+            Ok(DbMeminfo {
+                stamp, total_kb, available_kb
+            })
+        }).optional().unwrap().unwrap()
+    }
+}
+
+pub const CREATE_TABLE_HOSTMEM: &'static str = "\
+    CREATE TABLE IF NOT EXISTS stats_hostmem (
+        stamp TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        total_kb INTEGER NOT NULL,
+        available_kb INTEGER NOT NULL
+    );";
+
+pub struct DbMeminfo {
+    pub stamp: OffsetDateTime,
+    pub total_kb: usize,
+    pub available_kb: usize
+}
+
+impl DbMeminfo {
+    pub fn usage(&self) -> usize {
+        self.total_kb - self.available_kb
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index fa1ac5c..26e1cf0 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,12 +7,13 @@ mod settings;
 mod templated;
 mod timeparse;
 mod util;
+mod db;
 
-use std::{io::Write, os::unix::fs::MetadataExt, str::FromStr};
+use std::{fs::File, io::{BufRead, BufReader, Write}, os::unix::fs::MetadataExt, str::FromStr, sync::Arc, time::Duration};
 
 use axum::{
 	body::Body,
-	extract::Path,
+	extract::{Path, State},
 	http::{header, StatusCode},
 	response::Response,
 	routing::get,
@@ -21,9 +22,11 @@ use axum::{
 use bempline::{variables, Document, Options};
 use camino::Utf8PathBuf;
 use confindent::{Confindent, Node};
+use db::Database;
 pub use error::RuntimeError;
 use fs::Filesystem;
 use settings::Settings;
+use templated::Frontmatter;
 use tokio_util::io::ReaderStream;
 use tracing_subscriber::{fmt::time, prelude::*, EnvFilter};
 use util::{Referer, RemoteIp, SessionId};
@@ -33,6 +36,11 @@ use crate::{
 	templated::Templated,
 };
 
+#[derive(Clone)]
+pub struct AwakeState {
+	pub database: Arc<Database>
+}
+
 #[tokio::main]
 async fn main() {
 	match std::env::args().nth(1).as_deref() {
@@ -55,9 +63,27 @@ async fn main() {
 	let webroot: Utf8PathBuf = conf.child_parse("Webroot").unwrap();
 	let templates = conf.child_value("Templates").unwrap();
 	let hostname = conf.child_owned("Hostname").unwrap();
+	let dbpath = conf.child_owned("Database").unwrap();
+
+	let database = Database::new(dbpath.into());
+	database.create_tables();
 
 	let fs = Filesystem::new(&webroot);
 
+	let state = AwakeState {
+		database: Arc::new(database)
+	};
+
+	let mi_state = state.clone();
+	let meminfo_thread = std::thread::spawn(move || {
+		loop {
+			let meminfo = Meminfo::current();
+			mi_state.database.insert_host_meminfo(meminfo);
+
+			std::thread::sleep(Duration::from_secs(60));
+		}
+	});
+
 	let settings = Settings {
 		template_dir: Utf8PathBuf::from(webroot.join(templates))
 			.canonicalize_utf8()
@@ -69,29 +95,31 @@ async fn main() {
 		.route("/", get(index_handler))
 		.route("/*path", get(handler))
 		.layer(Extension(fs))
-		.layer(Extension(settings));
+		.layer(Extension(settings)).with_state(state);
 
 	let listener = tokio::net::TcpListener::bind("0.0.0.0:2560").await.unwrap();
 	axum::serve(listener, app).await.unwrap()
 }
 
 async fn index_handler(
+	state: State<AwakeState>,
 	fse: Extension<Filesystem>,
 	se: Extension<Settings>,
 	sid: SessionId,
 	rfr: Option<Referer>,
 ) -> Response {
-	handler(fse, se, sid, rfr, Path(String::from("/"))).await
+	handler(state, fse, se, sid, rfr, Path(String::from("/"))).await
 }
 
 async fn handler(
+	State(state): State<AwakeState>,
 	Extension(fs): Extension<Filesystem>,
 	Extension(settings): Extension<Settings>,
 	sid: SessionId,
 	rfr: Option<Referer>,
 	Path(path): Path<String>,
 ) -> Response {
-	match falible_handler(fs, settings, sid, rfr, path).await {
+	match falible_handler(state, fs, settings, sid, rfr, path).await {
 		Ok(resp) => resp,
 		Err(re) => Response::builder()
 			.body(Body::from(re.to_string()))
@@ -100,6 +128,7 @@ async fn handler(
 }
 
 async fn falible_handler(
+	state: AwakeState,
 	fs: Filesystem,
 	settings: Settings,
 	sid: SessionId,
@@ -134,7 +163,7 @@ async fn falible_handler(
 		let result = Templated::from_str(&content);
 
 		match result {
-			Ok(templated) => send_template(templated, resolve, webpath, settings, sid).await,
+			Ok(templated) => send_template(templated, resolve, webpath, state, settings, sid).await,
 			Err(e) => {
 				tracing::warn!("error sending template {e}");
 
@@ -203,6 +232,7 @@ async fn send_template(
 	templated: Templated,
 	resolve: PathResolution,
 	webpath: Webpath,
+	state: AwakeState,
 	settings: Settings,
 	sid: SessionId,
 ) -> Result<Response, RuntimeError> {
@@ -335,7 +365,11 @@ async fn send_template(
 	tracing::trace!("finished published block");
 
 	// insert the page content itself
-	let markedup = markup::process(&templated.content);
+	let mut markedup = markup::process(&templated.content);
+	if templated.frontmatter.get("use-template").is_some() {
+		markedup = template_content(state, &templated.frontmatter, markedup);
+	}
+
 	template.set("main", markedup);
 
 	Ok(Response::builder()
@@ -343,3 +377,66 @@ async fn send_template(
 		.body(Body::from(template.compile()))
 		.unwrap())
 }
+
+fn template_content(state: AwakeState, frontmatter: &Frontmatter, marked: String) -> String {
+	let Ok(mut doc) = Document::from_str(&marked, Options::default()) else {
+		return marked
+	};
+
+	if frontmatter.get("system-stats").is_some() {
+		let mem = state.database.get_last_host_meminfo();
+
+		doc.set("stats.mem.total", mem.total_kb / 1024);
+		doc.set("stats.mem.usage", mem.usage() / 1024);
+	}
+
+	doc.compile()
+}
+
+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
+	}
+}
\ No newline at end of file