xmtp_mls/subscriptions/
message_reader.rs

1//! Message delivery with shared network receipt and processing status.
2
3use parking_lot::Mutex;
4use std::sync::Arc;
5
6use super::{
7    incoming::{IncomingCoordinator, IncomingLease, IncomingScope, IncomingStatus},
8    local_delivery::{
9        DeliveryCursor, DeliveryScope, LocalDelivery, LocalDeliveryConfig, LocalDeliveryControl,
10        LocalDeliveryError, LocalDeliveryFilter, LocalDeliveryItem,
11    },
12};
13use crate::context::XmtpSharedContext;
14
15/// The network lease does not depend on application acknowledgement.
16pub struct MessageReader<C: XmtpSharedContext> {
17    delivery: LocalDelivery<C>,
18    control: MessageReaderControl,
19    configuration: crate::server_configuration::ServerConfigurationHandle,
20}
21
22/// Updates delivery selection and observes network progress without consuming messages.
23#[derive(Clone)]
24pub struct MessageReaderControl {
25    delivery: LocalDeliveryControl,
26    lease: Arc<Mutex<Option<Arc<IncomingLease>>>>,
27    changes: Arc<tokio::sync::Mutex<tokio::sync::watch::Receiver<u64>>>,
28    last_status: Arc<Mutex<IncomingStatus>>,
29    closed: tokio_util::sync::CancellationToken,
30}
31
32/// One independent status observer for a caller that snapshots before waiting.
33pub struct MessageReaderObserver {
34    changes: Option<tokio::sync::watch::Receiver<u64>>,
35    closed: tokio_util::sync::CancellationToken,
36}
37
38impl MessageReaderObserver {
39    pub async fn changed(&mut self) {
40        if let Some(changes) = &mut self.changes {
41            tokio::select! {
42                _ = self.closed.cancelled() => {},
43                _ = changes.changed() => {},
44            }
45        }
46    }
47}
48
49fn incoming_scope(scope: &DeliveryScope) -> IncomingScope {
50    match scope {
51        DeliveryScope::All => IncomingScope::AllGroups,
52        DeliveryScope::Groups(groups) => IncomingScope::Groups(groups.clone()),
53    }
54}
55
56impl<C: XmtpSharedContext + 'static> MessageReader<C> {
57    /// Hold network interest while opening default delivery or independent `from` replay.
58    pub fn new(
59        context: C,
60        scope: DeliveryScope,
61        filter: LocalDeliveryFilter,
62        from: Option<DeliveryCursor>,
63    ) -> Result<Self, LocalDeliveryError> {
64        context
65            .server_configuration()
66            .check()
67            .map_err(|error| LocalDeliveryError::Configuration(Box::new(error)))?;
68        let configuration = context.server_configuration().clone();
69        let config = LocalDeliveryConfig::from(context.incoming_runtime().policy());
70        let coordinator = IncomingCoordinator::for_context(&context);
71        let lease = Arc::new(coordinator.acquire_stream(incoming_scope(&scope)));
72        let delivery = LocalDelivery::new(context, scope, filter, from, config)?;
73        let control = MessageReaderControl {
74            delivery: delivery.control(),
75            last_status: Arc::new(Mutex::new(lease.snapshot())),
76            changes: Arc::new(tokio::sync::Mutex::new(lease.subscribe_changes())),
77            lease: Arc::new(Mutex::new(Some(lease))),
78            closed: tokio_util::sync::CancellationToken::new(),
79        };
80        Ok(Self {
81            delivery,
82            control,
83            configuration,
84        })
85    }
86
87    pub fn control(&self) -> MessageReaderControl {
88        self.control.clone()
89    }
90
91    /// Return one unacknowledged item; receipt can continue while its host handles it.
92    pub async fn next_delivery(
93        &mut self,
94    ) -> Result<Option<LocalDeliveryItem<C>>, LocalDeliveryError> {
95        let lease = self.control.lease.lock().clone();
96        let Some(lease) = lease else { return Ok(None) };
97        let configuration = self.configuration.clone();
98        let result = {
99            // Keep the pending read alive across status hints. Cancelling it on
100            // each hint could starve local storage work during a busy stream.
101            let delivery = self.delivery.next_delivery();
102            futures::pin_mut!(delivery);
103            loop {
104                if let Err(error) = configuration.check() {
105                    break Err(LocalDeliveryError::Configuration(Box::new(error)));
106                }
107                if let Err(error) = lease.check_recovery() {
108                    break Err(error.error());
109                }
110                tokio::select! {
111                    biased;
112                    _ = self.control.closed.cancelled() => {
113                        break configuration.check()
114                            .map(|()| None)
115                            .map_err(|error| LocalDeliveryError::Configuration(Box::new(error)));
116                    },
117                    item = &mut delivery => {
118                        if let Err(error) = configuration.check() {
119                            break Err(LocalDeliveryError::Configuration(Box::new(error)));
120                        }
121                        // Both futures can become ready while this task is suspended.
122                        // An ended lease must win before another item reaches the app.
123                        if matches!(&item, Ok(Some(_)))
124                            && let Err(error) = lease.check_recovery()
125                        {
126                            break Err(error.error());
127                        }
128                        break item;
129                    },
130                    _ = lease.changed() => {},
131                    _ = xmtp_common::time::sleep(super::recovery::RECOVERY_POLL) => {},
132                }
133            }
134        };
135        if !matches!(&result, Ok(Some(_))) {
136            // Release network interest and ownership on both error and EOF.
137            // The caller can immediately open a replacement stream.
138            self.close();
139        }
140        result
141    }
142
143    /// Report processing through fixed network heads, not application delivery progress.
144    pub fn catch_up_snapshot(&self) -> IncomingStatus {
145        self.control.catch_up_snapshot()
146    }
147
148    /// A Rust iterator acknowledges only when its consumer requests the next item.
149    pub fn into_stream(
150        self,
151    ) -> impl futures::Stream<Item = super::Result<xmtp_db::group_message::StoredGroupMessage>>
152    {
153        futures::stream::unfold(
154            Some((
155                self,
156                None::<super::local_delivery::DeliveryAcknowledgement<C>>,
157            )),
158            |state| async move {
159                let (mut reader, previous) = state?;
160                if let Some(previous) = previous
161                    && let Err(error) = previous.acknowledge()
162                {
163                    return Some((Err(error.into()), None));
164                }
165                loop {
166                    match reader.next_delivery().await {
167                        Ok(Some(item)) => match item.acknowledgement.check_owner() {
168                            Ok(()) => {
169                                return Some((
170                                    Ok(item.message),
171                                    Some((reader, Some(item.acknowledgement))),
172                                ));
173                            }
174                            Err(LocalDeliveryError::SelectionChanged) => continue,
175                            Err(error) => return Some((Err(error.into()), None)),
176                        },
177                        Ok(None) => return None,
178                        Err(error) => return Some((Err(error.into()), None)),
179                    }
180                }
181            },
182        )
183    }
184
185    /// Release network interest and delivery ownership without acknowledging the last item.
186    pub fn close(&mut self) {
187        self.delivery.close();
188        self.control.close();
189    }
190}
191
192impl MessageReaderControl {
193    #[cfg(any(test, feature = "test-utils"))]
194    pub fn lease_holder_count_for_test(&self) -> usize {
195        Arc::strong_count(&self.lease)
196    }
197
198    #[cfg(test)]
199    pub(crate) fn notify_change_for_test(&self) {
200        if let Some(lease) = self.lease.lock().as_ref() {
201            lease.notify_change_for_test();
202        }
203    }
204
205    /// Create a receiver before the caller reads its first status snapshot.
206    pub fn observer(&self) -> MessageReaderObserver {
207        MessageReaderObserver {
208            changes: self
209                .lease
210                .lock()
211                .as_ref()
212                .map(|lease| lease.subscribe_changes()),
213            closed: self.closed.clone(),
214        }
215    }
216
217    /// Replace network interest and local selection; excluded groups keep their saved D.
218    pub fn update_scope(&self, scope: DeliveryScope) {
219        self.delivery.update_scope(scope.clone());
220        if let Some(lease) = self.lease.lock().as_ref() {
221            lease.replace_scope(incoming_scope(&scope));
222        }
223    }
224
225    /// Apply filters to future selection without rewinding saved delivery progress.
226    pub fn update_filter(&self, filter: LocalDeliveryFilter) {
227        self.delivery.update_filter(filter);
228    }
229
230    /// Read current obligations, or the final cancelled snapshot after close.
231    pub fn catch_up_snapshot(&self) -> IncomingStatus {
232        if let Some(lease) = self.lease.lock().as_ref() {
233            let status = lease.snapshot();
234            *self.last_status.lock() = status.clone();
235            status
236        } else {
237            self.last_status.lock().clone()
238        }
239    }
240
241    /// Wait for a status hint or close; read a fresh snapshot after this returns.
242    pub async fn changed(&self) {
243        let mut changes = self.changes.lock().await;
244        tokio::select! {
245            _ = self.closed.cancelled() => {},
246            _ = changes.changed() => {},
247        }
248    }
249
250    /// Cancel the scope and fence delivery tokens without changing durable progress.
251    pub fn close(&self) {
252        self.delivery.close();
253        self.closed.cancel();
254        if let Some(lease) = self.lease.lock().take() {
255            let mut status = lease.snapshot();
256            status.cancel();
257            *self.last_status.lock() = status;
258            lease.close();
259        }
260    }
261}
262
263impl<C: XmtpSharedContext> Drop for MessageReader<C> {
264    fn drop(&mut self) {
265        self.control.close();
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272    use crate::{
273        groups::send_message_opts::SendMessageOpts, test::mock::generate_stored_msg, tester,
274    };
275    use xmtp_common::time::{Duration, Instant};
276    use xmtp_db::Store;
277    use xmtp_proto::types::Cursor;
278
279    // verifies: DMS-009
280    #[xmtp_common::test(unwrap_try = true)]
281    async fn one_dm_reader_receives_a_later_explicitly_joined_duplicate() {
282        tester!(alix, disable_workers);
283        tester!(bo, disable_workers);
284        let first = alix
285            .find_or_create_dm(bo.inbox_id().to_string(), None)
286            .await?;
287        let mut reader = MessageReader::new(
288            alix.context.clone(),
289            DeliveryScope::Groups(vec![first.group_id]),
290            LocalDeliveryFilter::default(),
291            None,
292        )?;
293        let second = bo
294            .find_or_create_dm(alix.inbox_id().to_string(), None)
295            .await?;
296        assert_ne!(first.group_id, second.group_id);
297        // This operation owns Welcome interest only until it returns.
298        alix.sync_welcomes().await?;
299        let installed: Option<xmtp_db::group::StoredGroup> =
300            xmtp_db::Fetch::fetch(&alix.context.db(), &second.group_id)?;
301        assert!(installed.is_some());
302        let message_id = second
303            .send_message(b"later duplicate", SendMessageOpts::default())
304            .await?;
305
306        let observed_group = xmtp_common::time::timeout(Duration::from_secs(15), async {
307            loop {
308                let item = reader
309                    .next_delivery()
310                    .await?
311                    .expect("the reader must remain open");
312                let matched = item.message.id == message_id;
313                let group_id = item.message.group_id;
314                item.acknowledgement.acknowledge()?;
315                if matched {
316                    break Ok::<_, LocalDeliveryError>(group_id);
317                }
318            }
319        })
320        .await??;
321        assert_eq!(observed_group, second.group_id);
322    }
323
324    // verifies: CONF-075
325    #[xmtp_common::test(unwrap_try = true)]
326    async fn blocked_connection_ends_a_pending_message_read_with_its_cause() {
327        use crate::{client::ClientError, server_configuration::BlockedConnection};
328        tester!(alix, disable_workers);
329        let mut reader = MessageReader::new(
330            alix.context.clone(),
331            DeliveryScope::All,
332            LocalDeliveryFilter::default(),
333            None,
334        )?;
335        let error = {
336            let pending = reader.next_delivery();
337            futures::pin_mut!(pending);
338            assert!(futures::poll!(&mut pending).is_pending());
339            alix.context.server_configuration().block_connection(
340                BlockedConnection::ClientVersionTooOld {
341                    client: "1.0.0".into(),
342                    minimum: "9999.0.0".into(),
343                },
344            );
345            alix.context.cancellation_token().cancel();
346            pending
347                .await
348                .err()
349                .expect("blocked connection must fail the read")
350        };
351        assert!(matches!(
352            error,
353            LocalDeliveryError::Configuration(error)
354                if matches!(*error, ClientError::ClientVersionTooOld { ref minimum, .. }
355                    if minimum == "9999.0.0")
356        ));
357        assert!(reader.next_delivery().await?.is_none());
358    }
359
360    // verifies: PROC-038
361    #[xmtp_common::test(unwrap_try = true)]
362    async fn exhaustion_wins_when_a_pending_local_item_is_ready_too() {
363        tester!(alix, disable_workers);
364        let group = alix.create_group(None, None)?;
365        for sequence in [100, 200] {
366            generate_stored_msg(Cursor(sequence), group.group_id).store(&alix.context.db())?;
367        }
368        let create = || {
369            MessageReader::new(
370                alix.context.clone(),
371                DeliveryScope::Groups(vec![group.group_id]),
372                LocalDeliveryFilter::default(),
373                None,
374            )
375        };
376        let mut reader = create()?;
377        let lease = reader.control.lease.lock().clone().unwrap();
378        let first = reader.next_delivery().await?.unwrap();
379        first.acknowledgement.check_owner()?;
380        let pending = reader.next_delivery();
381        futures::pin_mut!(pending);
382        assert!(futures::poll!(&mut pending).is_pending());
383        lease.fail_recovery_for_test(super::super::recovery::RecoveryFailure::Exhausted {
384            attempts: 10,
385            source: None,
386        });
387        // An already handed-off callback may finish. Its next local item is ready
388        // at the same time as the retained network failure.
389        first.acknowledgement.acknowledge()?;
390        assert!(matches!(
391            pending.await,
392            Err(LocalDeliveryError::NetworkRecoveryExhausted { .. })
393        ));
394        let mut replacement = create()?;
395        let second = replacement.next_delivery().await?.unwrap();
396        assert_ne!(second.message.id, first.message.id);
397    }
398
399    // verifies: PROC-039
400    #[xmtp_common::test(unwrap_try = true)]
401    async fn twice_exhausted_readers_release_ownership_before_replacement_on_same_client() {
402        tester!(alix, disable_workers);
403        let group = alix.create_group(None, None)?;
404        let create = || {
405            MessageReader::new(
406                alix.context.clone(),
407                DeliveryScope::Groups(vec![group.group_id]),
408                LocalDeliveryFilter::default(),
409                None,
410            )
411        };
412        for sequence in [100, 200] {
413            let message = generate_stored_msg(Cursor(sequence), group.group_id);
414            message.store(&alix.context.db())?;
415            let mut expired = create()?;
416            let lease = expired.control.lease.lock().clone().unwrap();
417            // Observe a future outage deadline without subtracting from the
418            // WASM clock. Publish the failure only to this consumer's lease.
419            let outage = super::super::recovery::RecoverySnapshot::default();
420            let now = Instant::now();
421            let mut recovery = super::super::recovery::RecoveryBudget::new(&outage, now);
422            let failure = recovery
423                .check(&outage, now + Duration::from_secs(3600))
424                .unwrap_err();
425            assert!(matches!(
426                failure,
427                super::super::recovery::RecoveryFailure::Exhausted { .. }
428            ));
429            lease.fail_recovery_for_test(failure);
430            assert!(matches!(
431                expired.next_delivery().await,
432                Err(LocalDeliveryError::NetworkRecoveryExhausted { .. })
433            ));
434            let mut replacement = create()?;
435            expired.close();
436            let delivered = replacement.next_delivery().await?.unwrap();
437            assert_eq!(delivered.message.id, message.id);
438            delivered.acknowledgement.check_owner()?;
439            delivered.acknowledgement.acknowledge()?;
440            replacement.close();
441        }
442    }
443}
xmtp_mls/subscriptions/message_reader.rs

xmtp_mls/subscriptions/
message_reader.rs
