Fix case where some thread messages were not updated in stable memory (…

…#6736)

backend/canisters/community/CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Consolidate `summary` and `c2c_summary` args ([#6723](https://github.com/open-chat-labs/open-chat/pull/6723))
+- Fix case where some thread messages were not updated in stable memory ([#6736](https://github.com/open-chat-labs/open-chat/pull/6736))
 
 ## [[2.0.1423](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1423-community)] - 2024-11-05
 

backend/canisters/group/CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Consolidate `summary` and `c2c_summary` args ([#6723](https://github.com/open-chat-labs/open-chat/pull/6723))
+- Fix case where some thread messages were not updated in stable memory ([#6736](https://github.com/open-chat-labs/open-chat/pull/6736))
 
 ## [[2.0.1424](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1424-group)] - 2024-11-05
 

backend/libraries/chat_events/src/chat_events.rs

@@ -13,7 +13,7 @@ use search::{Document, Query};
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
 use sha2::{Digest, Sha256};
-use std::cmp::max;
+use std::cmp::{max, min};
 use std::collections::hash_map::Entry::{Occupied, Vacant};
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::mem;
@@ -50,6 +50,8 @@ pub struct ChatEvents {
     search_index: SearchIndex,
     #[serde(default = "default_next_event_to_migrate_to_stable_memory")]
     next_event_to_migrate_to_stable_memory: Option<EventContext>,
+    #[serde(default)]
+    thread_messages_to_update_in_stable_memory: Vec<MessageIndex>,
 }
 
 fn default_next_event_to_migrate_to_stable_memory() -> Option<EventContext> {
@@ -81,11 +83,29 @@ impl ChatEvents {
     }
 
     pub fn migrate_next_batch_of_events_to_stable_storage(&mut self) -> bool {
+        let mut total_count = 0;
+        while !self.thread_messages_to_update_in_stable_memory.is_empty() {
+            if ic_cdk::api::instruction_counter() > 1_000_000_000 {
+                return false;
+            }
+
+            let batch: Vec<_> = self
+                .thread_messages_to_update_in_stable_memory
+                .drain(..min(100, self.thread_messages_to_update_in_stable_memory.len()))
+                .collect();
+
+            let count = batch.len();
+            for message_index in batch {
+                self.update_event_in_stable_memory(message_index.into());
+            }
+            info!(chat = ?self.chat, count, "Updated threads in stable memory");
+            total_count += count;
+        }
+
         if self.next_event_to_migrate_to_stable_memory.is_none() {
             return true;
         };
 
-        let mut total_count = 0;
         while ic_cdk::api::instruction_counter() < 1_000_000_000 {
             let EventContext {
                 thread_root_message_index: next_thread_root_message_index,
@@ -154,6 +174,7 @@ impl ChatEvents {
             anonymized_id: hex::encode(anonymized_id.to_be_bytes()),
             search_index: SearchIndex::default(),
             next_event_to_migrate_to_stable_memory: None,
+            thread_messages_to_update_in_stable_memory: Vec::new(),
         };
 
         events.push_event(None, ChatEventInternal::DirectChatCreated(DirectChatCreated {}), 0, now);
@@ -185,6 +206,7 @@ impl ChatEvents {
             anonymized_id: hex::encode(anonymized_id.to_be_bytes()),
             search_index: SearchIndex::default(),
             next_event_to_migrate_to_stable_memory: None,
+            thread_messages_to_update_in_stable_memory: Vec::new(),
         };
 
         events.push_event(