initial implementation

temperature code is haunted

.gitignore

@@ -1 +1,2 @@
 /target
+.env

Cargo.toml

@@ -5,3 +5,11 @@ authors = ["Robin Appelman <robin@icewind.nl>"]
 edition = "2018"
 
 [dependencies]
+libzetta = "0.2"
+color-eyre = "0.5"
+heim = { version = "0.1.0-rc.1", features = ["sensors", "cpu", "memory", "net", "host"] }
+warp = "0.3"
+tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
+ctrlc = { version = "3", features = ["termination"] }
+dotenv = "0.15"
+futures-util = "0.3"

src/heim.rs

@@ -0,0 +1,91 @@
+use color_eyre::Result;
+use futures_util::stream::{Stream, StreamExt};
+use futures_util::{future, TryStreamExt};
+use heim::units::{information, ratio, thermodynamic_temperature};
+use std::time::Duration;
+use tokio::time::sleep;
+
+#[derive(Debug, Clone)]
+pub enum TemperatureLabel {
+    CPU,
+}
+
+#[derive(Debug, Clone)]
+pub struct Temperature {
+    pub sensor: TemperatureLabel,
+    pub temperature: f32,
+}
+
+#[derive(Debug, Clone)]
+pub struct Memory {
+    pub total: u64,
+    pub free: u64,
+    pub available: u64,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct NetworkStats {
+    pub interface: String,
+    pub bytes_sent: u64,
+    pub bytes_received: u64,
+}
+
+#[derive(Default)]
+pub struct Heim {}
+
+impl Heim {
+    #[allow(dead_code)]
+    pub async fn temperatures() -> Result<Vec<Temperature>> {
+        let mut temperatures = Vec::new();
+        let results: Vec<_> = heim::sensors::temperatures().try_collect().await?;
+        // let results: Vec<TemperatureSensor> = Vec::new();
+        // pin_mut!(results);
+        for sensor in results {
+            if let Some(temp) = match (sensor.unit(), sensor.label()) {
+                ("k10temp", Some("Tdie")) => Some(Temperature {
+                    sensor: TemperatureLabel::CPU,
+                    temperature: sensor
+                        .current()
+                        .get::<thermodynamic_temperature::degree_celsius>(),
+                }),
+                _ => None,
+            } {
+                temperatures.push(temp);
+            }
+        }
+        Ok(temperatures)
+    }
+
+    pub async fn memory(&self) -> Result<Memory> {
+        let memory = heim::memory::memory().await?;
+        Ok(Memory {
+            total: memory.total().get::<information::byte>(),
+            free: memory.free().get::<information::byte>(),
+            available: memory.available().get::<information::byte>(),
+        })
+    }
+
+    pub async fn cpu_usage(&self) -> Result<f32> {
+        let cores = heim::cpu::logical_count().await?;
+        let measurement_1 = heim::cpu::usage().await?;
+        sleep(Duration::from_millis(100)).await;
+        let measurement_2 = heim::cpu::usage().await?;
+        Ok((measurement_2 - measurement_1).get::<ratio::percent>() / cores as f32)
+    }
+
+    pub async fn network_stats(&self) -> Result<impl Stream<Item = NetworkStats>> {
+        let networks = heim::net::io_counters().await?;
+        Ok(networks
+            .filter_map(|network| future::ready(network.ok()))
+            .filter(|network| future::ready(network.interface().starts_with("enp")))
+            .map(|network| NetworkStats {
+                interface: network.interface().into(),
+                bytes_sent: network.bytes_sent().get::<information::byte>(),
+                bytes_received: network.bytes_recv().get::<information::byte>(),
+            }))
+    }
+
+    pub async fn hostname(&self) -> Result<String> {
+        Ok(heim::host::platform().await?.hostname().to_string())
+    }
+}

src/main.rs

@@ -1,3 +1,123 @@
-fn main() {
+mod heim;
-    println!("Hello, world!");
+mod zfs;
+
+use crate::heim::{Heim, Memory, NetworkStats};
+use crate::zfs::{ZfsPool, ZFS};
+use color_eyre::{Report, Result};
+use futures_util::stream::StreamExt;
+use futures_util::{pin_mut, try_join};
+use std::fmt::Write;
+use warp::reject::Reject;
+use warp::{Filter, Rejection};
+
+#[derive(Debug)]
+struct ReportRejection(Report);
+
+impl From<Report> for ReportRejection {
+    fn from(report: Report) -> Self {
+        ReportRejection(report)
+    }
+}
+
+impl Reject for ReportRejection {}
+
+async fn get_metrics(heim: Heim, zfs: ZFS) -> Result<String, ReportRejection> {
+    let (hostname, pools, cpu, memory, network): (String, Vec<ZfsPool>, f32, Memory, _) = try_join! {
+        heim.hostname(),
+        zfs.pools(),
+        heim.cpu_usage(),
+        heim.memory(),
+        heim.network_stats(),
+    }?;
+    pin_mut!(network);
+    let mut result = String::with_capacity(256);
+    writeln!(&mut result, "cpu_usage{{host=\"{}\"}} {}", hostname, cpu).ok();
+    writeln!(
+        &mut result,
+        "memory_total{{host=\"{}\"}} {}",
+        hostname, memory.total
+    )
+    .ok();
+    writeln!(
+        &mut result,
+        "memory_available{{host=\"{}\"}} {}",
+        hostname, memory.available
+    )
+    .ok();
+    writeln!(
+        &mut result,
+        "memory_free{{host=\"{}\"}} {}",
+        hostname, memory.free
+    )
+    .ok();
+    for pool in pools {
+        writeln!(
+            &mut result,
+            "zfs_pool_size{{host=\"{}\", pool=\"{}\"}} {}",
+            hostname, pool.name, pool.size
+        )
+        .ok();
+        writeln!(
+            &mut result,
+            "zfs_pool_free{{host=\"{}\", pool=\"{}\"}} {}",
+            hostname, pool.name, pool.free
+        )
+        .ok();
+    }
+    while let Some(network) = network.next().await {
+        let network: NetworkStats = network;
+        writeln!(
+            &mut result,
+            "net_sent{{host=\"{}\", network=\"{}\"}} {}",
+            hostname, network.interface, network.bytes_sent
+        )
+        .ok();
+        writeln!(
+            &mut result,
+            "net_received{{host=\"{}\", network=\"{}\"}} {}",
+            hostname, network.interface, network.bytes_received
+        )
+        .ok();
+    }
+    // haunted ↓
+    // for temperature in Heim::temperatures().await? {
+    //     match temperature.sensor {
+    //         TemperatureLabel::CPU => writeln!(
+    //             &mut result,
+    //             "temperature{{host=\"{}\", sensor=\"cpu\"}} {}",
+    //             hostname, temperature.temperature
+    //         )
+    //         .ok(),
+    //     };
+    // }
+    Result::<_, ReportRejection>::Ok(result)
+}
+
+async fn serve_metrics(heim: Heim, zfs: ZFS) -> Result<String, Rejection> {
+    get_metrics(heim, zfs).await.map_err(warp::reject::custom)
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let host_port: u16 = dotenv::var("PORT")
+        .ok()
+        .map(|port| port.parse())
+        .transpose()?
+        .unwrap_or(80);
+
+    ctrlc::set_handler(move || {
+        std::process::exit(0);
+    })
+    .expect("Error setting Ctrl-C handler");
+
+    let heim = warp::any().map(|| Heim::default());
+    let zfs = warp::any().map(|| ZFS::default());
+
+    let metrics = warp::path!("metrics")
+        .and(heim)
+        .and(zfs)
+        .and_then(serve_metrics);
+
+    warp::serve(metrics).run(([0, 0, 0, 0], host_port)).await;
+    Ok(())
 }

src/zfs.rs

@@ -0,0 +1,38 @@
+use color_eyre::Result;
+use libzetta::zpool::{ZpoolEngine, ZpoolOpen3};
+
+#[derive(Clone, Debug)]
+pub struct ZfsPool {
+    pub name: String,
+    pub size: usize,
+    pub free: usize,
+}
+
+pub struct ZFS {
+    engine: ZpoolOpen3,
+}
+
+impl Default for ZFS {
+    fn default() -> Self {
+        ZFS {
+            engine: ZpoolOpen3::default(),
+        }
+    }
+}
+
+impl ZFS {
+    pub async fn pools(&self) -> Result<Vec<ZfsPool>> {
+        let pools = self.engine.all()?;
+        pools
+            .into_iter()
+            .map(|pool| {
+                let props = self.engine.read_properties(pool.name())?;
+                Ok(ZfsPool {
+                    name: pool.name().to_string(),
+                    size: *props.size(),
+                    free: *props.size() * (*props.capacity() as usize) / 100,
+                })
+            })
+            .collect()
+    }
+}