xmtp_mls/subscriptions/
incoming.rs

1//! One receiver and processing scheduler for each client context.
2
3mod controller;
4mod status;
5use super::connection_state::ConnectionStates;
6
7use crate::context::XmtpSharedContext;
8use controller::Controller;
9use parking_lot::Mutex;
10pub use status::*;
11use std::{
12    collections::{HashMap, HashSet},
13    sync::{
14        Arc,
15        atomic::{AtomicU64, Ordering},
16    },
17};
18use tokio::sync::{mpsc, watch};
19use xmtp_common::{BoxDynFuture, MaybeSend, MaybeSync, time::Instant};
20use xmtp_proto::{
21    api::NetworkError,
22    types::{Cursor, GroupId, IncomingBatchLimits, IncomingSubscription, Topic, TopicCursor},
23};
24
25pub(crate) type SubscriptionFuture =
26    BoxDynFuture<'static, Result<IncomingSubscription<NetworkError>, NetworkError>>;
27
28/// Exact method and wire inputs for a permanent server rejection.
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub(crate) enum RequestKey {
31    Bidi(TopicCursor),
32    QueryNewest(HashSet<Topic>),
33    /// Query can reduce its wire limit after a size error. Hold all limits at
34    /// this cursor because the caller cannot see the final attempted limit.
35    Query(Topic, Cursor),
36}
37
38pub(crate) type RejectedRequests = Arc<Mutex<Vec<(RequestKey, Arc<IncomingError>)>>>;
39
40pub(crate) trait SubscriptionFactory: MaybeSend + MaybeSync {
41    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture;
42
43    /// A suspended live receiver must not start automatic unary network work.
44    fn is_suspended(&self) -> bool {
45        false
46    }
47}
48
49/// Shared client runtime. Transport capability and limits are fixed at construction.
50/// The controller slot orders its final release with the next reader acquisition.
51#[derive(Default)]
52pub struct IncomingRuntime {
53    policy: super::policy::StreamPolicy,
54    pub(crate) factory: Option<Arc<dyn SubscriptionFactory>>,
55    #[cfg(any(test, feature = "test-utils"))]
56    pub(crate) original_factory: Option<Arc<dyn SubscriptionFactory>>,
57    pub(crate) coordinator: Mutex<Option<Arc<IncomingCoordinator>>>,
58    /// A new controller on this client cannot resend a rejected wire request.
59    pub(crate) rejected_requests: RejectedRequests,
60    /// A closed reader whose owner token must be released after storage repair.
61    pub(crate) retired_delivery_owner: Mutex<Option<xmtp_db::delivery::DeliveryOwner>>,
62}
63
64impl IncomingRuntime {
65    pub(crate) fn new(
66        policy: super::policy::StreamPolicy,
67        factory: Option<Arc<dyn SubscriptionFactory>>,
68    ) -> Self {
69        Self {
70            policy,
71            #[cfg(any(test, feature = "test-utils"))]
72            original_factory: factory.clone(),
73            factory,
74            coordinator: Mutex::new(None),
75            rejected_requests: Arc::new(Mutex::new(Vec::new())),
76            retired_delivery_owner: Mutex::new(None),
77        }
78    }
79
80    pub(crate) fn with_preflight<A: xmtp_proto::api_client::XmtpBackendClient + 'static>(
81        mut self,
82        api: xmtp_api::preflight::GuardedApi<A>,
83    ) -> Self {
84        self.factory = self
85            .factory
86            .map(|inner| Arc::new(GuardedFactory { inner, api }) as Arc<dyn SubscriptionFactory>);
87        self
88    }
89
90    pub(crate) fn policy(&self) -> &super::policy::StreamPolicy {
91        &self.policy
92    }
93
94    #[cfg(any(test, feature = "test-utils"))]
95    pub fn active_lease_count_for_test(&self) -> usize {
96        self.coordinator
97            .lock()
98            .as_ref()
99            .map_or(0, |coordinator| coordinator.state.statuses.lock().len())
100    }
101}
102
103struct GuardedFactory<A> {
104    inner: Arc<dyn SubscriptionFactory>,
105    api: xmtp_api::preflight::GuardedApi<A>,
106}
107impl<A: xmtp_proto::api_client::XmtpBackendClient + 'static> SubscriptionFactory
108    for GuardedFactory<A>
109{
110    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
111        let inner = self.inner.clone();
112        let api = self.api.clone();
113        Box::pin(async move {
114            api.check_preflight().await.map_err(NetworkError::new)?;
115            inner.open(cursors, limits).await
116        })
117    }
118    fn is_suspended(&self) -> bool {
119        self.inner.is_suspended()
120    }
121}
122
123impl<F> SubscriptionFactory for F
124where
125    F: Fn(TopicCursor, IncomingBatchLimits) -> SubscriptionFuture + MaybeSend + MaybeSync,
126{
127    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
128        self(cursors, limits)
129    }
130}
131
132xmtp_common::if_native! {
133pub(crate) struct BidiSubscriptionFactory<A> {
134    api: A,
135    transport: std::sync::OnceLock<Arc<xmtp_api_backend::BidiTransport<xmtp_api_backend::BackendBinding>>>,
136}
137
138impl<A> BidiSubscriptionFactory<A> {
139    pub(crate) fn new(api: A) -> Self {
140        Self { api, transport: std::sync::OnceLock::new() }
141    }
142}
143
144impl<A> SubscriptionFactory for BidiSubscriptionFactory<A>
145where
146    A: xmtp_proto::api_client::XmtpMlsBidiStreams
147        + super::router_callbacks::ApiClientIdentity
148        + Clone
149        + Send
150        + Sync
151        + 'static,
152    A::SubscribeStream: 'static,
153{
154    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
155        // Keep the shared transport alive between this client's receiving calls.
156        // The process registry must not keep credentials after factories and subscriptions drop.
157        let transport = self.transport.get_or_init(|| {
158            super::router_callbacks::shared_transport(self.api.clone())
159        }).clone();
160        Box::pin(async move {
161            transport
162                .lease_ordered(
163                    cursors
164                        .into_iter()
165                        .map(|(topic, cursor)| (topic, cursor.0))
166                        .collect(),
167                    xmtp_api_backend::DEFAULT_LEASE_DEPTH,
168                    limits,
169                )
170                .await
171                .map(|lease| {
172                    lease.into_incoming_subscription().map_error(move |error| {
173                        // The event stream retains this exact shared transport until it drops.
174                        let _owner = &transport;
175                        NetworkError::new(error)
176                    })
177                })
178                .map_err(NetworkError::new)
179        })
180    }
181
182    fn is_suspended(&self) -> bool {
183        super::router_callbacks::bidi_streams_suspended()
184    }
185}
186}
187
188/// When a fixed-target operation may query beyond durable receipt.
189#[derive(Clone, Copy, Debug, PartialEq, Eq)]
190pub enum IncomingReceivePolicy {
191    /// Give a healthy receiver one bounded wait before using Query fallback.
192    StreamFirst,
193    /// Query immediately when receipt is below the target, even with a healthy receiver.
194    ImmediateQuery,
195}
196
197/// Network interest held by one reader or bounded sync run.
198#[derive(Clone, Debug)]
199pub enum IncomingScope {
200    /// Fixed topic membership; other stored groups are not enrolled.
201    Topics(Vec<Topic>),
202    /// Live group interest. A removed group can recover through its next Welcome.
203    /// Welcome recovery does not add status targets or widen local delivery.
204    Groups(Vec<GroupId>),
205    /// Uses caller-sampled targets and one deadline across receiver changes.
206    Barrier {
207        /// Caller-sampled heads H; reconnects must not replace them.
208        targets: TopicCursor,
209        /// One absolute deadline, including target capture and processing.
210        deadline: Instant,
211        /// This operation's preference does not restrict reads required by other operations.
212        receive_policy: IncomingReceivePolicy,
213    },
214    /// Includes the installation's welcome topic and every stored group.
215    AllGroups,
216    /// Keeps device-sync groups and new installation Welcomes receiving without app delivery.
217    DeviceSyncGroups,
218}
219
220/// Shares receipt and ordered processing across all readers in one context.
221/// Database transactions provide the cross-process safety boundary.
222pub struct IncomingCoordinator {
223    commands: mpsc::UnboundedSender<Command>,
224    generations: AtomicU64,
225    state: Arc<SharedState>,
226}
227
228struct SharedState {
229    statuses: Mutex<HashMap<u64, IncomingStatus>>,
230    recovery: Arc<Mutex<super::recovery::RecoverySnapshot>>,
231    consumer_recovery: Mutex<HashMap<u64, super::recovery::RecoveryState>>,
232    changed: watch::Sender<u64>,
233    connection_states: ConnectionStates,
234}
235
236impl Default for SharedState {
237    fn default() -> Self {
238        let (changed, _) = watch::channel(0);
239        Self {
240            statuses: Mutex::new(HashMap::new()),
241            recovery: Arc::new(Mutex::new(Default::default())),
242            consumer_recovery: Mutex::new(HashMap::new()),
243            changed,
244            connection_states: ConnectionStates::default(),
245        }
246    }
247}
248
249impl SharedState {
250    fn notify(&self) {
251        self.changed
252            .send_modify(|revision| *revision = revision.wrapping_add(1));
253    }
254}
255
256enum Command {
257    Acquire {
258        id: u64,
259        scope: IncomingScope,
260    },
261    Replace {
262        id: u64,
263        generation: u64,
264        scope: IncomingScope,
265    },
266    Release(u64),
267    Wake,
268}
269
270impl IncomingCoordinator {
271    /// Reuse the context's live controller, or start one with an owned context handle.
272    pub fn for_context<C: XmtpSharedContext>(context: &C) -> Arc<Self> {
273        let mut slot = context.incoming_runtime().coordinator.lock();
274        if let Some(coordinator) = slot
275            .as_ref()
276            .filter(|coordinator| !coordinator.commands.is_closed())
277        {
278            return coordinator.clone();
279        }
280        let (commands, receiver) = mpsc::unbounded_channel();
281        let state = Arc::new(SharedState::default());
282        state
283            .connection_states
284            .set_writer(Arc::new(xmtp_events::PublicBusWriter::new(
285                context.events(),
286            )));
287        let coordinator = Arc::new(Self {
288            commands,
289            generations: AtomicU64::new(0),
290            state: state.clone(),
291        });
292        let context = context.context_ref().clone();
293        xmtp_common::spawn(None, Controller::new(context, receiver, state).run());
294        *slot = Some(coordinator.clone());
295        coordinator
296    }
297
298    /// Keep this scope active until the returned lease closes or drops.
299    pub fn acquire(self: &Arc<Self>, scope: IncomingScope) -> IncomingLease {
300        self.acquire_with_recovery(scope, false)
301    }
302
303    /// Application streams own a finite outage budget. Internal workers and
304    /// bounded sync operations keep their existing lifetime and deadline rules.
305    pub(crate) fn acquire_stream(self: &Arc<Self>, scope: IncomingScope) -> IncomingLease {
306        self.acquire_with_recovery(scope, true)
307    }
308
309    fn acquire_with_recovery(
310        self: &Arc<Self>,
311        scope: IncomingScope,
312        bounded: bool,
313    ) -> IncomingLease {
314        let id = self.generations.fetch_add(1, Ordering::Relaxed) + 1;
315        if bounded {
316            self.state.connection_states.open(id);
317        }
318        self.state
319            .statuses
320            .lock()
321            .insert(id, IncomingStatus::pending(id));
322        let recovery = self.state.recovery.lock().clone();
323        self.state.consumer_recovery.lock().insert(
324            id,
325            super::recovery::RecoveryState::new(recovery, bounded, Instant::now()),
326        );
327        let _ = self.commands.send(Command::Acquire { id, scope });
328        IncomingLease {
329            id,
330            coordinator: self.clone(),
331            changes: tokio::sync::Mutex::new(self.state.changed.subscribe()),
332            closed: std::sync::atomic::AtomicBool::new(false),
333        }
334    }
335
336    /// Request a fresh database check. This hint is not proof of processing.
337    pub fn wake(&self) {
338        let _ = self.commands.send(Command::Wake);
339    }
340}
341
342/// Holds network interest and scope status; application acknowledgement is separate.
343pub struct IncomingLease {
344    id: u64,
345    coordinator: Arc<IncomingCoordinator>,
346    changes: tokio::sync::Mutex<watch::Receiver<u64>>,
347    closed: std::sync::atomic::AtomicBool,
348}
349
350impl IncomingLease {
351    #[cfg(test)]
352    pub(crate) fn recovery_snapshot(&self) -> super::recovery::RecoverySnapshot {
353        self.coordinator
354            .state
355            .consumer_recovery
356            .lock()
357            .get(&self.id)
358            .map(|state| state.snapshot.clone())
359            .unwrap_or_default()
360    }
361
362    pub(crate) fn check_recovery(&self) -> Result<(), super::recovery::RecoveryFailure> {
363        let mut consumers = self.coordinator.state.consumer_recovery.lock();
364        match consumers.get_mut(&self.id) {
365            Some(state) => state.check(Instant::now()),
366            None => Ok(()),
367        }
368    }
369
370    #[cfg(test)]
371    pub(crate) fn fail_recovery_for_test(&self, failure: super::recovery::RecoveryFailure) {
372        self.coordinator
373            .state
374            .consumer_recovery
375            .lock()
376            .get_mut(&self.id)
377            .expect("test lease must have a recovery snapshot")
378            .snapshot
379            .terminal = Some(failure);
380        self.coordinator.state.notify();
381    }
382
383    /// Read the latest status for this lease's current scope generation.
384    pub fn snapshot(&self) -> IncomingStatus {
385        self.coordinator
386            .state
387            .statuses
388            .lock()
389            .get(&self.id)
390            .cloned()
391            .unwrap_or_else(|| IncomingStatus::cancelled(self.id))
392    }
393
394    /// Return the old obligations as cancelled before starting the new generation.
395    pub fn replace_scope(&self, scope: IncomingScope) -> IncomingStatus {
396        let mut statuses = self.coordinator.state.statuses.lock();
397        if self.closed.load(Ordering::Acquire) {
398            return statuses
399                .get(&self.id)
400                .cloned()
401                .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
402        }
403        let generation = self.coordinator.generations.fetch_add(1, Ordering::Relaxed) + 1;
404        let mut previous = statuses
405            .get(&self.id)
406            .cloned()
407            .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
408        previous.cancel();
409        let mut next = IncomingStatus::pending(generation);
410        next.previous = Some(Box::new(previous.without_previous()));
411        statuses.insert(self.id, next);
412        let _ = self.coordinator.commands.send(Command::Replace {
413            id: self.id,
414            generation,
415            scope,
416        });
417        drop(statuses);
418        self.coordinator.state.notify();
419        previous
420    }
421
422    pub fn replace_topics(&self, topics: Vec<Topic>) -> IncomingStatus {
423        self.replace_scope(IncomingScope::Topics(topics))
424    }
425
426    /// A notification is a hint. Read a fresh snapshot after it arrives.
427    pub async fn changed(&self) {
428        let mut changes = self.changes.lock().await;
429        if self.closed.load(Ordering::Acquire) {
430            return;
431        }
432        let _ = changes.changed().await;
433    }
434
435    /// Subscribe before reading a status so an update cannot be missed.
436    pub fn subscribe_changes(&self) -> watch::Receiver<u64> {
437        self.coordinator.state.changed.subscribe()
438    }
439
440    #[cfg(test)]
441    pub(crate) fn notify_change_for_test(&self) {
442        self.coordinator.state.notify();
443    }
444
445    /// Hold the shared observer to test that another reader uses its own observer.
446    #[cfg(any(test, feature = "test-utils"))]
447    pub async fn lock_change_receiver_for_test(
448        &self,
449    ) -> tokio::sync::MutexGuard<'_, watch::Receiver<u64>> {
450        self.changes.lock().await
451    }
452
453    /// Release interest even if another task still holds this lease to watch status.
454    pub fn close(&self) {
455        let mut statuses = self.coordinator.state.statuses.lock();
456        if self.closed.swap(true, Ordering::AcqRel) {
457            return;
458        }
459        self.coordinator.state.connection_states.close(self.id);
460        statuses.remove(&self.id);
461        self.coordinator
462            .state
463            .consumer_recovery
464            .lock()
465            .remove(&self.id);
466        let _ = self.coordinator.commands.send(Command::Release(self.id));
467        drop(statuses);
468        self.coordinator.state.notify();
469    }
470}
471
472impl Drop for IncomingLease {
473    fn drop(&mut self) {
474        self.close();
475    }
476}
477
478#[cfg(test)]
479mod lease_observer_tests {
480    use super::*;
481
482    // verifies: PROC-018
483    #[xmtp_common::test(unwrap_try = true)]
484    fn replacement_failure_keeps_cancelled_scope_generation() {
485        use crate::subscriptions::barrier::{
486            BarrierCause, BarrierError, BarrierFailure, BarrierTopic,
487        };
488        use xmtp_proto::types::Cursor;
489
490        let (commands, _receiver) = mpsc::unbounded_channel();
491        let coordinator = Arc::new(IncomingCoordinator {
492            commands,
493            generations: AtomicU64::new(0),
494            state: Arc::new(SharedState::default()),
495        });
496        let sibling = coordinator.acquire(IncomingScope::AllGroups);
497        let topic = Topic::new_group_message([7; 32]);
498        let lease = coordinator.acquire(IncomingScope::Topics(vec![topic.clone()]));
499        let old_generation = lease.snapshot().scope_generation;
500        coordinator
501            .state
502            .statuses
503            .lock()
504            .get_mut(&lease.id)?
505            .topics
506            .push(IncomingTopicStatus {
507                topic: topic.clone(),
508                scope_generation: old_generation,
509                registration: IncomingRegistration::Active,
510                target: Some(Cursor(9)),
511                received: Cursor(9),
512                processed: Cursor(8),
513                unresolved_welcomes: 0,
514                processing: IncomingProcessing::Pending,
515                blocked: None,
516                error: None,
517            });
518        let previous = lease.replace_scope(IncomingScope::Topics(vec![Topic::new_group_message(
519            [8; 32],
520        )]));
521        let current = lease.snapshot();
522        assert_eq!(previous.scope_generation, old_generation);
523        assert_eq!(previous.processing, IncomingProcessing::Cancelled);
524        assert_eq!(previous.topics[0].scope_generation, old_generation);
525        assert_eq!(previous.topics[0].processing, IncomingProcessing::Cancelled);
526        assert!(current.scope_generation > old_generation);
527        assert_eq!(current.previous.as_ref()?.scope_generation, old_generation);
528        let mut unfinished = BarrierTopic {
529            topic,
530            scope_generation: None,
531            target: Some(Cursor(9)),
532            received: Cursor(9),
533            processed: Cursor(8),
534            unresolved_welcomes: Vec::new(),
535            inactive: false,
536            cause: Some(BarrierCause::ProcessingPending),
537        };
538        unfinished.capture_scope_generation(&previous);
539        let details = crate::subscriptions::stream_failure::StreamBarrierFailure::from(
540            &BarrierError::Incomplete {
541                reason: BarrierFailure::Cancelled,
542                unfinished: vec![unfinished],
543            },
544        );
545        let reported = &details.unfinished[0];
546        assert_eq!(
547            reported.scope_generation.as_deref(),
548            Some(old_generation.to_string().as_str())
549        );
550        assert_ne!(
551            reported.scope_generation.as_deref(),
552            Some(current.scope_generation.to_string().as_str())
553        );
554        assert_eq!(reported.target.as_deref(), Some("9"));
555        assert_eq!(reported.received, "9");
556        assert_eq!(reported.processed, "8");
557        lease.close();
558        assert_eq!(sibling.snapshot().processing, IncomingProcessing::Pending);
559        assert!(coordinator.state.statuses.lock().contains_key(&sibling.id));
560    }
561
562    #[xmtp_common::test(unwrap_try = true)]
563    async fn lease_change_after_snapshot_is_not_lost() {
564        let (commands, _receiver) = mpsc::unbounded_channel();
565        let coordinator = Arc::new(IncomingCoordinator {
566            commands,
567            generations: AtomicU64::new(0),
568            state: Arc::new(SharedState::default()),
569        });
570        let lease = coordinator.acquire(IncomingScope::AllGroups);
571        let _status = lease.snapshot();
572        coordinator.state.notify();
573        xmtp_common::time::timeout(std::time::Duration::from_secs(1), lease.changed()).await?;
574    }
575
576    #[xmtp_common::test(unwrap_try = true)]
577    async fn lease_change_wakes_independent_observers() {
578        let (commands, _receiver) = mpsc::unbounded_channel();
579        let coordinator = Arc::new(IncomingCoordinator {
580            commands,
581            generations: AtomicU64::new(0),
582            state: Arc::new(SharedState::default()),
583        });
584        let lease = coordinator.acquire(IncomingScope::AllGroups);
585        let mut first_changes = lease.subscribe_changes();
586        let mut second_changes = lease.subscribe_changes();
587        let _status = lease.snapshot();
588        coordinator.state.notify();
589        let (first, second) =
590            xmtp_common::time::timeout(std::time::Duration::from_secs(1), async {
591                futures::join!(first_changes.changed(), second_changes.changed())
592            })
593            .await?;
594        first?;
595        second?;
596    }
597}
xmtp_mls/subscriptions/incoming.rs

xmtp_mls/subscriptions/
incoming.rs
