crazy refactor

src/voice.rs

@@ -10,7 +10,7 @@ use crate::{AppState, db};
 
 #[derive(Default)]
 pub struct VoiceHub {
-    rooms: RwLock<HashMap<Uuid, HashMap<Uuid, ClientHandle>>>,
+    rooms: RwLock<HashMap<Uuid, HashMap<Uuid, HashMap<Uuid, ClientHandle>>>>,
 }
 
 #[derive(Clone)]
@@ -72,7 +72,7 @@ enum ServerEvent {
     },
     PlaySound {
         user_id: Uuid,
-        sound_url: String,
+        media_url: String,
     },
 }
 
@@ -97,10 +97,20 @@ enum ClientEvent {
         is_muted: bool,
     },
     PlaySound {
-        sound_url: String,
+        sound_id: Uuid,
     },
 }
 
+#[derive(Default)]
+struct VoiceStateChange {
+    joined: Option<VoiceParticipant>,
+    left_user_id: Option<Uuid>,
+    video_changed: Option<(Uuid, bool)>,
+    screen_changed: Option<(Uuid, bool)>,
+    speaking_changed: Option<(Uuid, bool)>,
+    mute_changed: Option<(Uuid, bool)>,
+}
+
 impl VoiceHub {
     pub async fn participants(&self, room_id: Uuid) -> Vec<VoiceParticipant> {
         let rooms = self.rooms.read().await;
@@ -109,14 +119,7 @@ impl VoiceHub {
         };
 
         room.iter()
-            .map(|(user_id, handle)| VoiceParticipant {
-                user_id: *user_id,
-                display_name: handle.display_name.clone(),
-                is_sharing_video: handle.is_sharing_video,
-                is_sharing_screen: handle.is_sharing_screen,
-                is_speaking: handle.is_speaking,
-                is_muted: handle.is_muted,
-            })
+            .filter_map(|(user_id, connections)| aggregate_participant(Some(connections), *user_id))
             .collect()
     }
 
@@ -124,28 +127,23 @@ impl VoiceHub {
         &self,
         room_id: Uuid,
         user_id: Uuid,
+        connection_id: Uuid,
         display_name: String,
         tx: mpsc::UnboundedSender<ServerEvent>,
-    ) -> Vec<VoiceParticipant> {
+    ) -> (Vec<VoiceParticipant>, VoiceStateChange) {
         let mut rooms = self.rooms.write().await;
         let room = rooms.entry(room_id).or_default();
 
         let peers = room
             .iter()
-            .map(|(peer_id, peer)| VoiceParticipant {
-                user_id: *peer_id,
-                display_name: peer.display_name.clone(),
-                is_sharing_video: peer.is_sharing_video,
-                is_sharing_screen: peer.is_sharing_screen,
-                is_speaking: peer.is_speaking,
-                is_muted: peer.is_muted,
-            })
+            .filter_map(|(peer_id, connections)| aggregate_participant(Some(connections), *peer_id))
             .collect::<Vec<_>>();
 
-        room.insert(
-            user_id,
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        room.entry(user_id).or_default().insert(
+            connection_id,
             ClientHandle {
-                display_name: display_name.clone(),
+                display_name,
                 is_sharing_video: false,
                 is_sharing_screen: false,
                 is_speaking: false,
@@ -153,34 +151,31 @@ impl VoiceHub {
                 tx,
             },
         );
+        let current = aggregate_participant(room.get(&user_id), user_id);
 
-        for (peer_id, peer) in room.iter() {
-            if *peer_id != user_id {
-                let _ = peer.tx.send(ServerEvent::PeerJoined {
-                    user_id,
-                    display_name: display_name.clone(),
-                });
-            }
-        }
-
-        peers
+        (peers, diff_voice_state(previous, current))
     }
 
-    async fn leave(&self, room_id: Uuid, user_id: Uuid) {
+    async fn leave(&self, room_id: Uuid, user_id: Uuid, connection_id: Uuid) -> VoiceStateChange {
         let mut rooms = self.rooms.write().await;
         let Some(room) = rooms.get_mut(&room_id) else {
-            return;
+            return VoiceStateChange::default();
         };
 
-        room.remove(&user_id);
-
-        for peer in room.values() {
-            let _ = peer.tx.send(ServerEvent::PeerLeft { user_id });
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        if let Some(connections) = room.get_mut(&user_id) {
+            connections.remove(&connection_id);
+            if connections.is_empty() {
+                room.remove(&user_id);
+            }
         }
+        let current = aggregate_participant(room.get(&user_id), user_id);
 
         if room.is_empty() {
             rooms.remove(&room_id);
         }
+
+        diff_voice_state(previous, current)
     }
 
     pub async fn relay_signal(
@@ -196,7 +191,10 @@ impl VoiceHub {
             return;
         };
 
-        if let Some(target) = room.get(&to_user_id) {
+        if let Some(target) = room
+            .get(&to_user_id)
+            .and_then(|connections| connections.values().next())
+        {
             let _ = target.tx.send(ServerEvent::Signal {
                 from_user_id,
                 kind,
@@ -205,99 +203,185 @@ impl VoiceHub {
         }
     }
 
-    pub async fn set_video_status(&self, room_id: Uuid, user_id: Uuid, is_sharing_video: bool) {
+    async fn set_video_status(
+        &self,
+        room_id: Uuid,
+        user_id: Uuid,
+        connection_id: Uuid,
+        is_sharing_video: bool,
+    ) -> VoiceStateChange {
         let mut rooms = self.rooms.write().await;
         let Some(room) = rooms.get_mut(&room_id) else {
-            return;
+            return VoiceStateChange::default();
         };
 
-        if let Some(handle) = room.get_mut(&user_id) {
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        if let Some(handle) = room
+            .get_mut(&user_id)
+            .and_then(|connections| connections.get_mut(&connection_id))
+        {
             handle.is_sharing_video = is_sharing_video;
-
-            for (peer_id, peer) in room.iter() {
-                if *peer_id != user_id {
-                    let _ = peer.tx.send(ServerEvent::VideoStatusChanged {
-                        user_id,
-                        is_sharing_video,
-                    });
-                }
-            }
         }
+        let current = aggregate_participant(room.get(&user_id), user_id);
+        diff_voice_state(previous, current)
     }
 
-    pub async fn set_screen_status(&self, room_id: Uuid, user_id: Uuid, is_sharing_screen: bool) {
+    async fn set_screen_status(
+        &self,
+        room_id: Uuid,
+        user_id: Uuid,
+        connection_id: Uuid,
+        is_sharing_screen: bool,
+    ) -> VoiceStateChange {
         let mut rooms = self.rooms.write().await;
         let Some(room) = rooms.get_mut(&room_id) else {
-            return;
+            return VoiceStateChange::default();
         };
 
-        if let Some(handle) = room.get_mut(&user_id) {
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        if let Some(handle) = room
+            .get_mut(&user_id)
+            .and_then(|connections| connections.get_mut(&connection_id))
+        {
             handle.is_sharing_screen = is_sharing_screen;
-
-            for (peer_id, peer) in room.iter() {
-                if *peer_id != user_id {
-                    let _ = peer.tx.send(ServerEvent::ScreenStatusChanged {
-                        user_id,
-                        is_sharing_screen,
-                    });
-                }
-            }
         }
+        let current = aggregate_participant(room.get(&user_id), user_id);
+        diff_voice_state(previous, current)
     }
 
-    pub async fn set_speaking_status(&self, room_id: Uuid, user_id: Uuid, is_speaking: bool) {
+    async fn set_speaking_status(
+        &self,
+        room_id: Uuid,
+        user_id: Uuid,
+        connection_id: Uuid,
+        is_speaking: bool,
+    ) -> VoiceStateChange {
         let mut rooms = self.rooms.write().await;
         let Some(room) = rooms.get_mut(&room_id) else {
-            return;
+            return VoiceStateChange::default();
         };
 
-        if let Some(handle) = room.get_mut(&user_id) {
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        if let Some(handle) = room
+            .get_mut(&user_id)
+            .and_then(|connections| connections.get_mut(&connection_id))
+        {
             handle.is_speaking = is_speaking;
-
-            for (peer_id, peer) in room.iter() {
-                if *peer_id != user_id {
-                    let _ = peer.tx.send(ServerEvent::SpeakingStatusChanged {
-                        user_id,
-                        is_speaking,
-                    });
-                }
-            }
         }
+        let current = aggregate_participant(room.get(&user_id), user_id);
+        diff_voice_state(previous, current)
     }
 
-    pub async fn set_mute_status(&self, room_id: Uuid, user_id: Uuid, is_muted: bool) {
+    async fn set_mute_status(
+        &self,
+        room_id: Uuid,
+        user_id: Uuid,
+        connection_id: Uuid,
+        is_muted: bool,
+    ) -> VoiceStateChange {
         let mut rooms = self.rooms.write().await;
         let Some(room) = rooms.get_mut(&room_id) else {
-            return;
+            return VoiceStateChange::default();
         };
 
-        if let Some(handle) = room.get_mut(&user_id) {
+        let previous = aggregate_participant(room.get(&user_id), user_id);
+        if let Some(handle) = room
+            .get_mut(&user_id)
+            .and_then(|connections| connections.get_mut(&connection_id))
+        {
             handle.is_muted = is_muted;
-
-            for (peer_id, peer) in room.iter() {
-                if *peer_id != user_id {
-                    let _ = peer
-                        .tx
-                        .send(ServerEvent::MuteStatusChanged { user_id, is_muted });
-                }
-            }
         }
+        let current = aggregate_participant(room.get(&user_id), user_id);
+        diff_voice_state(previous, current)
     }
 
-    pub async fn play_sound(&self, room_id: Uuid, user_id: Uuid, sound_url: String) {
+    pub async fn play_sound(&self, room_id: Uuid, user_id: Uuid, media_url: String) {
         let rooms = self.rooms.read().await;
         let Some(room) = rooms.get(&room_id) else {
             return;
         };
 
-        for (peer_id, peer) in room.iter() {
+        for (peer_id, connections) in room.iter() {
             if *peer_id == user_id {
                 continue;
             }
-            let _ = peer.tx.send(ServerEvent::PlaySound {
+            for peer in connections.values() {
+                let _ = peer.tx.send(ServerEvent::PlaySound {
+                    user_id,
+                    media_url: media_url.clone(),
+                });
+            }
+        }
+    }
+
+    async fn emit_change(&self, room_id: Uuid, user_id: Uuid, change: VoiceStateChange) {
+        if is_voice_state_change_empty(&change) {
+            return;
+        }
+
+        let rooms = self.rooms.read().await;
+        let Some(room) = rooms.get(&room_id) else {
+            return;
+        };
+
+        if let Some(participant) = change.joined {
+            broadcast_voice_event(
+                room,
                 user_id,
-                sound_url: sound_url.clone(),
-            });
+                ServerEvent::PeerJoined {
+                    user_id: participant.user_id,
+                    display_name: participant.display_name,
+                },
+            );
+        }
+        if let Some(left_user_id) = change.left_user_id {
+            broadcast_voice_event(
+                room,
+                user_id,
+                ServerEvent::PeerLeft {
+                    user_id: left_user_id,
+                },
+            );
+        }
+        if let Some((changed_user_id, is_sharing_video)) = change.video_changed {
+            broadcast_voice_event(
+                room,
+                user_id,
+                ServerEvent::VideoStatusChanged {
+                    user_id: changed_user_id,
+                    is_sharing_video,
+                },
+            );
+        }
+        if let Some((changed_user_id, is_sharing_screen)) = change.screen_changed {
+            broadcast_voice_event(
+                room,
+                user_id,
+                ServerEvent::ScreenStatusChanged {
+                    user_id: changed_user_id,
+                    is_sharing_screen,
+                },
+            );
+        }
+        if let Some((changed_user_id, is_speaking)) = change.speaking_changed {
+            broadcast_voice_event(
+                room,
+                user_id,
+                ServerEvent::SpeakingStatusChanged {
+                    user_id: changed_user_id,
+                    is_speaking,
+                },
+            );
+        }
+        if let Some((changed_user_id, is_muted)) = change.mute_changed {
+            broadcast_voice_event(
+                room,
+                user_id,
+                ServerEvent::MuteStatusChanged {
+                    user_id: changed_user_id,
+                    is_muted,
+                },
+            );
         }
     }
 }
@@ -306,15 +390,26 @@ pub async fn handle_socket(state: AppState, socket: WebSocket, room_id: Uuid, us
     let Some(user) = db::get_user_by_id(&state.db, user_id).await.ok().flatten() else {
         return;
     };
+    let Ok(Some(guild_id)) = db::guild_id_for_channel(&state.db, room_id).await else {
+        return;
+    };
 
     let (mut ws_sender, mut ws_receiver) = socket.split();
     let (tx, mut rx) = mpsc::unbounded_channel::<ServerEvent>();
+    let connection_id = Uuid::new_v4();
 
-    let peers = state
+    let (peers, change) = state
         .voice
-        .join(room_id, user_id, user.display_name.clone(), tx.clone())
+        .join(
+            room_id,
+            user_id,
+            connection_id,
+            user.display_name.clone(),
+            tx.clone(),
+        )
         .await;
     let _ = tx.send(ServerEvent::Peers { peers });
+    state.voice.emit_change(room_id, user_id, change).await;
 
     let send_task = tokio::spawn(async move {
         while let Some(event) = rx.recv().await {
@@ -343,31 +438,42 @@ pub async fn handle_socket(state: AppState, socket: WebSocket, room_id: Uuid, us
                             .await;
                     }
                     Ok(ClientEvent::SetVideoStatus { is_sharing_video }) => {
-                        state
+                        let change = state
                             .voice
-                            .set_video_status(room_id, user_id, is_sharing_video)
+                            .set_video_status(room_id, user_id, connection_id, is_sharing_video)
                             .await;
+                        state.voice.emit_change(room_id, user_id, change).await;
                     }
                     Ok(ClientEvent::SetScreenStatus { is_sharing_screen }) => {
-                        state
+                        let change = state
                             .voice
-                            .set_screen_status(room_id, user_id, is_sharing_screen)
+                            .set_screen_status(room_id, user_id, connection_id, is_sharing_screen)
                             .await;
+                        state.voice.emit_change(room_id, user_id, change).await;
                     }
                     Ok(ClientEvent::SetSpeakingStatus { is_speaking }) => {
-                        state
+                        let change = state
                             .voice
-                            .set_speaking_status(room_id, user_id, is_speaking)
+                            .set_speaking_status(room_id, user_id, connection_id, is_speaking)
                             .await;
+                        state.voice.emit_change(room_id, user_id, change).await;
                     }
                     Ok(ClientEvent::SetMuteStatus { is_muted }) => {
-                        state
+                        let change = state
                             .voice
-                            .set_mute_status(room_id, user_id, is_muted)
+                            .set_mute_status(room_id, user_id, connection_id, is_muted)
                             .await;
+                        state.voice.emit_change(room_id, user_id, change).await;
                     }
-                    Ok(ClientEvent::PlaySound { sound_url }) => {
-                        state.voice.play_sound(room_id, user_id, sound_url).await;
+                    Ok(ClientEvent::PlaySound { sound_id }) => {
+                        if let Ok(Some(sound)) = db::get_sound_by_id(&state.db, sound_id).await
+                            && sound.guild_id == guild_id
+                        {
+                            state
+                                .voice
+                                .play_sound(room_id, user_id, sound.media_url)
+                                .await;
+                        }
                     }
                     Err(err) => {
                         let _ = tx.send(ServerEvent::Error {
@@ -382,5 +488,183 @@ pub async fn handle_socket(state: AppState, socket: WebSocket, room_id: Uuid, us
     }
 
     send_task.abort();
-    state.voice.leave(room_id, user_id).await;
+    let change = state.voice.leave(room_id, user_id, connection_id).await;
+    state.voice.emit_change(room_id, user_id, change).await;
+}
+
+fn aggregate_participant(
+    connections: Option<&HashMap<Uuid, ClientHandle>>,
+    user_id: Uuid,
+) -> Option<VoiceParticipant> {
+    let connections = connections?;
+    if connections.is_empty() {
+        return None;
+    }
+
+    let mut handles = connections.values();
+    let first = handles.next()?;
+    Some(VoiceParticipant {
+        user_id,
+        display_name: first.display_name.clone(),
+        is_sharing_video: connections.values().any(|handle| handle.is_sharing_video),
+        is_sharing_screen: connections.values().any(|handle| handle.is_sharing_screen),
+        is_speaking: connections.values().any(|handle| handle.is_speaking),
+        is_muted: connections.values().all(|handle| handle.is_muted),
+    })
+}
+
+fn diff_voice_state(
+    previous: Option<VoiceParticipant>,
+    current: Option<VoiceParticipant>,
+) -> VoiceStateChange {
+    match (previous, current) {
+        (None, None) => VoiceStateChange::default(),
+        (None, Some(current)) => VoiceStateChange {
+            joined: Some(current),
+            ..Default::default()
+        },
+        (Some(previous), None) => VoiceStateChange {
+            left_user_id: Some(previous.user_id),
+            ..Default::default()
+        },
+        (Some(previous), Some(current)) => VoiceStateChange {
+            video_changed: (previous.is_sharing_video != current.is_sharing_video)
+                .then_some((current.user_id, current.is_sharing_video)),
+            screen_changed: (previous.is_sharing_screen != current.is_sharing_screen)
+                .then_some((current.user_id, current.is_sharing_screen)),
+            speaking_changed: (previous.is_speaking != current.is_speaking)
+                .then_some((current.user_id, current.is_speaking)),
+            mute_changed: (previous.is_muted != current.is_muted)
+                .then_some((current.user_id, current.is_muted)),
+            ..Default::default()
+        },
+    }
+}
+
+fn broadcast_voice_event(
+    room: &HashMap<Uuid, HashMap<Uuid, ClientHandle>>,
+    source_user_id: Uuid,
+    event: ServerEvent,
+) {
+    for (peer_id, connections) in room {
+        if *peer_id == source_user_id {
+            continue;
+        }
+        for peer in connections.values() {
+            let _ = peer.tx.send(event.clone());
+        }
+    }
+}
+
+fn is_voice_state_change_empty(change: &VoiceStateChange) -> bool {
+    change.joined.is_none()
+        && change.left_user_id.is_none()
+        && change.video_changed.is_none()
+        && change.screen_changed.is_none()
+        && change.speaking_changed.is_none()
+        && change.mute_changed.is_none()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{VoiceHub, aggregate_participant};
+    use std::collections::HashMap;
+    use tokio::sync::mpsc;
+    use uuid::Uuid;
+
+    #[tokio::test]
+    async fn multiple_connections_keep_voice_participant_until_last_leave() {
+        let hub = VoiceHub::default();
+        let room_id = Uuid::new_v4();
+        let user_id = Uuid::new_v4();
+        let first_connection = Uuid::new_v4();
+        let second_connection = Uuid::new_v4();
+        let (tx1, _rx1) = mpsc::unbounded_channel();
+        let (tx2, _rx2) = mpsc::unbounded_channel();
+
+        let (_peers, first_join) = hub
+            .join(room_id, user_id, first_connection, "User".to_string(), tx1)
+            .await;
+        let (_peers, second_join) = hub
+            .join(room_id, user_id, second_connection, "User".to_string(), tx2)
+            .await;
+        let first_leave = hub.leave(room_id, user_id, first_connection).await;
+        let second_leave = hub.leave(room_id, user_id, second_connection).await;
+
+        assert!(first_join.joined.is_some());
+        assert!(second_join.joined.is_none());
+        assert!(first_leave.left_user_id.is_none());
+        assert_eq!(second_leave.left_user_id, Some(user_id));
+    }
+
+    #[tokio::test]
+    async fn voice_mute_state_only_flips_when_all_connections_are_muted() {
+        let hub = VoiceHub::default();
+        let room_id = Uuid::new_v4();
+        let user_id = Uuid::new_v4();
+        let first_connection = Uuid::new_v4();
+        let second_connection = Uuid::new_v4();
+        let (tx1, _rx1) = mpsc::unbounded_channel();
+        let (tx2, _rx2) = mpsc::unbounded_channel();
+
+        let _ = hub
+            .join(room_id, user_id, first_connection, "User".to_string(), tx1)
+            .await;
+        let _ = hub
+            .join(room_id, user_id, second_connection, "User".to_string(), tx2)
+            .await;
+
+        let first_mute = hub
+            .set_mute_status(room_id, user_id, first_connection, true)
+            .await;
+        let second_mute = hub
+            .set_mute_status(room_id, user_id, second_connection, true)
+            .await;
+        let unmute = hub
+            .set_mute_status(room_id, user_id, first_connection, false)
+            .await;
+
+        assert!(first_mute.mute_changed.is_none());
+        assert_eq!(second_mute.mute_changed, Some((user_id, true)));
+        assert_eq!(unmute.mute_changed, Some((user_id, false)));
+    }
+
+    #[test]
+    fn aggregate_participant_combines_connection_state() {
+        let user_id = Uuid::new_v4();
+        let first_connection = Uuid::new_v4();
+        let second_connection = Uuid::new_v4();
+        let mut connections = HashMap::new();
+        let (tx1, _rx1) = mpsc::unbounded_channel();
+        let (tx2, _rx2) = mpsc::unbounded_channel();
+
+        connections.insert(
+            first_connection,
+            super::ClientHandle {
+                display_name: "User".to_string(),
+                is_sharing_video: true,
+                is_sharing_screen: false,
+                is_speaking: false,
+                is_muted: true,
+                tx: tx1,
+            },
+        );
+        connections.insert(
+            second_connection,
+            super::ClientHandle {
+                display_name: "User".to_string(),
+                is_sharing_video: false,
+                is_sharing_screen: true,
+                is_speaking: true,
+                is_muted: false,
+                tx: tx2,
+            },
+        );
+
+        let participant = aggregate_participant(Some(&connections), user_id).unwrap();
+        assert!(participant.is_sharing_video);
+        assert!(participant.is_sharing_screen);
+        assert!(participant.is_speaking);
+        assert!(!participant.is_muted);
+    }
 }