Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

maintain channel subscriptions in RAM (#9512)

Browse source 
This avoids a giant database query on every leave/join event.

Release Notes:

- N/A
This commit is contained in:
Conrad Irwin 2024-03-18 15:14:16 -06:00 • committed by GitHub
parent 963618a4a6
commit cd9b865e0a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 222 additions and 165 deletions
Download patch file Download diff file Expand all files Collapse all files

14
crates/collab/src/db.rs
View file

@ -546,7 +546,7 @@ pub struct Channel {
} }
impl Channel { impl Channel {
fn from_model(value: channel::Model) -> Self { pub fn from_model(value: channel::Model) -> Self {
Channel { Channel {
id: value.id, id: value.id,
visibility: value.visibility, visibility: value.visibility,
@ -604,16 +604,14 @@ pub struct RejoinedChannelBuffer {
#[derive(Clone)] #[derive(Clone)]
pub struct JoinRoom { pub struct JoinRoom {
pub room: proto::Room, pub room: proto::Room,
pub channel_id: Option<ChannelId>, pub channel: Option<channel::Model>,
pub channel_members: Vec<UserId>,
} }
pub struct RejoinedRoom { pub struct RejoinedRoom {
pub room: proto::Room, pub room: proto::Room,
pub rejoined_projects: Vec<RejoinedProject>, pub rejoined_projects: Vec<RejoinedProject>,
pub reshared_projects: Vec<ResharedProject>, pub reshared_projects: Vec<ResharedProject>,
pub channel_id: Option<ChannelId>, pub channel: Option<channel::Model>,
pub channel_members: Vec<UserId>,
} }
pub struct ResharedProject { pub struct ResharedProject {
@ -649,8 +647,7 @@ pub struct RejoinedWorktree {
pub struct LeftRoom { pub struct LeftRoom {
pub room: proto::Room, pub room: proto::Room,
pub channel_id: Option<ChannelId>, pub channel: Option<channel::Model>,
pub channel_members: Vec<UserId>,
pub left_projects: HashMap<ProjectId, LeftProject>, pub left_projects: HashMap<ProjectId, LeftProject>,
pub canceled_calls_to_user_ids: Vec<UserId>, pub canceled_calls_to_user_ids: Vec<UserId>,
pub deleted: bool, pub deleted: bool,
@ -658,8 +655,7 @@ pub struct LeftRoom {
pub struct RefreshedRoom { pub struct RefreshedRoom {
pub room: proto::Room, pub room: proto::Room,
pub channel_id: Option<ChannelId>, pub channel: Option<channel::Model>,
pub channel_members: Vec<UserId>,
pub stale_participant_user_ids: Vec<UserId>, pub stale_participant_user_ids: Vec<UserId>,
pub canceled_calls_to_user_ids: Vec<UserId>, pub canceled_calls_to_user_ids: Vec<UserId>,
} }

4
crates/collab/src/db/ids.rs
View file

@ -91,7 +91,9 @@ id_type!(NotificationKindId);
id_type!(HostedProjectId); id_type!(HostedProjectId);
/// ChannelRole gives you permissions for both channels and calls. /// ChannelRole gives you permissions for both channels and calls.
#[derive(Eq, PartialEq, Copy, Clone, Debug, EnumIter, DeriveActiveEnum, Default, Hash)] #[derive(
Eq, PartialEq, Copy, Clone, Debug, EnumIter, DeriveActiveEnum, Default, Hash, Serialize,
)]
#[sea_orm(rs_type = "String", db_type = "String(None)")] #[sea_orm(rs_type = "String", db_type = "String(None)")]
pub enum ChannelRole { pub enum ChannelRole {
/// Admin can read/write and change permissions. /// Admin can read/write and change permissions.

53
crates/collab/src/db/queries/channels.rs
View file

@ -45,11 +45,7 @@ impl Database {
name: &str, name: &str,
parent_channel_id: Option<ChannelId>, parent_channel_id: Option<ChannelId>,
admin_id: UserId, admin_id: UserId,
) -> Result<( ) -> Result<(channel::Model, Option<channel_member::Model>)> {
Channel,
Option<channel_member::Model>,
Vec<channel_member::Model>,
)> {
let name = Self::sanitize_channel_name(name)?; let name = Self::sanitize_channel_name(name)?;
self.transaction(move |tx| async move { self.transaction(move |tx| async move {
let mut parent = None; let mut parent = None;
@ -90,12 +86,7 @@ impl Database {
); );
} }
let channel_members = channel_member::Entity::find() Ok((channel, membership))
.filter(channel_member::Column::ChannelId.eq(channel.root_id()))
.all(&*tx)
.await?;
Ok((Channel::from_model(channel), membership, channel_members))
}) })
.await .await
} }
@ -181,7 +172,7 @@ impl Database {
channel_id: ChannelId, channel_id: ChannelId,
visibility: ChannelVisibility, visibility: ChannelVisibility,
admin_id: UserId, admin_id: UserId,
) -> Result<(Channel, Vec<channel_member::Model>)> { ) -> Result<channel::Model> {
self.transaction(move |tx| async move { self.transaction(move |tx| async move {
let channel = self.get_channel_internal(channel_id, &tx).await?; let channel = self.get_channel_internal(channel_id, &tx).await?;
self.check_user_is_channel_admin(&channel, admin_id, &tx) self.check_user_is_channel_admin(&channel, admin_id, &tx)
@ -214,12 +205,7 @@ impl Database {
model.visibility = ActiveValue::Set(visibility); model.visibility = ActiveValue::Set(visibility);
let channel = model.update(&*tx).await?; let channel = model.update(&*tx).await?;
let channel_members = channel_member::Entity::find() Ok(channel)
.filter(channel_member::Column::ChannelId.eq(channel.root_id()))
.all(&*tx)
.await?;
Ok((Channel::from_model(channel), channel_members))
}) })
.await .await
} }
@ -245,21 +231,12 @@ impl Database {
&self, &self,
channel_id: ChannelId, channel_id: ChannelId,
user_id: UserId, user_id: UserId,
) -> Result<(Vec<ChannelId>, Vec<UserId>)> { ) -> Result<(ChannelId, Vec<ChannelId>)> {
self.transaction(move |tx| async move { self.transaction(move |tx| async move {
let channel = self.get_channel_internal(channel_id, &tx).await?; let channel = self.get_channel_internal(channel_id, &tx).await?;
self.check_user_is_channel_admin(&channel, user_id, &tx) self.check_user_is_channel_admin(&channel, user_id, &tx)
.await?; .await?;
let members_to_notify: Vec<UserId> = channel_member::Entity::find()
.filter(channel_member::Column::ChannelId.eq(channel.root_id()))
.select_only()
.column(channel_member::Column::UserId)
.distinct()
.into_values::<_, QueryUserIds>()
.all(&*tx)
.await?;
let channels_to_remove = self let channels_to_remove = self
.get_channel_descendants_excluding_self([&channel], &tx) .get_channel_descendants_excluding_self([&channel], &tx)
.await? .await?
@ -273,7 +250,7 @@ impl Database {
.exec(&*tx) .exec(&*tx)
.await?; .await?;
Ok((channels_to_remove, members_to_notify)) Ok((channel.root_id(), channels_to_remove))
}) })
.await .await
} }
@ -343,7 +320,7 @@ impl Database {
channel_id: ChannelId, channel_id: ChannelId,
admin_id: UserId, admin_id: UserId,
new_name: &str, new_name: &str,
) -> Result<(Channel, Vec<channel_member::Model>)> { ) -> Result<channel::Model> {
self.transaction(move |tx| async move { self.transaction(move |tx| async move {
let new_name = Self::sanitize_channel_name(new_name)?.to_string(); let new_name = Self::sanitize_channel_name(new_name)?.to_string();
@ -355,12 +332,7 @@ impl Database {
model.name = ActiveValue::Set(new_name.clone()); model.name = ActiveValue::Set(new_name.clone());
let channel = model.update(&*tx).await?; let channel = model.update(&*tx).await?;
let channel_members = channel_member::Entity::find() Ok(channel)
.filter(channel_member::Column::ChannelId.eq(channel.root_id()))
.all(&*tx)
.await?;
Ok((Channel::from_model(channel), channel_members))
}) })
.await .await
} }
@ -984,7 +956,7 @@ impl Database {
channel_id: ChannelId, channel_id: ChannelId,
new_parent_id: ChannelId, new_parent_id: ChannelId,
admin_id: UserId, admin_id: UserId,
) -> Result<(Vec<Channel>, Vec<channel_member::Model>)> { ) -> Result<(ChannelId, Vec<Channel>)> {
self.transaction(|tx| async move { self.transaction(|tx| async move {
let channel = self.get_channel_internal(channel_id, &tx).await?; let channel = self.get_channel_internal(channel_id, &tx).await?;
self.check_user_is_channel_admin(&channel, admin_id, &tx) self.check_user_is_channel_admin(&channel, admin_id, &tx)
@ -1039,12 +1011,7 @@ impl Database {
.map(|c| Channel::from_model(c)) .map(|c| Channel::from_model(c))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let channel_members = channel_member::Entity::find() Ok((root_id, channels))
.filter(channel_member::Column::ChannelId.eq(root_id))
.all(&*tx)
.await?;
Ok((channels, channel_members))
}) })
.await .await
} }

33
crates/collab/src/db/queries/rooms.rs
View file

@ -52,12 +52,7 @@ impl Database {
); );
let (channel, room) = self.get_channel_room(room_id, &tx).await?; let (channel, room) = self.get_channel_room(room_id, &tx).await?;
let channel_members; if channel.is_none() {
if let Some(channel) = &channel {
channel_members = self.get_channel_participants(channel, &tx).await?;
} else {
channel_members = Vec::new();
// Delete the room if it becomes empty. // Delete the room if it becomes empty.
if room.participants.is_empty() { if room.participants.is_empty() {
project::Entity::delete_many() project::Entity::delete_many()
@ -70,8 +65,7 @@ impl Database {
Ok(RefreshedRoom { Ok(RefreshedRoom {
room, room,
channel_id: channel.map(|channel| channel.id), channel,
channel_members,
stale_participant_user_ids, stale_participant_user_ids,
canceled_calls_to_user_ids, canceled_calls_to_user_ids,
}) })
@ -349,8 +343,7 @@ impl Database {
let room = self.get_room(room_id, &tx).await?; let room = self.get_room(room_id, &tx).await?;
Ok(JoinRoom { Ok(JoinRoom {
room, room,
channel_id: None, channel: None,
channel_members: vec![],
}) })
}) })
.await .await
@ -446,11 +439,9 @@ impl Database {
let (channel, room) = self.get_channel_room(room_id, &tx).await?; let (channel, room) = self.get_channel_room(room_id, &tx).await?;
let channel = channel.ok_or_else(|| anyhow!("no channel for room"))?; let channel = channel.ok_or_else(|| anyhow!("no channel for room"))?;
let channel_members = self.get_channel_participants(&channel, tx).await?;
Ok(JoinRoom { Ok(JoinRoom {
room, room,
channel_id: Some(channel.id), channel: Some(channel),
channel_members,
}) })
} }
@ -736,16 +727,10 @@ impl Database {
} }
let (channel, room) = self.get_channel_room(room_id, &tx).await?; let (channel, room) = self.get_channel_room(room_id, &tx).await?;
let channel_members = if let Some(channel) = &channel {
self.get_channel_participants(&channel, &tx).await?
} else {
Vec::new()
};
Ok(RejoinedRoom { Ok(RejoinedRoom {
room, room,
channel_id: channel.map(|channel| channel.id), channel,
channel_members,
rejoined_projects, rejoined_projects,
reshared_projects, reshared_projects,
}) })
@ -902,15 +887,9 @@ impl Database {
false false
}; };
let channel_members = if let Some(channel) = &channel {
self.get_channel_participants(channel, &tx).await?
} else {
Vec::new()
};
let left_room = LeftRoom { let left_room = LeftRoom {
room, room,
channel_id: channel.map(|channel| channel.id), channel,
channel_members,
left_projects, left_projects,
canceled_calls_to_user_ids, canceled_calls_to_user_ids,
deleted, deleted,

3
crates/collab/src/db/tests/channel_tests.rs
View file

@ -109,10 +109,9 @@ async fn test_channels(db: &Arc<Database>) {
assert!(db.get_channel(crdb_id, a_id).await.is_err()); assert!(db.get_channel(crdb_id, a_id).await.is_err());
// Remove a channel tree // Remove a channel tree
let (mut channel_ids, user_ids) = db.delete_channel(rust_id, a_id).await.unwrap(); let (_, mut channel_ids) = db.delete_channel(rust_id, a_id).await.unwrap();
channel_ids.sort(); channel_ids.sort();
assert_eq!(channel_ids, &[rust_id, cargo_id, cargo_ra_id]); assert_eq!(channel_ids, &[rust_id, cargo_id, cargo_ra_id]);
assert_eq!(user_ids, &[a_id]);
assert!(db.get_channel(rust_id, a_id).await.is_err()); assert!(db.get_channel(rust_id, a_id).await.is_err());
assert!(db.get_channel(cargo_id, a_id).await.is_err()); assert!(db.get_channel(cargo_id, a_id).await.is_err());

174
crates/collab/src/rpc.rs
View file

@ -3,10 +3,10 @@ mod connection_pool;
use crate::{ use crate::{
auth::{self, Impersonator}, auth::{self, Impersonator},
db::{ db::{
self, BufferId, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage, Database, self, BufferId, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage,
InviteMemberResult, MembershipUpdated, MessageId, NotificationId, Project, ProjectId, Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, Project,
RemoveChannelMemberResult, ReplicaId, RespondToChannelInvite, RoomId, ServerId, User, ProjectId, RemoveChannelMemberResult, ReplicaId, RespondToChannelInvite, RoomId, ServerId,
UserId, User, UserId,
}, },
executor::Executor, executor::Executor,
AppState, Error, Result, AppState, Error, Result,
@ -351,14 +351,8 @@ impl Server {
"refreshed room" "refreshed room"
); );
room_updated(&refreshed_room.room, &peer); room_updated(&refreshed_room.room, &peer);
if let Some(channel_id) = refreshed_room.channel_id { if let Some(channel) = refreshed_room.channel.as_ref() {
channel_updated( channel_updated(channel, &refreshed_room.room, &peer, &pool.lock());
channel_id,
&refreshed_room.room,
&refreshed_room.channel_members,
&peer,
&pool.lock(),
);
} }
contacts_to_update contacts_to_update
.extend(refreshed_room.stale_participant_user_ids.iter().copied()); .extend(refreshed_room.stale_participant_user_ids.iter().copied());
@ -699,6 +693,9 @@ impl Server {
{ {
let mut pool = self.connection_pool.lock(); let mut pool = self.connection_pool.lock();
pool.add_connection(connection_id, user.id, user.admin, zed_version); pool.add_connection(connection_id, user.id, user.admin, zed_version);
for membership in &channels_for_user.channel_memberships {
pool.subscribe_to_channel(user.id, membership.channel_id, membership.role)
}
self.peer.send( self.peer.send(
connection_id, connection_id,
build_initial_contacts_update(contacts, &pool), build_initial_contacts_update(contacts, &pool),
@ -1148,8 +1145,7 @@ async fn rejoin_room(
session: Session, session: Session,
) -> Result<()> { ) -> Result<()> {
let room; let room;
let channel_id; let channel;
let channel_members;
{ {
let mut rejoined_room = session let mut rejoined_room = session
.db() .db()
@ -1315,15 +1311,13 @@ async fn rejoin_room(
let rejoined_room = rejoined_room.into_inner(); let rejoined_room = rejoined_room.into_inner();
room = rejoined_room.room; room = rejoined_room.room;
channel_id = rejoined_room.channel_id; channel = rejoined_room.channel;
channel_members = rejoined_room.channel_members;
} }
if let Some(channel_id) = channel_id { if let Some(channel) = channel {
channel_updated( channel_updated(
channel_id, &channel,
&room, &room,
&channel_members,
&session.peer, &session.peer,
&*session.connection_pool().await, &*session.connection_pool().await,
); );
@ -2427,31 +2421,39 @@ async fn create_channel(
let db = session.db().await; let db = session.db().await;
let parent_id = request.parent_id.map(|id| ChannelId::from_proto(id)); let parent_id = request.parent_id.map(|id| ChannelId::from_proto(id));
let (channel, owner, channel_members) = db let (channel, membership) = db
.create_channel(&request.name, parent_id, session.user_id) .create_channel(&request.name, parent_id, session.user_id)
.await?; .await?;
let root_id = channel.root_id();
let channel = Channel::from_model(channel);
response.send(proto::CreateChannelResponse { response.send(proto::CreateChannelResponse {
channel: Some(channel.to_proto()), channel: Some(channel.to_proto()),
parent_id: request.parent_id, parent_id: request.parent_id,
})?; })?;
let connection_pool = session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
if let Some(owner) = owner { if let Some(membership) = membership {
connection_pool.subscribe_to_channel(
membership.user_id,
membership.channel_id,
membership.role,
);
let update = proto::UpdateUserChannels { let update = proto::UpdateUserChannels {
channel_memberships: vec![proto::ChannelMembership { channel_memberships: vec![proto::ChannelMembership {
channel_id: owner.channel_id.to_proto(), channel_id: membership.channel_id.to_proto(),
role: owner.role.into(), role: membership.role.into(),
}], }],
..Default::default() ..Default::default()
}; };
for connection_id in connection_pool.user_connection_ids(owner.user_id) { for connection_id in connection_pool.user_connection_ids(membership.user_id) {
session.peer.send(connection_id, update.clone())?; session.peer.send(connection_id, update.clone())?;
} }
} }
for channel_member in channel_members { for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
if !channel_member.role.can_see_channel(channel.visibility) { if !role.can_see_channel(channel.visibility) {
continue; continue;
} }
@ -2459,9 +2461,7 @@ async fn create_channel(
channels: vec![channel.to_proto()], channels: vec![channel.to_proto()],
..Default::default() ..Default::default()
}; };
for connection_id in connection_pool.user_connection_ids(channel_member.user_id) { session.peer.send(connection_id, update.clone())?;
session.peer.send(connection_id, update.clone())?;
}
} }
Ok(()) Ok(())
@ -2476,7 +2476,7 @@ async fn delete_channel(
let db = session.db().await; let db = session.db().await;
let channel_id = request.channel_id; let channel_id = request.channel_id;
let (removed_channels, member_ids) = db let (root_channel, removed_channels) = db
.delete_channel(ChannelId::from_proto(channel_id), session.user_id) .delete_channel(ChannelId::from_proto(channel_id), session.user_id)
.await?; .await?;
response.send(proto::Ack {})?; response.send(proto::Ack {})?;
@ -2488,10 +2488,8 @@ async fn delete_channel(
.extend(removed_channels.into_iter().map(|id| id.to_proto())); .extend(removed_channels.into_iter().map(|id| id.to_proto()));
let connection_pool = session.connection_pool().await; let connection_pool = session.connection_pool().await;
for member_id in member_ids { for (connection_id, _) in connection_pool.channel_connection_ids(root_channel) {
for connection_id in connection_pool.user_connection_ids(member_id) { session.peer.send(connection_id, update.clone())?;
session.peer.send(connection_id, update.clone())?;
}
} }
Ok(()) Ok(())
@ -2551,9 +2549,9 @@ async fn remove_channel_member(
.remove_channel_member(channel_id, member_id, session.user_id) .remove_channel_member(channel_id, member_id, session.user_id)
.await?; .await?;
let connection_pool = &session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
notify_membership_updated( notify_membership_updated(
&connection_pool, &mut connection_pool,
membership_update, membership_update,
member_id, member_id,
&session.peer, &session.peer,
@ -2588,25 +2586,33 @@ async fn set_channel_visibility(
let channel_id = ChannelId::from_proto(request.channel_id); let channel_id = ChannelId::from_proto(request.channel_id);
let visibility = request.visibility().into(); let visibility = request.visibility().into();
let (channel, channel_members) = db let channel_model = db
.set_channel_visibility(channel_id, visibility, session.user_id) .set_channel_visibility(channel_id, visibility, session.user_id)
.await?; .await?;
let root_id = channel_model.root_id();
let channel = Channel::from_model(channel_model);
let connection_pool = session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
for member in channel_members { for (user_id, role) in connection_pool
let update = if member.role.can_see_channel(channel.visibility) { .channel_user_ids(root_id)
.collect::<Vec<_>>()
.into_iter()
{
let update = if role.can_see_channel(channel.visibility) {
connection_pool.subscribe_to_channel(user_id, channel_id, role);
proto::UpdateChannels { proto::UpdateChannels {
channels: vec![channel.to_proto()], channels: vec![channel.to_proto()],
..Default::default() ..Default::default()
} }
} else { } else {
connection_pool.unsubscribe_from_channel(&user_id, &channel_id);
proto::UpdateChannels { proto::UpdateChannels {
delete_channels: vec![channel.id.to_proto()], delete_channels: vec![channel.id.to_proto()],
..Default::default() ..Default::default()
} }
}; };
for connection_id in connection_pool.user_connection_ids(member.user_id) { for connection_id in connection_pool.user_connection_ids(user_id) {
session.peer.send(connection_id, update.clone())?; session.peer.send(connection_id, update.clone())?;
} }
} }
@ -2635,9 +2641,9 @@ async fn set_channel_member_role(
match result { match result {
db::SetMemberRoleResult::MembershipUpdated(membership_update) => { db::SetMemberRoleResult::MembershipUpdated(membership_update) => {
let connection_pool = session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
notify_membership_updated( notify_membership_updated(
&connection_pool, &mut connection_pool,
membership_update, membership_update,
member_id, member_id,
&session.peer, &session.peer,
@ -2671,24 +2677,23 @@ async fn rename_channel(
) -> Result<()> { ) -> Result<()> {
let db = session.db().await; let db = session.db().await;
let channel_id = ChannelId::from_proto(request.channel_id); let channel_id = ChannelId::from_proto(request.channel_id);
let (channel, channel_members) = db let channel_model = db
.rename_channel(channel_id, session.user_id, &request.name) .rename_channel(channel_id, session.user_id, &request.name)
.await?; .await?;
let root_id = channel_model.root_id();
let channel = Channel::from_model(channel_model);
response.send(proto::RenameChannelResponse { response.send(proto::RenameChannelResponse {
channel: Some(channel.to_proto()), channel: Some(channel.to_proto()),
})?; })?;
let connection_pool = session.connection_pool().await; let connection_pool = session.connection_pool().await;
for channel_member in channel_members { let update = proto::UpdateChannels {
if !channel_member.role.can_see_channel(channel.visibility) { channels: vec![channel.to_proto()],
continue; ..Default::default()
} };
let update = proto::UpdateChannels { for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
channels: vec![channel.to_proto()], if role.can_see_channel(channel.visibility) {
..Default::default()
};
for connection_id in connection_pool.user_connection_ids(channel_member.user_id) {
session.peer.send(connection_id, update.clone())?; session.peer.send(connection_id, update.clone())?;
} }
} }
@ -2705,18 +2710,18 @@ async fn move_channel(
let channel_id = ChannelId::from_proto(request.channel_id); let channel_id = ChannelId::from_proto(request.channel_id);
let to = ChannelId::from_proto(request.to); let to = ChannelId::from_proto(request.to);
let (channels, channel_members) = session let (root_id, channels) = session
.db() .db()
.await .await
.move_channel(channel_id, to, session.user_id) .move_channel(channel_id, to, session.user_id)
.await?; .await?;
let connection_pool = session.connection_pool().await; let connection_pool = session.connection_pool().await;
for member in channel_members { for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
let channels = channels let channels = channels
.iter() .iter()
.filter_map(|channel| { .filter_map(|channel| {
if member.role.can_see_channel(channel.visibility) { if role.can_see_channel(channel.visibility) {
Some(channel.to_proto()) Some(channel.to_proto())
} else { } else {
None None
@ -2732,9 +2737,7 @@ async fn move_channel(
..Default::default() ..Default::default()
}; };
for connection_id in connection_pool.user_connection_ids(member.user_id) { session.peer.send(connection_id, update.clone())?;
session.peer.send(connection_id, update.clone())?;
}
} }
response.send(Ack {})?; response.send(Ack {})?;
@ -2771,10 +2774,10 @@ async fn respond_to_channel_invite(
.respond_to_channel_invite(channel_id, session.user_id, request.accept) .respond_to_channel_invite(channel_id, session.user_id, request.accept)
.await?; .await?;
let connection_pool = session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
if let Some(membership_update) = membership_update { if let Some(membership_update) = membership_update {
notify_membership_updated( notify_membership_updated(
&connection_pool, &mut connection_pool,
membership_update, membership_update,
session.user_id, session.user_id,
&session.peer, &session.peer,
@ -2866,14 +2869,17 @@ async fn join_channel_internal(
response.send(proto::JoinRoomResponse { response.send(proto::JoinRoomResponse {
room: Some(joined_room.room.clone()), room: Some(joined_room.room.clone()),
channel_id: joined_room.channel_id.map(|id| id.to_proto()), channel_id: joined_room
.channel
.as_ref()
.map(|channel| channel.id.to_proto()),
live_kit_connection_info, live_kit_connection_info,
})?; })?;
let connection_pool = session.connection_pool().await; let mut connection_pool = session.connection_pool().await;
if let Some(membership_updated) = membership_updated { if let Some(membership_updated) = membership_updated {
notify_membership_updated( notify_membership_updated(
&connection_pool, &mut connection_pool,
membership_updated, membership_updated,
session.user_id, session.user_id,
&session.peer, &session.peer,
@ -2886,9 +2892,10 @@ async fn join_channel_internal(
}; };
channel_updated( channel_updated(
channel_id, &joined_room
.channel
.ok_or_else(|| anyhow!("channel not returned"))?,
&joined_room.room, &joined_room.room,
&joined_room.channel_members,
&session.peer, &session.peer,
&*session.connection_pool().await, &*session.connection_pool().await,
); );
@ -3403,11 +3410,18 @@ fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
} }
fn notify_membership_updated( fn notify_membership_updated(
connection_pool: &ConnectionPool, connection_pool: &mut ConnectionPool,
result: MembershipUpdated, result: MembershipUpdated,
user_id: UserId, user_id: UserId,
peer: &Peer, peer: &Peer,
) { ) {
for membership in &result.new_channels.channel_memberships {
connection_pool.subscribe_to_channel(user_id, membership.channel_id, membership.role)
}
for channel_id in &result.removed_channels {
connection_pool.unsubscribe_from_channel(&user_id, channel_id)
}
let user_channels_update = proto::UpdateUserChannels { let user_channels_update = proto::UpdateUserChannels {
channel_memberships: result channel_memberships: result
.new_channels .new_channels
@ -3420,6 +3434,7 @@ fn notify_membership_updated(
.collect(), .collect(),
..Default::default() ..Default::default()
}; };
let mut update = build_channels_update(result.new_channels, vec![]); let mut update = build_channels_update(result.new_channels, vec![]);
update.delete_channels = result update.delete_channels = result
.removed_channels .removed_channels
@ -3533,9 +3548,8 @@ fn room_updated(room: &proto::Room, peer: &Peer) {
} }
fn channel_updated( fn channel_updated(
channel_id: ChannelId, channel: &db::channel::Model,
room: &proto::Room, room: &proto::Room,
channel_members: &[UserId],
peer: &Peer, peer: &Peer,
pool: &ConnectionPool, pool: &ConnectionPool,
) { ) {
@ -3547,15 +3561,16 @@ fn channel_updated(
broadcast( broadcast(
None, None,
channel_members pool.channel_connection_ids(channel.root_id())
.iter() .filter_map(|(channel_id, role)| {
.flat_map(|user_id| pool.user_connection_ids(*user_id)), role.can_see_channel(channel.visibility).then(|| channel_id)
}),
|peer_id| { |peer_id| {
peer.send( peer.send(
peer_id, peer_id,
proto::UpdateChannels { proto::UpdateChannels {
channel_participants: vec![proto::ChannelParticipants { channel_participants: vec![proto::ChannelParticipants {
channel_id: channel_id.to_proto(), channel_id: channel.id.to_proto(),
participant_user_ids: participants.clone(), participant_user_ids: participants.clone(),
}], }],
..Default::default() ..Default::default()
@ -3608,8 +3623,7 @@ async fn leave_room_for_session(session: &Session) -> Result<()> {
let live_kit_room; let live_kit_room;
let delete_live_kit_room; let delete_live_kit_room;
let room; let room;
let channel_members; let channel;
let channel_id;
if let Some(mut left_room) = session.db().await.leave_room(session.connection_id).await? { if let Some(mut left_room) = session.db().await.leave_room(session.connection_id).await? {
contacts_to_update.insert(session.user_id); contacts_to_update.insert(session.user_id);
@ -3623,19 +3637,17 @@ async fn leave_room_for_session(session: &Session) -> Result<()> {
live_kit_room = mem::take(&mut left_room.room.live_kit_room); live_kit_room = mem::take(&mut left_room.room.live_kit_room);
delete_live_kit_room = left_room.deleted; delete_live_kit_room = left_room.deleted;
room = mem::take(&mut left_room.room); room = mem::take(&mut left_room.room);
channel_members = mem::take(&mut left_room.channel_members); channel = mem::take(&mut left_room.channel);
channel_id = left_room.channel_id;
room_updated(&room, &session.peer); room_updated(&room, &session.peer);
} else { } else {
return Ok(()); return Ok(());
} }
if let Some(channel_id) = channel_id { if let Some(channel) = channel {
channel_updated( channel_updated(
channel_id, &channel,
&room, &room,
&channel_members,
&session.peer, &session.peer,
&*session.connection_pool().await, &*session.connection_pool().await,
); );

106
crates/collab/src/rpc/connection_pool.rs
View file

@ -1,6 +1,6 @@
use crate::db::UserId; use crate::db::{ChannelId, ChannelRole, UserId};
use anyhow::{anyhow, Result}; use anyhow::{anyhow, Result};
use collections::{BTreeMap, HashSet}; use collections::{BTreeMap, HashMap, HashSet};
use rpc::ConnectionId; use rpc::ConnectionId;
use serde::Serialize; use serde::Serialize;
use tracing::instrument; use tracing::instrument;
@ -10,6 +10,7 @@ use util::SemanticVersion;
pub struct ConnectionPool { pub struct ConnectionPool {
connections: BTreeMap<ConnectionId, Connection>, connections: BTreeMap<ConnectionId, Connection>,
connected_users: BTreeMap<UserId, ConnectedUser>, connected_users: BTreeMap<UserId, ConnectedUser>,
channels: ChannelPool,
} }
#[derive(Default, Serialize)] #[derive(Default, Serialize)]
@ -47,6 +48,7 @@ impl ConnectionPool {
pub fn reset(&mut self) { pub fn reset(&mut self) {
self.connections.clear(); self.connections.clear();
self.connected_users.clear(); self.connected_users.clear();
self.channels.clear();
} }
#[instrument(skip(self))] #[instrument(skip(self))]
@ -81,6 +83,7 @@ impl ConnectionPool {
connected_user.connection_ids.remove(&connection_id); connected_user.connection_ids.remove(&connection_id);
if connected_user.connection_ids.is_empty() { if connected_user.connection_ids.is_empty() {
self.connected_users.remove(&user_id); self.connected_users.remove(&user_id);
self.channels.remove_user(&user_id);
} }
self.connections.remove(&connection_id).unwrap(); self.connections.remove(&connection_id).unwrap();
Ok(()) Ok(())
@ -110,6 +113,38 @@ impl ConnectionPool {
.copied() .copied()
} }
pub fn channel_user_ids(
&self,
channel_id: ChannelId,
) -> impl Iterator<Item = (UserId, ChannelRole)> + '_ {
self.channels.users_to_notify(channel_id)
}
pub fn channel_connection_ids(
&self,
channel_id: ChannelId,
) -> impl Iterator<Item = (ConnectionId, ChannelRole)> + '_ {
self.channels
.users_to_notify(channel_id)
.flat_map(|(user_id, role)| {
self.user_connection_ids(user_id)
.map(move |connection_id| (connection_id, role))
})
}
pub fn subscribe_to_channel(
&mut self,
user_id: UserId,
channel_id: ChannelId,
role: ChannelRole,
) {
self.channels.subscribe(user_id, channel_id, role);
}
pub fn unsubscribe_from_channel(&mut self, user_id: &UserId, channel_id: &ChannelId) {
self.channels.unsubscribe(user_id, channel_id);
}
pub fn is_user_online(&self, user_id: UserId) -> bool { pub fn is_user_online(&self, user_id: UserId) -> bool {
!self !self
.connected_users .connected_users
@ -140,3 +175,70 @@ impl ConnectionPool {
} }
} }
} }
#[derive(Default, Serialize)]
pub struct ChannelPool {
by_user: HashMap<UserId, HashMap<ChannelId, ChannelRole>>,
by_channel: HashMap<ChannelId, HashSet<UserId>>,
}
impl ChannelPool {
pub fn clear(&mut self) {
self.by_user.clear();
self.by_channel.clear();
}
pub fn subscribe(&mut self, user_id: UserId, channel_id: ChannelId, role: ChannelRole) {
self.by_user
.entry(user_id)
.or_default()
.insert(channel_id, role);
self.by_channel
.entry(channel_id)
.or_default()
.insert(user_id);
}
pub fn unsubscribe(&mut self, user_id: &UserId, channel_id: &ChannelId) {
if let Some(channels) = self.by_user.get_mut(user_id) {
channels.remove(channel_id);
if channels.is_empty() {
self.by_user.remove(user_id);
}
}
if let Some(users) = self.by_channel.get_mut(channel_id) {
users.remove(user_id);
if users.is_empty() {
self.by_channel.remove(channel_id);
}
}
}
pub fn remove_user(&mut self, user_id: &UserId) {
if let Some(channels) = self.by_user.remove(&user_id) {
for channel_id in channels.keys() {
self.unsubscribe(user_id, &channel_id)
}
}
}
pub fn users_to_notify(
&self,
channel_id: ChannelId,
) -> impl '_ + Iterator<Item = (UserId, ChannelRole)> {
self.by_channel
.get(&channel_id)
.into_iter()
.flat_map(move |users| {
users.iter().flat_map(move |user_id| {
Some((
*user_id,
self.by_user
.get(user_id)
.and_then(|channels| channels.get(&channel_id))
.copied()?,
))
})
})
}
}