dedup mitemp

This commit is contained in:
Robin Appelman 2021-11-10 00:02:03 +01:00
commit b6624fd763
3 changed files with 120 additions and 80 deletions

View file

@ -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=\"{}\"}} {}",
device.hostname, addr, name, state.battery
"sensor_battery{{mac=\"{}\", name=\"{}\"}} {}",
addr, name, state.battery
)?;
}
if state.temperature > 0.0 {
writeln!(
writer,
"sensor_temperature{{tasmota_id=\"{}\", mac=\"{}\", name=\"{}\"}} {}",
device.hostname, addr, name, state.temperature
"sensor_temperature{{mac=\"{}\", name=\"{}\"}} {}",
addr, name, state.temperature
)?;
}
if state.humidity > 0.0 {
writeln!(
writer,
"sensor_humidity{{tasmota_id=\"{}\", mac=\"{}\", name=\"{}\"}} {}",
device.hostname, addr, name, state.humidity
"sensor_humidity{{mac=\"{}\", name=\"{}\"}} {}",
addr, name, state.humidity
)?;
}
Ok(())

View file

@ -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() {
format_device_state(
&mut response,
&device.key(),
&device.value(),
&mi_temp_names,
)
.unwrap();
for (device, state) in state.devices() {
format_device_state(&mut response, device, state).unwrap();
}
for (addr, state) in state.mi_temp() {
format_mi_temp_state(&mut response, *addr, &mi_temp_names, state).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();
device_state.update(json);
let mut device_states = device_states.lock().unwrap();
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();
device_state.update(json);
let mut device_states = device_states.lock().unwrap();
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| {
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
}
});
state
.lock()
.unwrap()
.retain(cleanup_time, ping_time, &client);
sleep(Duration::from_secs(60)).await;
}

View file

@ -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!(
Topic::Result(device.clone()),
Topic::from("stat/hostname/RESULT")
);
assert_eq!(Topic::Result(device), Topic::from("stat/hostname/RESULT"));
}