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::{collections::VecDeque, sync::Arc};
15use tokio::sync::{Notify, broadcast};
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::{LocalEvents, SubscribeError};
31use crate::context::XmtpSharedContext;
32
33type Result<T> = std::result::Result<T, LocalDeliveryError>;
34
35/// One stored message and its unacknowledged application handoff.
36pub struct LocalDeliveryItem<Context: XmtpSharedContext> {
37    /// Retained, published content selected from the local database.
38    pub message: StoredGroupMessage,
39    /// Resume strictly after this item. Reading the cursor does not acknowledge it.
40    pub cursor: DeliveryCursor,
41    /// Host-side token; acknowledge only after callback return or the next-item request.
42    pub acknowledgement: DeliveryAcknowledgement<Context>,
43}
44
45#[derive(Clone)]
46struct Selection {
47    scope: DeliveryScope,
48    filter: LocalDeliveryFilter,
49    revision: u64,
50}
51
52/// Change future selection while a callback or iterator holds its current item.
53#[derive(Clone)]
54pub struct LocalDeliveryControl {
55    selection: Arc<Mutex<Selection>>,
56    changed: Arc<Notify>,
57    closer: Arc<dyn CloseDelivery>,
58}
59
60trait CloseDelivery: xmtp_common::MaybeSend + xmtp_common::MaybeSync {
61    fn close_delivery(&self);
62}
63
64impl<C: XmtpSharedContext> CloseDelivery for DeliverySession<C> {
65    fn close_delivery(&self) {
66        self.close();
67    }
68}
69
70impl LocalDeliveryControl {
71    pub fn close(&self) {
72        self.closer.close_delivery();
73    }
74
75    /// Change selected groups without consuming excluded groups' backlog (STR-092).
76    pub fn update_scope(&self, scope: DeliveryScope) {
77        let mut selection = self.selection.lock();
78        selection.scope = scope;
79        selection.revision = selection.revision.wrapping_add(1);
80        self.changed.notify_one();
81    }
82
83    /// Change future selection; previously scanned rows stay consumed (STR-093).
84    pub fn update_filter(&self, filter: LocalDeliveryFilter) {
85        let mut selection = self.selection.lock();
86        selection.filter = filter;
87        selection.revision = selection.revision.wrapping_add(1);
88        self.changed.notify_one();
89    }
90
91    pub fn scope(&self) -> DeliveryScope {
92        self.selection.lock().scope.clone()
93    }
94}
95
96/// One local reader with at most one item awaiting application acknowledgement.
97pub struct LocalDelivery<Context: XmtpSharedContext> {
98    session: Arc<DeliverySession<Context>>,
99    config: LocalDeliveryConfig,
100    control: LocalDeliveryControl,
101    events: broadcast::Receiver<LocalEvents>,
102    candidates: VecDeque<DeliveryMessage>,
103    candidate_revision: u64,
104    /// At most one handoff can wait for acknowledgement (STR-094).
105    pending: Option<Arc<PendingAcknowledgement>>,
106    /// `Some` selects independent replay; it never reads or changes default D.
107    replay_position: Option<DeliveryCursor>,
108    renewal: Option<Box<dyn StreamHandle<StreamOutput = ()>>>,
109}
110
111impl<Context> LocalDelivery<Context>
112where
113    Context: XmtpSharedContext + 'static,
114{
115    /// Open a default owner, or independent replay when `from` is supplied.
116    pub(crate) fn new(
117        context: Context,
118        scope: DeliveryScope,
119        filter: LocalDeliveryFilter,
120        from: Option<DeliveryCursor>,
121        config: LocalDeliveryConfig,
122    ) -> Result<Self> {
123        config.validate()?;
124        if context.is_closed() {
125            return Err(LocalDeliveryError::Closed);
126        }
127        // Subscribe before the first database read. Polling also covers missed and cross-process writes.
128        let events = context.local_events().subscribe();
129        let now = now_ns();
130        let owner = if let Some(cursor) = from {
131            context
132                .db()
133                .replay_delivery_messages(cursor, &scope, now, 0)?;
134            None
135        } else {
136            let mut registered = context.delivery_owner().lock();
137            if context.is_closed() {
138                return Err(LocalDeliveryError::Closed);
139            }
140            let owner = context
141                .db()
142                .acquire_delivery_owner_with_clock(config.lease_duration_ns()?, now_ns)?;
143            *registered = Some(owner);
144            Some(owner)
145        };
146        let session = Arc::new(DeliverySession::new(
147            context,
148            owner,
149            from.map(|cursor| cursor.database_id),
150        ));
151        let renewal = if owner.is_some() {
152            let session = Arc::clone(&session);
153            Some(Box::new(xmtp_common::spawn(None, async move {
154                loop {
155                    tokio::select! {
156                        _ = session.cancel.cancelled() => { session.close(); break; },
157                        _ = sleep(config.renew_interval) => {}
158                    }
159                    let result = config
160                        .lease_duration_ns()
161                        .and_then(|lease_ns| session.renew(lease_ns).map_err(Into::into));
162                    if let Err(error) = result {
163                        session.fail(error);
164                        break;
165                    }
166                }
167            })) as Box<dyn StreamHandle<StreamOutput = ()>>)
168        } else {
169            None
170        };
171        Ok(Self {
172            session: Arc::clone(&session),
173            config,
174            control: LocalDeliveryControl {
175                selection: Arc::new(Mutex::new(Selection {
176                    scope,
177                    filter,
178                    revision: 0,
179                })),
180                changed: Arc::new(Notify::new()),
181                closer: session,
182            },
183            events,
184            candidates: VecDeque::new(),
185            candidate_revision: 0,
186            pending: None,
187            replay_position: from,
188            renewal,
189        })
190    }
191
192    pub fn control(&self) -> LocalDeliveryControl {
193        self.control.clone()
194    }
195
196    /// Close without acknowledging the last item. Its content remains eligible after restart.
197    pub fn close(&mut self) {
198        self.session.close();
199    }
200
201    /// Return one item with an explicit acknowledgement token for callback and host-queue adapters.
202    /// The next call waits for that token. It never acknowledges the previous item itself.
203    pub async fn next_delivery(&mut self) -> Result<Option<LocalDeliveryItem<Context>>> {
204        let result = self.next_inner().await;
205        if result.is_err() {
206            self.session.close();
207        }
208        result
209    }
210
211    async fn next_inner(&mut self) -> Result<Option<LocalDeliveryItem<Context>>> {
212        if let Some(pending) = self.pending.clone() {
213            loop {
214                let notified = pending.changed.notified();
215                {
216                    let mut state = pending.state.lock();
217                    if matches!(*state, AcknowledgementState::Waiting)
218                        && self.control.selection.lock().revision != pending.revision
219                    {
220                        *state = AcknowledgementState::Reselect;
221                    }
222                }
223                match *pending.state.lock() {
224                    AcknowledgementState::Acknowledged => {
225                        if self.replay_position.is_some() {
226                            self.replay_position = Some(pending.cursor);
227                        }
228                        self.pending = None;
229                        break;
230                    }
231                    AcknowledgementState::Reselect => {
232                        self.pending = None;
233                        self.candidates.clear();
234                        break;
235                    }
236                    AcknowledgementState::Rejected => {
237                        return Err(LocalDeliveryError::AcknowledgementRejected);
238                    }
239                    AcknowledgementState::Failed => {
240                        return Err(LocalDeliveryError::AcknowledgementFailed);
241                    }
242                    AcknowledgementState::Waiting | AcknowledgementState::Dispatched => {}
243                }
244                tokio::select! {
245                    _ = notified => {},
246                    _ = self.control.changed.notified() => {},
247                    _ = self.session.cancel.cancelled() => return self.session.end_result(),
248                }
249            }
250        }
251        loop {
252            if self.session.is_closed() {
253                return self.session.end_result();
254            }
255            let selection = self.control.selection.lock().clone();
256            if selection.revision != self.candidate_revision {
257                self.candidates.clear();
258                self.candidate_revision = selection.revision;
259            }
260            if self.candidates.is_empty() {
261                let now = now_ns();
262                self.candidates = {
263                    let db = self.session.context.db();
264                    if let Some(position) = self.replay_position {
265                        db.replay_delivery_messages_bounded(
266                            position,
267                            &selection.scope,
268                            now,
269                            self.config.batch_size,
270                            self.config.max_bytes,
271                        )?
272                    } else {
273                        db.default_delivery_messages_bounded(
274                            self.session.owner.ok_or(LocalDeliveryError::Closed)?,
275                            &selection.scope,
276                            now,
277                            self.config.batch_size,
278                            self.config.max_bytes,
279                        )?
280                    }
281                    .into()
282                };
283                if self.candidates.is_empty() {
284                    tokio::select! {
285                        _ = self.session.cancel.cancelled() => return self.session.end_result(),
286                        _ = self.control.changed.notified() => {},
287                        _ = self.events.recv() => {},
288                        _ = sleep(self.config.poll_interval) => {},
289                    }
290                    continue;
291                }
292            }
293            let mut scanned = 0;
294            while let Some(candidate) = self.candidates.pop_front() {
295                // Re-read selection for each item. A scope change never consumes excluded groups.
296                let selection = self.control.selection.lock().clone();
297                if selection.revision != self.candidate_revision {
298                    self.candidates.clear();
299                    self.candidate_revision = selection.revision;
300                    break;
301                }
302                self.session.check_owner()?;
303                let retained = self.session.context.db().delivery_message_is_retained(
304                    &candidate.message.id,
305                    candidate.cursor,
306                    now_ns(),
307                )?;
308                if !retained
309                    || !matches_filter(
310                        &self.session.context,
311                        candidate.message.group_id,
312                        &selection.filter,
313                    )?
314                {
315                    if !self.skip_candidate(&candidate, selection.revision)? {
316                        self.candidates.clear();
317                        break;
318                    }
319                    scanned += 1;
320                    if scanned >= self.config.batch_size {
321                        break;
322                    }
323                    continue;
324                }
325                // A read transaction or filter check can outlive a lease. Check again at handoff.
326                self.session.check_owner()?;
327                let pending = Arc::new(PendingAcknowledgement::new(
328                    candidate.cursor,
329                    selection.revision,
330                ));
331                let acknowledgement = DeliveryAcknowledgement::new(
332                    Arc::clone(&self.session),
333                    Arc::clone(&pending),
334                    candidate.message.group_id,
335                    self.control.clone(),
336                    candidate.message.id.clone(),
337                );
338                self.pending = Some(pending);
339                return Ok(Some(LocalDeliveryItem {
340                    message: candidate.message,
341                    cursor: candidate.cursor,
342                    acknowledgement,
343                }));
344            }
345            tokio::task::yield_now().await;
346        }
347    }
348
349    /// Order filtered progress with scope changes. Excluded scopes never advance D.
350    fn skip_candidate(&mut self, candidate: &DeliveryMessage, revision: u64) -> Result<bool> {
351        let selection = self.control.selection.lock();
352        if selection.revision != revision {
353            return Ok(false);
354        }
355        self.session.check_owner()?;
356        if let Some(owner) = self.session.owner {
357            self.session.context.db().acknowledge_delivery_with_clock(
358                owner,
359                candidate.message.group_id,
360                candidate.cursor,
361                now_ns,
362            )?;
363        } else {
364            self.replay_position = Some(candidate.cursor);
365        }
366        Ok(true)
367    }
368
369    /// Rust iterator adapter. A new poll acknowledges the previously returned item.
370    /// Dropping this stream while it is suspended at `yield` leaves that item unacknowledged.
371    pub fn into_stream(self) -> impl Stream<Item = super::Result<StoredGroupMessage>> {
372        self.into_cursor_stream()
373            .map(|item| item.map(|item| item.message).map_err(SubscribeError::from))
374    }
375
376    /// Cursor-bearing Rust replay/default iterator with the same acknowledgement boundary.
377    pub fn into_cursor_stream(self) -> impl Stream<Item = Result<DeliveryMessage>> {
378        futures::stream::unfold(
379            Some((self, None::<DeliveryAcknowledgement<Context>>)),
380            |state| async move {
381                let (mut reader, previous) = state?;
382                if let Some(previous) = previous
383                    && let Err(error) = previous.acknowledge()
384                {
385                    return Some((Err(error), None));
386                }
387                loop {
388                    match reader.next_delivery().await {
389                        Ok(Some(LocalDeliveryItem {
390                            message,
391                            cursor,
392                            acknowledgement,
393                        })) => match acknowledgement.check_owner() {
394                            Ok(()) => {
395                                return Some((
396                                    Ok(DeliveryMessage { message, cursor }),
397                                    Some((reader, Some(acknowledgement))),
398                                ));
399                            }
400                            Err(LocalDeliveryError::SelectionChanged) => continue,
401                            Err(error) => return Some((Err(error), None)),
402                        },
403                        Ok(None) => return None,
404                        Err(error) => return Some((Err(error), None)),
405                    }
406                }
407            },
408        )
409    }
410
411    /// Read history and its resume boundary in one database snapshot (STR-099).
412    pub fn history_snapshot(
413        context: &Context,
414        scope: &DeliveryScope,
415        filter: &LocalDeliveryFilter,
416        limit: u32,
417    ) -> Result<DeliverySnapshot> {
418        let settings = context.incoming_runtime().policy();
419        Ok(context.db().delivery_history_snapshot_filtered(
420            scope,
421            filter,
422            now_ns(),
423            limit.min(settings.max_local_read_rows),
424            settings.max_local_read_bytes,
425        )?)
426    }
427}
428
429impl<Context: XmtpSharedContext> Drop for LocalDelivery<Context> {
430    fn drop(&mut self) {
431        self.session.close();
432        if let Some(task) = &self.renewal {
433            task.end();
434        }
435    }
436}
437
438fn matches_filter<Context: XmtpSharedContext>(
439    context: &Context,
440    group_id: GroupId,
441    filter: &LocalDeliveryFilter,
442) -> Result<bool> {
443    let db = context.db();
444    if let Some(kind) = filter.conversation_type {
445        let group: Option<StoredGroup> = db.fetch(&group_id)?;
446        if group.is_none_or(|group| group.conversation_type != kind) {
447            return Ok(false);
448        }
449    }
450    if let Some(states) = &filter.consent_states {
451        let consent = db
452            .get_consent_record(hex::encode(group_id), ConsentType::ConversationId)
453            .map_err(StorageError::from)?
454            .map_or(ConsentState::Unknown, |record| record.state);
455        if !states.contains(&consent) {
456            return Ok(false);
457        }
458    }
459    Ok(true)
460}
461
462#[cfg(test)]
463mod tests;
xmtp_mls/subscriptions/local_delivery.rs

xmtp_mls/subscriptions/
local_delivery.rs
