1mod controller;
4mod status;
5pub mod stream_api;
6use super::connection_state::ConnectionStates;
7
8use crate::context::XmtpSharedContext;
9use controller::Controller;
10use parking_lot::Mutex;
11pub use status::*;
12use std::{
13 collections::{HashMap, HashSet},
14 sync::{
15 Arc,
16 atomic::{AtomicU64, Ordering},
17 },
18};
19use tokio::sync::{mpsc, watch};
20use xmtp_common::{BoxDynFuture, MaybeSend, MaybeSync, time::Instant};
21use xmtp_proto::{
22 api::NetworkError,
23 types::{Cursor, GroupId, IncomingBatchLimits, IncomingSubscription, Topic, TopicCursor},
24};
25
26pub(crate) type SubscriptionFuture =
27 BoxDynFuture<'static, Result<IncomingSubscription<NetworkError>, NetworkError>>;
28
29#[derive(Clone, Debug, PartialEq, Eq)]
31pub(crate) enum RequestKey {
32 Bidi(TopicCursor),
33 QueryNewest(HashSet<Topic>),
34 Query(Topic, Cursor),
37}
38
39pub(crate) type RejectedRequests = Arc<Mutex<Vec<(RequestKey, Arc<IncomingError>)>>>;
40
41pub(crate) trait SubscriptionFactory: MaybeSend + MaybeSync {
42 fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture;
43
44 fn bind_limits(
46 &self,
47 _limits: &xmtp_configuration::LimitsConfiguration,
48 ) -> Option<Arc<dyn SubscriptionFactory>> {
49 None
50 }
51
52 fn is_suspended(&self) -> bool {
54 false
55 }
56}
57
58#[derive(Debug, Default)]
63pub(crate) struct RateLimitHold {
64 opens: Option<Instant>,
65}
66
67pub(crate) type SharedRateLimitHold = Arc<Mutex<RateLimitHold>>;
68
69impl RateLimitHold {
70 fn until(delay: xmtp_common::time::Duration) -> Instant {
71 xmtp_common::time::saturating_after(Instant::now(), delay)
72 }
73
74 pub(crate) fn hold_opens(&mut self, delay: xmtp_common::time::Duration) {
76 let until = Self::until(delay);
77 self.opens = Some(self.opens.map_or(until, |at| at.max(until)));
78 }
79
80 pub(crate) fn opens_until(&self) -> Option<Instant> {
82 self.opens.filter(|until| Instant::now() < *until)
83 }
84}
85
86#[derive(Default)]
89pub struct IncomingRuntime {
90 policy: super::policy::StreamPolicy,
91 pub(crate) factory: Option<Arc<dyn SubscriptionFactory>>,
92 #[cfg(any(test, feature = "test-utils"))]
93 pub(crate) original_factory: Option<Arc<dyn SubscriptionFactory>>,
94 pub(crate) coordinator: Mutex<Option<Arc<IncomingCoordinator>>>,
95 pub(crate) rejected_requests: RejectedRequests,
97 pub(crate) rate_limit_hold: SharedRateLimitHold,
99}
100
101impl IncomingRuntime {
102 pub(crate) fn new(
103 policy: super::policy::StreamPolicy,
104 factory: Option<Arc<dyn SubscriptionFactory>>,
105 ) -> Self {
106 Self {
107 policy,
108 #[cfg(any(test, feature = "test-utils"))]
109 original_factory: factory.clone(),
110 factory,
111 coordinator: Mutex::new(None),
112 rejected_requests: Arc::new(Mutex::new(Vec::new())),
113 rate_limit_hold: Arc::default(),
114 }
115 }
116
117 pub(crate) fn with_preflight<A: xmtp_proto::api_client::XmtpBackendClient + 'static>(
118 mut self,
119 api: xmtp_api::preflight::GuardedApi<A>,
120 ) -> Self {
121 self.factory = self
122 .factory
123 .map(|inner| Arc::new(GuardedFactory { inner, api }) as Arc<dyn SubscriptionFactory>);
124 self
125 }
126
127 pub(crate) fn policy(&self) -> &super::policy::StreamPolicy {
128 &self.policy
129 }
130
131 #[cfg(any(test, feature = "test-utils"))]
132 pub fn active_lease_count_for_test(&self) -> usize {
133 self.coordinator
134 .lock()
135 .as_ref()
136 .map_or(0, |coordinator| coordinator.state.statuses.lock().len())
137 }
138}
139
140struct GuardedFactory<A> {
141 inner: Arc<dyn SubscriptionFactory>,
142 api: xmtp_api::preflight::GuardedApi<A>,
143}
144impl<A: xmtp_proto::api_client::XmtpBackendClient + 'static> SubscriptionFactory
145 for GuardedFactory<A>
146{
147 fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
148 let inner = self.inner.clone();
149 let api = self.api.clone();
150 Box::pin(async move {
151 api.check_preflight().await.map_err(NetworkError::new)?;
152 inner.open(cursors, limits).await
153 })
154 }
155 fn is_suspended(&self) -> bool {
156 self.inner.is_suspended()
157 }
158}
159
160xmtp_common::if_wasm! {
161pub(crate) fn streaming_factory<A>(api: A) -> Arc<dyn SubscriptionFactory>
164where
165 A: xmtp_proto::api_client::XmtpMlsStreams + Clone + 'static,
166{
167 Arc::new(move |cursors: TopicCursor, limits| -> SubscriptionFuture {
168 let api = api.clone();
169 Box::pin(async move {
170 api.subscribe_envelopes_with_cursors(&cursors, limits)
171 .await
172 .map(|subscription| subscription.map_error(NetworkError::new))
173 .map_err(NetworkError::new)
174 })
175 })
176}
177}
178
179impl<F> SubscriptionFactory for F
180where
181 F: Fn(TopicCursor, IncomingBatchLimits) -> SubscriptionFuture + MaybeSend + MaybeSync,
182{
183 fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
184 self(cursors, limits)
185 }
186}
187
188xmtp_common::if_native! {
189pub(crate) struct BidiSubscriptionFactory<A> {
190 api: A,
191 mutate: Option<xmtp_api_backend::MutateLimits>,
192 transport: std::sync::OnceLock<Arc<xmtp_api_backend::BidiTransport<xmtp_api_backend::BackendBinding>>>,
193}
194
195impl<A> BidiSubscriptionFactory<A> {
196 pub(crate) fn new(api: A) -> Self {
197 Self { api, mutate: None, transport: std::sync::OnceLock::new() }
198 }
199}
200
201pub(crate) fn streaming_factory<A>(api: A) -> Arc<dyn SubscriptionFactory>
204where
205 BidiSubscriptionFactory<A>: SubscriptionFactory + 'static,
206{
207 Arc::new(BidiSubscriptionFactory::new(api))
208}
209
210impl<A> SubscriptionFactory for BidiSubscriptionFactory<A>
211where
212 A: xmtp_proto::api_client::XmtpMlsBidiStreams
213 + stream_api::shared_wires::ApiClientIdentity
214 + Clone
215 + Send
216 + Sync
217 + 'static,
218 A::SubscribeStream: 'static,
219{
220 fn open(&self, cursors: TopicCursor, limits: IncomingBatchLimits) -> SubscriptionFuture {
221 let transport = self.transport.get_or_init(|| {
224 let mutate = self.mutate.unwrap_or_else(|| {
225 xmtp_api_backend::MutateLimits::from_limits(&self.api.bidi_limits())
226 });
227 stream_api::shared_wires::shared_transport_within(self.api.clone(), mutate)
228 }).clone();
229 Box::pin(async move {
230 transport
231 .lease_ordered(
232 cursors
233 .into_iter()
234 .map(|(topic, cursor)| (topic, cursor.0))
235 .collect(),
236 xmtp_api_backend::DEFAULT_LEASE_DEPTH,
237 limits,
238 )
239 .await
240 .map(|lease| {
241 lease.into_incoming_subscription().map_error(move |error| {
242 let _owner = &transport;
244 NetworkError::new(error)
245 })
246 })
247 .map_err(NetworkError::new)
248 })
249 }
250
251 fn bind_limits(
252 &self,
253 limits: &xmtp_configuration::LimitsConfiguration,
254 ) -> Option<Arc<dyn SubscriptionFactory>> {
255 Some(Arc::new(Self {
256 api: self.api.clone(),
257 mutate: Some(xmtp_api_backend::MutateLimits::from_limits(limits)),
258 transport: std::sync::OnceLock::new(),
259 }))
260 }
261
262 fn is_suspended(&self) -> bool {
263 stream_api::shared_wires::bidi_streams_suspended()
264 }
265}
266}
267
268#[derive(Clone, Copy, Debug, PartialEq, Eq)]
270pub enum IncomingReceivePolicy {
271 StreamFirst,
273 ImmediateQuery,
275}
276
277#[derive(Clone, Debug)]
279pub enum IncomingScope {
280 Topics(Vec<Topic>),
282 Groups(Vec<GroupId>),
285 Barrier {
287 targets: TopicCursor,
289 deadline: Instant,
291 receive_policy: IncomingReceivePolicy,
293 },
294 AllGroups,
296 DeviceSyncGroups,
298}
299
300pub struct IncomingCoordinator {
303 commands: mpsc::UnboundedSender<Command>,
304 generations: AtomicU64,
305 state: Arc<SharedState>,
306}
307
308struct SharedState {
309 statuses: Mutex<HashMap<u64, IncomingStatus>>,
310 recovery: Arc<Mutex<super::recovery::RecoverySnapshot>>,
311 consumer_recovery: Mutex<HashMap<u64, super::recovery::RecoveryState>>,
312 changed: watch::Sender<u64>,
313 connection_states: ConnectionStates,
314}
315
316impl Default for SharedState {
317 fn default() -> Self {
318 let (changed, _) = watch::channel(0);
319 Self {
320 statuses: Mutex::new(HashMap::new()),
321 recovery: Arc::new(Mutex::new(Default::default())),
322 consumer_recovery: Mutex::new(HashMap::new()),
323 changed,
324 connection_states: ConnectionStates::default(),
325 }
326 }
327}
328
329impl SharedState {
330 fn notify(&self) {
331 self.changed
332 .send_modify(|revision| *revision = revision.wrapping_add(1));
333 }
334}
335
336enum Command {
337 Acquire {
338 id: u64,
339 scope: IncomingScope,
340 },
341 Replace {
342 id: u64,
343 generation: u64,
344 scope: IncomingScope,
345 },
346 Release(u64),
347 Wake,
348}
349
350impl IncomingCoordinator {
351 pub fn for_context<C: XmtpSharedContext>(context: &C) -> Arc<Self> {
353 let mut slot = context.incoming_runtime().coordinator.lock();
354 if let Some(coordinator) = slot
355 .as_ref()
356 .filter(|coordinator| !coordinator.commands.is_closed())
357 {
358 return coordinator.clone();
359 }
360 let (commands, receiver) = mpsc::unbounded_channel();
361 let state = Arc::new(SharedState::default());
362 state
363 .connection_states
364 .set_writer(Arc::new(xmtp_events::PublicBusWriter::new(
365 context.events(),
366 )));
367 let coordinator = Arc::new(Self {
368 commands,
369 generations: AtomicU64::new(0),
370 state: state.clone(),
371 });
372 let context = context.context_ref().clone();
373 xmtp_common::spawn(None, Controller::new(context, receiver, state).run());
374 *slot = Some(coordinator.clone());
375 coordinator
376 }
377
378 pub fn acquire(self: &Arc<Self>, scope: IncomingScope) -> IncomingLease {
380 self.acquire_with_recovery(scope, false)
381 }
382
383 fn acquire_with_recovery(
384 self: &Arc<Self>,
385 scope: IncomingScope,
386 bounded: bool,
387 ) -> IncomingLease {
388 let id = self.generations.fetch_add(1, Ordering::Relaxed) + 1;
389 if bounded {
390 self.state.connection_states.open(id);
391 }
392 self.state
393 .statuses
394 .lock()
395 .insert(id, IncomingStatus::pending(id));
396 let recovery = self.state.recovery.lock().clone();
397 self.state.consumer_recovery.lock().insert(
398 id,
399 super::recovery::RecoveryState::new(recovery, bounded, Instant::now()),
400 );
401 let _ = self.commands.send(Command::Acquire { id, scope });
402 IncomingLease {
403 id,
404 coordinator: self.clone(),
405 changes: tokio::sync::Mutex::new(self.state.changed.subscribe()),
406 closed: std::sync::atomic::AtomicBool::new(false),
407 }
408 }
409
410 pub fn wake(&self) {
412 let _ = self.commands.send(Command::Wake);
413 }
414}
415
416pub struct IncomingLease {
418 id: u64,
419 coordinator: Arc<IncomingCoordinator>,
420 changes: tokio::sync::Mutex<watch::Receiver<u64>>,
421 closed: std::sync::atomic::AtomicBool,
422}
423
424impl IncomingLease {
425 #[cfg(test)]
426 pub(crate) fn recovery_snapshot(&self) -> super::recovery::RecoverySnapshot {
427 self.coordinator
428 .state
429 .consumer_recovery
430 .lock()
431 .get(&self.id)
432 .map(|state| state.snapshot.clone())
433 .unwrap_or_default()
434 }
435
436 #[cfg(any(test, feature = "test-hooks"))]
441 #[doc(hidden)]
442 pub fn exhaust_recovery_for_test(
443 &self,
444 failures: u64,
445 outage: xmtp_common::time::Duration,
446 ) -> bool {
447 let now = Instant::now();
448 let mut budget = super::recovery::RecoveryBudget::new(
449 &super::recovery::RecoverySnapshot::default(),
450 now,
451 );
452 let snapshot = super::recovery::RecoverySnapshot {
453 failures,
454 ..Default::default()
455 };
456 let Err(failure) = budget.check(&snapshot, now + outage) else {
457 return false;
458 };
459 let exhausted = matches!(failure, super::recovery::RecoveryFailure::Exhausted { .. });
460 self.coordinator
461 .state
462 .consumer_recovery
463 .lock()
464 .get_mut(&self.id)
465 .expect("test lease must have a recovery snapshot")
466 .snapshot
467 .terminal = Some(failure);
468 self.coordinator.state.notify();
469 exhausted
470 }
471
472 pub fn snapshot(&self) -> IncomingStatus {
474 self.coordinator
475 .state
476 .statuses
477 .lock()
478 .get(&self.id)
479 .cloned()
480 .unwrap_or_else(|| IncomingStatus::cancelled(self.id))
481 }
482
483 pub fn replace_scope(&self, scope: IncomingScope) -> IncomingStatus {
485 let mut statuses = self.coordinator.state.statuses.lock();
486 if self.closed.load(Ordering::Acquire) {
487 return statuses
488 .get(&self.id)
489 .cloned()
490 .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
491 }
492 let generation = self.coordinator.generations.fetch_add(1, Ordering::Relaxed) + 1;
493 let mut previous = statuses
494 .get(&self.id)
495 .cloned()
496 .unwrap_or_else(|| IncomingStatus::cancelled(self.id));
497 previous.cancel();
498 let mut next = IncomingStatus::pending(generation);
499 next.previous = Some(Box::new(previous.without_previous()));
500 statuses.insert(self.id, next);
501 let _ = self.coordinator.commands.send(Command::Replace {
502 id: self.id,
503 generation,
504 scope,
505 });
506 drop(statuses);
507 self.coordinator.state.notify();
508 previous
509 }
510
511 pub fn replace_topics(&self, topics: Vec<Topic>) -> IncomingStatus {
512 self.replace_scope(IncomingScope::Topics(topics))
513 }
514
515 pub async fn changed(&self) {
517 let mut changes = self.changes.lock().await;
518 if self.closed.load(Ordering::Acquire) {
519 return;
520 }
521 let _ = changes.changed().await;
522 }
523
524 pub fn subscribe_changes(&self) -> watch::Receiver<u64> {
526 self.coordinator.state.changed.subscribe()
527 }
528
529 #[cfg(any(test, feature = "test-hooks"))]
531 pub fn notify_change_for_test(&self) {
532 self.coordinator.state.notify();
533 }
534
535 #[cfg(any(test, feature = "test-utils"))]
537 pub async fn lock_change_receiver_for_test(
538 &self,
539 ) -> tokio::sync::MutexGuard<'_, watch::Receiver<u64>> {
540 self.changes.lock().await
541 }
542
543 pub fn close(&self) {
545 let mut statuses = self.coordinator.state.statuses.lock();
546 if self.closed.swap(true, Ordering::AcqRel) {
547 return;
548 }
549 self.coordinator.state.connection_states.close(self.id);
550 statuses.remove(&self.id);
551 self.coordinator
552 .state
553 .consumer_recovery
554 .lock()
555 .remove(&self.id);
556 let _ = self.coordinator.commands.send(Command::Release(self.id));
557 drop(statuses);
558 self.coordinator.state.notify();
559 }
560}
561
562impl Drop for IncomingLease {
563 fn drop(&mut self) {
564 self.close();
565 }
566}
567
568#[cfg(test)]
569mod lease_observer_tests {
570 use super::*;
571
572 #[xmtp_common::test(unwrap_try = true)]
574 fn replacement_failure_keeps_cancelled_scope_generation() {
575 use crate::subscriptions::barrier::{
576 BarrierCause, BarrierError, BarrierFailure, BarrierTopic,
577 };
578 use xmtp_proto::types::Cursor;
579
580 let (commands, _receiver) = mpsc::unbounded_channel();
581 let coordinator = Arc::new(IncomingCoordinator {
582 commands,
583 generations: AtomicU64::new(0),
584 state: Arc::new(SharedState::default()),
585 });
586 let sibling = coordinator.acquire(IncomingScope::AllGroups);
587 let topic = Topic::new_group_message([7; 32]);
588 let lease = coordinator.acquire(IncomingScope::Topics(vec![topic.clone()]));
589 let old_generation = lease.snapshot().scope_generation;
590 coordinator
591 .state
592 .statuses
593 .lock()
594 .get_mut(&lease.id)?
595 .topics
596 .push(IncomingTopicStatus {
597 topic: topic.clone(),
598 scope_generation: old_generation,
599 registration: IncomingRegistration::Active,
600 target: Some(Cursor(9)),
601 received: Cursor(9),
602 processed: Cursor(8),
603 unresolved_welcomes: 0,
604 processing: IncomingProcessing::Pending,
605 blocked: None,
606 error: None,
607 });
608 let previous = lease.replace_scope(IncomingScope::Topics(vec![Topic::new_group_message(
609 [8; 32],
610 )]));
611 let current = lease.snapshot();
612 assert_eq!(previous.scope_generation, old_generation);
613 assert_eq!(previous.processing, IncomingProcessing::Cancelled);
614 assert_eq!(previous.topics[0].scope_generation, old_generation);
615 assert_eq!(previous.topics[0].processing, IncomingProcessing::Cancelled);
616 assert!(current.scope_generation > old_generation);
617 assert_eq!(current.previous.as_ref()?.scope_generation, old_generation);
618 let mut unfinished = BarrierTopic {
619 topic,
620 scope_generation: None,
621 target: Some(Cursor(9)),
622 received: Cursor(9),
623 processed: Cursor(8),
624 unresolved_welcomes: Vec::new(),
625 inactive: false,
626 cause: Some(BarrierCause::ProcessingPending),
627 };
628 unfinished.capture_scope_generation(&previous);
629 let details = crate::subscriptions::stream_failure::StreamBarrierFailure::from(
630 &BarrierError::Incomplete {
631 reason: BarrierFailure::Cancelled,
632 unfinished: vec![unfinished],
633 },
634 );
635 let reported = &details.unfinished[0];
636 assert_eq!(
637 reported.scope_generation.as_deref(),
638 Some(old_generation.to_string().as_str())
639 );
640 assert_ne!(
641 reported.scope_generation.as_deref(),
642 Some(current.scope_generation.to_string().as_str())
643 );
644 assert_eq!(reported.target.as_deref(), Some("9"));
645 assert_eq!(reported.received, "9");
646 assert_eq!(reported.processed, "8");
647 lease.close();
648 assert_eq!(sibling.snapshot().processing, IncomingProcessing::Pending);
649 assert!(coordinator.state.statuses.lock().contains_key(&sibling.id));
650 }
651
652 #[xmtp_common::test(unwrap_try = true)]
653 async fn lease_change_after_snapshot_is_not_lost() {
654 let (commands, _receiver) = mpsc::unbounded_channel();
655 let coordinator = Arc::new(IncomingCoordinator {
656 commands,
657 generations: AtomicU64::new(0),
658 state: Arc::new(SharedState::default()),
659 });
660 let lease = coordinator.acquire(IncomingScope::AllGroups);
661 let _status = lease.snapshot();
662 coordinator.state.notify();
663 xmtp_common::time::timeout(std::time::Duration::from_secs(1), lease.changed()).await?;
664 }
665
666 #[xmtp_common::test(unwrap_try = true)]
667 async fn lease_change_wakes_independent_observers() {
668 let (commands, _receiver) = mpsc::unbounded_channel();
669 let coordinator = Arc::new(IncomingCoordinator {
670 commands,
671 generations: AtomicU64::new(0),
672 state: Arc::new(SharedState::default()),
673 });
674 let lease = coordinator.acquire(IncomingScope::AllGroups);
675 let mut first_changes = lease.subscribe_changes();
676 let mut second_changes = lease.subscribe_changes();
677 let _status = lease.snapshot();
678 coordinator.state.notify();
679 let (first, second) =
680 xmtp_common::time::timeout(std::time::Duration::from_secs(1), async {
681 futures::join!(first_changes.changed(), second_changes.changed())
682 })
683 .await?;
684 first?;
685 second?;
686 }
687}
688
689#[cfg(all(test, not(target_arch = "wasm32")))]
690mod tests {
691 mod preflight;
692}