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

xmtp_mls/subscriptions/
message_reader.rs
