xmtp_db/encrypted_store/
refresh_state.rs

1use std::collections::HashMap;
2
3use diesel::{
4    backend::Backend,
5    deserialize::{self, FromSql, FromSqlRow},
6    expression::AsExpression,
7    prelude::*,
8    serialize::{self, IsNull, Output, ToSql},
9    sql_types::Integer,
10};
11use xmtp_proto::types::Cursor;
12
13use super::{ConnectionExt, Sqlite, db_connection::DbConnection, schema::refresh_state};
14use crate::{StorageError, StoreOrIgnore, impl_store_or_ignore};
15
16#[repr(i32)]
17#[derive(Debug, Clone, Copy, PartialEq, Eq, AsExpression, Hash, FromSqlRow)]
18#[diesel(sql_type = Integer)]
19pub enum EntityKind {
20    Welcome = 1,
21    ApplicationMessage = 2,       // All group envelopes, including commits
22    CommitLogUpload = 3, // Rowid of the last local entry we uploaded to the remote commit log
23    CommitLogDownload = 4, // Server log sequence id of last remote entry we downloaded from the remote commit log
24    CommitLogForkCheckLocal = 5, // Last rowid verified in local commit log
25    CommitLogForkCheckRemote = 6, // Last rowid verified in remote commit log
26    Identity = 8,
27    DeliveryAllocator = 9,
28    Delivery = 10,
29}
30
31pub trait HasEntityKind {
32    fn entity_kind(&self) -> EntityKind;
33}
34
35impl HasEntityKind for xmtp_proto::types::GroupMessage {
36    fn entity_kind(&self) -> EntityKind {
37        EntityKind::ApplicationMessage
38    }
39}
40
41impl HasEntityKind for xmtp_proto::types::WelcomeMessage {
42    fn entity_kind(&self) -> EntityKind {
43        EntityKind::Welcome
44    }
45}
46
47impl std::fmt::Display for EntityKind {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        use EntityKind::*;
50        match self {
51            Welcome => write!(f, "welcome"),
52            ApplicationMessage => write!(f, "group"),
53            CommitLogUpload => write!(f, "commit_log_upload"),
54            CommitLogDownload => write!(f, "commit_log_download"),
55            CommitLogForkCheckLocal => write!(f, "commit_log_fork_check_local"),
56            CommitLogForkCheckRemote => write!(f, "commit_log_fork_check_remote"),
57            Identity => write!(f, "identity"),
58            DeliveryAllocator => write!(f, "delivery_allocator"),
59            Delivery => write!(f, "delivery"),
60        }
61    }
62}
63
64impl ToSql<Integer, Sqlite> for EntityKind
65where
66    i32: ToSql<Integer, Sqlite>,
67{
68    fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Sqlite>) -> serialize::Result {
69        out.set_value(*self as i32);
70        Ok(IsNull::No)
71    }
72}
73
74impl FromSql<Integer, Sqlite> for EntityKind
75where
76    i32: FromSql<Integer, Sqlite>,
77{
78    fn from_sql(bytes: <Sqlite as Backend>::RawValue<'_>) -> deserialize::Result<Self> {
79        match i32::from_sql(bytes)? {
80            1 => Ok(EntityKind::Welcome),
81            2 => Ok(EntityKind::ApplicationMessage),
82            3 => Ok(EntityKind::CommitLogUpload),
83            4 => Ok(EntityKind::CommitLogDownload),
84            5 => Ok(EntityKind::CommitLogForkCheckLocal),
85            6 => Ok(EntityKind::CommitLogForkCheckRemote),
86            8 => Ok(EntityKind::Identity),
87            9 => Ok(EntityKind::DeliveryAllocator),
88            10 => Ok(EntityKind::Delivery),
89            x => Err(format!("Unrecognized variant {}", x).into()),
90        }
91    }
92}
93
94#[derive(Insertable, Identifiable, Queryable, Selectable, Debug, Clone)]
95#[diesel(table_name = refresh_state)]
96#[diesel(primary_key(entity_id, entity_kind))]
97pub struct RefreshState {
98    pub entity_id: Vec<u8>,
99    pub entity_kind: EntityKind,
100    /// Network kinds store P here; local delivery kinds store their separate D position.
101    pub sequence_id: i64,
102    /// F for ordered network admission. None is not proof of a received network prefix.
103    pub received_sequence_id: Option<i64>,
104}
105
106impl_store_or_ignore!(RefreshState, refresh_state);
107
108pub trait QueryRefreshState {
109    fn get_refresh_state<Id: AsRef<[u8]>>(
110        &self,
111        entity_id: Id,
112        entity_kind: EntityKind,
113    ) -> Result<Option<RefreshState>, StorageError>;
114
115    /// Read one ledger position. Create a zero position when it is absent.
116    fn get_last_cursor<Id: AsRef<[u8]>>(
117        &self,
118        id: Id,
119        entity_kind: EntityKind,
120    ) -> Result<Cursor, StorageError>;
121
122    /// Return the minimum position across the requested kinds for each stored id.
123    /// An absent kind has position zero. Ids with no rows are absent from the map.
124    fn get_last_cursor_for_ids<Id: AsRef<[u8]>>(
125        &self,
126        ids: &[Id],
127        entities: &[EntityKind],
128    ) -> Result<HashMap<Vec<u8>, Cursor>, StorageError>;
129
130    /// Advance a ledger position only when the new position is greater.
131    fn update_cursor<Id: AsRef<[u8]>>(
132        &self,
133        entity_id: Id,
134        entity_kind: EntityKind,
135        cursor: Cursor,
136    ) -> Result<bool, StorageError>;
137
138    fn latest_cursor_for_id<Id: AsRef<[u8]>>(
139        &self,
140        entity_id: Id,
141        entities: &[EntityKind],
142    ) -> Result<Cursor, StorageError> {
143        Ok(self
144            .get_last_cursor_for_ids(&[entity_id.as_ref()], entities)?
145            .remove(entity_id.as_ref())
146            .unwrap_or_default())
147    }
148
149    fn get_remote_log_cursors(
150        &self,
151        conversation_ids: &[&[u8]],
152    ) -> Result<HashMap<Vec<u8>, Cursor>, StorageError> {
153        conversation_ids
154            .iter()
155            .map(|id| {
156                self.get_last_cursor(id, EntityKind::CommitLogDownload)
157                    .map(|cursor| (id.to_vec(), cursor))
158            })
159            .collect()
160    }
161}
162
163impl<T: QueryRefreshState> QueryRefreshState for &T {
164    fn get_refresh_state<Id: AsRef<[u8]>>(
165        &self,
166        entity_id: Id,
167        entity_kind: EntityKind,
168    ) -> Result<Option<RefreshState>, StorageError> {
169        (**self).get_refresh_state(entity_id, entity_kind)
170    }
171
172    fn get_last_cursor<Id: AsRef<[u8]>>(
173        &self,
174        id: Id,
175        entity_kind: EntityKind,
176    ) -> Result<Cursor, StorageError> {
177        (**self).get_last_cursor(id, entity_kind)
178    }
179
180    fn get_last_cursor_for_ids<Id: AsRef<[u8]>>(
181        &self,
182        ids: &[Id],
183        entities: &[EntityKind],
184    ) -> Result<HashMap<Vec<u8>, Cursor>, StorageError> {
185        (**self).get_last_cursor_for_ids(ids, entities)
186    }
187
188    fn update_cursor<Id: AsRef<[u8]>>(
189        &self,
190        entity_id: Id,
191        entity_kind: EntityKind,
192        cursor: Cursor,
193    ) -> Result<bool, StorageError> {
194        (**self).update_cursor(entity_id, entity_kind, cursor)
195    }
196}
197
198impl<C: ConnectionExt> QueryRefreshState for DbConnection<C> {
199    #[xmtp_common::db_span]
200    fn get_refresh_state<Id: AsRef<[u8]>>(
201        &self,
202        entity_id: Id,
203        entity_kind: EntityKind,
204    ) -> Result<Option<RefreshState>, StorageError> {
205        Ok(self.raw_query(|conn| {
206            refresh_state::table
207                .find((entity_id.as_ref(), entity_kind))
208                .select(RefreshState::as_select())
209                .first(conn)
210                .optional()
211        })?)
212    }
213
214    #[xmtp_common::db_span]
215    fn get_last_cursor<Id: AsRef<[u8]>>(
216        &self,
217        id: Id,
218        entity_kind: EntityKind,
219    ) -> Result<Cursor, StorageError> {
220        RefreshState {
221            entity_id: id.as_ref().to_vec(),
222            entity_kind,
223            sequence_id: 0,
224            received_sequence_id: None,
225        }
226        .store_or_ignore(self)?;
227        Ok(Cursor(
228            self.get_refresh_state(id, entity_kind)?
229                .ok_or(StorageError::DbDeserialize)?
230                .sequence_id as u64,
231        ))
232    }
233
234    #[xmtp_common::db_span]
235    fn get_last_cursor_for_ids<Id: AsRef<[u8]>>(
236        &self,
237        ids: &[Id],
238        entities: &[EntityKind],
239    ) -> Result<HashMap<Vec<u8>, Cursor>, StorageError> {
240        use super::schema::refresh_state::dsl;
241        use diesel::dsl::{count, min};
242        use std::collections::HashSet;
243
244        // Leave room for the kind filters under SQLite's bind parameter limit.
245        const IDS_PER_QUERY: usize = 900;
246        let entities: HashSet<_> = entities.iter().copied().collect();
247        if entities.is_empty() {
248            return Ok(HashMap::new());
249        }
250        Ok(self.raw_query(|conn| {
251            let mut result = HashMap::new();
252            for chunk in ids.chunks(IDS_PER_QUERY) {
253                let ids: Vec<_> = chunk.iter().map(AsRef::as_ref).collect();
254                let rows = dsl::refresh_state
255                    .filter(dsl::entity_kind.eq_any(&entities))
256                    .filter(dsl::entity_id.eq_any(ids))
257                    .group_by(dsl::entity_id)
258                    .select((
259                        dsl::entity_id,
260                        min(dsl::sequence_id),
261                        count(dsl::entity_kind),
262                    ))
263                    .load::<(Vec<u8>, Option<i64>, i64)>(conn)?;
264                for (id, sequence, kinds) in rows {
265                    let sequence = if kinds as usize == entities.len() {
266                        sequence.unwrap_or_default() as u64
267                    } else {
268                        0
269                    };
270                    result.insert(id, Cursor(sequence));
271                }
272            }
273            Ok(result)
274        })?)
275    }
276
277    #[xmtp_common::db_span]
278    fn update_cursor<Id: AsRef<[u8]>>(
279        &self,
280        entity_id: Id,
281        entity_kind: EntityKind,
282        cursor: Cursor,
283    ) -> Result<bool, StorageError> {
284        use super::schema::refresh_state::dsl;
285        use diesel::{query_dsl::methods::FilterDsl, upsert::excluded};
286        let state = RefreshState {
287            entity_id: entity_id.as_ref().to_vec(),
288            entity_kind,
289            sequence_id: i64::try_from(cursor.0).map_err(|_| StorageError::DbSerialize)?,
290            received_sequence_id: None,
291        };
292        Ok(self.raw_query(|conn| {
293            diesel::insert_into(dsl::refresh_state)
294                .values(&state)
295                .on_conflict((dsl::entity_id, dsl::entity_kind))
296                .do_update()
297                .set(dsl::sequence_id.eq(excluded(dsl::sequence_id)))
298                .filter(dsl::sequence_id.lt(excluded(dsl::sequence_id)))
299                .execute(conn)
300        })? > 0)
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::test_utils::with_connection;
308
309    #[xmtp_common::test(unwrap_try = true)]
310    fn cursor_defaults_and_advances_only_in_order() {
311        with_connection(|conn| {
312            let id = [1, 2, 3];
313            let kind = EntityKind::ApplicationMessage;
314            assert!(conn.get_refresh_state(id, kind).unwrap().is_none());
315            assert_eq!(conn.get_last_cursor(id, kind).unwrap(), Cursor(0));
316            assert!(conn.get_refresh_state(id, kind).unwrap().is_some());
317            assert!(conn.update_cursor(id, kind, Cursor(123)).unwrap());
318            assert!(!conn.update_cursor(id, kind, Cursor(122)).unwrap());
319            assert!(!conn.update_cursor(id, kind, Cursor(123)).unwrap());
320            assert!(conn.update_cursor(id, kind, Cursor(124)).unwrap());
321            assert_eq!(conn.get_last_cursor(id, kind).unwrap(), Cursor(124));
322        });
323    }
324
325    #[rstest::rstest]
326    #[case(Some(500), Some(250), 250)]
327    #[case(Some(100), Some(200), 100)]
328    #[case(Some(500), None, 0)]
329    #[case(None, Some(250), 0)]
330    #[case(None, None, 0)]
331    #[xmtp_common::test(unwrap_try = true)]
332    async fn cursor_meets_requested_kinds(
333        #[case] application: Option<u64>,
334        #[case] identity: Option<u64>,
335        #[case] expected: u64,
336    ) {
337        with_connection(|conn| {
338            let id = [1, 2, 3];
339            for (kind, value) in [
340                (EntityKind::ApplicationMessage, application),
341                (EntityKind::Identity, identity),
342            ] {
343                if let Some(value) = value {
344                    conn.update_cursor(id, kind, Cursor(value)).unwrap();
345                }
346            }
347            conn.update_cursor(id, EntityKind::Welcome, Cursor(999))
348                .unwrap();
349            let kinds = [EntityKind::ApplicationMessage, EntityKind::Identity];
350            assert_eq!(
351                conn.latest_cursor_for_id(id, &kinds).unwrap(),
352                Cursor(expected)
353            );
354            assert_eq!(
355                conn.latest_cursor_for_id(id, &[EntityKind::Welcome])
356                    .unwrap(),
357                Cursor(999)
358            );
359        });
360    }
361
362    #[rstest::rstest]
363    #[case(0)]
364    #[case(1)]
365    #[case(900)]
366    #[case(1000)]
367    #[case(2000)]
368    #[xmtp_common::test(unwrap_try = true)]
369    async fn cursor_queries_batch_ids(#[case] count: u64) {
370        with_connection(|conn| {
371            let ids: Vec<_> = (0..count).map(u64::to_be_bytes).collect();
372            for (index, id) in ids.iter().enumerate() {
373                conn.update_cursor(id, EntityKind::ApplicationMessage, Cursor(index as u64))
374                    .unwrap();
375            }
376            let found = conn
377                .get_last_cursor_for_ids(&ids, &[EntityKind::ApplicationMessage])
378                .unwrap();
379            assert_eq!(found.len(), ids.len());
380            for (index, id) in ids.iter().enumerate() {
381                assert_eq!(found.get(id.as_slice()), Some(&Cursor(index as u64)));
382            }
383            let missing = conn
384                .get_last_cursor_for_ids(&[[255; 8]], &[EntityKind::ApplicationMessage])
385                .unwrap();
386            assert!(missing.is_empty());
387        });
388    }
389}
xmtp_db/encrypted_store/refresh_state.rs

xmtp_db/encrypted_store/
refresh_state.rs
