Skip to main content

xmtp_api_backend/queries/
bidi_transport.rs

1//! One backend subscription connection serves multiple topic leases.
2//!
3//! Each registration has a fixed catch-up target from Applied. Each lease
4//! completes after its targets are met. Updates that lack acknowledgements
5//! keep resume callers waiting, including updates that only remove topics.
6//!
7//! Messages carry no registration tag. Route them by their metadata topic.
8//! Each lease keeps a monotonic delivery position above its requested floor.
9//! A lease that needs older messages removes and re-adds the held topic.
10//! Until the remove acknowledgement, only old holders receive queued messages.
11//! At that boundary, current interest determines the re-add cursor and holders.
12//! After the add acknowledgement, those holders receive the new registration.
13//! The delivery positions discard overlap without storing message buffers.
14//!
15//! A connection failure keeps leases alive. Reconnect uses the current topic
16//! set and the minimum durable receipt position of its leases. Suspend releases
17//! the connection and keeps this state. Resume waits for all acknowledgements
18//! and targets. A slow lease is closed so its consumer can recover from storage.
19//!
20//! Ordered leases raise their floors only after durable receipt. A reconnect
21//! resets delivery positions to these floors so uncommitted batches replay.
22//! The consumer owns durable progress. This module does not decode MLS data
23//! or promise exactly-once callbacks across a process crash.
24
25use 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/// Keep the cool-down wait short in tests. Match the middleware's test value,
43/// so the wire waits exactly one cool-down and no longer.
44#[cfg(test)]
45const AUTH_LOCKOUT_COOLDOWN: std::time::Duration = std::time::Duration::from_millis(500);
46
47use super::bidi::{BidiBinding, Connection, Event, TryMutateError};
48
49/// Default event capacity for one lease. Slow consumers must recover from storage.
50pub const DEFAULT_LEASE_DEPTH: usize = 64;
51/// Retry a queued update when the connection is quiet and capacity may be free.
52const 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
61/// Maximum adds per backend update. The wire also has a separate topic cap.
62pub(crate) const MAX_MUTATE_TOPICS: usize = xmtp_configuration::BACKEND_DEFAULT_MAX_UPDATE_ADDS;
63/// Leave space for the request wrapper within the backend byte limit.
64pub(crate) const MAX_MUTATE_BYTES: usize =
65    xmtp_configuration::BACKEND_DEFAULT_MAX_REQUEST_BYTES - PER_ENTRY_OVERHEAD;
66/// Conservative protobuf overhead per topic, including cursor and length fields.
67const PER_ENTRY_OVERHEAD: usize = 64;
68
69/// The interest-update frame shapes one deployment accepts (CFG-064). Adds and
70/// removes have separate caps because the backend enforces them separately.
71#[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    /// Take the caps a deployment published, leaving the same room for the
90    /// request wrapper the compiled default leaves.
91    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
107/// Keep input order while splitting at either the topic count or byte budget.
108fn 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
132/// Split an add set into updates within the backend count and byte limits.
133/// The caller must also enforce the total topic limit for its connection.
134pub 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);
143/// Bound the background drain after request half-close.
144const GRACEFUL_CLOSE_BUDGET: std::time::Duration = std::time::Duration::from_secs(5);
145
146/// Keep the source error and its retry decision after type erasure.
147#[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    /// True while an authentication cool-down runs. The open cannot succeed now
167    /// but will become possible again without any caller action.
168    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    /// Capture the lockout state of a concrete error before it is erased.
193    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
218/// Map backend envelopes to topics and positions for the lease ledger.
219pub trait TransportBinding: BidiBinding
220where
221    Self::GroupMessage: Clone,
222    Self::WelcomeMessage: Clone,
223{
224    /// A position that can represent each fixed backend target.
225    type Cursor: Copy + Send + std::fmt::Debug + From<u64> + Into<u64> + 'static;
226
227    /// Build an update with a nonzero ID. IDs increase on each connection.
228    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    /// Return true when the position is at or above the delivered cursor.
243    fn covers(position: &Self::Cursor, delivered: &Self::Cursor) -> bool;
244    /// Return the greatest cursor covered by both positions.
245    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    /// The lease exceeds the topic limit. This error is not retryable.
257    #[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    /// The receive queue is full. Reopen from durable receipt positions.
264    #[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>>>;
271/// One bounded wire frame remains one queue item even when it covers many topics.
272type 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
289/// Encrypted envelopes and initial catch-up completion for one lease.
290pub enum LeaseEvent<B: TransportBinding>
291where
292    B::GroupMessage: Clone,
293    B::WelcomeMessage: Clone,
294{
295    /// Every topic has met its initial target. This is emitted once per lease.
296    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
304/// A cloneable handle to the connection and topic ledger task.
305pub 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
325/// Local topic interest. Dropping the handle removes its interest immediately.
326pub 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    /// Read the next event. A closed channel requires recovery from storage.
346    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    /// Read raw ordered events, including a terminal receive-capacity error.
355    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    /// Call this only after the supplied positions commit to local storage.
369    pub fn acknowledge_received(&self, cursors: TopicCursor) {
370        let _ = self.cmds.send(Cmd::Received {
371            id: self.id,
372            cursors,
373        });
374    }
375
376    /// Keep this lease alive until its owned event stream is dropped.
377    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    /// Open lazily through `opener`. A suspended transport waits for resume.
405    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    /// Open lazily, chunking interest updates to what the deployment published
414    /// (CFG-064).
415    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    /// Register topics with exclusive floors and a bounded event channel.
470    /// Refuse empty leases and leases that would exceed the wire topic limit.
471    /// While suspended, register locally without opening a connection.
472    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    /// Register an ordered raw receiver. Receipt acknowledgements raise its resume floors.
488    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    /// Half-close the connection. Keep leases and positions for resume.
513    #[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    /// Queue suspend synchronously so callers can order lifecycle commands.
521    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    /// Reopen and wait until every update is acknowledged and every target is met.
530    /// A later suspend or removal of the last lease also releases this waiter.
531    #[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    /// Queue resume without waiting for catch-up. Return its completion receiver.
539    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
599/// An update that still needs its acknowledgement.
600struct PendingUpdate<C> {
601    adds: Vec<(Topic, C)>,
602    removes: Vec<Topic>,
603}
604
605/// One topic registration on the current connection.
606struct 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    /// Old holders receive queued frames until the remove acknowledgement.
617    Removing,
618}
619
620/// The enclosing lease and topic map identify this obligation.
621struct 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    /// Raise reconnect floors from committed receipt, never from stream delivery.
684    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    /// Emit the accepted read positions and fixed targets before raw data.
706    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            // A full channel is detected before the next payload copy.
730            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    /// Require exactly one fixed target for each topic added by this update.
741    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    /// Validate the full frame before copying it into any raw receive channel.
754    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            // Reserve one slot for the complete validated frame before copying.
821            // One frame can contain more topics than the queue has slots.
822            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    /// Create local interest before opening or changing a connection.
855    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    /// Forget connection state. Keep each lease's delivery guard.
886    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    /// Reopen at the meet of the observed position and all requested floors.
915    ///
916    /// A floor never rises, so this meet is the floor in practice. See the
917    /// module header for why that makes reopen cost grow with lease age.
918    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    /// Record acknowledgements before the updates enter the connection queue.
934    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    /// Join an active registration, or remove it before requesting older data.
997    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, &registration.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    /// Apply ordered acknowledgement boundaries and return topics to re-add.
1026    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                // An absent target leaves an existing registration unchanged.
1047                if matches!(registration.state, RegistrationState::Adding) {
1048                    registration.state = RegistrationState::Active;
1049                }
1050                for holder in &registration.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    /// Check changed topics only. Send completion after their message batches.
1062    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 &registration.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(&registration.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    /// Remove local interest immediately. The caller queues wire removals.
1113    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    /// Route one frame by topic. A lease's position only moves forward.
1135    /// Reserve channel capacity before copying a lease's payload batch.
1136    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            // A new registration cannot deliver before its add acknowledgement.
1153            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 &registration.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
1208/// The task owns both the wire and the ledger. Send updates with try_mutate.
1209/// Never wait for command capacity here: this task must keep reading events.
1210enum 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            // A finite snapshot lets queued leases share an update without
1290            // letting a continuous command producer starve wire events.
1291            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        // The opener already sent the first Update on this connection.
1367        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        // Cancel an unsent add only when none of its topics has a holder.
1489        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                    // Keep a pending removal until Applied. A replacement must
1507                    // not add this topic before that old boundary is consumed.
1508                    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    /// Resume from suspension immediately. During an outage, join the pending
1545    /// catch-up without replacing the scheduled reconnect backoff.
1546    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                    // The cool-down clears on its own, so the wire must wait for
1750                    // it. A shutdown here would lose every subscription for the
1751                    // life of the process because nothing restarts this task.
1752                    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    /// Keep receiving commands during a dial. Suspend cancels the dial at once.
1786    /// Other commands keep their order and run after the dial ends.
1787    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    /// Merge concurrent resume calls into the same attempt and backoff.
1808    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    /// Combine a same-kind prefix without changing queued frames or ack state.
1864    /// Duplicate topics and add/remove boundaries stay in separate updates.
1865    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        // CFG-064: the deployment caps adds and removes separately, so the
1877        // merged frame is bounded by the cap for the kind it carries. Merging a
1878        // removes-only prefix up to `add_cap` would build a frame the backend
1879        // rejects with INVALID_ARGUMENT wherever a deployment publishes a
1880        // smaller `max_update_removes`.
1881        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    /// Commit a coalesced acknowledgement ID only after the wire accepts it.
1936    /// Backpressure leaves the original queue and pending IDs intact.
1937    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
2020/// Release the connection without blocking the ledger on command capacity.
2021/// Discard remaining events and abort the connection when the drain budget ends.
2022fn 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; // actor already gone — nothing to drain
2032            }
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;