Skip to main content

xmtp_mls/subscriptions/
incoming.rs

1//! One receiver and processing scheduler for each client context.
2
3mod controller;
4mod status;
5
6use crate::context::XmtpSharedContext;
7use controller::Controller;
8use parking_lot::Mutex;
9pub use status::*;
10use std::{
11    collections::HashMap,
12    sync::{
13        Arc,
14        atomic::{AtomicU64, Ordering},
15    },
16};
17use tokio::sync::{mpsc, watch};
18use xmtp_common::{BoxDynFuture, MaybeSend, MaybeSync, time::Instant};
19use xmtp_proto::{
20    api::NetworkError,
21    types::{GroupId, IncomingBatchLimits, IncomingSubscription, Topic, TopicCursor},
22};
23
24pub(crate) type SubscriptionFuture =
25    BoxDynFuture<'static, Result<IncomingSubscription<NetworkError>, NetworkError>>;
26
27pub(crate) trait SubscriptionFactory: MaybeSend + MaybeSync {
28    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture;
29
30    /// A suspended live receiver must not start automatic unary network work.
31    fn is_suspended(&self) -> bool {
32        false
33    }
34}
35
36/// Shared client runtime. Transport capability and limits are fixed at construction.
37/// The controller slot orders its final release with the next reader acquisition.
38#[derive(Default)]
39pub struct IncomingRuntime {
40    policy: super::policy::StreamPolicy,
41    pub(crate) factory: Option<Arc<dyn SubscriptionFactory>>,
42    pub(crate) coordinator: Mutex<Option<Arc<IncomingCoordinator>>>,
43}
44
45impl IncomingRuntime {
46    pub(crate) fn new(
47        policy: super::policy::StreamPolicy,
48        factory: Option<Arc<dyn SubscriptionFactory>>,
49    ) -> Self {
50        Self {
51            policy,
52            factory,
53            coordinator: Mutex::new(None),
54        }
55    }
56
57    pub(crate) fn policy(&self) -> &super::policy::StreamPolicy {
58        &self.policy
59    }
60}
61
62impl<F> SubscriptionFactory for F
63where
64    F: Fn(TopicCursor, IncomingBatchLimits) -> SubscriptionFuture + MaybeSend + MaybeSync,
65{
66    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
67        self(cursors, limits)
68    }
69}
70
71xmtp_common::if_native! {
72pub(crate) struct BidiSubscriptionFactory<A> {
73    pub(crate) api: A,
74}
75
76impl<A> SubscriptionFactory for BidiSubscriptionFactory<A>
77where
78    A: xmtp_proto::api_client::XmtpMlsBidiStreams
79        + super::router_callbacks::ApiClientIdentity
80        + Clone
81        + Send
82        + Sync
83        + 'static,
84    A::SubscribeStream: 'static,
85{
86    fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
87        let api = self.api.clone();
88        Box::pin(async move {
89            super::router_callbacks::shared_transport(api)
90                .lease_ordered(
91                    cursors
92                        .into_iter()
93                        .map(|(topic, cursor)| (topic, cursor.0))
94                        .collect(),
95                    xmtp_api_backend::DEFAULT_LEASE_DEPTH,
96                    limits,
97                )
98                .await
99                .map(|lease| {
100                    lease
101                        .into_incoming_subscription()
102                        .map_error(NetworkError::new)
103                })
104                .map_err(NetworkError::new)
105        })
106    }
107
108    fn is_suspended(&self) -> bool {
109        super::router_callbacks::bidi_streams_suspended()
110    }
111}
112}
113
114/// When a fixed-target operation may query beyond durable receipt.
115#[derive(Clone, Copy, Debug, PartialEq, Eq)]
116pub enum IncomingReceivePolicy {
117    /// Give a healthy receiver one bounded wait before using Query fallback.
118    StreamFirst,
119    /// Query immediately when receipt is below the target, even with a healthy receiver.
120    ImmediateQuery,
121}
122
123/// Network interest held by one reader or bounded sync run.
124#[derive(Clone, Debug)]
125pub enum IncomingScope {
126    /// Fixed topic membership; other stored groups are not enrolled.
127    Topics(Vec<Topic>),
128    /// Live group interest. A removed group can recover through its next Welcome.
129    /// Welcome recovery does not add status targets or widen local delivery.
130    Groups(Vec<GroupId>),
131    /// Uses caller-sampled targets and one deadline across receiver changes.
132    Barrier {
133        /// Caller-sampled heads H; reconnects must not replace them.
134        targets: TopicCursor,
135        /// One absolute deadline, including target capture and processing.
136        deadline: Instant,
137        /// This operation's preference does not restrict reads required by other operations.
138        receive_policy: IncomingReceivePolicy,
139    },
140    /// Includes the installation's welcome topic and every stored group.
141    AllGroups,
142    /// Keeps device-sync groups and new installation Welcomes receiving without app delivery.
143    DeviceSyncGroups,
144}
145
146/// Shares receipt and ordered processing across all readers in one context.
147/// Database transactions provide the cross-process safety boundary.
148pub struct IncomingCoordinator {
149    commands: mpsc::UnboundedSender<Command>,
150    generations: AtomicU64,
151    state: Arc<SharedState>,
152}
153
154struct SharedState {
155    statuses: Mutex<HashMap<u64, IncomingStatus>>,
156    changed: watch::Sender<u64>,
157}
158
159impl Default for SharedState {
160    fn default() -> Self {
161        let (changed, _) = watch::channel(0);
162        Self {
163            statuses: Mutex::new(HashMap::new()),
164            changed,
165        }
166    }
167}
168
169impl SharedState {
170    fn notify(&self) {
171        self.changed
172            .send_modify(|revision| *revision = revision.wrapping_add(1));
173    }
174}
175
176enum Command {
177    Acquire {
178        id: u64,
179        scope: IncomingScope,
180    },
181    Replace {
182        id: u64,
183        generation: u64,
184        scope: IncomingScope,
185    },
186    Release(u64),
187    Wake,
188}
189
190impl IncomingCoordinator {
191    /// Reuse the context's live controller, or start one with an owned context handle.
192    pub fn for_context<C: XmtpSharedContext>(context: &C) -> Arc<Self> {
193        let mut slot = context.incoming_runtime().coordinator.lock();
194        if let Some(coordinator) = slot
195            .as_ref()
196            .filter(|coordinator| !coordinator.commands.is_closed())
197        {
198            return coordinator.clone();
199        }
200        let (commands, receiver) = mpsc::unbounded_channel();
201        let state = Arc::new(SharedState::default());
202        let coordinator = Arc::new(Self {
203            commands,
204            generations: AtomicU64::new(0),
205            state: state.clone(),
206        });
207        let context = context.context_ref().clone();
208        xmtp_common::spawn(None, Controller::new(context, receiver, state).run());
209        *slot = Some(coordinator.clone());
210        coordinator
211    }
212
213    /// Keep this scope active until the returned lease closes or drops.
214    pub fn acquire(self: &Arc<Self>, scope: IncomingScope) -> IncomingLease {
215        let id = self.generations.fetch_add(1, Ordering::Relaxed) + 1;
216        self.state
217            .statuses
218            .lock()
219            .insert(id, IncomingStatus::pending(id));
220        let _ = self.commands.send(Command::Acquire { id, scope });
221        IncomingLease {
222            id,
223            coordinator: self.clone(),
224            changes: tokio::sync::Mutex::new(self.state.changed.subscribe()),
225            closed: std::sync::atomic::AtomicBool::new(false),
226        }
227    }
228
229    /// Request a fresh database check. This hint is not proof of processing.
230    pub fn wake(&self) {
231        let _ = self.commands.send(Command::Wake);
232    }
233}
234
235/// Holds network interest and scope status; application acknowledgement is separate.
236pub struct IncomingLease {
237    id: u64,
238    coordinator: Arc<IncomingCoordinator>,
239    changes: tokio::sync::Mutex<watch::Receiver<u64>>,
240    closed: std::sync::atomic::AtomicBool,
241}
242
243impl IncomingLease {
244    /// Read the latest status for this lease's current scope generation.
245    pub fn snapshot(&self) -> IncomingStatus {
246        self.coordinator
247            .state
248            .statuses
249            .lock()
250            .get(&self.id)
251            .cloned()
252            .unwrap_or_else(|| IncomingStatus::cancelled(self.id))
253    }
254
255    /// Return the old obligations as cancelled before starting the new generation.
256    pub fn replace_scope(&self, scope: IncomingScope) -> IncomingStatus {
257        let mut statuses = self.coordinator.state.statuses.lock();
258        if self.closed.load(Ordering::Acquire) {
259            return statuses
260                .get(&self.id)
261                .cloned()
262                .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
263        }
264        let generation = self.coordinator.generations.fetch_add(1, Ordering::Relaxed) + 1;
265        let mut previous = statuses
266            .get(&self.id)
267            .cloned()
268            .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
269        previous.cancel();
270        let mut next = IncomingStatus::pending(generation);
271        next.previous = Some(Box::new(previous.without_previous()));
272        statuses.insert(self.id, next);
273        let _ = self.coordinator.commands.send(Command::Replace {
274            id: self.id,
275            generation,
276            scope,
277        });
278        drop(statuses);
279        self.coordinator.state.notify();
280        previous
281    }
282
283    pub fn replace_topics(&self, topics: Vec<Topic>) -> IncomingStatus {
284        self.replace_scope(IncomingScope::Topics(topics))
285    }
286
287    /// A notification is a hint. Read a fresh snapshot after it arrives.
288    pub async fn changed(&self) {
289        let mut changes = self.changes.lock().await;
290        if self.closed.load(Ordering::Acquire) {
291            return;
292        }
293        let _ = changes.changed().await;
294    }
295
296    /// Release interest even if another task still holds this lease to watch status.
297    pub fn close(&self) {
298        let mut statuses = self.coordinator.state.statuses.lock();
299        if self.closed.swap(true, Ordering::AcqRel) {
300            return;
301        }
302        statuses.remove(&self.id);
303        let _ = self.coordinator.commands.send(Command::Release(self.id));
304        drop(statuses);
305        self.coordinator.state.notify();
306    }
307}
308
309impl Drop for IncomingLease {
310    fn drop(&mut self) {
311        self.close();
312    }
313}