Make StableMemoryMap use strongly typed keys (#6937)

backend/canisters/community/CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Implement `MembersStableStorage` which stores members in stable memory ([#6931](https://github.com/open-chat-labs/open-chat/pull/6931))
 - Migrate chat members to stable memory using timer job ([#6933](https://github.com/open-chat-labs/open-chat/pull/6933))
 - Write group members to stable memory when importing group into community ([#6935](https://github.com/open-chat-labs/open-chat/pull/6935))
+- Make `StableMemoryMap` use strongly typed keys ([#6937](https://github.com/open-chat-labs/open-chat/pull/6937))
 
 ## [[2.0.1479](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1479-community)] - 2024-11-28
 

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

@@ -7,7 +7,7 @@ use activity_notification_state::ActivityNotificationState;
 use candid::Principal;
 use canister_state_macros::canister_state;
 use canister_timer_jobs::TimerJobs;
-use chat_events::{ChannelThreadKeyPrefix, ChatMetricsInternal, KeyPrefix};
+use chat_events::ChatMetricsInternal;
 use community_canister::EventsResponse;
 use constants::MINUTE_IN_MS;
 use event_store_producer::{EventStoreClient, EventStoreClientBuilder, EventStoreClientInfo};
@@ -28,6 +28,7 @@ use rand::rngs::StdRng;
 use rand::RngCore;
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
+use stable_memory_map::{ChatEventKeyPrefix, KeyPrefix};
 use std::cell::RefCell;
 use std::ops::Deref;
 use std::time::Duration;
@@ -241,9 +242,12 @@ impl RuntimeState {
             }
             final_prize_payments.extend(result.final_prize_payments);
             for thread in result.threads {
-                self.data.stable_memory_keys_to_garbage_collect.push(
-                    KeyPrefix::ChannelThread(ChannelThreadKeyPrefix::new(channel.id, thread.root_message_index)).to_vec(),
-                );
+                self.data
+                    .stable_memory_keys_to_garbage_collect
+                    .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_channel(
+                        channel.id,
+                        Some(thread.root_message_index),
+                    )));
             }
         }
         jobs::garbage_collect_stable_memory::start_job_if_required(self);
@@ -366,7 +370,7 @@ struct Data {
     user_cache: UserCache,
     user_event_sync_queue: GroupedTimerJobQueue<UserEventBatch>,
     #[serde(default)]
-    stable_memory_keys_to_garbage_collect: Vec<Vec<u8>>,
+    stable_memory_keys_to_garbage_collect: Vec<KeyPrefix>,
     #[serde(default)]
     members_migrated_to_stable_memory: bool,
 }

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

@@ -11,7 +11,7 @@ use ic_cdk::post_upgrade;
 use instruction_counts_log::InstructionCountFunctionId;
 use stable_memory::get_reader;
 use tracing::info;
-use types::MultiUserChat;
+use types::{Chat, MultiUserChat};
 
 #[post_upgrade]
 #[trace]
@@ -26,6 +26,7 @@ fn post_upgrade(args: Args) {
 
     let community_id = ic_cdk::id().into();
     for channel in data.channels.iter_mut() {
+        channel.chat.events.set_chat(Chat::Channel(community_id, channel.id));
         channel.chat.members.set_member_default_timestamps();
         channel
             .chat

backend/canisters/community/impl/src/updates/delete_channel.rs

@@ -4,8 +4,8 @@ use crate::{
 };
 use canister_api_macros::update;
 use canister_tracing_macros::trace;
-use chat_events::{ChannelKeyPrefix, ChannelThreadKeyPrefix, KeyPrefix};
 use community_canister::delete_channel::{Response::*, *};
+use stable_memory_map::{ChatEventKeyPrefix, KeyPrefix, MemberKeyPrefix};
 use types::{ChannelDeleted, ChannelId};
 
 #[update(candid = true, msgpack = true)]
@@ -54,19 +54,22 @@ fn delete_channel_impl(channel_id: ChannelId, state: &mut RuntimeState) -> Respo
     state
         .data
         .stable_memory_keys_to_garbage_collect
-        .push(KeyPrefix::Channel(ChannelKeyPrefix::new(channel_id)).to_vec());
+        .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_channel(channel_id, None)));
 
     for message_index in channel.chat.events.thread_keys() {
         state
             .data
             .stable_memory_keys_to_garbage_collect
-            .push(KeyPrefix::ChannelThread(ChannelThreadKeyPrefix::new(channel_id, message_index)).to_vec());
+            .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_channel(
+                channel_id,
+                Some(message_index),
+            )));
     }
 
     state
         .data
         .stable_memory_keys_to_garbage_collect
-        .push(group_chat_core::MembersKeyPrefix::Channel(channel_id.as_u32()).to_vec());
+        .push(KeyPrefix::from(MemberKeyPrefix::new_from_channel(channel_id)));
 
     crate::jobs::garbage_collect_stable_memory::start_job_if_required(state);
 

backend/canisters/group/CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Implement `MembersStableStorage` which stores members in stable memory ([#6931](https://github.com/open-chat-labs/open-chat/pull/6931))
 - Migrate chat members to stable memory using timer job ([#6933](https://github.com/open-chat-labs/open-chat/pull/6933))
 - Export members from stable memory when importing group into community ([#6935](https://github.com/open-chat-labs/open-chat/pull/6935))
+- Make `StableMemoryMap` use strongly typed keys ([#6937](https://github.com/open-chat-labs/open-chat/pull/6937))
 
 ## [[2.0.1480](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1480-group)] - 2024-11-28
 

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

@@ -7,7 +7,7 @@ use activity_notification_state::ActivityNotificationState;
 use candid::Principal;
 use canister_state_macros::canister_state;
 use canister_timer_jobs::TimerJobs;
-use chat_events::{GroupChatThreadKeyPrefix, KeyPrefix, Reader};
+use chat_events::Reader;
 use constants::{DAY_IN_MS, HOUR_IN_MS, MINUTE_IN_MS, OPENCHAT_BOT_USER_ID};
 use event_store_producer::{EventStoreClient, EventStoreClientBuilder, EventStoreClientInfo};
 use event_store_producer_cdk_runtime::CdkRuntime;
@@ -24,6 +24,7 @@ use msgpack::serialize_then_unwrap;
 use notifications_canister::c2c_push_notification;
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
+use stable_memory_map::{ChatEventKeyPrefix, KeyPrefix};
 use std::cell::RefCell;
 use std::collections::hash_map::Entry::{Occupied, Vacant};
 use std::collections::{HashMap, HashSet};
@@ -361,7 +362,9 @@ impl RuntimeState {
         for thread in result.threads {
             self.data
                 .stable_memory_keys_to_garbage_collect
-                .push(KeyPrefix::GroupChatThread(GroupChatThreadKeyPrefix::new(thread.root_message_index)).to_vec());
+                .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_group_chat(Some(
+                    thread.root_message_index,
+                ))));
         }
         jobs::garbage_collect_stable_memory::start_job_if_required(self);
     }
@@ -474,7 +477,7 @@ struct Data {
     user_event_sync_queue: GroupedTimerJobQueue<UserEventBatch>,
     #[serde(default)]
     members_migrated_to_stable_memory: bool,
-    stable_memory_keys_to_garbage_collect: Vec<Vec<u8>>,
+    stable_memory_keys_to_garbage_collect: Vec<KeyPrefix>,
 }
 
 fn init_instruction_counts_log() -> InstructionCountsLog {

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

@@ -10,7 +10,7 @@ use ic_cdk::post_upgrade;
 use instruction_counts_log::InstructionCountFunctionId;
 use stable_memory::get_reader;
 use tracing::info;
-use types::MultiUserChat;
+use types::{Chat, MultiUserChat};
 
 #[post_upgrade]
 #[trace]
@@ -23,8 +23,10 @@ fn post_upgrade(args: Args) {
     let (mut data, errors, logs, traces): (Data, Vec<LogEntry>, Vec<LogEntry>, Vec<LogEntry>) =
         msgpack::deserialize(reader).unwrap();
 
+    let chat_id = ic_cdk::id().into();
+    data.chat.events.set_chat(Chat::Group(chat_id));
     data.chat.members.set_member_default_timestamps();
-    data.chat.members.set_chat(MultiUserChat::Group(ic_cdk::id().into()));
+    data.chat.members.set_chat(MultiUserChat::Group(chat_id));
 
     canister_logger::init_with_logs(data.test_mode, errors, logs, traces);
 

backend/canisters/user/CHANGELOG.md

@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Extract stable memory map so it can store additional datasets ([#6876](https://github.com/open-chat-labs/open-chat/pull/6876))
 - Make `ChannelId` comparisons use their 32bit representation ([#6885](https://github.com/open-chat-labs/open-chat/pull/6885))
 - Remove chat event updates after 31 days ([#6916](https://github.com/open-chat-labs/open-chat/pull/6916))
+- Make `StableMemoryMap` use strongly typed keys ([#6937](https://github.com/open-chat-labs/open-chat/pull/6937))
 
 ### Removed
 

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

@@ -25,6 +25,7 @@ use model::streak::Streak;
 use notifications_canister::c2c_push_notification;
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
+use stable_memory_map::KeyPrefix;
 use std::cell::RefCell;
 use std::collections::HashSet;
 use std::ops::Deref;
@@ -255,7 +256,7 @@ struct Data {
     pub referred_by: Option<UserId>,
     pub referrals: Referrals,
     pub message_activity_events: MessageActivityEvents,
-    pub stable_memory_keys_to_garbage_collect: Vec<Vec<u8>>,
+    pub stable_memory_keys_to_garbage_collect: Vec<KeyPrefix>,
 }
 
 impl Data {

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

@@ -1,12 +1,13 @@
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::{get_stable_memory_map_memory, get_upgrades_memory};
 use crate::{mutate_state, Data};
+use candid::Principal;
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
 use stable_memory::get_reader;
 use tracing::info;
-use types::CanisterId;
+use types::{CanisterId, Chat};
 use user_canister::post_upgrade::Args;
 
 #[post_upgrade]
@@ -38,6 +39,7 @@ fn post_upgrade(args: Args) {
     mutate_state(|state| {
         let now = state.env.now();
         for chat in state.data.direct_chats.iter_mut() {
+            chat.events.set_chat(Chat::Direct(Principal::from(chat.them).into()));
             chat.events.remove_spurious_video_call_in_progress(now);
 
             let count_removed = chat.events.prune_updated_events(now);

backend/canisters/user/impl/src/updates/delete_direct_chat.rs

@@ -2,7 +2,7 @@ use crate::guards::caller_is_owner;
 use crate::{mutate_state, run_regular_jobs, RuntimeState};
 use canister_api_macros::update;
 use canister_tracing_macros::trace;
-use chat_events::{DirectChatKeyPrefix, DirectChatThreadKeyPrefix, KeyPrefix};
+use stable_memory_map::{ChatEventKeyPrefix, KeyPrefix};
 use user_canister::delete_direct_chat::{Response::*, *};
 
 #[update(guard = "caller_is_owner", msgpack = true)]
@@ -19,16 +19,22 @@ fn delete_direct_chat_impl(args: Args, state: &mut RuntimeState) -> Response {
         if args.block_user {
             state.data.block_user(args.user_id, now);
         }
+
         state
             .data
             .stable_memory_keys_to_garbage_collect
-            .push(KeyPrefix::DirectChat(DirectChatKeyPrefix::new(args.user_id)).to_vec());
+            .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_direct_chat(args.user_id, None)));
+
         for message_index in chat.events.thread_keys() {
             state
                 .data
                 .stable_memory_keys_to_garbage_collect
-                .push(KeyPrefix::DirectChatThread(DirectChatThreadKeyPrefix::new(args.user_id, message_index)).to_vec());
+                .push(KeyPrefix::from(ChatEventKeyPrefix::new_from_direct_chat(
+                    args.user_id,
+                    Some(message_index),
+                )));
         }
+
         crate::jobs::garbage_collect_stable_memory::start_job_if_required(state);
         Success
     } else {

backend/libraries/chat_events/src/lib.rs

@@ -17,4 +17,3 @@ pub use crate::chat_events_list::*;
 pub use crate::events_map::*;
 pub use crate::message_content_internal::*;
 pub use crate::metrics::*;
-pub use crate::stable_memory::key::*;