1mod 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 fn is_suspended(&self) -> bool {
32 false
33 }
34}
35
36#[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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
116pub enum IncomingReceivePolicy {
117 StreamFirst,
119 ImmediateQuery,
121}
122
123#[derive(Clone, Debug)]
125pub enum IncomingScope {
126 Topics(Vec<Topic>),
128 Groups(Vec<GroupId>),
131 Barrier {
133 targets: TopicCursor,
135 deadline: Instant,
137 receive_policy: IncomingReceivePolicy,
139 },
140 AllGroups,
142 DeviceSyncGroups,
144}
145
146pub 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 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 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 pub fn wake(&self) {
231 let _ = self.commands.send(Command::Wake);
232 }
233}
234
235pub 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 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 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 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 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}