dedup mitemp

src/device.rs

@@ -1,9 +1,79 @@
 use color_eyre::{eyre::WrapErr, Report, Result};
 use json::JsonValue;
-use std::collections::BTreeMap;
+use rumqttc::{AsyncClient, QoS};
+use std::collections::{BTreeMap, HashMap};
 use std::convert::TryFrom;
 use std::fmt::{self, Debug, Display, Formatter, Write};
 use std::time::Instant;
+use tokio::task::spawn;
+
+#[derive(Default)]
+pub struct DeviceStates {
+    pub devices: HashMap<Device, DeviceState>,
+    pub mi_temp_devices: BTreeMap<BDAddr, MiTempState>,
+}
+
+impl DeviceStates {
+    pub fn devices(&self) -> impl Iterator<Item = (&Device, &DeviceState)> {
+        self.devices.iter()
+    }
+
+    pub fn update(&mut self, device: Device, json: JsonValue) {
+        let device = self.devices.entry(device).or_default();
+
+        for (key, value) in json.entries() {
+            if let Some(addr) = key.strip_prefix("MJ_HT_V1") {
+                let addr = addr.trim_start_matches('-');
+                match BDAddr::from_mi_temp_mac_part(addr) {
+                    Ok(addr) => {
+                        let state = self.mi_temp_devices.entry(addr).or_default();
+                        state.update(value);
+                    }
+                    Err(e) => eprintln!("Failed to parse mitemp mac: {:#}", e),
+                }
+            }
+        }
+
+        device.update(json);
+    }
+
+    pub fn mi_temp(&self) -> impl Iterator<Item = (&BDAddr, &MiTempState)> {
+        self.mi_temp_devices.iter()
+    }
+
+    pub fn retain(&mut self, cleanup_time: Instant, ping_time: Instant, client: &AsyncClient) {
+        self.devices.retain(|device, state| {
+            if state.last_seen < cleanup_time {
+                println!("{} hasn't been seen for 15m, removing", device.hostname);
+                false
+            } else if state.last_seen < ping_time || state.name.is_empty() {
+                println!(
+                    "{} hasn't been seen for 10m or has no name set, pinging",
+                    device.hostname
+                );
+                let send_client = client.clone();
+                let topic = device.get_topic("cmnd", "DeviceName");
+                spawn(async move {
+                    if let Err(e) = send_client.publish(topic, QoS::AtMostOnce, false, "").await {
+                        eprintln!("Failed to ping device: {:#}", e);
+                    }
+                });
+                true
+            } else {
+                true
+            }
+        });
+
+        self.mi_temp_devices.retain(|device, state| {
+            if state.last_seen < cleanup_time {
+                println!("{} hasn't been seen for 15m, removing", device);
+                false
+            } else {
+                true
+            }
+        });
+    }
+}
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Device {
@@ -24,7 +94,6 @@ pub struct DeviceState {
     pub power_yesterday: Option<f32>,
     pub power_today: Option<f32>,
     pub co2: Option<f32>,
-    pub mi_temp_devices: BTreeMap<BDAddr, MiTempState>,
     pub pms_state: Option<PMSState>,
     pub last_seen: Instant,
 }
@@ -38,7 +107,6 @@ impl Default for DeviceState {
             power_yesterday: Default::default(),
             power_today: Default::default(),
             co2: Default::default(),
-            mi_temp_devices: Default::default(),
             pms_state: Default::default(),
             last_seen: Instant::now(),
         }
@@ -86,32 +154,33 @@ impl DeviceState {
             let pms = self.pms_state.get_or_insert(PMSState::default());
             pms.update(&json["PMS5003"]);
         }
-
-        for (key, value) in json.entries() {
-            if let Some(addr) = key.strip_prefix("MJ_HT_V1") {
-                let addr = addr.trim_start_matches('-');
-                match BDAddr::from_mi_temp_mac_part(addr) {
-                    Ok(addr) => {
-                        let state = self.mi_temp_devices.entry(addr).or_default();
-                        state.update(value);
-                    }
-                    Err(e) => eprintln!("Failed to parse mitemp mac: {:#}", e),
-                }
-            }
-        }
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct MiTempState {
     temperature: f32,
     humidity: f32,
     dew_point: f32,
     battery: u8,
+    pub last_seen: Instant,
+}
+
+impl Default for MiTempState {
+    fn default() -> Self {
+        MiTempState {
+            temperature: 0.0,
+            humidity: 0.0,
+            dew_point: 0.0,
+            battery: 0,
+            last_seen: Instant::now(),
+        }
+    }
 }
 
 impl MiTempState {
     pub fn update(&mut self, json: &JsonValue) {
+        self.last_seen = Instant::now();
         if let Some(temperature) = json["Temperature"]
             .as_number()
             .and_then(|num| f32::try_from(num).ok())
@@ -143,7 +212,6 @@ pub fn format_device_state<W: Write>(
     mut writer: W,
     device: &Device,
     state: &DeviceState,
-    mi_temp_names: &BTreeMap<BDAddr, String>,
 ) -> std::fmt::Result {
     if state.name.is_empty() {
         println!("{} has no name set, skipping", device.hostname);
@@ -196,10 +264,6 @@ pub fn format_device_state<W: Write>(
         )?;
     }
 
-    for (addr, state) in state.mi_temp_devices.iter() {
-        format_mi_temp_state(&mut writer, device, *addr, mi_temp_names, state)?;
-    }
-
     if let Some(pms) = state.pms_state.as_ref() {
         format_pms_state(&mut writer, device, state, pms)?;
     }
@@ -209,7 +273,6 @@ pub fn format_device_state<W: Write>(
 
 pub fn format_mi_temp_state<W: Write>(
     mut writer: W,
-    device: &Device,
     addr: BDAddr,
     names: &BTreeMap<BDAddr, String>,
     state: &MiTempState,
@@ -227,24 +290,24 @@ pub fn format_mi_temp_state<W: Write>(
     if state.battery > 0 {
         writeln!(
             writer,
-            "sensor_battery{{tasmota_id=\"{}\", mac=\"{}\", name=\"{}\"}} {}",
+            "sensor_battery{{mac=\"{}\", name=\"{}\"}} {}",
-            device.hostname, addr, name, state.battery
+            addr, name, state.battery
         )?;
     }
 
     if state.temperature > 0.0 {
         writeln!(
             writer,
-            "sensor_temperature{{tasmota_id=\"{}\", mac=\"{}\", name=\"{}\"}} {}",
+            "sensor_temperature{{mac=\"{}\", name=\"{}\"}} {}",
-            device.hostname, addr, name, state.temperature
+            addr, name, state.temperature
         )?;
     }
 
     if state.humidity > 0.0 {
         writeln!(
             writer,
-            "sensor_humidity{{tasmota_id=\"{}\", mac=\"{}\", name=\"{}\"}} {}",
+            "sensor_humidity{{mac=\"{}\", name=\"{}\"}} {}",
-            device.hostname, addr, name, state.humidity
+            addr, name, state.humidity
         )?;
     }
     Ok(())

src/main.rs

@@ -4,29 +4,28 @@ mod mqtt;
 mod topic;
 
 use crate::config::Config;
-use crate::device::{format_device_state, Device, DeviceState};
+use crate::device::{format_device_state, format_mi_temp_state, Device, DeviceStates};
 use crate::mqtt::mqtt_stream;
 use crate::topic::Topic;
 use color_eyre::{eyre::WrapErr, Result};
-use dashmap::DashMap;
+
 use pin_utils::pin_mut;
 use rumqttc::{AsyncClient, Publish, QoS};
+
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::time::Instant;
 use tokio::task::spawn;
 use tokio::time::{sleep, Duration};
 use tokio_stream::{Stream, StreamExt};
 use warp::Filter;
 
-type DeviceStates = Arc<DashMap<Device, DeviceState>>;
-
 #[tokio::main]
 async fn main() -> Result<()> {
     let config = Config::from_env()?;
     let host_port = config.host_port;
 
-    let device_states = DeviceStates::default();
+    let device_states = <Arc<Mutex<DeviceStates>>>::default();
 
     ctrlc::set_handler(move || {
         std::process::exit(0);
@@ -48,16 +47,14 @@ async fn main() -> Result<()> {
 
     let metrics = warp::path!("metrics")
         .and(state)
-        .map(move |state: DeviceStates| {
+        .map(move |state: Arc<Mutex<DeviceStates>>| {
+            let state = state.lock().unwrap();
             let mut response = String::new();
-            for device in state.iter() {
+            for (device, state) in state.devices() {
-                format_device_state(
+                format_device_state(&mut response, device, state).unwrap();
-                    &mut response,
+            }
-                    &device.key(),
+            for (addr, state) in state.mi_temp() {
-                    &device.value(),
+                format_mi_temp_state(&mut response, *addr, &mi_temp_names, state).unwrap()
-                    &mi_temp_names,
-                )
-                .unwrap();
             }
             response
         });
@@ -69,7 +66,7 @@ async fn main() -> Result<()> {
 async fn mqtt_loop(
     client: AsyncClient,
     stream: impl Stream<Item = Result<Publish>>,
-    states: DeviceStates,
+    states: Arc<Mutex<DeviceStates>>,
 ) {
     pin_mut!(stream);
     loop {
@@ -96,7 +93,7 @@ async fn command(client: &AsyncClient, device: &Device, command: &str) -> Result
 async fn mqtt_client<S: Stream<Item = Result<Publish>>>(
     client: AsyncClient,
     stream: &mut Pin<&mut S>,
-    device_states: DeviceStates,
+    device_states: Arc<Mutex<DeviceStates>>,
 ) -> Result<()> {
     while let Some(message) = stream.next().await {
         let message = message?;
@@ -108,7 +105,7 @@ async fn mqtt_client<S: Stream<Item = Result<Publish>>>(
         let topic = Topic::from(message.topic.as_str());
 
         match topic {
-            Topic::LWT(device) => {
+            Topic::Lwt(device) => {
                 // on discovery, ask the device for it's power state and name
                 let send_client = client.clone();
                 spawn(async move {
@@ -124,15 +121,15 @@ async fn mqtt_client<S: Stream<Item = Result<Publish>>>(
             Topic::Result(device) => {
                 let payload = std::str::from_utf8(message.payload.as_ref()).unwrap_or_default();
                 if let Ok(json) = json::parse(payload) {
-                    let mut device_state = device_states.entry(device).or_default();
+                    let mut device_states = device_states.lock().unwrap();
-                    device_state.update(json);
+                    device_states.update(device, json);
                 }
             }
             Topic::Sensor(device) => {
                 let payload = std::str::from_utf8(message.payload.as_ref()).unwrap_or_default();
                 if let Ok(json) = json::parse(payload) {
-                    let mut device_state = device_states.entry(device).or_default();
+                    let mut device_states = device_states.lock().unwrap();
-                    device_state.update(json);
+                    device_states.update(device, json);
                 }
             }
             _ => {}
@@ -141,32 +138,15 @@ async fn mqtt_client<S: Stream<Item = Result<Publish>>>(
     Ok(())
 }
 
-async fn cleanup(client: AsyncClient, devices: DeviceStates) {
+async fn cleanup(client: AsyncClient, state: Arc<Mutex<DeviceStates>>) {
     loop {
         let ping_time = Instant::now() - Duration::from_secs(10 * 60);
         let cleanup_time = Instant::now() - Duration::from_secs(15 * 60);
 
-        devices.retain(|device, state| {
+        state
-            if state.last_seen < cleanup_time {
+            .lock()
-                println!("{} hasn't been seen for 15m, removing", device.hostname);
+            .unwrap()
-                false
+            .retain(cleanup_time, ping_time, &client);
-            } else if state.last_seen < ping_time || state.name.is_empty() {
-                println!(
-                    "{} hasn't been seen for 10m or has no name set, pinging",
-                    device.hostname
-                );
-                let send_client = client.clone();
-                let topic = device.get_topic("cmnd", "DeviceName");
-                spawn(async move {
-                    if let Err(e) = send_client.publish(topic, QoS::AtMostOnce, false, "").await {
-                        eprintln!("Failed to ping device: {:#}", e);
-                    }
-                });
-                true
-            } else {
-                true
-            }
-        });
 
         sleep(Duration::from_secs(60)).await;
     }

src/topic.rs

@@ -2,7 +2,7 @@ use crate::device::Device;
 
 #[derive(Debug, Eq, PartialEq)]
 pub enum Topic {
-    LWT(Device),
+    Lwt(Device),
     Power(Device),
     State(Device),
     Sensor(Device),
@@ -20,7 +20,7 @@ impl From<&str> for Topic {
                 hostname: hostname.to_string(),
             };
             match (prefix, cmd) {
-                ("tele", "LWT") => Topic::LWT(device),
+                ("tele", "LWT") => Topic::Lwt(device),
                 ("tele", "STATE") => Topic::State(device),
                 ("stat", "POWER") => Topic::Power(device),
                 ("tele", "SENSOR") => Topic::Sensor(device),
@@ -38,7 +38,7 @@ fn parse_topic() {
     let device = Device {
         hostname: "hostname".to_string(),
     };
-    assert_eq!(Topic::LWT(device.clone()), Topic::from("tele/hostname/LWT"));
+    assert_eq!(Topic::Lwt(device.clone()), Topic::from("tele/hostname/LWT"));
     assert_eq!(
         Topic::Power(device.clone()),
         Topic::from("stat/hostname/POWER")
@@ -51,8 +51,5 @@ fn parse_topic() {
         Topic::Sensor(device.clone()),
         Topic::from("tele/hostname/SENSOR")
     );
-    assert_eq!(
+    assert_eq!(Topic::Result(device), Topic::from("stat/hostname/RESULT"));
-        Topic::Result(device.clone()),
-        Topic::from("stat/hostname/RESULT")
-    );
 }