WIP: heartbeat

src/heartbeat.rs

@@ -0,0 +1,94 @@
+use js_int::UInt;
+use matrix_sdk::{
+    room::Joined,
+    ruma::{
+        events::EmptyStateKey, exports::ruma_macros::EventContent, OwnedRoomId, OwnedServerName,
+        UserId,
+    },
+};
+use serde::{Deserialize, Serialize};
+use std::time::Duration;
+use std::{collections::HashMap, env, process::exit, sync::Arc};
+use tokio::sync::RwLock;
+
+type HeartBeatStats = Arc<RwLock<HashMap<OwnedRoomId, HashMap<OwnedServerName, Millis>>>>;
+
+#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
+#[ruma_event(type = "xyz.dandellion.matrix-heartbeat.settings", kind = State, state_key_type = EmptyStateKey)]
+pub struct HeartBeatSettingsContent {
+    pub periode: UInt, // in seconds between each heartbeat
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
+#[ruma_event(type = "xyz.dandellion.matrix-heartbeat.notification-settings", kind = State, state_key_type = String)]
+pub struct NotificationSettingsContent {
+    // Bot to notify: Bot to count: How many counters to miss before sending a notification
+    /* {
+        '@heartbeat:dodsorf.as': { '@bot_heartbeat:pvv.ntnu.no': 5 },
+        '@bot_heartbeat:pvv.ntnu.no': { '@heartbeat:dodsorf.as': 5 },
+    } */
+    // This way you can get an alert when incoming or outgoing federation is broken.
+    // Or you can montitor a different server than your own
+    notification_rules: HashMap<String, HashMap<String, UInt>>,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
+#[ruma_event(type = "xyz.dandellion.matrix-heartbeat.heartbeat", kind = MessageLike)]
+pub struct HeartbeatContent {
+    // Server to different pong stats from last heartbeat
+    pongs: HashMap<OwnedServerName, Millis>,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct Millis {
+    pub last: UInt,
+}
+
+pub async fn get_heartbeat_settings(
+    room: &Joined,
+) -> anyhow::Result<Option<HeartBeatSettingsContent>> {
+    let settings = match room
+        .get_state_event_static::<HeartBeatSettingsContent>()
+        .await?
+    {
+        Some(s) => s,
+        None => return Ok(None),
+    };
+    let settings = match settings.deserialize() {
+        Ok(r) => r,
+        Err(e) => return Err(anyhow::Error::new(e)),
+    };
+    let settings = match settings {
+        matrix_sdk::ruma::events::SyncStateEvent::Original(s) => s,
+        matrix_sdk::ruma::events::SyncStateEvent::Redacted(_) => return Ok(None),
+    };
+    let settings = settings.content;
+    Ok(Some(settings))
+}
+
+pub async fn start_loop(room: Joined, d: Duration, heartbeat_stats: HeartBeatStats) {
+    let mut interval = tokio::time::interval(d);
+    println!("Sending pings in {} every {:?}", &room.room_id(), d);
+    loop {
+        let heartbeat_stats = heartbeat_stats.clone();
+        send_heartbeat(&room, heartbeat_stats).await;
+        interval.tick().await;
+    }
+}
+
+async fn send_heartbeat(room: &Joined, heartbeat_stats: HeartBeatStats) -> anyhow::Result<()> {
+    let room_id = room.room_id().to_owned();
+    println!("Sending ping in {}", &room_id);
+    let default_pong = HashMap::new();
+    {
+        let heartbeat_stats_r = heartbeat_stats.read().await;
+        let pongs = heartbeat_stats_r.get(&room_id).unwrap_or(&default_pong);
+        let heartbeat = HeartbeatContent {
+            pongs: pongs.clone(),
+        };
+        room.send(heartbeat, None).await.unwrap();
+    }
+    let mut heartbeat_stats_w = heartbeat_stats.write().await;
+    heartbeat_stats_w.insert(room_id, default_pong);
+    Ok(())
+}

src/main.rs

@@ -1,18 +1,27 @@
 // Parts of this file was lifted from https://github.com/matrix-org/matrix-rust-sdk/commit/3db90fbe026c222167000fcfdee8b35b44ae5694
 // And are thus licensed under Apache
 
+use heartbeat::HeartBeatSettingsContent;
 use matrix_sdk::{
     config::SyncSettings,
     room::Room,
     ruma::events::room::message::{MessageType, OriginalSyncRoomMessageEvent},
-    ruma::events::{macros::EventContent, room::member::StrippedRoomMemberEvent},
+    ruma::{
+        api::client::membership::join_room_by_id_or_alias,
+        events::{macros::EventContent, room::member::StrippedRoomMemberEvent},
+        OwnedRoomId, OwnedServerName, OwnedServerSigningKeyId,
+    },
     Client,
 };
 use serde::{Deserialize, Serialize};
-use std::{env, process::exit};
-use tokio::time::{sleep, Duration};
+use std::{collections::HashMap, env, process::exit, sync::Arc};
+use tokio::sync::RwLock;
+use tokio::time::{self, sleep, Duration};
 
+mod heartbeat;
 mod ping;
+use crate::heartbeat::start_loop;
+use crate::heartbeat::{get_heartbeat_settings, Millis};
 
 async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) {
     if let Room::Joined(room) = room {
@@ -20,7 +29,7 @@ async fn on_room_message(event: OriginalSyncRoomMessageEvent, room: Room) {
             return;
         };
 
-        if text_content.body.starts_with("!ping ") {
+        if text_content.body.starts_with("!ping") {
             ping::ping(event, room).await;
         }
     }
@@ -89,9 +98,68 @@ async fn login_and_sync(
 
     println!("Finished initial sync");
 
+    let heartbeat_senders = Arc::new(RwLock::new(HashMap::new()));
+    let heartbeat_stats = Arc::new(RwLock::new(HashMap::<
+        OwnedRoomId,
+        HashMap<OwnedServerName, heartbeat::Millis>,
+    >::new()));
+
+    let joined_rooms = client.joined_rooms();
+    for room in joined_rooms {
+        let hbs = get_heartbeat_settings(&room).await;
+        match hbs {
+            Ok(Some(a)) => {
+                let room_id = room.room_id().to_owned();
+                let mut heartbeat_senders = heartbeat_senders.write().await;
+                heartbeat_senders.insert(
+                    room_id,
+                    tokio::spawn({
+                        let heartbeat_stats = heartbeat_stats.clone();
+                        async move {
+                            start_loop(
+                                room,
+                                std::time::Duration::from_secs(a.periode.into()),
+                                heartbeat_stats.clone(),
+                            )
+                            .await
+                        }
+                    }),
+                );
+            }
+            Err(e) => println!("{}", e),
+            Ok(None) => {}
+        };
+    }
+
     client.add_event_handler(on_room_message);
     client.add_event_handler(on_stripped_state_member);
 
+    client.add_event_handler({
+        let heartbeat_stats = heartbeat_stats.clone();
+        move |event: heartbeat::OriginalSyncHeartbeat, room: Room| {
+            let heartbeat_stats = heartbeat_stats.clone();
+            async move {
+                let room_id = room.room_id().to_owned();
+                println!("Recieved a heartbeat in {}", &room_id);
+                let mut heartbeat_stats = heartbeat_stats.write().await;
+
+                let millis = Millis {
+                    last: ping::duration_to_jsint(&ping::get_event_time_delta(&event)),
+                };
+
+                let sender: OwnedServerName = event.sender.server_name().to_owned();
+
+                if let Some(hm) = heartbeat_stats.get_mut(&room_id) {
+                    hm.insert(sender, millis);
+                } else {
+                    let mut hm = HashMap::new();
+                    hm.insert(sender, millis);
+                    heartbeat_stats.insert(room_id, hm);
+                }
+            }
+        }
+    });
+
     // since we called `sync_once` before we entered our sync loop we must pass
     // that sync token to `sync`
     let settings = SyncSettings::default().token(response.next_batch);

src/ping.rs

@@ -3,8 +3,8 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use matrix_sdk::{
     room::Joined,
     ruma::{
-        events::room::message::OriginalSyncRoomMessageEvent, exports::ruma_macros::EventContent,
-        OwnedEventId, UInt,
+        events::room::message::OriginalSyncRoomMessageEvent, events::OriginalSyncMessageLikeEvent,
+        exports::ruma_macros::EventContent, OwnedEventId, UInt,
     },
 };
 use serde::{Deserialize, Serialize};
@@ -42,20 +42,24 @@ struct PongMixin {
 
 impl PongMixin {
     fn new(event: &OriginalSyncRoomMessageEvent, diff: &Duration) -> Self {
-        let millis = diff.as_millis();
-        let clamped_ms: u64 = if millis < js_int::MAX_SAFE_UINT as u128 {
-            millis as u64
-        } else {
-            js_int::MAX_SAFE_UINT
-        };
         PongMixin {
             ping: event.event_id.to_owned(),
             from: event.sender.server_name().to_string(),
-            ms: UInt::new_saturating(clamped_ms),
+            ms: duration_to_jsint(diff),
         }
     }
 }
 
+pub fn duration_to_jsint(d: &Duration) -> UInt {
+    let millis = d.as_millis();
+    let clamped_ms: u64 = if millis < js_int::MAX_SAFE_UINT as u128 {
+        millis as u64
+    } else {
+        js_int::MAX_SAFE_UINT
+    };
+    UInt::new_saturating(clamped_ms)
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
 #[ruma_event(type = "m.room.message", kind = MessageLike)]
 struct AckEventContent {
@@ -79,7 +83,9 @@ impl AckEventContent {
     }
 }
 
-pub async fn ping(event: OriginalSyncRoomMessageEvent, room: Joined) {
+pub fn get_event_time_delta<T: matrix_sdk::ruma::events::MessageLikeEventContent>(
+    event: &OriginalSyncMessageLikeEvent<T>,
+) -> Duration {
     let now = SystemTime::now()
         .duration_since(UNIX_EPOCH)
         .expect("timetravel");
@@ -89,9 +95,12 @@ pub async fn ping(event: OriginalSyncRoomMessageEvent, room: Joined) {
         .expect("timestamp was not a real time")
         .duration_since(UNIX_EPOCH)
         .expect("timetravel");
-    let diff = now - then;
 
+    now - then
+}
+
+pub async fn ping(event: OriginalSyncRoomMessageEvent, room: Joined) {
+    let diff = get_event_time_delta(&event);
     let content = AckEventContent::new(&event, &diff);
-
     room.send(content, None).await.unwrap();
 }