1use std::collections::{HashMap, HashSet};
26
27use prost::Message;
28use tokio::sync::{mpsc, oneshot};
29use xmtp_common::rate_limit::Bucket;
30use xmtp_common::{BoxDynFuture, MaybeSend, MaybeSync, RetryableError};
31#[cfg(not(test))]
32use xmtp_configuration::AUTH_LOCKOUT_COOLDOWN;
33use xmtp_proto::api::ApiClientError;
34use xmtp_proto::{
35 backend_v1::ServerEnvelope,
36 types::{
37 Cursor, IncomingBatchLimits, IncomingEvent, IncomingSubscription, OrderedEnvelopeBatch,
38 Topic, TopicCursor,
39 },
40};
41
42#[cfg(test)]
45const AUTH_LOCKOUT_COOLDOWN: std::time::Duration = std::time::Duration::from_millis(500);
46
47use super::bidi::{BidiBinding, Connection, Event, TryMutateError};
48
49pub const DEFAULT_LEASE_DEPTH: usize = 64;
51const OUTBOX_RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(25);
53
54fn update_budget() -> Bucket {
55 Bucket::new(
56 xmtp_configuration::BACKEND_DEFAULT_MAX_UPDATE_FRAMES_PER_SECOND,
57 xmtp_configuration::BACKEND_DEFAULT_MAX_UPDATE_BURST,
58 )
59}
60
61pub(crate) const MAX_MUTATE_TOPICS: usize = xmtp_configuration::BACKEND_DEFAULT_MAX_UPDATE_ADDS;
63pub(crate) const MAX_MUTATE_BYTES: usize =
65 xmtp_configuration::BACKEND_DEFAULT_MAX_REQUEST_BYTES - PER_ENTRY_OVERHEAD;
66const PER_ENTRY_OVERHEAD: usize = 64;
68
69#[derive(Clone, Copy, Debug, PartialEq, Eq)]
72pub struct MutateLimits {
73 pub add_cap: usize,
74 pub remove_cap: usize,
75 pub byte_cap: usize,
76}
77
78impl Default for MutateLimits {
79 fn default() -> Self {
80 Self {
81 add_cap: MAX_MUTATE_TOPICS,
82 remove_cap: xmtp_configuration::BACKEND_DEFAULT_MAX_UPDATE_REMOVES,
83 byte_cap: MAX_MUTATE_BYTES,
84 }
85 }
86}
87
88impl MutateLimits {
89 pub fn from_limits(limits: &xmtp_configuration::LimitsConfiguration) -> Self {
92 Self {
93 add_cap: limits.max_update_adds.max(1),
94 remove_cap: limits.max_update_removes.max(1),
95 byte_cap: limits
96 .max_request_bytes
97 .saturating_sub(PER_ENTRY_OVERHEAD)
98 .max(1),
99 }
100 }
101}
102
103fn topic_wire_cost(topic: &Topic) -> usize {
104 1 + topic.identifier().len() + PER_ENTRY_OVERHEAD
105}
106
107fn chunk_by_budget<T>(
109 items: Vec<T>,
110 max_count: usize,
111 max_bytes: usize,
112 wire_cost: impl Fn(&T) -> usize,
113) -> Vec<Vec<T>> {
114 let mut chunks: Vec<Vec<T>> = Vec::new();
115 let mut current: Vec<T> = Vec::new();
116 let mut current_bytes = 0usize;
117 for item in items {
118 let cost = wire_cost(&item);
119 if !current.is_empty() && (current.len() >= max_count || current_bytes + cost > max_bytes) {
120 chunks.push(std::mem::take(&mut current));
121 current_bytes = 0;
122 }
123 current_bytes += cost;
124 current.push(item);
125 }
126 if !current.is_empty() {
127 chunks.push(current);
128 }
129 chunks
130}
131
132pub fn chunk_mutate_adds<C>(adds: Vec<(Topic, C)>) -> Vec<Vec<(Topic, C)>> {
135 chunk_by_budget(adds, MAX_MUTATE_TOPICS, MAX_MUTATE_BYTES, |(topic, _)| {
136 topic_wire_cost(topic)
137 })
138}
139
140const RECONNECT_INITIAL_DELAY: std::time::Duration = std::time::Duration::from_millis(100);
141const RECONNECT_MAX_DELAY: std::time::Duration = std::time::Duration::from_secs(30);
142const MIN_STABLE_UPTIME: std::time::Duration = std::time::Duration::from_secs(10);
143const GRACEFUL_CLOSE_BUDGET: std::time::Duration = std::time::Duration::from_secs(5);
145
146#[derive(Debug)]
148pub struct OpenError {
149 retryable: bool,
150 locked_out: bool,
151 source: Box<dyn std::error::Error + Send + Sync + 'static>,
152}
153
154impl OpenError {
155 pub fn new<E>(e: E) -> Self
156 where
157 E: std::error::Error + xmtp_common::RetryableError + Send + Sync + 'static,
158 {
159 Self {
160 retryable: e.is_retryable(),
161 locked_out: Self::locked_out_of(&e),
162 source: Box::new(e),
163 }
164 }
165
166 fn is_locked_out(&self) -> bool {
169 self.locked_out
170 }
171
172 #[cfg(test)]
173 pub fn retryable(e: impl Into<Box<dyn std::error::Error + Send + Sync + 'static>>) -> Self {
174 Self {
175 retryable: true,
176 locked_out: false,
177 source: e.into(),
178 }
179 }
180
181 #[cfg(test)]
182 pub fn unretryable(e: impl Into<Box<dyn std::error::Error + Send + Sync + 'static>>) -> Self {
183 Self {
184 retryable: false,
185 locked_out: false,
186 source: e.into(),
187 }
188 }
189}
190
191impl OpenError {
192 fn locked_out_of<E: 'static>(error: &E) -> bool {
194 (error as &dyn std::any::Any)
195 .downcast_ref::<ApiClientError>()
196 .is_some_and(ApiClientError::is_locked_out)
197 }
198}
199
200impl std::fmt::Display for OpenError {
201 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202 self.source.fmt(f)
203 }
204}
205
206impl std::error::Error for OpenError {
207 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
208 Some(&*self.source)
209 }
210}
211
212impl xmtp_common::RetryableError for OpenError {
213 fn is_retryable(&self) -> bool {
214 self.retryable
215 }
216}
217
218pub trait TransportBinding: BidiBinding
220where
221 Self::GroupMessage: Clone,
222 Self::WelcomeMessage: Clone,
223{
224 type Cursor: Copy + Send + std::fmt::Debug + From<u64> + Into<u64> + 'static;
226
227 fn build_mutate(
229 adds: impl IntoIterator<Item = (Topic, Self::Cursor)>,
230 removes: impl IntoIterator<Item = Topic>,
231 mutate_id: u64,
232 ) -> Self::Mutate;
233
234 fn group_topic(msg: &Self::GroupMessage) -> Option<Topic>;
235 fn welcome_topic(msg: &Self::WelcomeMessage) -> Option<Topic>;
236
237 fn group_cursor(msg: &Self::GroupMessage) -> Option<Self::Cursor>;
238 fn welcome_cursor(msg: &Self::WelcomeMessage) -> Option<Self::Cursor>;
239 fn group_envelope(msg: &Self::GroupMessage) -> &ServerEnvelope;
240 fn welcome_envelope(msg: &Self::WelcomeMessage) -> &ServerEnvelope;
241 fn advance(position: &mut Self::Cursor, delivered: Self::Cursor);
242 fn covers(position: &Self::Cursor, delivered: &Self::Cursor) -> bool;
244 fn meet(a: Self::Cursor, b: Self::Cursor) -> Self::Cursor;
246}
247
248#[derive(Debug, thiserror::Error)]
249pub enum TransportError {
250 #[error("the bidi transport is closed")]
251 Closed,
252 #[error("opening the bidi wire failed: {0}")]
253 Open(#[source] OpenError),
254 #[error("a lease must name at least one topic")]
255 Empty,
256 #[error("the bidi wire topic limit would be exceeded")]
258 TooManyTopics,
259 #[error("invalid subscription frame: {0}")]
260 Protocol(&'static str),
261 #[error("incoming subscription delivery exceeds its receive limit")]
262 Capacity,
263 #[error("incoming subscription delivery queue is full")]
265 Backpressure,
266 #[error(transparent)]
267 Wire(#[from] std::sync::Arc<super::bidi::ConnectionFailure>),
268}
269
270type IncomingFailure = std::sync::Arc<parking_lot::Mutex<Option<TransportError>>>;
271type IncomingFrame = Result<Vec<IncomingEvent>, TransportError>;
273
274impl xmtp_common::RetryableError for TransportError {
275 fn is_retryable(&self) -> bool {
276 match self {
277 Self::Backpressure => true,
278 Self::Open(e) => e.is_retryable(),
279 Self::Wire(e) => e.is_retryable(),
280 Self::Closed
281 | Self::Empty
282 | Self::TooManyTopics
283 | Self::Protocol(_)
284 | Self::Capacity => false,
285 }
286 }
287}
288
289pub enum LeaseEvent<B: TransportBinding>
291where
292 B::GroupMessage: Clone,
293 B::WelcomeMessage: Clone,
294{
295 CatchUpComplete,
297 GroupMessages(Vec<B::GroupMessage>),
298 WelcomeMessages(Vec<B::WelcomeMessage>),
299}
300
301#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
302struct LeaseId(u64);
303
304pub struct BidiTransport<B: TransportBinding>
306where
307 B::GroupMessage: Clone,
308 B::WelcomeMessage: Clone,
309{
310 cmds: mpsc::UnboundedSender<Cmd<B>>,
311}
312
313impl<B: TransportBinding> Clone for BidiTransport<B>
314where
315 B::GroupMessage: Clone,
316 B::WelcomeMessage: Clone,
317{
318 fn clone(&self) -> Self {
319 Self {
320 cmds: self.cmds.clone(),
321 }
322 }
323}
324
325pub struct TopicLease<B: TransportBinding>
327where
328 B::GroupMessage: Clone,
329 B::WelcomeMessage: Clone,
330{
331 id: LeaseId,
332 topics: Vec<Topic>,
333 events: mpsc::Receiver<LeaseEvent<B>>,
334 incoming: Option<mpsc::Receiver<IncomingFrame>>,
335 incoming_pending: std::vec::IntoIter<IncomingEvent>,
336 incoming_failure: IncomingFailure,
337 cmds: mpsc::UnboundedSender<Cmd<B>>,
338}
339
340impl<B: TransportBinding> TopicLease<B>
341where
342 B::GroupMessage: Clone,
343 B::WelcomeMessage: Clone,
344{
345 pub async fn next(&mut self) -> Option<LeaseEvent<B>> {
347 self.events.recv().await
348 }
349
350 pub fn topics(&self) -> &[Topic] {
351 &self.topics
352 }
353
354 pub async fn next_incoming(&mut self) -> Option<Result<IncomingEvent, TransportError>> {
356 loop {
357 if let Some(event) = self.incoming_pending.next() {
358 return Some(Ok(event));
359 }
360 match self.incoming.as_mut()?.recv().await {
361 Some(Ok(events)) => self.incoming_pending = events.into_iter(),
362 Some(Err(error)) => return Some(Err(error)),
363 None => return self.incoming_failure.lock().take().map(Err),
364 }
365 }
366 }
367
368 pub fn acknowledge_received(&self, cursors: TopicCursor) {
370 let _ = self.cmds.send(Cmd::Received {
371 id: self.id,
372 cursors,
373 });
374 }
375
376 pub fn into_incoming_subscription(self) -> IncomingSubscription<TransportError> {
378 let cmds = self.cmds.clone();
379 let id = self.id;
380 let events = futures::stream::unfold(self, |mut lease| async move {
381 lease.next_incoming().await.map(|event| (event, lease))
382 });
383 IncomingSubscription::new(Box::pin(events), move |cursors| {
384 let _ = cmds.send(Cmd::Received { id, cursors });
385 })
386 }
387}
388
389impl<B: TransportBinding> Drop for TopicLease<B>
390where
391 B::GroupMessage: Clone,
392 B::WelcomeMessage: Clone,
393{
394 fn drop(&mut self) {
395 let _ = self.cmds.send(Cmd::Deref(self.id));
396 }
397}
398
399impl<B: TransportBinding> BidiTransport<B>
400where
401 B::GroupMessage: Clone,
402 B::WelcomeMessage: Clone,
403{
404 pub fn new<O, Fut>(opener: O, initially_suspended: bool) -> Self
406 where
407 O: Fn(B::Mutate) -> Fut + MaybeSend + MaybeSync + 'static,
408 Fut: Future<Output = Result<Connection<B>, OpenError>> + MaybeSend + 'static,
409 {
410 Self::spawn(opener, initially_suspended, MutateLimits::default())
411 }
412
413 pub fn new_within<O, Fut>(opener: O, initially_suspended: bool, mutate: MutateLimits) -> Self
416 where
417 O: Fn(B::Mutate) -> Fut + MaybeSend + MaybeSync + 'static,
418 Fut: Future<Output = Result<Connection<B>, OpenError>> + MaybeSend + 'static,
419 {
420 Self::spawn(opener, initially_suspended, mutate)
421 }
422
423 #[cfg(test)]
424 pub(crate) fn new_with_chunk_limits<O, Fut>(
425 opener: O,
426 initially_suspended: bool,
427 chunk_cap: usize,
428 chunk_bytes: usize,
429 ) -> Self
430 where
431 O: Fn(B::Mutate) -> Fut + MaybeSend + MaybeSync + 'static,
432 Fut: Future<Output = Result<Connection<B>, OpenError>> + MaybeSend + 'static,
433 {
434 Self::spawn(
435 opener,
436 initially_suspended,
437 MutateLimits {
438 add_cap: chunk_cap,
439 remove_cap: chunk_cap,
440 byte_cap: chunk_bytes,
441 },
442 )
443 }
444
445 fn spawn<O, Fut>(opener: O, initially_suspended: bool, mutate: MutateLimits) -> Self
446 where
447 O: Fn(B::Mutate) -> Fut + MaybeSend + MaybeSync + 'static,
448 Fut: Future<Output = Result<Connection<B>, OpenError>> + MaybeSend + 'static,
449 {
450 let (cmds, cmds_rx) = mpsc::unbounded_channel();
451 let opener: Opener<B> = Box::new(
452 move |initial| -> BoxDynFuture<'static, Result<Connection<B>, OpenError>> {
453 Box::pin(opener(initial))
454 },
455 );
456 xmtp_common::spawn(
457 None,
458 run_ledger::<B>(
459 opener,
460 cmds_rx,
461 cmds.clone().downgrade(),
462 initially_suspended,
463 mutate,
464 ),
465 );
466 Self { cmds }
467 }
468
469 pub async fn lease(
473 &self,
474 subs: Vec<(Topic, B::Cursor)>,
475 depth: usize,
476 ) -> Result<TopicLease<B>, TransportError> {
477 if subs.is_empty() {
478 return Err(TransportError::Empty);
479 }
480 let (reply, response) = oneshot::channel();
481 self.cmds
482 .send(Cmd::Lease { subs, depth, reply })
483 .map_err(|_| TransportError::Closed)?;
484 response.await.map_err(|_| TransportError::Closed)?
485 }
486
487 pub async fn lease_ordered(
489 &self,
490 subs: Vec<(Topic, B::Cursor)>,
491 depth: usize,
492 limits: IncomingBatchLimits,
493 ) -> Result<TopicLease<B>, TransportError> {
494 if subs.is_empty() {
495 return Err(TransportError::Empty);
496 }
497 if limits.max_rows == 0 || limits.max_bytes == 0 {
498 return Err(TransportError::Capacity);
499 }
500 let (reply, response) = oneshot::channel();
501 self.cmds
502 .send(Cmd::LeaseOrdered {
503 subs,
504 depth,
505 limits,
506 reply,
507 })
508 .map_err(|_| TransportError::Closed)?;
509 response.await.map_err(|_| TransportError::Closed)?
510 }
511
512 #[xmtp_common::span(prefix = "bidi")]
514 pub async fn suspend(&self) -> Result<(), TransportError> {
515 self.enqueue_suspend()?
516 .await
517 .map_err(|_| TransportError::Closed)
518 }
519
520 pub fn enqueue_suspend(&self) -> Result<oneshot::Receiver<()>, TransportError> {
522 let (reply, response) = oneshot::channel();
523 self.cmds
524 .send(Cmd::Suspend { reply })
525 .map_err(|_| TransportError::Closed)?;
526 Ok(response)
527 }
528
529 #[xmtp_common::span(prefix = "bidi")]
532 pub async fn resume(&self) -> Result<(), TransportError> {
533 self.enqueue_resume()?
534 .await
535 .map_err(|_| TransportError::Closed)
536 }
537
538 pub fn enqueue_resume(&self) -> Result<oneshot::Receiver<()>, TransportError> {
540 let (reply, response) = oneshot::channel();
541 self.cmds
542 .send(Cmd::Resume { reply })
543 .map_err(|_| TransportError::Closed)?;
544 Ok(response)
545 }
546}
547
548enum Cmd<B: TransportBinding>
549where
550 B::GroupMessage: Clone,
551 B::WelcomeMessage: Clone,
552{
553 Lease {
554 subs: Vec<(Topic, B::Cursor)>,
555 depth: usize,
556 reply: oneshot::Sender<Result<TopicLease<B>, TransportError>>,
557 },
558 LeaseOrdered {
559 subs: Vec<(Topic, B::Cursor)>,
560 depth: usize,
561 limits: IncomingBatchLimits,
562 reply: oneshot::Sender<Result<TopicLease<B>, TransportError>>,
563 },
564 Received {
565 id: LeaseId,
566 cursors: TopicCursor,
567 },
568 Deref(LeaseId),
569 Suspend {
570 reply: oneshot::Sender<()>,
571 },
572 Resume {
573 reply: oneshot::Sender<()>,
574 },
575}
576
577trait OpenWire<B: BidiBinding>: MaybeSend + MaybeSync {
578 fn open(&self, initial: B::Mutate) -> BoxDynFuture<'static, Result<Connection<B>, OpenError>>;
579}
580
581impl<B: BidiBinding, F> OpenWire<B> for F
582where
583 F: Fn(B::Mutate) -> BoxDynFuture<'static, Result<Connection<B>, OpenError>>
584 + MaybeSend
585 + MaybeSync,
586{
587 fn open(&self, initial: B::Mutate) -> BoxDynFuture<'static, Result<Connection<B>, OpenError>> {
588 self(initial)
589 }
590}
591
592type Opener<B> = Box<dyn OpenWire<B>>;
593#[derive(Debug, Clone, Copy)]
594enum DeliveryKind {
595 Group,
596 Welcome,
597}
598
599struct PendingUpdate<C> {
601 adds: Vec<(Topic, C)>,
602 removes: Vec<Topic>,
603}
604
605struct TopicRegistration<C> {
607 target: Option<u64>,
608 delivered: C,
609 holders: HashSet<LeaseId>,
610 state: RegistrationState,
611}
612
613enum RegistrationState {
614 Adding,
615 Active,
616 Removing,
618}
619
620struct LeaseObligation {
622 satisfied: bool,
623}
624
625struct LeaseState<B: TransportBinding>
626where
627 B::GroupMessage: Clone,
628 B::WelcomeMessage: Clone,
629{
630 floors: HashMap<Topic, B::Cursor>,
631 delivered: HashMap<Topic, B::Cursor>,
632 obligations: HashMap<Topic, LeaseObligation>,
633 unmet: usize,
634 notified: bool,
635 events: mpsc::Sender<LeaseEvent<B>>,
636 incoming: Option<(mpsc::Sender<IncomingFrame>, IncomingBatchLimits)>,
637 incoming_failure: IncomingFailure,
638}
639
640struct Ledger<B: TransportBinding>
641where
642 B::GroupMessage: Clone,
643 B::WelcomeMessage: Clone,
644{
645 leases: HashMap<LeaseId, LeaseState<B>>,
646 by_topic: HashMap<Topic, HashSet<LeaseId>>,
647 last_seen: HashMap<Topic, B::Cursor>,
648 registrations: HashMap<Topic, TopicRegistration<B::Cursor>>,
649 pending_updates: HashMap<u64, PendingUpdate<B::Cursor>>,
650 dirty_topics: HashSet<Topic>,
651 failed_incoming: HashSet<LeaseId>,
652 next_lease: u64,
653 next_update: u64,
654 mutate: MutateLimits,
655}
656
657impl<B: TransportBinding> Default for Ledger<B>
658where
659 B::GroupMessage: Clone,
660 B::WelcomeMessage: Clone,
661{
662 fn default() -> Self {
663 Self {
664 leases: HashMap::new(),
665 by_topic: HashMap::new(),
666 last_seen: HashMap::new(),
667 registrations: HashMap::new(),
668 pending_updates: HashMap::new(),
669 dirty_topics: HashSet::new(),
670 failed_incoming: HashSet::new(),
671 next_lease: 0,
672 next_update: 0,
673 mutate: MutateLimits::default(),
674 }
675 }
676}
677
678impl<B: TransportBinding> Ledger<B>
679where
680 B::GroupMessage: Clone,
681 B::WelcomeMessage: Clone,
682{
683 fn received(&mut self, id: LeaseId, cursors: TopicCursor) {
685 let Some(lease) = self.leases.get_mut(&id) else {
686 return;
687 };
688 if lease.incoming.is_none() {
689 return;
690 }
691 for (topic, cursor) in cursors {
692 if cursor.0 > i64::MAX as u64 {
693 continue;
694 }
695 if let Some(floor) = lease.floors.get_mut(&topic) {
696 B::advance(floor, cursor.0.into());
697 if let Some(delivered) = lease.delivered.get_mut(&topic) {
698 B::advance(delivered, cursor.0.into());
699 }
700 self.dirty_topics.insert(topic);
701 }
702 }
703 }
704
705 fn incoming_registered(&mut self, id: LeaseId, topics: impl IntoIterator<Item = Topic>) {
707 let Some(lease) = self.leases.get(&id) else {
708 return;
709 };
710 let Some((sender, _)) = &lease.incoming else {
711 return;
712 };
713 let mut starts = TopicCursor::new();
714 let mut targets = TopicCursor::new();
715 for topic in topics {
716 let Some(registration) = self.registrations.get(&topic) else {
717 continue;
718 };
719 let Some(target) = registration.target else {
720 continue;
721 };
722 let Some(start) = lease.delivered.get(&topic) else {
723 continue;
724 };
725 starts.insert(topic.clone(), Cursor((*start).into()));
726 targets.insert(topic, Cursor(target));
727 }
728 if !starts.is_empty() {
729 if sender
731 .try_send(Ok(vec![IncomingEvent::Registered { starts, targets }]))
732 .is_err()
733 {
734 *lease.incoming_failure.lock() = Some(TransportError::Backpressure);
735 self.failed_incoming.insert(id);
736 }
737 }
738 }
739
740 fn valid_targets(&self, id: u64, targets: &[(Topic, u64)]) -> bool {
742 let Some(update) = self.pending_updates.get(&id) else {
743 return false;
744 };
745 let unique: HashSet<_> = targets.iter().map(|(topic, _)| topic).collect();
746 targets.len() == update.adds.len()
747 && unique.len() == targets.len()
748 && targets.iter().all(|(topic, target)| {
749 *target <= i64::MAX as u64 && update.adds.iter().any(|(added, _)| added == topic)
750 })
751 }
752
753 fn demux_incoming<M>(
755 &mut self,
756 messages: &[M],
757 envelope_of: impl Fn(&M) -> &ServerEnvelope,
758 ) -> Result<Vec<LeaseId>, &'static str> {
759 let mut positions = HashMap::new();
760 for message in messages {
761 let envelope = envelope_of(message);
762 let meta = envelope.meta.as_ref().ok_or("metadata")?;
763 let topic =
764 Topic::parse(&meta.topic.as_ref().ok_or("topic")?.topic).map_err(|_| "topic")?;
765 let (_, cursor, _) =
766 crate::envelope::metadata(meta, topic.kind()).map_err(|_| "metadata")?;
767 let Some(registration) = self.registrations.get(&topic) else {
768 continue;
769 };
770 if registration.target.is_none() {
771 return Err("messages before Applied");
772 }
773 let previous = positions
774 .entry(topic)
775 .or_insert_with(|| Cursor(registration.delivered.into()));
776 if cursor <= *previous {
777 return Err("cursor order");
778 }
779 *previous = cursor;
780 }
781 let mut dropped = Vec::new();
782 for (id, lease) in &mut self.leases {
783 let Some((sender, limits)) = &lease.incoming else {
784 continue;
785 };
786 let mut selected: HashMap<Topic, Vec<&ServerEnvelope>> = HashMap::new();
787 let mut rows = 0usize;
788 let mut bytes = 0usize;
789 for message in messages {
790 let envelope = envelope_of(message);
791 let meta = envelope.meta.as_ref().ok_or("metadata")?;
792 let topic = Topic::parse(&meta.topic.as_ref().ok_or("topic")?.topic)
793 .map_err(|_| "topic")?;
794 let Some(registration) = self.registrations.get(&topic) else {
795 continue;
796 };
797 if !registration.holders.contains(id) {
798 continue;
799 }
800 let Some(delivered) = lease.delivered.get(&topic) else {
801 continue;
802 };
803 if meta.cursor.as_ref().ok_or("cursor")?.sequence_id <= (*delivered).into() {
804 continue;
805 }
806 rows += 1;
807 bytes = bytes
808 .checked_add(envelope.encoded_len())
809 .ok_or("byte count")?;
810 selected.entry(topic).or_default().push(envelope);
811 }
812 if rows > limits.max_rows || bytes > limits.max_bytes {
813 *lease.incoming_failure.lock() = Some(TransportError::Capacity);
814 dropped.push(*id);
815 continue;
816 }
817 if selected.is_empty() {
818 continue;
819 }
820 let Ok(permit) = sender.try_reserve() else {
823 *lease.incoming_failure.lock() = Some(TransportError::Backpressure);
824 dropped.push(*id);
825 continue;
826 };
827 let mut batches = Vec::with_capacity(selected.len());
828 for (topic, envelopes) in selected {
829 let delivered = lease.delivered.get_mut(&topic).ok_or("lease topic")?;
830 let after = Cursor((*delivered).into());
831 let last = envelopes
832 .last()
833 .and_then(|envelope| envelope.meta.as_ref()?.cursor.as_ref())
834 .ok_or("cursor")?
835 .sequence_id;
836 let batch = OrderedEnvelopeBatch {
837 topic,
838 after,
839 envelopes: envelopes.into_iter().cloned().collect(),
840 };
841 batches.push(IncomingEvent::OrderedBatch(batch));
842 B::advance(delivered, last.into());
843 }
844 permit.send(Ok(batches));
845 }
846 Ok(dropped)
847 }
848
849 fn next_update_id(&mut self) -> u64 {
850 self.next_update += 1;
851 self.next_update
852 }
853
854 fn register(
856 &mut self,
857 subs: &[(Topic, B::Cursor)],
858 events: mpsc::Sender<LeaseEvent<B>>,
859 ) -> LeaseId {
860 self.next_lease += 1;
861 let id = LeaseId(self.next_lease);
862 let floors: HashMap<_, _> = subs.iter().cloned().collect();
863 for topic in floors.keys() {
864 self.by_topic.entry(topic.clone()).or_default().insert(id);
865 }
866 self.leases.insert(
867 id,
868 LeaseState {
869 delivered: floors.clone(),
870 obligations: floors
871 .keys()
872 .map(|topic| (topic.clone(), LeaseObligation { satisfied: false }))
873 .collect(),
874 unmet: floors.len(),
875 floors,
876 notified: false,
877 events,
878 incoming: None,
879 incoming_failure: IncomingFailure::default(),
880 },
881 );
882 id
883 }
884
885 fn reset_wire(&mut self) {
887 self.registrations.clear();
888 self.pending_updates.clear();
889 self.dirty_topics.clear();
890 for lease in self.leases.values_mut() {
891 lease.obligations.clear();
892 lease.unmet = lease.floors.len();
893 if lease.incoming.is_some() {
894 lease.delivered.clone_from(&lease.floors);
895 }
896 }
897 }
898
899 fn reset_obligation(&mut self, id: LeaseId, topic: &Topic) {
900 if let Some(lease) = self.leases.get_mut(&id) {
901 if let Some(obligation) = lease.obligations.get_mut(topic) {
902 if obligation.satisfied {
903 obligation.satisfied = false;
904 lease.unmet += 1;
905 }
906 } else {
907 lease
908 .obligations
909 .insert(topic.clone(), LeaseObligation { satisfied: false });
910 }
911 }
912 }
913
914 fn resume_cursor(&self, topic: &Topic) -> Option<B::Cursor> {
919 let holders = self.by_topic.get(topic)?;
920 holders
921 .iter()
922 .filter_map(|id| self.leases.get(id)?.floors.get(topic).copied())
923 .reduce(B::meet)
924 }
925
926 fn resume_adds(&self) -> Vec<(Topic, B::Cursor)> {
927 self.by_topic
928 .keys()
929 .filter_map(|topic| Some((topic.clone(), self.resume_cursor(topic)?)))
930 .collect()
931 }
932
933 fn prepare_adds(&mut self, adds: Vec<(Topic, B::Cursor)>) -> Vec<(u64, B::Mutate)> {
935 let chunks = chunk_by_budget(
936 adds,
937 self.mutate.add_cap,
938 self.mutate.byte_cap,
939 |(topic, _)| topic_wire_cost(topic),
940 );
941 chunks
942 .into_iter()
943 .map(|adds| {
944 for (topic, cursor) in &adds {
945 self.dirty_topics.insert(topic.clone());
946 let holders = self.by_topic.get(topic).cloned().unwrap_or_default();
947 for id in &holders {
948 self.reset_obligation(*id, topic);
949 }
950 self.registrations
951 .entry(topic.clone())
952 .or_insert_with(|| TopicRegistration {
953 target: None,
954 delivered: *cursor,
955 holders,
956 state: RegistrationState::Adding,
957 });
958 }
959 let id = self.next_update_id();
960 let update = B::build_mutate(adds.clone(), [], id);
961 self.pending_updates.insert(
962 id,
963 PendingUpdate {
964 adds,
965 removes: vec![],
966 },
967 );
968 (id, update)
969 })
970 .collect()
971 }
972
973 fn prepare_removes(&mut self, removes: Vec<Topic>) -> Vec<(u64, B::Mutate)> {
974 chunk_by_budget(
975 removes,
976 self.mutate.remove_cap,
977 self.mutate.byte_cap,
978 topic_wire_cost,
979 )
980 .into_iter()
981 .map(|removes| {
982 let id = self.next_update_id();
983 let update = B::build_mutate([], removes.clone(), id);
984 self.pending_updates.insert(
985 id,
986 PendingUpdate {
987 adds: vec![],
988 removes,
989 },
990 );
991 (id, update)
992 })
993 .collect()
994 }
995
996 fn join(&mut self, id: LeaseId, subs: Vec<(Topic, B::Cursor)>) -> Vec<(u64, B::Mutate)> {
998 let mut adds = Vec::new();
999 let mut removes = Vec::new();
1000 let mut joined = Vec::new();
1001 for (topic, floor) in subs {
1002 self.dirty_topics.insert(topic.clone());
1003 let Some(registration) = self.registrations.get_mut(&topic) else {
1004 adds.push((topic, floor));
1005 continue;
1006 };
1007 match &mut registration.state {
1008 RegistrationState::Removing => {}
1009 _ if !B::covers(&floor, ®istration.delivered) => {
1010 registration.state = RegistrationState::Removing;
1011 removes.push(topic);
1012 }
1013 _ => {
1014 registration.holders.insert(id);
1015 joined.push(topic);
1016 }
1017 }
1018 }
1019 let mut updates = self.prepare_removes(removes);
1020 updates.extend(self.prepare_adds(adds));
1021 self.incoming_registered(id, joined);
1022 updates
1023 }
1024
1025 fn applied(&mut self, id: u64, targets: Vec<(Topic, u64)>) -> Vec<(Topic, B::Cursor)> {
1027 let Some(update) = self.pending_updates.remove(&id) else {
1028 tracing::warn!(id, "received Applied for an unknown update");
1029 return Vec::new();
1030 };
1031 let targets: HashMap<_, _> = targets.into_iter().collect();
1032 let mut readds = Vec::new();
1033 for topic in update.removes {
1034 self.registrations.remove(&topic);
1035 if let Some(cursor) = self.resume_cursor(&topic) {
1036 readds.push((topic, cursor));
1037 }
1038 }
1039 let mut registered: HashMap<LeaseId, Vec<Topic>> = HashMap::new();
1040 for (topic, _) in update.adds {
1041 self.dirty_topics.insert(topic.clone());
1042 if let Some(registration) = self.registrations.get_mut(&topic) {
1043 if let Some(target) = targets.get(&topic) {
1044 registration.target = Some(*target);
1045 }
1046 if matches!(registration.state, RegistrationState::Adding) {
1048 registration.state = RegistrationState::Active;
1049 }
1050 for holder in ®istration.holders {
1051 registered.entry(*holder).or_default().push(topic.clone());
1052 }
1053 }
1054 }
1055 for (holder, topics) in registered {
1056 self.incoming_registered(holder, topics);
1057 }
1058 readds
1059 }
1060
1061 fn recheck(&mut self) -> Vec<LeaseId> {
1063 let failed: Vec<_> = self.failed_incoming.drain().collect();
1064 let mut candidates = HashSet::new();
1065 for topic in self.dirty_topics.drain() {
1066 let Some(registration) = self.registrations.get(&topic) else {
1067 continue;
1068 };
1069 let Some(target) = registration.target else {
1070 continue;
1071 };
1072 let target_cursor = B::Cursor::from(target);
1073 for id in ®istration.holders {
1074 candidates.insert(*id);
1075 let Some(lease) = self.leases.get_mut(id) else {
1076 continue;
1077 };
1078 let Some(obligation) = lease.obligations.get_mut(&topic) else {
1079 continue;
1080 };
1081 if !obligation.satisfied
1082 && (target == 0
1083 || B::covers(®istration.delivered, &target_cursor)
1084 || lease
1085 .floors
1086 .get(&topic)
1087 .is_some_and(|floor| B::covers(floor, &target_cursor)))
1088 {
1089 obligation.satisfied = true;
1090 lease.unmet -= 1;
1091 }
1092 }
1093 }
1094 candidates
1095 .into_iter()
1096 .filter_map(|id| {
1097 let lease = self.leases.get_mut(&id)?;
1098 if !lease.notified && lease.unmet == 0 {
1099 if lease.incoming.is_none()
1100 && lease.events.try_send(LeaseEvent::CatchUpComplete).is_err()
1101 {
1102 return Some(id);
1103 }
1104 lease.notified = true;
1105 }
1106 None
1107 })
1108 .chain(failed)
1109 .collect()
1110 }
1111
1112 fn deref(&mut self, id: LeaseId) -> Vec<Topic> {
1114 let Some(lease) = self.leases.remove(&id) else {
1115 return vec![];
1116 };
1117 let mut removes = Vec::new();
1118 for topic in lease.floors.keys() {
1119 if let Some(holders) = self.by_topic.get_mut(topic) {
1120 holders.remove(&id);
1121 if holders.is_empty() {
1122 self.by_topic.remove(topic);
1123 self.last_seen.remove(topic);
1124 removes.push(topic.clone());
1125 }
1126 }
1127 if let Some(registration) = self.registrations.get_mut(topic) {
1128 registration.holders.remove(&id);
1129 }
1130 }
1131 removes
1132 }
1133
1134 fn demux<M: Clone>(
1137 &mut self,
1138 messages: Vec<M>,
1139 kind: DeliveryKind,
1140 topic_of: impl Fn(&M) -> Option<Topic>,
1141 cursor_of: impl Fn(&M) -> Option<B::Cursor>,
1142 event: impl Fn(Vec<M>) -> LeaseEvent<B>,
1143 ) -> Vec<LeaseId> {
1144 let mut batches: HashMap<LeaseId, Vec<&M>> = HashMap::new();
1145 for message in &messages {
1146 let Some(topic) = topic_of(message) else {
1147 continue;
1148 };
1149 let Some(registration) = self.registrations.get_mut(&topic) else {
1150 continue;
1151 };
1152 if registration.target.is_none() {
1154 continue;
1155 }
1156 self.dirty_topics.insert(topic.clone());
1157 let cursor = cursor_of(message);
1158 if let Some(cursor) = cursor {
1159 B::advance(&mut registration.delivered, cursor);
1160 B::advance(
1161 self.last_seen.entry(topic.clone()).or_insert(cursor),
1162 cursor,
1163 );
1164 }
1165 for id in ®istration.holders {
1166 let Some(lease) = self.leases.get_mut(id) else {
1167 continue;
1168 };
1169 if lease.incoming.is_some() {
1170 continue;
1171 }
1172 if let Some(cursor) = cursor {
1173 let Some(position) = lease.delivered.get_mut(&topic) else {
1174 continue;
1175 };
1176 if B::covers(position, &cursor) {
1177 continue;
1178 }
1179 B::advance(position, cursor);
1180 }
1181 batches.entry(*id).or_default().push(message);
1182 }
1183 }
1184 batches
1185 .into_iter()
1186 .filter_map(|(id, batch)| {
1187 let lease = self.leases.get(&id)?;
1188 if let Ok(permit) = lease.events.try_reserve() {
1189 permit.send(event(batch.into_iter().cloned().collect()));
1190 None
1191 } else {
1192 tracing::warn!(
1193 lease = id.0,
1194 ?kind,
1195 "closing a lease whose delivery channel is full"
1196 );
1197 Some(id)
1198 }
1199 })
1200 .collect()
1201 }
1202
1203 fn caught_up(&self) -> bool {
1204 self.pending_updates.is_empty() && self.leases.values().all(|lease| lease.unmet == 0)
1205 }
1206}
1207
1208enum Step<B: TransportBinding>
1211where
1212 B::GroupMessage: Clone,
1213 B::WelcomeMessage: Clone,
1214{
1215 Cmd(Option<Cmd<B>>),
1216 Wire(Option<Event<B::GroupMessage, B::WelcomeMessage>>),
1217 Retry,
1218 Reconnect,
1219}
1220
1221async fn run_ledger<B: TransportBinding>(
1222 opener: Opener<B>,
1223 cmds: mpsc::UnboundedReceiver<Cmd<B>>,
1224 lease_cmds: mpsc::WeakUnboundedSender<Cmd<B>>,
1225 initially_suspended: bool,
1226 mutate: MutateLimits,
1227) where
1228 B::GroupMessage: Clone,
1229 B::WelcomeMessage: Clone,
1230{
1231 LedgerTask {
1232 opener,
1233 cmds,
1234 lease_cmds,
1235 ledger: Ledger {
1236 mutate,
1237 ..Ledger::default()
1238 },
1239 conn: None,
1240 reconnect_delay: RECONNECT_INITIAL_DELAY,
1241 reconnect_at: tokio::time::Instant::now(),
1242 wire_opened_at: None,
1243 wire_span: None,
1244 suspended: initially_suspended,
1245 wire_opens: 0,
1246 resume_notify: Vec::new(),
1247 outbox: Outbox::default(),
1248 update_budget: update_budget(),
1249 deferred: std::collections::VecDeque::new(),
1250 }
1251 .run()
1252 .await
1253}
1254
1255enum Flow {
1256 Continue,
1257 Shutdown,
1258}
1259
1260struct LedgerTask<B: TransportBinding>
1261where
1262 B::GroupMessage: Clone,
1263 B::WelcomeMessage: Clone,
1264{
1265 opener: Opener<B>,
1266 cmds: mpsc::UnboundedReceiver<Cmd<B>>,
1267 lease_cmds: mpsc::WeakUnboundedSender<Cmd<B>>,
1268 ledger: Ledger<B>,
1269 conn: Option<Connection<B>>,
1270 reconnect_delay: std::time::Duration,
1271 reconnect_at: tokio::time::Instant,
1272 wire_opened_at: Option<tokio::time::Instant>,
1273 wire_span: Option<tracing::Span>,
1274 suspended: bool,
1275 wire_opens: u64,
1276 resume_notify: Vec<oneshot::Sender<()>>,
1277 outbox: Outbox<B::Mutate>,
1278 update_budget: Bucket,
1279 deferred: std::collections::VecDeque<Cmd<B>>,
1280}
1281
1282impl<B: TransportBinding> LedgerTask<B>
1283where
1284 B::GroupMessage: Clone,
1285 B::WelcomeMessage: Clone,
1286{
1287 async fn run(mut self) {
1288 loop {
1289 let queued = self.deferred.len() + self.cmds.len();
1292 for _ in 0..queued {
1293 let cmd = self
1294 .deferred
1295 .pop_front()
1296 .or_else(|| self.cmds.try_recv().ok());
1297 let Some(cmd) = cmd else { break };
1298 if let Flow::Shutdown = self.command(cmd).await {
1299 return;
1300 }
1301 }
1302 self.flush_outbox();
1303 let flow = match self.next_step().await {
1304 Step::Retry => Flow::Continue,
1305 Step::Cmd(None) => self.shutdown(),
1306 Step::Cmd(Some(cmd)) => self.command(cmd).await,
1307 Step::Wire(Some(event)) => self.wire_event(event),
1308 Step::Wire(None) => self.wire_died(),
1309 Step::Reconnect => self.reconnect().await,
1310 };
1311 if let Flow::Shutdown = flow {
1312 return;
1313 }
1314 }
1315 }
1316
1317 async fn command(&mut self, cmd: Cmd<B>) -> Flow {
1318 match cmd {
1319 Cmd::Lease { subs, depth, reply } => self.lease(subs, depth, None, reply).await,
1320 Cmd::LeaseOrdered {
1321 subs,
1322 depth,
1323 limits,
1324 reply,
1325 } => self.lease(subs, depth, Some(limits), reply).await,
1326 Cmd::Received { id, cursors } => {
1327 self.ledger.received(id, cursors);
1328 Flow::Continue
1329 }
1330 Cmd::Deref(id) => self.deref(id),
1331 Cmd::Suspend { reply } => self.suspend(reply),
1332 Cmd::Resume { reply } => self.resume(reply).await,
1333 }
1334 }
1335
1336 async fn next_step(&mut self) -> Step<B> {
1337 if let Some(cmd) = self.deferred.pop_front() {
1338 return Step::Cmd(Some(cmd));
1339 }
1340 let retry_after = self.update_budget.wait().max(OUTBOX_RETRY_INTERVAL);
1341 match self.conn.as_mut() {
1342 Some(wire) => tokio::select! {
1343 cmd = self.cmds.recv() => Step::Cmd(cmd),
1344 event = wire.next() => Step::Wire(event),
1345 _ = xmtp_common::time::sleep(retry_after), if !self.outbox.is_empty() => Step::Retry,
1346 },
1347 None if !self.ledger.leases.is_empty() && !self.suspended => tokio::select! {
1348 cmd = self.cmds.recv() => Step::Cmd(cmd),
1349 _ = tokio::time::sleep_until(self.reconnect_at) => Step::Reconnect,
1350 },
1351 None => Step::Cmd(self.cmds.recv().await),
1352 }
1353 }
1354
1355 fn shutdown(&mut self) -> Flow {
1356 self.outbox.clear();
1357 self.close_wire_span("shutdown");
1358 if let Some(wire) = self.conn.take() {
1359 close_gracefully(wire);
1360 }
1361 Flow::Shutdown
1362 }
1363
1364 fn open_wire_span(&mut self) {
1365 self.update_budget = update_budget();
1366 self.update_budget.take();
1368 self.wire_opened_at = Some(tokio::time::Instant::now());
1369 self.wire_span = Some(tracing::info_span!(
1370 parent: None,
1371 "bidi_wire",
1372 operation = "bidi.wire_session",
1373 reason = tracing::field::Empty,
1374 ));
1375 }
1376
1377 fn close_wire_span(&mut self, reason: &'static str) {
1378 if let Some(span) = self.wire_span.take() {
1379 span.record("reason", reason);
1380 }
1381 }
1382
1383 async fn lease(
1384 &mut self,
1385 subs: Vec<(Topic, B::Cursor)>,
1386 depth: usize,
1387 incoming_limits: Option<IncomingBatchLimits>,
1388 reply: oneshot::Sender<Result<TopicLease<B>, TransportError>>,
1389 ) -> Flow {
1390 let mut positions = HashMap::new();
1391 let mut subs_unique: Vec<(Topic, B::Cursor)> = Vec::new();
1392 for (topic, floor) in subs {
1393 if let Some(index) = positions.get(&topic).copied() {
1394 let (_, cursor): &mut (Topic, B::Cursor) = &mut subs_unique[index];
1395 *cursor = B::meet(*cursor, floor);
1396 } else {
1397 positions.insert(topic.clone(), subs_unique.len());
1398 subs_unique.push((topic, floor));
1399 }
1400 }
1401 let subs = subs_unique;
1402 let added = subs
1403 .iter()
1404 .filter(|(topic, _)| !self.ledger.by_topic.contains_key(topic))
1405 .count();
1406 if self.ledger.by_topic.len() + added
1407 > xmtp_configuration::BACKEND_DEFAULT_MAX_STREAM_TOPICS
1408 {
1409 let _ = reply.send(Err(TransportError::TooManyTopics));
1410 return Flow::Continue;
1411 }
1412 let cold = self.conn.is_none() && self.ledger.leases.is_empty() && !self.suspended;
1413 let topics = subs.iter().map(|(topic, _)| topic.clone()).collect();
1414 let (tx, events) = mpsc::channel(depth.max(1));
1415 let id = self.ledger.register(&subs, tx);
1416 let incoming = incoming_limits.map(|limits| {
1417 let (sender, receiver) = mpsc::channel(depth.max(1));
1418 if let Some(lease) = self.ledger.leases.get_mut(&id) {
1419 lease.incoming = Some((sender, limits));
1420 }
1421 receiver
1422 });
1423 let incoming_failure = self
1424 .ledger
1425 .leases
1426 .get(&id)
1427 .map(|lease| lease.incoming_failure.clone())
1428 .unwrap_or_default();
1429 if cold {
1430 self.outbox.updates.extend(self.ledger.prepare_adds(subs));
1431 let Some((_, initial)) = self.outbox.updates.pop_front() else {
1432 return Flow::Shutdown;
1433 };
1434 match self.open_preemptibly(initial).await {
1435 OpenOutcome::Opened(wire) => {
1436 self.conn = Some(wire);
1437 self.open_wire_span();
1438 self.wire_opens += 1;
1439 }
1440 OpenOutcome::Failed(error) => {
1441 self.ledger.deref(id);
1442 self.ledger.reset_wire();
1443 self.outbox.clear();
1444 let _ = reply.send(Err(TransportError::Open(error)));
1445 return Flow::Continue;
1446 }
1447 OpenOutcome::Suspended(ack) => {
1448 self.ledger.reset_wire();
1449 self.outbox.clear();
1450 self.suspend_preempted(ack);
1451 }
1452 OpenOutcome::Shutdown => return Flow::Shutdown,
1453 }
1454 } else if self.conn.is_some() {
1455 self.outbox.updates.extend(self.ledger.join(id, subs));
1456 let dropped = self.ledger.recheck();
1457 let removes = self.drop_leases(dropped);
1458 self.retire(removes);
1459 }
1460 let Some(cmds) = self.lease_cmds.upgrade() else {
1461 return Flow::Continue;
1462 };
1463 let _ = reply.send(Ok(TopicLease {
1464 id,
1465 topics,
1466 events,
1467 incoming,
1468 incoming_pending: Vec::new().into_iter(),
1469 incoming_failure,
1470 cmds,
1471 }));
1472 Flow::Continue
1473 }
1474
1475 fn deref(&mut self, id: LeaseId) -> Flow {
1476 let removes = self.drop_leases(vec![id]);
1477 self.retire(removes);
1478 self.settle_idle_waiters();
1479 self.settle_caught_up_waiters();
1480 Flow::Continue
1481 }
1482
1483 fn drop_leases(&mut self, dropped: Vec<LeaseId>) -> Vec<Topic> {
1484 let mut removes = Vec::new();
1485 for lease in dropped {
1486 removes.extend(self.ledger.deref(lease));
1487 }
1488 let unused: HashSet<_> = self
1490 .ledger
1491 .pending_updates
1492 .iter()
1493 .filter_map(|(id, update)| {
1494 (!update.adds.is_empty()
1495 && update
1496 .adds
1497 .iter()
1498 .all(|(topic, _)| !self.ledger.by_topic.contains_key(topic)))
1499 .then_some(*id)
1500 })
1501 .collect();
1502 for id in self.outbox.purge(&unused) {
1503 if let Some(update) = self.ledger.pending_updates.remove(&id) {
1504 for (topic, _) in update.adds {
1505 self.ledger.dirty_topics.insert(topic.clone());
1506 if !self
1509 .ledger
1510 .registrations
1511 .get(&topic)
1512 .is_some_and(|registration| {
1513 matches!(registration.state, RegistrationState::Removing)
1514 })
1515 {
1516 self.ledger.registrations.remove(&topic);
1517 }
1518 }
1519 }
1520 }
1521 removes
1522 }
1523
1524 fn suspend(&mut self, reply: oneshot::Sender<()>) -> Flow {
1525 tracing::info!(
1526 leases = self.ledger.leases.len(),
1527 had_wire = self.conn.is_some(),
1528 "bidi transport: suspending — going off the network"
1529 );
1530 self.suspended = true;
1531 self.outbox.clear();
1532 self.ledger.reset_wire();
1533 self.close_wire_span("suspend");
1534 if let Some(wire) = self.conn.take() {
1535 close_gracefully(wire);
1536 }
1537 for waiter in self.resume_notify.drain(..) {
1538 let _ = waiter.send(());
1539 }
1540 let _ = reply.send(());
1541 Flow::Continue
1542 }
1543
1544 async fn resume(&mut self, reply: oneshot::Sender<()>) -> Flow {
1547 tracing::info!(
1548 leases = self.ledger.leases.len(),
1549 "bidi transport: resuming — catch up, then done"
1550 );
1551 let was_suspended = self.suspended;
1552 self.suspended = false;
1553 self.resume_notify.push(reply);
1554 if self.conn.is_some() {
1555 self.settle_caught_up_waiters();
1556 return Flow::Continue;
1557 }
1558 if self.ledger.leases.is_empty() {
1559 self.settle_idle_waiters();
1560 return Flow::Continue;
1561 }
1562 if was_suspended {
1563 self.reconnect_delay = RECONNECT_INITIAL_DELAY;
1564 self.reconnect().await
1565 } else {
1566 Flow::Continue
1567 }
1568 }
1569
1570 fn wire_event(&mut self, event: Event<B::GroupMessage, B::WelcomeMessage>) -> Flow {
1571 let has_incoming = self
1572 .ledger
1573 .leases
1574 .values()
1575 .any(|lease| lease.incoming.is_some());
1576 let incoming_dropped = if has_incoming {
1577 match &event {
1578 Event::Applied { id, targets } if !self.ledger.valid_targets(*id, targets) => {
1579 return self.fail_incoming("Applied targets");
1580 }
1581 Event::GroupMessages { messages } => {
1582 self.ledger.demux_incoming(messages, B::group_envelope)
1583 }
1584 Event::WelcomeMessages { messages } => {
1585 self.ledger.demux_incoming(messages, B::welcome_envelope)
1586 }
1587 _ => Ok(Vec::new()),
1588 }
1589 } else {
1590 Ok(Vec::new())
1591 };
1592 let incoming_dropped = match incoming_dropped {
1593 Ok(dropped) => dropped,
1594 Err(reason) => return self.fail_incoming(reason),
1595 };
1596 let mut dropped = match event {
1597 Event::Started { .. } => Vec::new(),
1598 Event::Applied { id, targets } => {
1599 let readds = self.ledger.applied(id, targets);
1600 self.outbox.updates.extend(self.ledger.prepare_adds(readds));
1601 Vec::new()
1602 }
1603 Event::GroupMessages { messages } => self.ledger.demux(
1604 messages,
1605 DeliveryKind::Group,
1606 B::group_topic,
1607 B::group_cursor,
1608 LeaseEvent::GroupMessages,
1609 ),
1610 Event::WelcomeMessages { messages } => self.ledger.demux(
1611 messages,
1612 DeliveryKind::Welcome,
1613 B::welcome_topic,
1614 B::welcome_cursor,
1615 LeaseEvent::WelcomeMessages,
1616 ),
1617 };
1618 dropped.extend(incoming_dropped);
1619 dropped.extend(self.ledger.recheck());
1620 let removes = self.drop_leases(dropped);
1621 self.retire(removes);
1622 self.settle_idle_waiters();
1623 self.settle_caught_up_waiters();
1624 Flow::Continue
1625 }
1626
1627 fn fail_incoming(&mut self, reason: &'static str) -> Flow {
1628 let dropped = self
1629 .ledger
1630 .leases
1631 .iter()
1632 .filter_map(|(id, lease)| {
1633 lease.incoming.as_ref()?;
1634 *lease.incoming_failure.lock() = Some(TransportError::Protocol(reason));
1635 Some(*id)
1636 })
1637 .collect();
1638 let removes = self.drop_leases(dropped);
1639 self.retire(removes);
1640 self.settle_idle_waiters();
1641 Flow::Continue
1642 }
1643
1644 fn arm_reconnect(&mut self) -> std::time::Duration {
1645 let delay = self.reconnect_delay + xmtp_common::time::rand_offset(self.reconnect_delay);
1646 self.reconnect_at = tokio::time::Instant::now() + delay;
1647 delay
1648 }
1649
1650 fn wire_died(&mut self) -> Flow {
1651 self.outbox.clear();
1652 let failure = self.conn.as_ref().and_then(Connection::failure);
1653 drop(self.conn.take());
1654 let dropped = self
1655 .ledger
1656 .leases
1657 .iter()
1658 .filter_map(|(id, lease)| {
1659 let (sender, _) = lease.incoming.as_ref()?;
1660 let event = failure
1661 .as_ref()
1662 .map_or(Ok(vec![IncomingEvent::Disconnected]), |error| {
1663 Err(TransportError::Wire(error.clone()))
1664 });
1665 let full = match sender.try_send(event) {
1666 Ok(()) => false,
1667 Err(error) => {
1668 *lease.incoming_failure.lock() = Some(match error.into_inner() {
1669 Err(error) => error,
1670 Ok(_) => TransportError::Backpressure,
1671 });
1672 true
1673 }
1674 };
1675 (full || failure.as_ref().is_some_and(|error| !error.is_retryable())).then_some(*id)
1676 })
1677 .collect();
1678 self.drop_leases(dropped);
1679 self.ledger.reset_wire();
1680 self.close_wire_span("wire_end");
1681 if failure.is_some_and(|error| !error.is_retryable()) {
1682 return Flow::Shutdown;
1683 }
1684 let stable = self
1685 .wire_opened_at
1686 .take()
1687 .is_some_and(|opened| opened.elapsed() >= MIN_STABLE_UPTIME);
1688 self.reconnect_delay = if stable {
1689 RECONNECT_INITIAL_DELAY
1690 } else {
1691 (self.reconnect_delay * 2).min(RECONNECT_MAX_DELAY)
1692 };
1693 let retry_in = self.arm_reconnect();
1694 if !self.ledger.leases.is_empty() {
1695 tracing::warn!(
1696 leases = self.ledger.leases.len(),
1697 pending_updates = self.ledger.pending_updates.len(),
1698 retry_in_ms = retry_in.as_millis() as u64,
1699 "bidi transport: wire died; reopening from lease floors"
1700 );
1701 }
1702 self.settle_idle_waiters();
1703 Flow::Continue
1704 }
1705
1706 #[tracing::instrument(skip_all, fields(operation = "bidi.reconnect"))]
1707 async fn reconnect(&mut self) -> Flow {
1708 match self.reopen().await {
1709 AfterReopen::Proceed => Flow::Continue,
1710 AfterReopen::Suspended(ack) => {
1711 self.suspend_preempted(ack);
1712 Flow::Continue
1713 }
1714 AfterReopen::Shutdown => Flow::Shutdown,
1715 }
1716 }
1717
1718 fn suspend_preempted(&mut self, ack: oneshot::Sender<()>) {
1719 self.suspended = true;
1720 self.park_deferred_resumes();
1721 for waiter in self.resume_notify.drain(..) {
1722 let _ = waiter.send(());
1723 }
1724 let _ = ack.send(());
1725 }
1726
1727 async fn reopen(&mut self) -> AfterReopen {
1728 self.ledger.reset_wire();
1729 let adds = self.ledger.resume_adds();
1730 if adds.is_empty() {
1731 self.settle_idle_waiters();
1732 return AfterReopen::Proceed;
1733 }
1734 self.outbox.updates.extend(self.ledger.prepare_adds(adds));
1735 let Some((_, initial)) = self.outbox.updates.pop_front() else {
1736 return AfterReopen::Proceed;
1737 };
1738 match self.open_preemptibly(initial).await {
1739 OpenOutcome::Opened(wire) => {
1740 self.conn = Some(wire);
1741 self.open_wire_span();
1742 self.wire_opens += 1;
1743 AfterReopen::Proceed
1744 }
1745 OpenOutcome::Failed(error) => {
1746 self.outbox.clear();
1747 self.ledger.reset_wire();
1748 if error.is_locked_out() {
1749 tracing::warn!("bidi reopen waits for the auth cool-down: {error}");
1753 self.reconnect_delay = AUTH_LOCKOUT_COOLDOWN;
1754 self.arm_reconnect();
1755 self.park_deferred_resumes();
1756 return AfterReopen::Proceed;
1757 }
1758 if !error.is_retryable() {
1759 tracing::error!("bidi reconnect failed permanently: {error}");
1760 let error = std::sync::Arc::new(super::bidi::ConnectionFailure::Wire(
1761 xmtp_proto::api::NetworkError::new(error),
1762 ));
1763 for lease in self.ledger.leases.values() {
1764 if lease.incoming.is_some() {
1765 *lease.incoming_failure.lock() =
1766 Some(TransportError::Wire(error.clone()));
1767 }
1768 }
1769 return AfterReopen::Shutdown;
1770 }
1771 self.reconnect_delay = (self.reconnect_delay * 2).min(RECONNECT_MAX_DELAY);
1772 self.arm_reconnect();
1773 self.park_deferred_resumes();
1774 AfterReopen::Proceed
1775 }
1776 OpenOutcome::Suspended(ack) => {
1777 self.outbox.clear();
1778 self.ledger.reset_wire();
1779 AfterReopen::Suspended(ack)
1780 }
1781 OpenOutcome::Shutdown => AfterReopen::Shutdown,
1782 }
1783 }
1784
1785 async fn open_preemptibly(&mut self, mutate: B::Mutate) -> OpenOutcome<B> {
1788 let open = self.opener.open(mutate);
1789 tokio::pin!(open);
1790 loop {
1791 tokio::select! {
1792 result = &mut open => {
1793 return match result {
1794 Ok(wire) => OpenOutcome::Opened(wire),
1795 Err(e) => OpenOutcome::Failed(e),
1796 };
1797 }
1798 cmd = self.cmds.recv() => match cmd {
1799 Some(Cmd::Suspend { reply }) => return OpenOutcome::Suspended(reply),
1800 Some(cmd) => self.deferred.push_back(cmd),
1801 None => return OpenOutcome::Shutdown,
1802 },
1803 }
1804 }
1805 }
1806
1807 fn park_deferred_resumes(&mut self) {
1809 for cmd in std::mem::take(&mut self.deferred) {
1810 match cmd {
1811 Cmd::Resume { reply } => self.resume_notify.push(reply),
1812 other => self.deferred.push_back(other),
1813 }
1814 }
1815 }
1816
1817 fn settle_idle_waiters(&mut self) {
1818 if !self.ledger.leases.is_empty() {
1819 return;
1820 }
1821 self.ledger.reset_wire();
1822 for waiter in self.resume_notify.drain(..) {
1823 let _ = waiter.send(());
1824 }
1825 }
1826
1827 fn settle_caught_up_waiters(&mut self) {
1828 if self.conn.is_none() || !self.ledger.caught_up() || !self.outbox.is_empty() {
1829 return;
1830 }
1831 for waiter in self.resume_notify.drain(..) {
1832 let _ = waiter.send(());
1833 }
1834 }
1835
1836 fn retire(&mut self, removes: Vec<Topic>) {
1837 if self.ledger.leases.is_empty() {
1838 self.outbox.clear();
1839 self.ledger.reset_wire();
1840 self.close_wire_span("idle");
1841 if let Some(wire) = self.conn.take() {
1842 close_gracefully(wire);
1843 }
1844 return;
1845 }
1846 if self.conn.is_none() {
1847 return;
1848 }
1849 let mut to_remove = Vec::new();
1850 for topic in removes {
1851 if let Some(registration) = self.ledger.registrations.get_mut(&topic)
1852 && !matches!(registration.state, RegistrationState::Removing)
1853 {
1854 registration.state = RegistrationState::Removing;
1855 to_remove.push(topic);
1856 }
1857 }
1858 self.outbox
1859 .updates
1860 .extend(self.ledger.prepare_removes(to_remove));
1861 }
1862
1863 fn coalesced_prefix(&self) -> Option<(usize, PendingUpdate<B::Cursor>)> {
1866 if self.outbox.updates.len() < 2 {
1867 return None;
1868 }
1869 let (first_id, _) = self.outbox.updates.front()?;
1870 let first = self.ledger.pending_updates.get(first_id)?;
1871 let adds_only = !first.adds.is_empty() && first.removes.is_empty();
1872 let removes_only = first.adds.is_empty() && !first.removes.is_empty();
1873 if !adds_only && !removes_only {
1874 return None;
1875 }
1876 let topic_cap = if adds_only {
1882 self.ledger.mutate.add_cap
1883 } else {
1884 self.ledger.mutate.remove_cap
1885 };
1886 let mut topics: HashSet<_> = first
1887 .adds
1888 .iter()
1889 .map(|(topic, _)| topic)
1890 .chain(first.removes.iter())
1891 .collect();
1892 let mut bytes: usize = topics.iter().map(|topic| topic_wire_cost(topic)).sum();
1893 let mut count = 1;
1894 for (id, _) in self.outbox.updates.iter().skip(1) {
1895 let Some(next) = self.ledger.pending_updates.get(id) else {
1896 break;
1897 };
1898 if (adds_only && (next.adds.is_empty() || !next.removes.is_empty()))
1899 || (removes_only && (next.removes.is_empty() || !next.adds.is_empty()))
1900 {
1901 break;
1902 }
1903 let next_topics: Vec<_> = next
1904 .adds
1905 .iter()
1906 .map(|(topic, _)| topic)
1907 .chain(next.removes.iter())
1908 .collect();
1909 let next_bytes: usize = next_topics.iter().map(|topic| topic_wire_cost(topic)).sum();
1910 if topics.len() + next_topics.len() > topic_cap
1911 || bytes + next_bytes > self.ledger.mutate.byte_cap
1912 || next_topics.iter().any(|topic| topics.contains(topic))
1913 {
1914 break;
1915 }
1916 topics.extend(next_topics);
1917 bytes += next_bytes;
1918 count += 1;
1919 }
1920 if count == 1 {
1921 return None;
1922 }
1923 let mut merged = PendingUpdate {
1924 adds: Vec::new(),
1925 removes: Vec::new(),
1926 };
1927 for (id, _) in self.outbox.updates.iter().take(count) {
1928 let update = self.ledger.pending_updates.get(id)?;
1929 merged.adds.extend(update.adds.iter().cloned());
1930 merged.removes.extend(update.removes.iter().cloned());
1931 }
1932 Some((count, merged))
1933 }
1934
1935 fn flush_outbox(&mut self) {
1938 let Some(wire) = self.conn.as_ref() else {
1939 return;
1940 };
1941 while let Some((id, _)) = self.outbox.updates.front() {
1942 if !self.update_budget.take() {
1943 return;
1944 }
1945 let id = *id;
1946 if let Some((count, merged)) = self.coalesced_prefix() {
1947 let update = B::build_mutate(merged.adds.clone(), merged.removes.clone(), id);
1948 if wire.try_mutate(update).is_err() {
1949 self.update_budget.refund();
1950 return;
1951 }
1952 for (old_id, _) in self.outbox.updates.drain(..count) {
1953 self.ledger.pending_updates.remove(&old_id);
1954 }
1955 self.ledger.pending_updates.insert(id, merged);
1956 continue;
1957 }
1958 let Some((id, update)) = self.outbox.updates.pop_front() else {
1959 break;
1960 };
1961 match wire.try_mutate(update) {
1962 Ok(()) => {}
1963 Err(TryMutateError::Full(update)) | Err(TryMutateError::Closed(update)) => {
1964 self.update_budget.refund();
1965 self.outbox.updates.push_front((id, update));
1966 return;
1967 }
1968 }
1969 }
1970 }
1971}
1972
1973enum OpenOutcome<B: BidiBinding> {
1974 Opened(Connection<B>),
1975 Failed(OpenError),
1976 Suspended(oneshot::Sender<()>),
1977 Shutdown,
1978}
1979
1980enum AfterReopen {
1981 Proceed,
1982 Suspended(oneshot::Sender<()>),
1983 Shutdown,
1984}
1985
1986struct Outbox<M> {
1987 updates: std::collections::VecDeque<(u64, M)>,
1988}
1989
1990impl<M> Default for Outbox<M> {
1991 fn default() -> Self {
1992 Self {
1993 updates: std::collections::VecDeque::new(),
1994 }
1995 }
1996}
1997
1998impl<M> Outbox<M> {
1999 fn is_empty(&self) -> bool {
2000 self.updates.is_empty()
2001 }
2002 fn clear(&mut self) {
2003 self.updates.clear();
2004 }
2005
2006 fn purge(&mut self, ids: &HashSet<u64>) -> Vec<u64> {
2007 let mut purged = Vec::new();
2008 self.updates.retain(|(id, _)| {
2009 if ids.contains(id) {
2010 purged.push(*id);
2011 false
2012 } else {
2013 true
2014 }
2015 });
2016 purged
2017 }
2018}
2019
2020fn close_gracefully<B: TransportBinding>(wire: Connection<B>)
2023where
2024 B::GroupMessage: Clone,
2025 B::WelcomeMessage: Clone,
2026{
2027 xmtp_common::spawn(None, async move {
2028 let mut wire = wire;
2029 let drained = xmtp_common::time::timeout(GRACEFUL_CLOSE_BUDGET, async {
2030 if wire.finish().await.is_err() {
2031 return; }
2033 while wire.next().await.is_some() {}
2034 })
2035 .await;
2036 if drained.is_err() {
2037 tracing::debug!(
2038 budget_ms = GRACEFUL_CLOSE_BUDGET.as_millis() as u64,
2039 "bidi transport: graceful-close drain budget expired; dropping the wire"
2040 );
2041 }
2042 });
2043}
2044
2045#[cfg(test)]
2046mod tests;