From 1c50587cad882e7fe751720fd4d7aee1b3dc4740 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 11 Sep 2023 17:37:05 -0700
Subject: [PATCH] Remove channel chat participant when connection is lost

---
 crates/collab/src/db/queries/buffers.rs  | 23 ++++++++
 crates/collab/src/db/queries/messages.rs | 19 ++++++
 crates/collab/src/db/queries/rooms.rs    | 73 ++++++++++--------------
 crates/collab/src/rpc.rs                 |  3 +-
 4 files changed, 74 insertions(+), 44 deletions(-)

diff --git a/crates/collab/src/db/queries/buffers.rs b/crates/collab/src/db/queries/buffers.rs
index 00de20140320640d98517556a5321433ad891f86..62ead11932af5cdab9049d719eef7b1238b73ffc 100644
--- a/crates/collab/src/db/queries/buffers.rs
+++ b/crates/collab/src/db/queries/buffers.rs
@@ -249,6 +249,29 @@ impl Database {
         .await
     }
 
+    pub async fn channel_buffer_connection_lost(
+        &self,
+        connection: ConnectionId,
+        tx: &DatabaseTransaction,
+    ) -> Result<()> {
+        channel_buffer_collaborator::Entity::update_many()
+            .filter(
+                Condition::all()
+                    .add(channel_buffer_collaborator::Column::ConnectionId.eq(connection.id as i32))
+                    .add(
+                        channel_buffer_collaborator::Column::ConnectionServerId
+                            .eq(connection.owner_id as i32),
+                    ),
+            )
+            .set(channel_buffer_collaborator::ActiveModel {
+                connection_lost: ActiveValue::set(true),
+                ..Default::default()
+            })
+            .exec(&*tx)
+            .await?;
+        Ok(())
+    }
+
     pub async fn leave_channel_buffers(
         &self,
         connection: ConnectionId,
diff --git a/crates/collab/src/db/queries/messages.rs b/crates/collab/src/db/queries/messages.rs
index 4593d8f2f3717ce6e9894220095b43a688e3a379..60bb245207718ceb2a1744d060bbb57e20831e20 100644
--- a/crates/collab/src/db/queries/messages.rs
+++ b/crates/collab/src/db/queries/messages.rs
@@ -25,6 +25,25 @@ impl Database {
         .await
     }
 
+    pub async fn channel_chat_connection_lost(
+        &self,
+        connection_id: ConnectionId,
+        tx: &DatabaseTransaction,
+    ) -> Result<()> {
+        channel_chat_participant::Entity::delete_many()
+            .filter(
+                Condition::all()
+                    .add(
+                        channel_chat_participant::Column::ConnectionServerId
+                            .eq(connection_id.owner_id),
+                    )
+                    .add(channel_chat_participant::Column::ConnectionId.eq(connection_id.id)),
+            )
+            .exec(tx)
+            .await?;
+        Ok(())
+    }
+
     pub async fn leave_channel_chat(
         &self,
         channel_id: ChannelId,
diff --git a/crates/collab/src/db/queries/rooms.rs b/crates/collab/src/db/queries/rooms.rs
index e348b50beeedb47c609d455ad759f59adea01adf..fb81fef176ab8bfa38ece1b4beb284088a96e0bb 100644
--- a/crates/collab/src/db/queries/rooms.rs
+++ b/crates/collab/src/db/queries/rooms.rs
@@ -890,54 +890,43 @@ impl Database {
 
     pub async fn connection_lost(&self, connection: ConnectionId) -> Result<()> {
         self.transaction(|tx| async move {
-            let participant = room_participant::Entity::find()
-                .filter(
-                    Condition::all()
-                        .add(
-                            room_participant::Column::AnsweringConnectionId
-                                .eq(connection.id as i32),
-                        )
-                        .add(
-                            room_participant::Column::AnsweringConnectionServerId
-                                .eq(connection.owner_id as i32),
-                        ),
-                )
-                .one(&*tx)
-                .await?;
-
-            if let Some(participant) = participant {
-                room_participant::Entity::update(room_participant::ActiveModel {
-                    answering_connection_lost: ActiveValue::set(true),
-                    ..participant.into_active_model()
-                })
-                .exec(&*tx)
+            self.room_connection_lost(connection, &*tx).await?;
+            self.channel_buffer_connection_lost(connection, &*tx)
                 .await?;
-            }
-
-            channel_buffer_collaborator::Entity::update_many()
-                .filter(
-                    Condition::all()
-                        .add(
-                            channel_buffer_collaborator::Column::ConnectionId
-                                .eq(connection.id as i32),
-                        )
-                        .add(
-                            channel_buffer_collaborator::Column::ConnectionServerId
-                                .eq(connection.owner_id as i32),
-                        ),
-                )
-                .set(channel_buffer_collaborator::ActiveModel {
-                    connection_lost: ActiveValue::set(true),
-                    ..Default::default()
-                })
-                .exec(&*tx)
-                .await?;
-
+            self.channel_chat_connection_lost(connection, &*tx).await?;
             Ok(())
         })
         .await
     }
 
+    pub async fn room_connection_lost(
+        &self,
+        connection: ConnectionId,
+        tx: &DatabaseTransaction,
+    ) -> Result<()> {
+        let participant = room_participant::Entity::find()
+            .filter(
+                Condition::all()
+                    .add(room_participant::Column::AnsweringConnectionId.eq(connection.id as i32))
+                    .add(
+                        room_participant::Column::AnsweringConnectionServerId
+                            .eq(connection.owner_id as i32),
+                    ),
+            )
+            .one(&*tx)
+            .await?;
+
+        if let Some(participant) = participant {
+            room_participant::Entity::update(room_participant::ActiveModel {
+                answering_connection_lost: ActiveValue::set(true),
+                ..participant.into_active_model()
+            })
+            .exec(&*tx)
+            .await?;
+        }
+        Ok(())
+    }
+
     fn build_incoming_call(
         room: &proto::Room,
         called_user_id: UserId,
diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index b508c45dc53f0a1017d44d1df96643165fc59b53..5ecd05ae1bb411f4333a61742115f619e582d231 100644
--- a/crates/collab/src/rpc.rs
+++ b/crates/collab/src/rpc.rs
@@ -904,9 +904,8 @@ async fn connection_lost(
                     room_updated(&room, &session.peer);
                 }
             }
-            update_user_contacts(session.user_id, &session).await?;
-
 
+            update_user_contacts(session.user_id, &session).await?;
         }
         _ = teardown.changed().fuse() => {}
     }