xmtp_mls/subscriptions/
local_delivery.rs

1//! Local message delivery. Network receipt and MLS processing run independently.
2
3mod acknowledgement;
4mod types;
5
6pub use acknowledgement::DeliveryAcknowledgement;
7pub(crate) use types::LocalDeliveryConfig;
8pub use types::{LocalDeliveryError, LocalDeliveryFilter};
9pub use xmtp_db::delivery::{DeliveryCursor, DeliveryScope, DeliverySnapshot};
10
11use acknowledgement::{AcknowledgementState, DeliverySession, PendingAcknowledgement};
12use futures::{Stream, StreamExt};
13use parking_lot::Mutex;
14use std::{borrow::Cow, collections::VecDeque, sync::Arc};
15use tokio::sync::Notify;
16use xmtp_common::{
17    StreamHandle,
18    time::{now_ns, sleep},
19};
20use xmtp_db::{
21    Fetch, StorageError,
22    consent_record::{ConsentState, ConsentType},
23    delivery::{DeliveryMessage, QueryDelivery},
24    group::StoredGroup,
25    group_message::StoredGroupMessage,
26    prelude::*,
27};
28use xmtp_proto::types::GroupId;
29
30use super::{DEFAULT_STREAM_CONSENT_STATES, SubscribeError, internal::InternalEvent};
31use crate::context::XmtpSharedContext;
32use xmtp_events::{EventFilter, Subscription};
33
34type Result<T> = std::result::Result<T, LocalDeliveryError>;
35
36/// One stored message and its unacknowledged application handoff.
37pub struct LocalDeliveryItem<Context: XmtpSharedContext> {
38    /// Retained, published content selected from the local database.
39    pub message: StoredGroupMessage,
40    /// Resume strictly after this item. Reading the cursor does not acknowledge it.
41    pub cursor: DeliveryCursor,
42    /// Host-side token; acknowledge only after callback return or the next-item request.
43    pub acknowledgement: DeliveryAcknowledgement<Context>,
44}
45
46#[derive(Clone)]
47struct Selection {
48    scope: DeliveryScope,
49    filter: LocalDeliveryFilter,
50    revision: u64,
51}
52
53/// Change future selection while a callback or iterator holds its current item.
54#[derive(Clone)]
55pub struct LocalDeliveryControl {
56    selection: Arc<Mutex<Selection>>,
57    changed: Arc<Notify>,
58    closer: Arc<dyn CloseDelivery>,
59}
60
61trait CloseDelivery: xmtp_common::MaybeSend + xmtp_common::MaybeSync {
62    fn close_delivery(&self);
63}
64
65impl<C: XmtpSharedContext> CloseDelivery for DeliverySession<C> {
66    fn close_delivery(&self) {
67        self.close();
68    }
69}
70
71impl LocalDeliveryControl {
72    pub fn close(&self) {
73        self.closer.close_delivery();
74    }
75
76    /// Change selected groups without consuming excluded groups' backlog.
77    pub fn update_scope(&self, scope: DeliveryScope) {
78        let mut selection = self.selection.lock();
79        selection.scope = scope;
80        selection.revision = selection.revision.wrapping_add(1);
81        self.changed.notify_one();
82    }
83
84    /// Change future selection; previously scanned rows stay consumed.
85    pub fn update_filter(&self, filter: LocalDeliveryFilter) {
86        let mut selection = self.selection.lock();
87        selection.filter = filter;
88        selection.revision = selection.revision.wrapping_add(1);
89        self.changed.notify_one();
90    }
91
92    pub fn scope(&self) -> DeliveryScope {
93        self.selection.lock().scope.clone()
94    }
95}
96
97/// One local reader with at most one item awaiting application acknowledgement.
98pub struct LocalDelivery<Context: XmtpSharedContext> {
99    session: Arc<DeliverySession<Context>>,
100    config: LocalDeliveryConfig,
101    control: LocalDeliveryControl,
102    events: Subscription<InternalEvent>,
103    candidates: VecDeque<DeliveryMessage>,
104    candidate_revision: u64,
105    /// At most one handoff can wait for acknowledgement.
106    pending: Option<Arc<PendingAcknowledgement>>,
107    /// `Some` selects independent replay; it never reads or changes default D.
108    replay_position: Option<DeliveryCursor>,
109    renewal: Option<Box<dyn StreamHandle<StreamOutput = ()>>>,
110}
111
112impl<Context> LocalDelivery<Context>
113where
114    Context: XmtpSharedContext + 'static,
115{
116    /// Open a default owner, or independent replay when `from` is supplied.
117    pub(crate) fn new(
118        context: Context,
119        scope: DeliveryScope,
120        filter: LocalDeliveryFilter,
121        from: Option<DeliveryCursor>,
122        config: LocalDeliveryConfig,
123    ) -> Result<Self> {
124        config.validate()?;
125        if context.is_closed() {
126            return Err(LocalDeliveryError::Closed);
127        }
128        // Subscribe before the first database read. Polling also covers missed and cross-process writes.
129        let events = context.events().subscribe(
130            EventFilter::default().with_internal(|event| {
131                matches!(event, InternalEvent::MessageStored { is_sync: false, .. })
132            }),
133            Some(10),
134        );
135        let now = now_ns();
136        let database_id = context.db().stream_database_id()?;
137        let owner = if let Some(cursor) = from {
138            context
139                .db()
140                .replay_delivery_messages(cursor, &scope, now, 0)?;
141            None
142        } else {
143            let mut registered = context.delivery_owner().lock();
144            if context.is_closed() {
145                return Err(LocalDeliveryError::Closed);
146            }
147            let mut retired = context.incoming_runtime().retired_delivery_owner.lock();
148            if let Some(owner) = *retired {
149                // Release only the closed reader's token. SQL fences a different owner.
150                context.db().release_delivery_owner(owner)?;
151                if *registered == Some(owner) {
152                    *registered = None;
153                }
154                *retired = None;
155            }
156            let owner = context
157                .db()
158                .acquire_delivery_owner_with_clock(config.lease_duration_ns()?, now_ns)?;
159            *registered = Some(owner);
160            Some(owner)
161        };
162        let session = Arc::new(DeliverySession::new(context, owner, database_id));
163        let renewal = if owner.is_some() {
164            let session = Arc::clone(&session);
165            Some(Box::new(xmtp_common::spawn(None, async move {
166                loop {
167                    tokio::select! {
168                        _ = session.cancel.cancelled() => { session.close(); break; },
169                        _ = sleep(config.renew_interval) => {}
170                    }
171                    let result = config
172                        .lease_duration_ns()
173                        .and_then(|lease_ns| session.renew(lease_ns).map_err(Into::into));
174                    if let Err(error) = result {
175                        session.fail(error);
176                        break;
177                    }
178                }
179            })) as Box<dyn StreamHandle<StreamOutput = ()>>)
180        } else {
181            None
182        };
183        Ok(Self {
184            session: Arc::clone(&session),
185            config,
186            control: LocalDeliveryControl {
187                selection: Arc::new(Mutex::new(Selection {
188                    scope,
189                    filter,
190                    revision: 0,
191                })),
192                changed: Arc::new(Notify::new()),
193                closer: session,
194            },
195            events,
196            candidates: VecDeque::new(),
197            candidate_revision: 0,
198            pending: None,
199            replay_position: from,
200            renewal,
201        })
202    }
203
204    pub fn control(&self) -> LocalDeliveryControl {
205        self.control.clone()
206    }
207
208    /// Close without acknowledging the last item. Its content remains eligible after restart.
209    pub fn close(&mut self) {
210        self.session.close();
211    }
212
213    /// Return one item with an explicit acknowledgement token for callback and host-queue adapters.
214    /// The next call waits for that token. It never acknowledges the previous item itself.
215    // implements: PROC-040
216    pub async fn next_delivery(&mut self) -> Result<Option<LocalDeliveryItem<Context>>> {
217        let result = self.next_inner().await;
218        if result.is_err() {
219            self.session.close();
220        }
221        result
222    }
223
224    // implements: PROC-026, PROC-052, PROC-034
225    async fn next_inner(&mut self) -> Result<Option<LocalDeliveryItem<Context>>> {
226        if let Some(pending) = self.pending.clone() {
227            loop {
228                let notified = pending.changed.notified();
229                {
230                    let mut state = pending.state.lock();
231                    // Read the cause under the same lock as acknowledgement and rejection.
232                    if let Some(error) = self.session.background_error() {
233                        return Err(error);
234                    }
235                    if matches!(*state, AcknowledgementState::Waiting)
236                        && self.control.selection.lock().revision != pending.revision
237                    {
238                        *state = AcknowledgementState::Reselect;
239                    }
240                    match *state {
241                        AcknowledgementState::Acknowledged => {
242                            if self.replay_position.is_some() {
243                                self.replay_position = Some(pending.cursor);
244                            }
245                            self.pending = None;
246                            break;
247                        }
248                        AcknowledgementState::Reselect => {
249                            self.pending = None;
250                            self.candidates.clear();
251                            break;
252                        }
253                        AcknowledgementState::Rejected => {
254                            return Err(LocalDeliveryError::AcknowledgementRejected);
255                        }
256                        AcknowledgementState::Failed => {
257                            return Err(LocalDeliveryError::AcknowledgementFailed);
258                        }
259                        AcknowledgementState::Waiting | AcknowledgementState::Dispatched => {}
260                    }
261                }
262                tokio::select! {
263                    _ = notified => {},
264                    _ = self.control.changed.notified() => {},
265                    _ = self.session.cancel.cancelled() => return self.session.end_result(),
266                }
267            }
268        }
269        loop {
270            if self.session.is_closed() {
271                return self.session.end_result();
272            }
273            let selection = self.control.selection.lock().clone();
274            if selection.revision != self.candidate_revision {
275                self.candidates.clear();
276                self.candidate_revision = selection.revision;
277            }
278            if self.candidates.is_empty() {
279                let now = now_ns();
280                self.candidates = {
281                    let db = self.session.context.db();
282                    if let Some(position) = self.replay_position {
283                        db.replay_delivery_messages_bounded(
284                            position,
285                            &selection.scope,
286                            now,
287                            self.config.batch_size,
288                            self.config.max_bytes,
289                        )?
290                    } else {
291                        db.default_delivery_messages_bounded(
292                            self.session.owner().ok_or(LocalDeliveryError::Closed)?,
293                            &selection.scope,
294                            now,
295                            self.config.batch_size,
296                            self.config.max_bytes,
297                        )?
298                    }
299                    .into()
300                };
301                if self.candidates.is_empty() {
302                    tokio::select! {
303                        _ = self.session.cancel.cancelled() => return self.session.end_result(),
304                        _ = self.control.changed.notified() => {},
305                        _ = self.events.next() => {},
306                        _ = sleep(self.config.poll_interval) => {},
307                    }
308                    continue;
309                }
310            }
311            let mut scanned = 0;
312            while let Some(candidate) = self.candidates.pop_front() {
313                // Re-read selection for each item. A scope change never consumes excluded groups.
314                let selection = self.control.selection.lock().clone();
315                if selection.revision != self.candidate_revision {
316                    self.candidates.clear();
317                    self.candidate_revision = selection.revision;
318                    break;
319                }
320                self.session.check_owner()?;
321                let retained = self.session.context.db().delivery_message_is_retained(
322                    &candidate.message.id,
323                    candidate.cursor,
324                    now_ns(),
325                )?;
326                if !retained
327                    || !matches_filter(
328                        &self.session.context,
329                        candidate.message.group_id,
330                        &selection.scope,
331                        &selection.filter,
332                    )?
333                {
334                    if !self.skip_candidate(&candidate, selection.revision)? {
335                        self.candidates.clear();
336                        break;
337                    }
338                    scanned += 1;
339                    if scanned >= self.config.batch_size {
340                        break;
341                    }
342                    continue;
343                }
344                // A read transaction or filter check can outlive a lease. Check again at handoff.
345                self.session.check_owner()?;
346                let pending = Arc::new(PendingAcknowledgement::new(
347                    candidate.cursor,
348                    selection.revision,
349                ));
350                let acknowledgement = DeliveryAcknowledgement::new(
351                    Arc::clone(&self.session),
352                    Arc::clone(&pending),
353                    candidate.message.group_id,
354                    self.control.clone(),
355                    candidate.message.id.clone(),
356                );
357                self.pending = Some(pending);
358                return Ok(Some(LocalDeliveryItem {
359                    message: candidate.message,
360                    cursor: candidate.cursor,
361                    acknowledgement,
362                }));
363            }
364            tokio::task::yield_now().await;
365        }
366    }
367
368    /// Order filtered progress with scope changes. Excluded scopes never advance D.
369    // implements: PROC-047
370    fn skip_candidate(&mut self, candidate: &DeliveryMessage, revision: u64) -> Result<bool> {
371        let selection = self.control.selection.lock();
372        if selection.revision != revision {
373            return Ok(false);
374        }
375        self.session.check_owner()?;
376        if let Some(owner) = self.session.owner() {
377            self.session.context.db().acknowledge_delivery_with_clock(
378                owner,
379                candidate.message.group_id,
380                candidate.cursor,
381                now_ns,
382            )?;
383        } else {
384            self.replay_position = Some(candidate.cursor);
385        }
386        Ok(true)
387    }
388
389    /// Rust iterator adapter. A new poll acknowledges the previously returned item.
390    /// Dropping this stream while it is suspended at `yield` leaves that item unacknowledged.
391    pub fn into_stream(self) -> impl Stream<Item = super::Result<StoredGroupMessage>> {
392        self.into_cursor_stream()
393            .map(|item| item.map(|item| item.message).map_err(SubscribeError::from))
394    }
395
396    /// Cursor-bearing Rust replay/default iterator with the same acknowledgement boundary.
397    // implements: PROC-052
398    pub fn into_cursor_stream(self) -> impl Stream<Item = Result<DeliveryMessage>> {
399        futures::stream::unfold(
400            Some((self, None::<DeliveryAcknowledgement<Context>>)),
401            |state| async move {
402                let (mut reader, previous) = state?;
403                if let Some(previous) = previous
404                    && let Err(error) = previous.acknowledge()
405                {
406                    return Some((Err(error), None));
407                }
408                loop {
409                    match reader.next_delivery().await {
410                        Ok(Some(LocalDeliveryItem {
411                            message,
412                            cursor,
413                            acknowledgement,
414                        })) => match acknowledgement.check_owner() {
415                            Ok(()) => {
416                                return Some((
417                                    Ok(DeliveryMessage { message, cursor }),
418                                    Some((reader, Some(acknowledgement))),
419                                ));
420                            }
421                            Err(LocalDeliveryError::SelectionChanged) => continue,
422                            Err(error) => return Some((Err(error), None)),
423                        },
424                        Ok(None) => return None,
425                        Err(error) => return Some((Err(error), None)),
426                    }
427                }
428            },
429        )
430    }
431
432    /// Read history, relations, deletions, and its boundary in one database snapshot.
433    pub fn enriched_history_snapshot(
434        context: &Context,
435        scope: &DeliveryScope,
436        filter: &LocalDeliveryFilter,
437        limit: u32,
438    ) -> Result<EnrichedDeliverySnapshot> {
439        let settings = context.incoming_runtime().policy();
440        let filter = effective_filter(scope, filter);
441        Ok(context
442            .db()
443            .delivery_history_snapshot_projected_with_clock(
444                scope,
445                &filter,
446                now_ns,
447                limit.min(settings.max_local_read_rows),
448                settings.max_local_read_bytes,
449                |conn, snapshot| {
450                    let storage = conn.key_store();
451                    enrich_delivery_snapshot(storage.db(), snapshot)
452                },
453            )?)
454    }
455
456    /// Read history and its resume boundary in one database snapshot.
457    pub fn history_snapshot(
458        context: &Context,
459        scope: &DeliveryScope,
460        filter: &LocalDeliveryFilter,
461        limit: u32,
462    ) -> Result<DeliverySnapshot> {
463        let settings = context.incoming_runtime().policy();
464        let filter = effective_filter(scope, filter);
465        Ok(context
466            .db()
467            .delivery_history_snapshot_projected_with_clock(
468                scope,
469                &filter,
470                now_ns,
471                limit.min(settings.max_local_read_rows),
472                settings.max_local_read_bytes,
473                |_, snapshot| Ok(snapshot),
474            )?)
475    }
476}
477
478impl<Context: XmtpSharedContext> Drop for LocalDelivery<Context> {
479    fn drop(&mut self) {
480        self.session.close();
481        if let Some(task) = &self.renewal {
482            task.end();
483        }
484    }
485}
486
487// implements: CONS-042, CONS-043
488fn effective_consent_states<'a>(
489    scope: &DeliveryScope,
490    filter: &'a LocalDeliveryFilter,
491) -> Option<&'a [ConsentState]> {
492    match (&filter.consent_states, scope) {
493        (Some(states), _) => Some(states),
494        (None, DeliveryScope::All) => Some(&DEFAULT_STREAM_CONSENT_STATES),
495        (None, _) => None,
496    }
497}
498
499fn effective_filter<'a>(
500    scope: &DeliveryScope,
501    filter: &'a LocalDeliveryFilter,
502) -> Cow<'a, LocalDeliveryFilter> {
503    match effective_consent_states(scope, filter) {
504        Some(states) if filter.consent_states.is_none() => Cow::Owned(LocalDeliveryFilter {
505            conversation_type: filter.conversation_type,
506            consent_states: Some(states.to_vec()),
507        }),
508        _ => Cow::Borrowed(filter),
509    }
510}
511
512fn matches_filter<Context: XmtpSharedContext>(
513    context: &Context,
514    group_id: GroupId,
515    scope: &DeliveryScope,
516    filter: &LocalDeliveryFilter,
517) -> Result<bool> {
518    let db = context.db();
519    if let Some(kind) = filter.conversation_type {
520        let group: Option<StoredGroup> = db.fetch(&group_id)?;
521        if group.is_none_or(|group| group.conversation_type != kind) {
522            return Ok(false);
523        }
524    }
525    // Borrow the effective states; this runs once per candidate.
526    if let Some(states) = effective_consent_states(scope, filter) {
527        let consent = db
528            .get_consent_record(hex::encode(group_id), ConsentType::ConversationId)
529            .map_err(StorageError::from)?
530            .map_or(ConsentState::Unknown, |record| record.state);
531        if !states.contains(&consent) {
532            return Ok(false);
533        }
534    }
535    Ok(true)
536}
537
538#[cfg(test)]
539mod tests;
540
541/// Selected rows with canonical relations and deletion placeholders.
542pub struct EnrichedDeliverySnapshot {
543    pub messages: Vec<crate::messages::enrichment::EnrichedStoredMessage>,
544    pub cursor: DeliveryCursor,
545}
546
547fn enrich_delivery_snapshot(
548    db: impl xmtp_db::DbQuery,
549    snapshot: DeliverySnapshot,
550) -> std::result::Result<EnrichedDeliverySnapshot, StorageError> {
551    use crate::messages::enrichment::{EnrichMessageError, enrich_messages_with_stored};
552    let mut groups = std::collections::HashMap::new();
553    let mut projected: Vec<_> = (0..snapshot.messages.len()).map(|_| None).collect();
554    for (index, item) in snapshot.messages.into_iter().enumerate() {
555        groups
556            .entry(item.message.group_id)
557            .or_insert_with(Vec::new)
558            .push(((index, item.cursor), item.message));
559    }
560    for (group_id, items) in groups {
561        let (positions, rows): (Vec<_>, Vec<_>) = items.into_iter().unzip();
562        let enriched =
563            enrich_messages_with_stored(&db, &group_id, rows).map_err(|error| match error {
564                EnrichMessageError::Storage(error) => error,
565                EnrichMessageError::DbConnection(error) => StorageError::from(error),
566            })?;
567        if enriched.len() != positions.len() {
568            return Err(StorageError::DbDeserialize);
569        }
570        for ((index, cursor), mut row) in positions.into_iter().zip(enriched) {
571            row.delivery_cursor = Some(cursor);
572            projected[index] = Some(row);
573        }
574    }
575    Ok(EnrichedDeliverySnapshot {
576        messages: projected
577            .into_iter()
578            .map(|row| row.ok_or(StorageError::DbDeserialize))
579            .collect::<std::result::Result<_, _>>()?,
580        cursor: snapshot.cursor,
581    })
582}
xmtp_mls/subscriptions/local_delivery.rs

xmtp_mls/subscriptions/
local_delivery.rs
