xmtp_db/encrypted_store/
delivery.rs

1//! Database-local message order and fenced default-consumer progress.
2
3use diesel::{
4    prelude::*,
5    sql_types::{BigInt, Integer, Nullable},
6};
7use serde::{Deserialize, Serialize};
8use xmtp_proto::types::GroupId;
9
10use super::{
11    consent_record::{ConsentState, ConsentType},
12    group::ConversationType,
13    group_message::{DeliveryStatus, StoredGroupMessage},
14    refresh_state::EntityKind,
15    schema::{
16        group_messages as messages, groups, refresh_state as progress,
17        user_preferences as preferences,
18    },
19    stream_storage::{StreamStorageError, stream_transaction},
20};
21use crate::{ConnectionExt, NotFound, StorageError};
22
23const PREFERENCES_ID: i32 = 0;
24const ALLOCATOR_ID: &[u8] = &[];
25
26/// An exclusive replay position in one client database, not a network cursor.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
28pub struct DeliveryCursor {
29    /// Stable across reopen. Whole-database restore rotates it to reject stale cursors.
30    pub database_id: [u8; 16],
31    /// Immutable local order. Zero precedes all messages; deleted values are never reused.
32    pub delivery_sequence: u64,
33}
34
35/// Random lease token that fences every default-delivery progress write.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct DeliveryOwner(pub [u8; 16]);
38
39/// Groups eligible for selection. Default progress outside this scope stays unchanged.
40#[derive(Debug, Clone, Default)]
41pub enum DeliveryScope {
42    #[default]
43    All,
44    Groups(Vec<GroupId>),
45}
46
47/// Current selection filters; excluded rows can advance D only within the active scope.
48#[derive(Debug, Clone, Default)]
49pub struct DeliveryFilter {
50    pub conversation_type: Option<ConversationType>,
51    pub consent_states: Option<Vec<ConsentState>>,
52}
53
54/// One retained, published message and its immutable local replay position.
55#[derive(Debug, Clone)]
56pub struct DeliveryMessage {
57    pub message: StoredGroupMessage,
58    pub cursor: DeliveryCursor,
59}
60
61/// History and its resume cursor read from one database snapshot.
62#[derive(Debug, Clone)]
63pub struct DeliverySnapshot {
64    pub messages: Vec<DeliveryMessage>,
65    /// Covers the snapshot, even when filters or limits omit some older messages.
66    pub cursor: DeliveryCursor,
67}
68
69/// Local message order, retained-history reads, and fenced per-group D positions.
70pub trait QueryDelivery: ConnectionExt + Sized {
71    /// Read the database identity used to reject foreign or pre-restore cursors.
72    fn stream_database_id(&self) -> Result<[u8; 16], StorageError> {
73        self.raw_query(|conn| Ok(database_id(conn)))?
74    }
75
76    /// Rotate only under exclusive restore lifecycle access. Existing consumer tokens are fenced.
77    fn rotate_stream_database_id(&self) -> Result<[u8; 16], StorageError> {
78        let identity = xmtp_common::rand_array::<16>();
79        stream_transaction(self, |conn| {
80            diesel::update(preferences::table.find(PREFERENCES_ID))
81                .set((
82                    preferences::stream_database_id.eq(identity.as_slice()),
83                    preferences::delivery_owner.eq(None::<Vec<u8>>),
84                    preferences::delivery_owner_until_ns.eq(None::<i64>),
85                ))
86                .execute(conn)?;
87            Ok(identity)
88        })
89    }
90
91    /// Allocate within the transaction that makes a message deliverable.
92    /// A duplicate keeps its number. An optimistic unpublished row has no number.
93    fn assign_delivery_sequence(&self, message_id: &[u8]) -> Result<Option<u64>, StorageError> {
94        stream_transaction(self, |conn| assign_sequence(conn, message_id))
95    }
96
97    /// Read the persistent allocator, not the maximum remaining message row.
98    fn current_delivery_cursor(&self) -> Result<DeliveryCursor, StorageError> {
99        self.raw_query(|conn| Ok(conn.transaction::<_, StorageError, _>(current_cursor)))?
100    }
101
102    /// Acquire the sole default-consumer lease at a supplied time; fail if one is active.
103    fn acquire_delivery_owner(
104        &self,
105        now_ns: i64,
106        until_ns: i64,
107    ) -> Result<DeliveryOwner, StorageError> {
108        let lease_ns = until_ns
109            .checked_sub(now_ns)
110            .ok_or(StreamStorageError::InvalidDeliveryPosition)?;
111        self.acquire_delivery_owner_with_clock(lease_ns, || now_ns)
112    }
113
114    /// Read the clock after the writer is acquired, not before a possible lock wait.
115    fn acquire_delivery_owner_with_clock(
116        &self,
117        lease_ns: i64,
118        clock: impl FnOnce() -> i64,
119    ) -> Result<DeliveryOwner, StorageError> {
120        if lease_ns <= 0 {
121            return Err(StreamStorageError::InvalidDeliveryPosition.into());
122        }
123        let owner = DeliveryOwner(xmtp_common::rand_array::<16>());
124        stream_transaction(self, |conn| {
125            let now_ns = clock();
126            let until_ns = now_ns
127                .checked_add(lease_ns)
128                .ok_or(StreamStorageError::InvalidDeliveryPosition)?;
129            let changed = diesel::update(
130                preferences::table.find(PREFERENCES_ID).filter(
131                    preferences::delivery_owner
132                        .is_null()
133                        .or(preferences::delivery_owner_until_ns.le(now_ns)),
134                ),
135            )
136            .set((
137                preferences::delivery_owner.eq(owner.0.as_slice()),
138                preferences::delivery_owner_until_ns.eq(until_ns),
139            ))
140            .execute(conn)?;
141            if changed == 0 {
142                return Err(StreamStorageError::AlreadyActive.into());
143            }
144            Ok(owner)
145        })
146    }
147
148    /// An expired token cannot be renewed. The caller must acquire a new token.
149    fn renew_delivery_owner(
150        &self,
151        owner: DeliveryOwner,
152        now_ns: i64,
153        until_ns: i64,
154    ) -> Result<(), StorageError> {
155        let lease_ns = until_ns
156            .checked_sub(now_ns)
157            .ok_or(StreamStorageError::InvalidDeliveryPosition)?;
158        self.renew_delivery_owner_with_clock(owner, lease_ns, || now_ns)
159    }
160
161    /// Extend only the current, unexpired token using time read after the writer lock.
162    fn renew_delivery_owner_with_clock(
163        &self,
164        owner: DeliveryOwner,
165        lease_ns: i64,
166        clock: impl FnOnce() -> i64,
167    ) -> Result<(), StorageError> {
168        if lease_ns <= 0 {
169            return Err(StreamStorageError::InvalidDeliveryPosition.into());
170        }
171        stream_transaction(self, |conn| {
172            let now_ns = clock();
173            let until_ns = now_ns
174                .checked_add(lease_ns)
175                .ok_or(StreamStorageError::InvalidDeliveryPosition)?;
176            let changed = diesel::update(
177                preferences::table
178                    .find(PREFERENCES_ID)
179                    .filter(preferences::delivery_owner.eq(owner.0.as_slice()))
180                    .filter(preferences::delivery_owner_until_ns.gt(now_ns)),
181            )
182            .set(preferences::delivery_owner_until_ns.eq(until_ns))
183            .execute(conn)?;
184            if changed == 0 {
185                return Err(StreamStorageError::NotCurrentOwner.into());
186            }
187            Ok(())
188        })
189    }
190
191    /// Reject expired or replaced tokens before handing a message to the app.
192    fn check_delivery_owner(&self, owner: DeliveryOwner, now_ns: i64) -> Result<(), StorageError> {
193        self.check_delivery_owner_with_clock(owner, || now_ns)
194    }
195
196    /// Check ownership using fresh time after the database connection is available.
197    fn check_delivery_owner_with_clock(
198        &self,
199        owner: DeliveryOwner,
200        clock: impl FnOnce() -> i64,
201    ) -> Result<(), StorageError> {
202        self.raw_query(|conn| Ok(check_owner(conn, owner, clock())))?
203    }
204
205    /// Release this token only; a stale consumer cannot release its replacement.
206    fn release_delivery_owner(&self, owner: DeliveryOwner) -> Result<(), StorageError> {
207        self.raw_query(|conn| {
208            diesel::update(
209                preferences::table
210                    .find(PREFERENCES_ID)
211                    .filter(preferences::delivery_owner.eq(owner.0.as_slice())),
212            )
213            .set((
214                preferences::delivery_owner.eq(None::<Vec<u8>>),
215                preferences::delivery_owner_until_ns.eq(None::<i64>),
216            ))
217            .execute(conn)
218        })?;
219        Ok(())
220    }
221
222    /// A buffered candidate can expire or be deleted while its previous item is held.
223    fn delivery_message_is_retained(
224        &self,
225        message_id: &[u8],
226        cursor: DeliveryCursor,
227        now_ns: i64,
228    ) -> Result<bool, StorageError> {
229        self.raw_query(|conn| {
230            Ok(conn.transaction::<_, StorageError, _>(|conn| {
231                validate_cursor(conn, cursor)?;
232                Ok(diesel::select(diesel::dsl::exists(
233                    messages::table
234                        .find(message_id)
235                        .filter(messages::delivery_sequence.eq(cursor.delivery_sequence as i64))
236                        .filter(messages::delivery_status.eq(DeliveryStatus::Published))
237                        .filter(
238                            messages::expire_at_ns
239                                .is_null()
240                                .or(messages::expire_at_ns.gt(now_ns)),
241                        ),
242                ))
243                .get_result::<bool>(conn)?)
244            }))
245        })?
246    }
247
248    /// Read retained candidates above each group's default position.
249    /// The caller applies consent/type filters and acknowledges scanned rows by the same owner.
250    fn default_delivery_messages(
251        &self,
252        owner: DeliveryOwner,
253        scope: &DeliveryScope,
254        now_ns: i64,
255        limit: u32,
256    ) -> Result<Vec<DeliveryMessage>, StorageError> {
257        self.default_delivery_messages_bounded(owner, scope, now_ns, limit, u64::MAX)
258    }
259
260    /// Bound rows and bytes before loading message bodies; this read never advances D.
261    fn default_delivery_messages_bounded(
262        &self,
263        owner: DeliveryOwner,
264        scope: &DeliveryScope,
265        now_ns: i64,
266        limit: u32,
267        max_bytes: u64,
268    ) -> Result<Vec<DeliveryMessage>, StorageError> {
269        self.raw_query(|conn| {
270            Ok(conn.transaction::<_, StorageError, _>(|conn| {
271                check_owner(conn, owner, now_ns)?;
272                read_messages(
273                    conn,
274                    scope,
275                    DeliveryReadOptions {
276                        after: None,
277                        filter: None,
278                        now_ns,
279                        limit,
280                        max_bytes,
281                        descending: false,
282                    },
283                )
284            }))
285        })?
286    }
287
288    /// Explicit replay does not read or write default delivery positions.
289    fn replay_delivery_messages(
290        &self,
291        after: DeliveryCursor,
292        scope: &DeliveryScope,
293        now_ns: i64,
294        limit: u32,
295    ) -> Result<Vec<DeliveryMessage>, StorageError> {
296        self.replay_delivery_messages_bounded(after, scope, now_ns, limit, u64::MAX)
297    }
298
299    /// Read a bounded retained prefix strictly after the cursor, without an owner or D writes.
300    fn replay_delivery_messages_bounded(
301        &self,
302        after: DeliveryCursor,
303        scope: &DeliveryScope,
304        now_ns: i64,
305        limit: u32,
306        max_bytes: u64,
307    ) -> Result<Vec<DeliveryMessage>, StorageError> {
308        self.raw_query(|conn| {
309            Ok(conn.transaction::<_, StorageError, _>(|conn| {
310                validate_cursor(conn, after)?;
311                read_messages(
312                    conn,
313                    scope,
314                    DeliveryReadOptions {
315                        after: Some(after.delivery_sequence),
316                        filter: None,
317                        now_ns,
318                        limit,
319                        max_bytes,
320                        descending: false,
321                    },
322                )
323            }))
324        })?
325    }
326
327    /// Return recent retained history and a cursor from the same database snapshot.
328    fn delivery_history_snapshot(
329        &self,
330        scope: &DeliveryScope,
331        now_ns: i64,
332        limit: u32,
333    ) -> Result<DeliverySnapshot, StorageError> {
334        self.delivery_history_snapshot_bounded(scope, now_ns, limit, u64::MAX)
335    }
336
337    /// Bound history allocation while capturing its resume cursor in the same snapshot.
338    fn delivery_history_snapshot_bounded(
339        &self,
340        scope: &DeliveryScope,
341        now_ns: i64,
342        limit: u32,
343        max_bytes: u64,
344    ) -> Result<DeliverySnapshot, StorageError> {
345        self.delivery_history_snapshot_filtered(
346            scope,
347            &DeliveryFilter::default(),
348            now_ns,
349            limit,
350            max_bytes,
351        )
352    }
353
354    /// Apply history filters before its limit, with history and cursor in one read transaction.
355    fn delivery_history_snapshot_filtered(
356        &self,
357        scope: &DeliveryScope,
358        filter: &DeliveryFilter,
359        now_ns: i64,
360        limit: u32,
361        max_bytes: u64,
362    ) -> Result<DeliverySnapshot, StorageError> {
363        self.raw_query(|conn| {
364            Ok(conn.transaction::<_, StorageError, _>(|conn| {
365                let cursor = current_cursor(conn)?;
366                let mut messages = read_messages(
367                    conn,
368                    scope,
369                    DeliveryReadOptions {
370                        after: Some(0),
371                        filter: Some(filter),
372                        now_ns,
373                        limit,
374                        max_bytes,
375                        descending: true,
376                    },
377                )?;
378                messages.reverse();
379                Ok(DeliverySnapshot { messages, cursor })
380            }))
381        })?
382    }
383
384    /// Acknowledge after callback return or the next iterator request, never on queue insertion.
385    /// This also fences progress for rows excluded by a consent/type filter.
386    fn acknowledge_delivery(
387        &self,
388        owner: DeliveryOwner,
389        group_id: GroupId,
390        cursor: DeliveryCursor,
391        now_ns: i64,
392    ) -> Result<(), StorageError> {
393        self.acknowledge_delivery_with_clock(owner, group_id, cursor, || now_ns)
394    }
395
396    /// Advance this group's D only after a fresh owner check under the state writer.
397    fn acknowledge_delivery_with_clock(
398        &self,
399        owner: DeliveryOwner,
400        group_id: GroupId,
401        cursor: DeliveryCursor,
402        clock: impl FnOnce() -> i64,
403    ) -> Result<(), StorageError> {
404        stream_transaction(self, |conn| {
405            check_owner(conn, owner, clock())?;
406            validate_cursor(conn, cursor)?;
407            use diesel::{query_dsl::methods::FilterDsl, upsert::excluded};
408            diesel::insert_into(progress::table)
409                .values((
410                    progress::entity_id.eq(group_id.as_ref()),
411                    progress::entity_kind.eq(EntityKind::Delivery),
412                    progress::sequence_id.eq(cursor.delivery_sequence as i64),
413                ))
414                .on_conflict((progress::entity_id, progress::entity_kind))
415                .do_update()
416                .set(progress::sequence_id.eq(excluded(progress::sequence_id)))
417                .filter(progress::sequence_id.lt(excluded(progress::sequence_id)))
418                .execute(conn)?;
419            Ok(())
420        })
421    }
422}
423
424impl<C: ConnectionExt> QueryDelivery for C {}
425
426/// The caller holds the writer and has made the message visible in this transaction.
427pub(crate) fn assign_sequence(
428    conn: &mut diesel::SqliteConnection,
429    message_id: &[u8],
430) -> Result<Option<u64>, StorageError> {
431    let (sequence, status) = messages::table
432        .find(message_id)
433        .select((messages::delivery_sequence, messages::delivery_status))
434        .first::<(Option<i64>, DeliveryStatus)>(conn)
435        .optional()?
436        .ok_or_else(|| NotFound::MessageById(message_id.to_vec()))?;
437    if sequence.is_some() || status != DeliveryStatus::Published {
438        return Ok(sequence.map(|value| value as u64));
439    }
440    let sequence = diesel::update(
441        progress::table
442            .find((ALLOCATOR_ID, EntityKind::DeliveryAllocator))
443            .filter(progress::sequence_id.lt(i64::MAX)),
444    )
445    .set(progress::sequence_id.eq(progress::sequence_id + 1))
446    .returning(progress::sequence_id)
447    .get_result::<i64>(conn)
448    .optional()?
449    .ok_or(StreamStorageError::DeliveryExhausted)?;
450    diesel::update(
451        messages::table
452            .find(message_id)
453            .filter(messages::delivery_sequence.is_null()),
454    )
455    .set(messages::delivery_sequence.eq(sequence))
456    .execute(conn)?;
457    Ok(Some(sequence as u64))
458}
459
460fn database_id(conn: &mut diesel::SqliteConnection) -> Result<[u8; 16], StorageError> {
461    preferences::table
462        .find(PREFERENCES_ID)
463        .select(preferences::stream_database_id)
464        .first::<Vec<u8>>(conn)?
465        .try_into()
466        .map_err(|_| StorageError::DbDeserialize)
467}
468
469fn current_cursor(conn: &mut diesel::SqliteConnection) -> Result<DeliveryCursor, StorageError> {
470    let sequence = progress::table
471        .find((ALLOCATOR_ID, EntityKind::DeliveryAllocator))
472        .select(progress::sequence_id)
473        .first::<i64>(conn)?;
474    Ok(DeliveryCursor {
475        database_id: database_id(conn)?,
476        delivery_sequence: sequence as u64,
477    })
478}
479
480fn validate_cursor(
481    conn: &mut diesel::SqliteConnection,
482    cursor: DeliveryCursor,
483) -> Result<(), StorageError> {
484    let current = current_cursor(conn)?;
485    if cursor.database_id != current.database_id {
486        return Err(StreamStorageError::ForeignCursor.into());
487    }
488    if cursor.delivery_sequence > current.delivery_sequence {
489        return Err(StreamStorageError::InvalidDeliveryPosition.into());
490    }
491    Ok(())
492}
493
494fn check_owner(
495    conn: &mut diesel::SqliteConnection,
496    owner: DeliveryOwner,
497    now_ns: i64,
498) -> Result<(), StorageError> {
499    let valid = preferences::table
500        .find(PREFERENCES_ID)
501        .filter(preferences::delivery_owner.eq(owner.0.as_slice()))
502        .filter(preferences::delivery_owner_until_ns.gt(now_ns))
503        .select(preferences::id)
504        .first::<i32>(conn)
505        .optional()?
506        .is_some();
507    if !valid {
508        return Err(StreamStorageError::NotCurrentOwner.into());
509    }
510    Ok(())
511}
512
513/// Selection and allocation bounds for one retained-message read.
514struct DeliveryReadOptions<'a> {
515    /// An absent replay position selects each group's saved default progress.
516    after: Option<u64>,
517    filter: Option<&'a DeliveryFilter>,
518    now_ns: i64,
519    limit: u32,
520    max_bytes: u64,
521    descending: bool,
522}
523
524/// Select scalar sizes first so message bodies cannot exceed the read budget.
525fn read_messages(
526    conn: &mut diesel::SqliteConnection,
527    scope: &DeliveryScope,
528    options: DeliveryReadOptions<'_>,
529) -> Result<Vec<DeliveryMessage>, StorageError> {
530    let DeliveryReadOptions {
531        after,
532        filter,
533        now_ns,
534        limit,
535        max_bytes,
536        descending,
537    } = options;
538    // Read only scalar lengths first. Large blobs never enter the candidate buffer before the budget check.
539    let row_bytes = diesel::dsl::sql::<BigInt>(
540        "length(group_messages.id) + length(group_messages.group_id) + length(group_messages.decrypted_message_bytes) + \
541        length(group_messages.sender_installation_id) + length(CAST(group_messages.sender_inbox_id AS BLOB)) + \
542        length(CAST(group_messages.authority_id AS BLOB)) + COALESCE(length(group_messages.reference_id), 0) + \
543        COALESCE(length(group_messages.envelope_hash), 0) + length(CAST(group_messages.idempotency_key AS BLOB)) + ")
544        .bind::<BigInt, _>(std::mem::size_of::<DeliveryMessage>() as i64);
545    let mut query = messages::table
546        .inner_join(groups::table)
547        .filter(groups::conversation_type.ne_all(ConversationType::virtual_types()))
548        .filter(messages::delivery_sequence.is_not_null())
549        .filter(messages::delivery_status.eq(DeliveryStatus::Published))
550        .filter(
551            messages::expire_at_ns
552                .is_null()
553                .or(messages::expire_at_ns.gt(now_ns)),
554        )
555        .select((messages::delivery_sequence.assume_not_null(), row_bytes))
556        .into_boxed();
557    if let DeliveryScope::Groups(groups) = scope {
558        query = query.filter(messages::group_id.eq_any(groups));
559    }
560    if let Some(filter) = filter {
561        if let Some(kind) = filter.conversation_type {
562            query = query.filter(groups::conversation_type.eq(kind));
563        }
564        if let Some(states) = &filter.consent_states {
565            let consent = diesel::dsl::sql::<Integer>(
566                "COALESCE((SELECT state FROM consent_records WHERE entity_type = ",
567            )
568            .bind::<Integer, _>(ConsentType::ConversationId as i32)
569            .sql(" AND entity = lower(hex(group_messages.group_id))), ")
570            .bind::<Integer, _>(ConsentState::Unknown as i32)
571            .sql(")");
572            query = query.filter(
573                consent.eq_any(states.iter().map(|state| *state as i32).collect::<Vec<_>>()),
574            );
575        }
576    }
577    if let Some(after) = after {
578        query = query.filter(messages::delivery_sequence.gt(after as i64));
579    } else {
580        query = query.filter(
581            messages::delivery_sequence.gt(diesel::dsl::sql::<Nullable<BigInt>>(
582                "COALESCE((SELECT sequence_id FROM refresh_state WHERE entity_kind = ",
583            )
584            .bind::<Integer, _>(EntityKind::Delivery as i32)
585            .sql(" AND entity_id = group_messages.group_id), 0)")),
586        );
587    }
588    query = if descending {
589        query.order(messages::delivery_sequence.desc())
590    } else {
591        query.order(messages::delivery_sequence.asc())
592    };
593    let sizes = query.limit(i64::from(limit)).load::<(i64, i64)>(conn)?;
594    let mut sequences = Vec::with_capacity(sizes.len());
595    let mut total_bytes = 0_u64;
596    for (sequence, bytes) in sizes {
597        let bytes = u64::try_from(bytes).map_err(|_| StorageError::DbDeserialize)?;
598        if bytes > max_bytes.saturating_sub(total_bytes) {
599            if sequences.is_empty() {
600                return Err(StreamStorageError::LocalReadCapacity {
601                    bytes,
602                    limit: max_bytes,
603                }
604                .into());
605            }
606            break;
607        }
608        total_bytes += bytes;
609        sequences.push(sequence);
610    }
611    let mut query = messages::table
612        .filter(messages::delivery_sequence.eq_any(sequences))
613        .select((
614            StoredGroupMessage::as_select(),
615            messages::delivery_sequence.assume_not_null(),
616        ))
617        .into_boxed();
618    query = if descending {
619        query.order(messages::delivery_sequence.desc())
620    } else {
621        query.order(messages::delivery_sequence.asc())
622    };
623    let database_id = database_id(conn)?;
624    let rows = query.load::<(StoredGroupMessage, i64)>(conn)?;
625    Ok(rows
626        .into_iter()
627        .map(|(message, sequence)| DeliveryMessage {
628            message,
629            cursor: DeliveryCursor {
630                database_id,
631                delivery_sequence: sequence as u64,
632            },
633        })
634        .collect())
635}
636
637#[cfg(test)]
638mod tests;
xmtp_db/encrypted_store/delivery.rs

xmtp_db/encrypted_store/
delivery.rs
