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, conversation_consent},
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.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                // The buffered row can be marked deleted after prefetch: hand
322                // out its current row.
323                // implements: PROC-025, PROC-060
324                let retained = self.session.context.db().retained_delivery_message(
325                    &candidate.message.id,
326                    candidate.cursor,
327                    now_ns(),
328                )?;
329                let current = match retained {
330                    Some(current)
331                        if matches_filter(
332                            &self.session.context,
333                            candidate.message.group_id,
334                            &selection.scope,
335                            &selection.filter,
336                        )? =>
337                    {
338                        Some(current)
339                    }
340                    _ => None,
341                };
342                let Some(current) = current else {
343                    if !self.skip_candidate(&candidate, selection.revision)? {
344                        self.candidates.clear();
345                        break;
346                    }
347                    scanned += 1;
348                    if scanned >= self.config.batch_size {
349                        break;
350                    }
351                    continue;
352                };
353                // A read transaction or filter check can outlive a lease. Check again at handoff.
354                self.session.check_owner()?;
355                let pending = Arc::new(PendingAcknowledgement::new(
356                    candidate.cursor,
357                    selection.revision,
358                ));
359                let acknowledgement = DeliveryAcknowledgement::new(
360                    Arc::clone(&self.session),
361                    Arc::clone(&pending),
362                    candidate.message.group_id,
363                    self.control.clone(),
364                    candidate.message.id.clone(),
365                );
366                self.pending = Some(pending);
367                return Ok(Some(LocalDeliveryItem {
368                    message: current,
369                    cursor: candidate.cursor,
370                    acknowledgement,
371                }));
372            }
373            tokio::task::yield_now().await;
374        }
375    }
376
377    /// Order filtered progress with scope changes. Excluded scopes never advance D.
378    // implements: PROC-047
379    fn skip_candidate(&mut self, candidate: &DeliveryMessage, revision: u64) -> Result<bool> {
380        let selection = self.control.selection.lock();
381        if selection.revision != revision {
382            return Ok(false);
383        }
384        self.session.check_owner()?;
385        if let Some(owner) = self.session.owner() {
386            self.session.context.db().acknowledge_delivery_with_clock(
387                owner,
388                candidate.message.group_id,
389                candidate.cursor,
390                now_ns,
391            )?;
392        } else {
393            self.replay_position = Some(candidate.cursor);
394        }
395        Ok(true)
396    }
397
398    /// Rust iterator adapter. A new poll acknowledges the previously returned item.
399    /// Dropping this stream while it is suspended at `yield` leaves that item unacknowledged.
400    pub fn into_stream(self) -> impl Stream<Item = super::Result<StoredGroupMessage>> {
401        self.into_cursor_stream()
402            .map(|item| item.map(|item| item.message).map_err(SubscribeError::from))
403    }
404
405    /// Cursor-bearing Rust replay/default iterator with the same acknowledgement boundary.
406    // implements: PROC-052
407    pub fn into_cursor_stream(self) -> impl Stream<Item = Result<DeliveryMessage>> {
408        futures::stream::unfold(
409            Some((self, None::<DeliveryAcknowledgement<Context>>)),
410            |state| async move {
411                let (mut reader, previous) = state?;
412                if let Some(previous) = previous
413                    && let Err(error) = previous.acknowledge()
414                {
415                    return Some((Err(error), None));
416                }
417                loop {
418                    match reader.next_delivery().await {
419                        Ok(Some(LocalDeliveryItem {
420                            message,
421                            cursor,
422                            acknowledgement,
423                        })) => match acknowledgement.check_owner() {
424                            Ok(()) => {
425                                return Some((
426                                    Ok(DeliveryMessage { message, cursor }),
427                                    Some((reader, Some(acknowledgement))),
428                                ));
429                            }
430                            Err(LocalDeliveryError::SelectionChanged) => continue,
431                            Err(error) => return Some((Err(error), None)),
432                        },
433                        Ok(None) => return None,
434                        Err(error) => return Some((Err(error), None)),
435                    }
436                }
437            },
438        )
439    }
440
441    /// Read history, relations, deletions, and its boundary in one database snapshot.
442    pub fn enriched_history_snapshot(
443        context: &Context,
444        scope: &DeliveryScope,
445        filter: &LocalDeliveryFilter,
446        limit: u32,
447    ) -> Result<EnrichedDeliverySnapshot> {
448        let settings = context.stream_policy();
449        let filter = effective_filter(scope, filter);
450        Ok(context
451            .db()
452            .delivery_history_snapshot_projected_with_clock(
453                scope,
454                &filter,
455                now_ns,
456                limit.min(settings.max_local_read_rows),
457                settings.max_local_read_bytes,
458                |conn, snapshot| {
459                    let storage = conn.key_store();
460                    enrich_delivery_snapshot(storage.db(), snapshot)
461                },
462            )?)
463    }
464
465    /// Read history and its resume boundary in one database snapshot.
466    pub fn history_snapshot(
467        context: &Context,
468        scope: &DeliveryScope,
469        filter: &LocalDeliveryFilter,
470        limit: u32,
471    ) -> Result<DeliverySnapshot> {
472        let settings = context.stream_policy();
473        let filter = effective_filter(scope, filter);
474        Ok(context
475            .db()
476            .delivery_history_snapshot_projected_with_clock(
477                scope,
478                &filter,
479                now_ns,
480                limit.min(settings.max_local_read_rows),
481                settings.max_local_read_bytes,
482                |_, snapshot| Ok(snapshot),
483            )?)
484    }
485}
486
487impl<Context: XmtpSharedContext> Drop for LocalDelivery<Context> {
488    fn drop(&mut self) {
489        self.session.close();
490        if let Some(task) = &self.renewal {
491            task.end();
492        }
493    }
494}
495
496// implements: CONS-042, CONS-043
497fn effective_consent_states<'a>(
498    scope: &DeliveryScope,
499    filter: &'a LocalDeliveryFilter,
500) -> Option<&'a [ConsentState]> {
501    match (&filter.consent_states, scope) {
502        (Some(states), _) => Some(states),
503        (None, DeliveryScope::All) => Some(&DEFAULT_STREAM_CONSENT_STATES),
504        (None, _) => None,
505    }
506}
507
508fn effective_filter<'a>(
509    scope: &DeliveryScope,
510    filter: &'a LocalDeliveryFilter,
511) -> Cow<'a, LocalDeliveryFilter> {
512    match effective_consent_states(scope, filter) {
513        Some(states) if filter.consent_states.is_none() => Cow::Owned(LocalDeliveryFilter {
514            conversation_type: filter.conversation_type,
515            consent_states: Some(states.to_vec()),
516        }),
517        _ => Cow::Borrowed(filter),
518    }
519}
520
521fn matches_filter<Context: XmtpSharedContext>(
522    context: &Context,
523    group_id: GroupId,
524    scope: &DeliveryScope,
525    filter: &LocalDeliveryFilter,
526) -> Result<bool> {
527    let db = context.db();
528    if let Some(kind) = filter.conversation_type {
529        let group: Option<StoredGroup> = db.fetch(&group_id)?;
530        if group.is_none_or(|group| group.conversation_type != kind) {
531            return Ok(false);
532        }
533    }
534    // Borrow the effective states; this runs once per candidate.
535    if let Some(states) = effective_consent_states(scope, filter) {
536        let consent = conversation_consent(&db, context.inbox_id(), &group_id)?;
537        if !states.contains(&consent) {
538            return Ok(false);
539        }
540    }
541    Ok(true)
542}
543
544#[cfg(test)]
545mod tests;
546
547/// Selected rows with canonical relations and deletion placeholders.
548pub struct EnrichedDeliverySnapshot {
549    pub messages: Vec<crate::messages::enrichment::EnrichedStoredMessage>,
550    pub cursor: DeliveryCursor,
551}
552
553fn enrich_delivery_snapshot(
554    db: impl xmtp_db::DbQuery,
555    snapshot: DeliverySnapshot,
556) -> std::result::Result<EnrichedDeliverySnapshot, StorageError> {
557    use crate::messages::enrichment::{EnrichMessageError, enrich_messages_with_stored};
558    let mut groups = std::collections::HashMap::new();
559    let mut projected: Vec<_> = (0..snapshot.messages.len()).map(|_| None).collect();
560    for (index, item) in snapshot.messages.into_iter().enumerate() {
561        groups
562            .entry(item.message.group_id)
563            .or_insert_with(Vec::new)
564            .push(((index, item.cursor), item.message));
565    }
566    for (group_id, items) in groups {
567        let (positions, rows): (Vec<_>, Vec<_>) = items.into_iter().unzip();
568        let enriched =
569            enrich_messages_with_stored(&db, &group_id, rows).map_err(|error| match error {
570                EnrichMessageError::Storage(error) => error,
571                EnrichMessageError::DbConnection(error) => StorageError::from(error),
572            })?;
573        if enriched.len() != positions.len() {
574            return Err(StorageError::DbDeserialize);
575        }
576        for ((index, cursor), mut row) in positions.into_iter().zip(enriched) {
577            row.delivery_cursor = Some(cursor);
578            projected[index] = Some(row);
579        }
580    }
581    Ok(EnrichedDeliverySnapshot {
582        messages: projected
583            .into_iter()
584            .map(|row| row.ok_or(StorageError::DbDeserialize))
585            .collect::<std::result::Result<_, _>>()?,
586        cursor: snapshot.cursor,
587    })
588}
xmtp_mls/subscriptions/local_delivery.rs

xmtp_mls/subscriptions/
local_delivery.rs
