feat: heartbeat

Author: fractalbeauty

player/src/connection.rs

@@ -235,6 +235,8 @@ impl ConnectionActor {
     }
 
     async fn run(&mut self) {
+        let mut heartbeat_interval = tokio::time::interval(tokio::time::Duration::from_secs(60));
+
         loop {
             tokio::select! {
                 // received a command with a message to send to the server
@@ -323,6 +325,11 @@ impl ConnectionActor {
                         }
                     }
                 },
+
+                // time to send a heartbeat
+                _ = heartbeat_interval.tick() => {
+                    self.stream.send(&Message::Heartbeat).await.unwrap();
+                },
             }
         }
 
@@ -399,6 +406,7 @@ impl ConnectionActor {
             Message::RefreshRoomList => println!("unexpected message: {:?}", message),
             Message::Text(_) => println!("unexpected message: {:?}", message),
             Message::Request { .. } => println!("unexpected message: {:?}", message),
+            Message::Heartbeat => println!("unexpected message: {:?}", message),
         }
     }
 

player/src/main.rs

@@ -37,7 +37,7 @@ async fn main() -> std::io::Result<()> {
     let args: Vec<String> = std::env::args().collect();
     let is_temp = match args.get(1).map(|s| s.as_str()) {
         Some("-t") => true, // use temporary identity for testing
-        _ => false, // len is 1 but index is 1 :3
+        _ => false,         // len is 1 but index is 1 :3
     };
 
     let mut t = MainThread::setup(is_temp).await;

protocol/src/lib.rs

@@ -62,6 +62,8 @@ pub enum Message {
     AudioData(AudioData),
     Text(String),
 
+    Heartbeat,
+
     Notification(Notification),
 
     Request { request_id: u32, data: Request },

server/Cargo.toml

@@ -8,4 +8,4 @@ edition.workspace = true
 bincode = "1.3.3"
 futures = "0.3.28"
 protocol = { path = "../protocol" }
-tokio = { version = "1.28.1", features = ["macros", "net", "rt-multi-thread", "io-util", "sync"] }
+tokio = { version = "1.28.1", features = ["macros", "net", "rt-multi-thread", "io-util", "sync", "time"] }

server/src/main.rs

@@ -2,6 +2,7 @@ use std::collections::{HashMap, VecDeque};
 use std::error::Error;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 
 use futures::StreamExt;
 use tokio::io;
@@ -163,6 +164,7 @@ struct PeerHandle {
 struct Peer {
     id: String,
     room_id: Option<u32>,
+    last_heartbeat: Instant,
     addr: SocketAddr,
     stream: FrameStream,
     rx: Rx, // mpsc channel used to receive messages from peers
@@ -189,6 +191,7 @@ impl Peer {
         Ok(Peer {
             id,
             room_id: None,
+            last_heartbeat: Instant::now(),
             addr,
             stream,
             rx,
@@ -277,10 +280,19 @@ async fn handle_handshaken_connection(
                     return handle_authenticated_connection(state, stream, id, name).await;
                 }
 
-                _ => return Err("not allowed".into()),
+                _ => {
+                    println!("not allowed request was {:?}", request);
+                    return Err("not allowed".into());
+                }
             },
 
-            _ => return Err("not allowed".into()),
+            // sure
+            Message::Heartbeat => {}
+
+            _ => {
+                println!("not allowed message was {:?}", msg);
+                return Err("not allowed".into());
+            }
         }
     }
 
@@ -362,6 +374,8 @@ async fn handle_authenticated_connection(
     // create a new Peer
     let mut peer = Peer::new(state.clone(), stream, id, name).await?;
 
+    let mut heartbeat_check_interval = tokio::time::interval(tokio::time::Duration::from_secs(10));
+
     // process incoming messages until disconnected
     loop {
         tokio::select! {
@@ -387,6 +401,14 @@ async fn handle_authenticated_connection(
                 // socket disconnected
                 None => break
             },
+
+            // if we haven't received a heartbeat in the last 2 minutes, disconnect them
+            _ = heartbeat_check_interval.tick() => {
+                let elapsed = peer.last_heartbeat.elapsed();
+                if elapsed > Duration::from_secs(120) {
+                    break
+                }
+            }
         }
     }
 
@@ -487,6 +509,11 @@ async fn handle_message(
             Ok(())
         }
 
+        Message::Heartbeat => {
+            peer.last_heartbeat = Instant::now();
+            Ok(())
+        }
+
         // should not be reached
         Message::QueryRoomList => Ok(()),
         Message::Notification(_) => Ok(()),