xmtp_mls/utils/
cleanup_duplicate_updates.rs

1use std::collections::HashSet;
2use std::hash::{DefaultHasher, Hash, Hasher};
3
4use tracing::info;
5use xmtp_db::diesel::prelude::*;
6use xmtp_db::user_preferences::StoredUserPreferences;
7use xmtp_db::{ConnectionExt, DbConnection};
8use xmtp_db::{
9    group::{ConversationType, GroupQueryArgs, QueryGroup},
10    group_message::{ContentType, MsgQueryArgs},
11    prelude::QueryGroupMessage,
12};
13
14use crate::groups::mls_sync::GroupMessageProcessingError;
15use crate::messages::decoded_message::MessageBody;
16use crate::messages::enrichment::enrich_messages;
17
18const BATCH_SIZE: i64 = 100;
19
20pub async fn perform<C>(db: DbConnection<C>)
21where
22    C: ConnectionExt,
23{
24    if let Err(err) = perform_inner(db).await {
25        tracing::error!("Duplicate cleanup task failed: {err:?}");
26    }
27}
28
29async fn perform_inner<C>(db: DbConnection<C>) -> Result<(), GroupMessageProcessingError>
30where
31    C: ConnectionExt,
32{
33    let prefs = StoredUserPreferences::load(&db)?;
34    if prefs.dm_group_updates_migrated {
35        info!("DM group updates migration has already been performed. Skipping.");
36        return Ok(());
37    }
38
39    let mut group_offset = 0;
40    let mut groups;
41    loop {
42        groups = db.find_groups_by_id_paged(
43            GroupQueryArgs {
44                conversation_type: Some(ConversationType::Dm),
45                limit: Some(BATCH_SIZE),
46                ..Default::default()
47            },
48            group_offset,
49        )?;
50
51        if groups.is_empty() {
52            break;
53        }
54
55        for group in groups {
56            let mut sent_after_ns = None;
57            let mut msgs;
58            let mut originals: HashSet<u64> = HashSet::default();
59
60            loop {
61                msgs = db.get_group_messages(
62                    &group.id,
63                    &MsgQueryArgs {
64                        content_types: Some(vec![ContentType::GroupUpdated]),
65                        sent_after_ns,
66                        limit: Some(BATCH_SIZE),
67                        ..Default::default()
68                    },
69                )?;
70
71                {
72                    let Some(msg) = msgs.last() else {
73                        break;
74                    };
75                    sent_after_ns = Some(msg.sent_at_ns);
76                }
77
78                let msgs = enrich_messages(&db, &group.id, msgs)?;
79
80                for msg in msgs {
81                    let MessageBody::GroupUpdated(update) = msg.content else {
82                        continue;
83                    };
84
85                    let mut hasher = DefaultHasher::new();
86                    update.hash(&mut hasher);
87                    if originals.insert(hasher.finish()) {
88                        continue;
89                    }
90
91                    db.raw_query(|conn| {
92                        xmtp_db::diesel::delete(xmtp_db::schema::group_messages::table)
93                            .filter(xmtp_db::schema::group_messages::id.eq(&msg.metadata.id))
94                            .execute(conn)
95                    })?;
96
97                    tokio::task::yield_now().await;
98                }
99            }
100
101            tokio::task::yield_now().await;
102        }
103
104        group_offset += BATCH_SIZE;
105    }
106
107    db.raw_query(|conn| {
108        xmtp_db::diesel::update(xmtp_db::schema::user_preferences::table)
109            .set(xmtp_db::schema::user_preferences::dm_group_updates_migrated.eq(true))
110            .execute(conn)
111    })?;
112
113    Ok(())
114}
115
116#[cfg(test)]
117mod tests {
118    use crate::tester;
119    use prost::Message;
120
121    use super::*;
122    use xmtp_common::{rand_vec, time::now_ns};
123    use xmtp_content_types::{ContentCodec, encryption::sha256, group_updated::GroupUpdatedCodec};
124    use xmtp_db::{
125        Store,
126        group_message::{DeliveryStatus, GroupMessageKind, StoredGroupMessage},
127    };
128    use xmtp_proto::xmtp::mls::message_contents::{
129        GroupUpdated,
130        group_updated::{Inbox, MetadataFieldChange},
131    };
132
133    #[xmtp_common::test(unwrap_try = true)]
134    async fn test_cleanup_works_as_expected() {
135        tester!(alix);
136        tester!(bo);
137        let mut sequence_id = 0;
138
139        let mut gen_update_msg = |group_id, payload| {
140            let mut encoded_payload = Vec::new();
141            GroupUpdatedCodec::encode(payload)?.encode(&mut encoded_payload)?;
142            sequence_id += 1;
143
144            StoredGroupMessage {
145                id: sha256(&rand_vec::<12>()),
146                group_id,
147                decrypted_message_bytes: encoded_payload,
148                sent_at_ns: now_ns(),
149                kind: GroupMessageKind::MembershipChange,
150                sender_installation_id: vec![1, 2, 3],
151                sender_inbox_id: "123".to_string(),
152                delivery_status: DeliveryStatus::Published,
153                content_type: ContentType::GroupUpdated,
154                version_major: 0,
155                version_minor: 0,
156                authority_id: "unknown".to_string(),
157                reference_id: None,
158                sequence_id,
159                envelope_hash: None,
160                expiry_ns: None,
161                expire_at_ns: None,
162                inserted_at_ns: 0,
163                should_push: true,
164                idempotency_key: String::new(),
165            }
166        };
167
168        let (dm, _) = alix.test_talk_in_dm_with(&bo).await?;
169        dm.sync().await?;
170        // Let startup cleanup finish before preparing this test's unmigrated state.
171        xmtp_common::wait_for_eq(
172            || async {
173                StoredUserPreferences::load(alix.db())
174                    .unwrap()
175                    .dm_group_updates_migrated
176            },
177            true,
178        )
179        .await?;
180        alix.db().raw_query(|conn| {
181            xmtp_db::diesel::update(xmtp_db::schema::user_preferences::table)
182                .set(xmtp_db::schema::user_preferences::dm_group_updates_migrated.eq(false))
183                .execute(conn)
184        })?;
185        let old_updates = dm.find_messages_v2(&MsgQueryArgs {
186            content_types: Some(vec![ContentType::GroupUpdated]),
187            ..Default::default()
188        })?;
189
190        // Insert some duplicate group_updated messages
191        let payload1 = GroupUpdated {
192            added_inboxes: vec![Inbox {
193                inbox_id: "123".to_string(),
194            }],
195            ..Default::default()
196        };
197        let payload2 = GroupUpdated {
198            metadata_field_changes: vec![MetadataFieldChange {
199                field_name: "expiration_setting".to_string(),
200                old_value: None,
201                new_value: Some("123".to_string()),
202            }],
203
204            ..Default::default()
205        };
206
207        let mut duplicates = vec![];
208
209        for i in 0..3 {
210            let msg1 = gen_update_msg(dm.group_id, payload1.clone());
211            msg1.store(&alix.db())?;
212            let msg2 = gen_update_msg(dm.group_id, payload2.clone());
213            msg2.store(&alix.db())?;
214
215            if i > 0 {
216                duplicates.push(msg1.id);
217                duplicates.push(msg2.id);
218            }
219        }
220
221        perform_inner(alix.db()).await?;
222
223        let msgs = dm.find_messages_v2(&MsgQueryArgs {
224            content_types: Some(vec![ContentType::GroupUpdated]),
225            ..Default::default()
226        })?;
227
228        for msg in &msgs {
229            assert!(
230                !duplicates.contains(&msg.metadata.id),
231                "A duplicate has remained in the db {:?}",
232                msg.metadata.id
233            );
234        }
235        // Make sure the old update didn't get deleted. The +2 is for the 2 dummy updates.
236        assert_eq!(msgs.len(), old_updates.len() + 2);
237
238        // Let's insert another duplicate and make sure it stays this time.
239        // We don't want the perform to run more than once.
240        let msg = gen_update_msg(dm.group_id, payload1.clone());
241        msg.store(&alix.db())?;
242        perform_inner(alix.db()).await?;
243
244        // The duplicate should remain because perform will only clean up once.
245        let msgs = dm.find_messages_v2(&MsgQueryArgs {
246            content_types: Some(vec![ContentType::GroupUpdated]),
247            ..Default::default()
248        })?;
249        assert!(msgs.iter().any(|m| m.metadata.id == msg.id));
250    }
251}
xmtp_mls/utils/cleanup_duplicate_updates.rs

xmtp_mls/utils/
cleanup_duplicate_updates.rs
