Merge member Ids set with channel links map (#7027)

backend/canisters/community/CHANGELOG.md

@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Allow bots to send a subset of message types ([#7016](https://github.com/open-chat-labs/open-chat/pull/7016))
 - Switch to using `PrincipalToStableMemoryMap` ([#7023](https://github.com/open-chat-labs/open-chat/pull/7023))
 - Reduce size by grouping member -> channel links per userId ([#7025](https://github.com/open-chat-labs/open-chat/pull/7025))
+- Merge member Ids set with channel links map ([#7027](https://github.com/open-chat-labs/open-chat/pull/7027))
 
 ### Removed
 

backend/canisters/community/impl/src/activity_notifications.rs

@@ -49,7 +49,7 @@ pub(crate) fn extract_activity(now: TimestampMillis, data: &Data) -> PublicCommu
 
     let mut activity = PublicCommunityActivity {
         timestamp: now,
-        member_count: data.members.len(),
+        member_count: data.members.len() as u32,
         channel_count: data.channels.public_channels().len() as u32,
         ..Default::default()
     };

backend/canisters/community/impl/src/jobs/import_groups.rs

@@ -359,7 +359,7 @@ fn add_community_members_to_channel_if_public(channel_id: ChannelId, state: &mut
         // If this is a public channel, add all community members to it
         if channel.chat.is_public.value && channel.chat.gate_config.value.is_none() {
             let now = state.env.now();
-            let user_ids: Vec<_> = state.data.members.member_ids().iter().copied().collect();
+            let user_ids: Vec<_> = state.data.members.iter_member_ids().collect();
             add_members_to_public_channel_unchecked(
                 user_ids.into_iter(),
                 channel,

backend/canisters/community/impl/src/lib.rs

@@ -217,7 +217,7 @@ impl RuntimeState {
             banner_id: Document::id(&data.banner),
             is_public: data.is_public.value,
             latest_event_index: data.events.latest_event_index(),
-            member_count: data.members.len(),
+            member_count: data.members.len() as u32,
             permissions: data.permissions.value.clone(),
             frozen: data.frozen.value.clone(),
             gate: data.gate_config.value.as_ref().map(|gc| gc.gate.clone()),
@@ -288,7 +288,7 @@ impl RuntimeState {
             git_commit_id: utils::git::git_commit_id().to_string(),
             public: self.data.is_public.value,
             date_created: self.data.date_created,
-            members: self.data.members.len(),
+            members: self.data.members.len() as u32,
             admins: self.data.members.admins().len() as u32,
             owners: self.data.members.owners().len() as u32,
             blocked: self.data.members.blocked().len() as u32,

backend/canisters/community/impl/src/lifecycle/post_upgrade.rs

@@ -18,10 +18,11 @@ fn post_upgrade(args: Args) {
     let memory = get_upgrades_memory();
     let reader = get_reader(&memory);
 
-    let (data, errors, logs, traces): (Data, Vec<LogEntry>, Vec<LogEntry>, Vec<LogEntry>) =
+    let (mut data, errors, logs, traces): (Data, Vec<LogEntry>, Vec<LogEntry>, Vec<LogEntry>) =
         msgpack::deserialize(reader).unwrap();
 
     canister_logger::init_with_logs(data.test_mode, errors, logs, traces);
+    data.members.move_member_ids_into_channel_links_map();
 
     let env = init_env(data.rng_seed);
     init_state(env, data, args.wasm_version);

backend/canisters/community/impl/src/model/members.rs

@@ -7,6 +7,7 @@ use principal_to_user_id_map::{deserialize_principal_to_user_id_map_from_heap, P
 use rand::RngCore;
 use serde::de::{SeqAccess, Visitor};
 use serde::{Deserialize, Deserializer, Serialize};
+use std::collections::btree_map::Entry::Vacant;
 use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::Formatter;
 use types::{
@@ -23,14 +24,15 @@ const MAX_MEMBERS_PER_COMMUNITY: u32 = 100_000;
 #[derive(Serialize, Deserialize)]
 pub struct CommunityMembers {
     members_map: MembersStableStorage,
-    #[serde(deserialize_with = "deserialize_member_channel_links")]
-    member_channel_links: BTreeMap<UserId, Vec<ChannelId>>,
+    #[serde(alias = "member_channel_links", deserialize_with = "deserialize_members_and_channels")]
+    members_and_channels: BTreeMap<UserId, Vec<ChannelId>>,
     member_channel_links_removed: BTreeMap<(UserId, ChannelId), TimestampMillis>,
     user_groups: UserGroups,
-    // This includes the userIds of community members and also users invited to the community
     #[serde(deserialize_with = "deserialize_principal_to_user_id_map_from_heap")]
     principal_to_user_id_map: PrincipalToUserIdMap,
-    member_ids: BTreeSet<UserId>,
+    // This includes the userIds of community members and also users invited to the community
+    #[deprecated]
+    member_ids: Vec<UserId>,
     owners: BTreeSet<UserId>,
     admins: BTreeSet<UserId>,
     bots: BTreeMap<UserId, UserType>,
@@ -44,6 +46,15 @@ pub struct CommunityMembers {
 }
 
 impl CommunityMembers {
+    pub fn move_member_ids_into_channel_links_map(&mut self) {
+        #[allow(deprecated)]
+        for user_id in std::mem::take(&mut self.member_ids) {
+            if let Vacant(e) = self.members_and_channels.entry(user_id) {
+                e.insert(Vec::new());
+            }
+        }
+    }
+
     pub fn new(
         creator_principal: Principal,
         creator_user_id: UserId,
@@ -68,13 +79,14 @@ impl CommunityMembers {
         let mut principal_to_user_id_map = PrincipalToUserIdMap::default();
         principal_to_user_id_map.insert(creator_principal, creator_user_id);
 
+        #[allow(deprecated)]
         CommunityMembers {
             members_map: MembersStableStorage::new(member),
-            member_channel_links: [(creator_user_id, public_channels)].into_iter().collect(),
+            members_and_channels: [(creator_user_id, public_channels)].into_iter().collect(),
             member_channel_links_removed: BTreeMap::new(),
             user_groups: UserGroups::default(),
             principal_to_user_id_map,
-            member_ids: [creator_user_id].into_iter().collect(),
+            member_ids: Vec::new(),
             owners: [creator_user_id].into_iter().collect(),
             admins: BTreeSet::new(),
             bots: if creator_user_type.is_bot() {
@@ -102,9 +114,9 @@ impl CommunityMembers {
     ) -> AddResult {
         if self.blocked.contains(&user_id) {
             AddResult::Blocked
-        } else if self.member_ids.contains(&user_id) {
-            AddResult::AlreadyInCommunity
-        } else {
+        } else if let Vacant(e) = self.members_and_channels.entry(user_id) {
+            e.insert(Vec::new());
+
             if referred_by == Some(user_id) {
                 referred_by = None;
             }
@@ -137,9 +149,9 @@ impl CommunityMembers {
                     self.members_with_referrals.insert(referrer);
                 }
             }
-            self.member_ids.insert(user_id);
-
             AddResult::Success(member)
+        } else {
+            AddResult::AlreadyInCommunity
         }
     }
 
@@ -198,13 +210,12 @@ impl CommunityMembers {
                 self.members_with_referrals.remove(&referrer);
             }
         }
-        self.member_channel_links.remove(&user_id);
+        self.members_and_channels.remove(&user_id);
         let channels_removed: Vec<_> = self.channels_removed_for_member(user_id).map(|(c, _)| c).collect();
         for channel_id in channels_removed {
             self.member_channel_links_removed.remove(&(user_id, channel_id));
         }
         self.user_groups.remove_user_from_all(&member.user_id, now);
-        self.member_ids.remove(&user_id);
         self.prune_then_insert_member_update(user_id, MemberUpdate::Removed, now);
 
         Some(member)
@@ -315,7 +326,7 @@ impl CommunityMembers {
         rng: &mut R,
         now: TimestampMillis,
     ) -> Option<u32> {
-        users.retain(|u| self.member_ids.contains(u));
+        users.retain(|u| self.members_and_channels.contains_key(u));
 
         self.user_groups.create(name, users, rng, now)
     }
@@ -328,7 +339,7 @@ impl CommunityMembers {
         users_to_remove: Vec<UserId>,
         now: TimestampMillis,
     ) -> bool {
-        users_to_add.retain(|u| self.member_ids.contains(u));
+        users_to_add.retain(|u| self.members_and_channels.contains_key(u));
 
         self.user_groups
             .update(user_group_id, name, users_to_add, users_to_remove, now)
@@ -361,11 +372,8 @@ impl CommunityMembers {
     }
 
     pub fn mark_member_joined_channel(&mut self, user_id: UserId, channel_id: ChannelId) {
-        if self.member_ids.contains(&user_id) {
-            self.member_channel_links
-                .entry(user_id)
-                .or_default()
-                .push_if_not_contains(channel_id);
+        if let Some(channel_ids) = self.members_and_channels.get_mut(&user_id) {
+            channel_ids.push_if_not_contains(channel_id);
             self.member_channel_links_removed.remove(&(user_id, channel_id));
         }
     }
@@ -377,10 +385,8 @@ impl CommunityMembers {
         channel_deleted: bool,
         now: TimestampMillis,
     ) {
-        if self.member_ids.contains(&user_id) {
-            if let Some(channel_ids) = self.member_channel_links.get_mut(&user_id) {
-                channel_ids.retain(|id| *id != channel_id);
-            }
+        if let Some(channel_ids) = self.members_and_channels.get_mut(&user_id) {
+            channel_ids.retain(|id| *id != channel_id);
             if !channel_deleted {
                 self.member_channel_links_removed.insert((user_id, channel_id), now);
             }
@@ -410,7 +416,7 @@ impl CommunityMembers {
     }
 
     pub fn user_limit_reached(&self) -> Option<u32> {
-        if self.member_ids.len() >= MAX_MEMBERS_PER_COMMUNITY as usize {
+        if self.members_and_channels.len() >= MAX_MEMBERS_PER_COMMUNITY as usize {
             Some(MAX_MEMBERS_PER_COMMUNITY)
         } else {
             None
@@ -428,7 +434,7 @@ impl CommunityMembers {
     pub fn lookup_user_id(&self, user_id_or_principal: Principal) -> Option<UserId> {
         self.principal_to_user_id_map.get(&user_id_or_principal).or_else(|| {
             let user_id: UserId = user_id_or_principal.into();
-            self.member_ids.contains(&user_id).then_some(user_id)
+            self.members_and_channels.contains_key(&user_id).then_some(user_id)
         })
     }
 
@@ -445,19 +451,19 @@ impl CommunityMembers {
     }
 
     pub fn contains(&self, user_id: &UserId) -> bool {
-        self.member_ids.contains(user_id)
+        self.members_and_channels.contains_key(user_id)
     }
 
-    pub fn len(&self) -> u32 {
-        self.member_ids.len() as u32
+    pub fn len(&self) -> usize {
+        self.members_and_channels.len()
     }
 
-    pub fn member_ids(&self) -> &BTreeSet<UserId> {
-        &self.member_ids
+    pub fn iter_member_ids(&self) -> impl Iterator<Item = UserId> + '_ {
+        self.members_and_channels.keys().copied()
     }
 
     pub fn channels_for_member(&self, user_id: UserId) -> &[ChannelId] {
-        self.member_channel_links
+        self.members_and_channels
             .get(&user_id)
             .map(|v| v.as_slice())
             .unwrap_or_default()
@@ -643,7 +649,7 @@ impl CommunityMembers {
             }
         }
 
-        assert_eq!(member_ids, self.member_ids);
+        assert_eq!(member_ids, self.members_and_channels.keys().copied().collect::<BTreeSet<_>>());
         assert_eq!(owners, self.owners);
         assert_eq!(admins, self.admins);
         assert_eq!(lapsed, self.lapsed);
@@ -661,15 +667,13 @@ impl Members for CommunityMembers {
     }
 
     fn can_member_lapse(&self, user_id: &UserId) -> bool {
-        self.member_ids.contains(user_id) && !self.owners.contains(user_id) && !self.lapsed.contains(user_id)
+        self.members_and_channels.contains_key(user_id) && !self.owners.contains(user_id) && !self.lapsed.contains(user_id)
     }
 
     fn iter_members_who_can_lapse(&self) -> Box<dyn Iterator<Item = UserId> + '_> {
         Box::new(
-            self.member_ids
-                .iter()
-                .filter(|id| !self.owners.contains(id) && !self.lapsed.contains(id))
-                .copied(),
+            self.iter_member_ids()
+                .filter(|id| !self.owners.contains(id) && !self.lapsed.contains(id)),
         )
     }
 }
@@ -856,7 +860,7 @@ impl<'de> Visitor<'de> for MemberChannelLinksVisitor {
     }
 }
 
-fn deserialize_member_channel_links<'de, D: Deserializer<'de>>(d: D) -> Result<BTreeMap<UserId, Vec<ChannelId>>, D::Error> {
+fn deserialize_members_and_channels<'de, D: Deserializer<'de>>(d: D) -> Result<BTreeMap<UserId, Vec<ChannelId>>, D::Error> {
     d.deserialize_seq(MemberChannelLinksVisitor)
 }
 

backend/canisters/community/impl/src/model/members/proptests.rs

@@ -5,7 +5,7 @@ use ic_stable_structures::DefaultMemoryImpl;
 use proptest::collection::vec as pvec;
 use proptest::prelude::*;
 use proptest::prop_oneof;
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 use test_strategy::proptest;
 use types::{CommunityPermissions, CommunityRole, TimestampMillis, UserId, UserType};
 
@@ -88,10 +88,10 @@ fn execute_operation(members: &mut CommunityMembers, op: Operation, timestamp: T
             referred_by_index,
         } => {
             let referred_by = referred_by_index.and_then(|i| {
-                if members.member_ids.is_empty() {
+                if members.members_and_channels.is_empty() {
                     None
                 } else {
-                    Some(get(&members.member_ids, i))
+                    Some(get_from_map(&members.members_and_channels, i))
                 }
             });
             members.add(user_id, user_id.into(), UserType::User, referred_by, timestamp);
@@ -101,8 +101,8 @@ fn execute_operation(members: &mut CommunityMembers, op: Operation, timestamp: T
             user_index,
             role,
         } => {
-            let owner = get(&members.owners, owner_index);
-            let user_id = get(&members.member_ids, user_index);
+            let owner = get_from_set(&members.owners, owner_index);
+            let user_id = get_from_map(&members.members_and_channels, user_index);
             members.change_role(
                 owner,
                 user_id,
@@ -114,35 +114,35 @@ fn execute_operation(members: &mut CommunityMembers, op: Operation, timestamp: T
             );
         }
         Operation::SetDisplayName { user_index, value } => {
-            let user_id = get(&members.member_ids, user_index);
+            let user_id = get_from_map(&members.members_and_channels, user_index);
             members.set_display_name(user_id, value, timestamp);
         }
         Operation::Remove { user_index } => {
-            let user_id = get(&members.member_ids, user_index);
+            let user_id = get_from_map(&members.members_and_channels, user_index);
             if members.owners.len() != 1 || members.owners.first() != Some(&user_id) {
                 members.remove(user_id, None, timestamp);
             }
         }
         Operation::Block { user_index } => {
-            let user_id = get(&members.member_ids, user_index);
+            let user_id = get_from_map(&members.members_and_channels, user_index);
             if members.owners.len() != 1 || members.owners.first() != Some(&user_id) {
                 members.remove(user_id, None, timestamp);
                 members.block(user_id, timestamp);
             }
         }
         Operation::Unblock { user_index } => {
             if !members.blocked.is_empty() {
-                let user_id = get(&members.blocked, user_index);
+                let user_id = get_from_set(&members.blocked, user_index);
                 members.unblock(user_id, timestamp);
             }
         }
         Operation::Lapse { user_index } => {
-            let user_id = get(&members.member_ids, user_index);
+            let user_id = get_from_map(&members.members_and_channels, user_index);
             members.update_lapsed(user_id, true, timestamp);
         }
         Operation::Unlapse { user_index } => {
             if !members.lapsed.is_empty() {
-                let user_id = get(&members.lapsed, user_index);
+                let user_id = get_from_set(&members.lapsed, user_index);
                 members.update_lapsed(user_id, false, timestamp);
             }
         }
@@ -151,19 +151,24 @@ fn execute_operation(members: &mut CommunityMembers, op: Operation, timestamp: T
         }
         Operation::SetSuspended { user_index, suspended } => {
             if suspended {
-                let user_id = get(&members.member_ids, user_index);
+                let user_id = get_from_map(&members.members_and_channels, user_index);
                 members.set_suspended(user_id, true, timestamp);
             } else if !members.suspended.is_empty() {
-                let user_id = get(&members.suspended, user_index);
+                let user_id = get_from_set(&members.suspended, user_index);
                 members.set_suspended(user_id, false, timestamp);
             }
         }
     };
 }
 
-fn get(set: &BTreeSet<UserId>, index: usize) -> UserId {
+fn get_from_map<V>(map: &BTreeMap<UserId, V>, index: usize) -> UserId {
+    let index = index % map.len();
+    map.iter().nth(index).map(|(u, _)| *u).unwrap()
+}
+
+fn get_from_set(set: &BTreeSet<UserId>, index: usize) -> UserId {
     let index = index % set.len();
-    *set.iter().nth(index).unwrap()
+    set.iter().nth(index).copied().unwrap()
 }
 
 fn principal(index: usize) -> Principal {