xmtp_mls/subscriptions/incoming/stream_api/
shared_wires.rs

1//! The process registry of native shared Bidi transports, and the suspension
2//! intent that every new transport inherits.
3//!
4//! Clients share a wire only with the same API client Arc, host, and admitted limits.
5//! Each client factory owns its transport. The registry keeps only weak handles.
6//! A live transport retains the API client, so its address cannot be reused.
7
8use std::collections::HashMap;
9use std::sync::{Arc, LazyLock, Weak};
10
11use parking_lot::Mutex;
12
13use xmtp_api_backend::{
14    BackendBinding, BidiConnection, BidiTransport, MutateLimits, OpenError, TransportError,
15};
16use xmtp_proto::api_client::XmtpMlsBidiStreams;
17
18/// Stable identity of the Arc that owns a backend API client.
19pub trait ApiClientIdentity {
20    /// The address of the owning Arc, the same for every clone of it.
21    fn api_client_identity(&self) -> usize;
22}
23
24impl<C: ?Sized> ApiClientIdentity for Arc<C> {
25    fn api_client_identity(&self) -> usize {
26        Arc::as_ptr(self).cast::<()>() as usize
27    }
28}
29
30type WireKey = (String, usize, MutateLimits);
31
32struct SharedWires {
33    transports: HashMap<WireKey, Weak<BidiTransport<BackendBinding>>>,
34    suspend_requested: bool,
35}
36
37static SHARED_WIRES: LazyLock<Mutex<SharedWires>> = LazyLock::new(|| {
38    Mutex::new(SharedWires {
39        transports: HashMap::new(),
40        suspend_requested: false,
41    })
42});
43
44/// Prevent automatic unary receipt from bypassing native stream suspension.
45pub(in crate::subscriptions::incoming) fn bidi_streams_suspended() -> bool {
46    SHARED_WIRES.lock().suspend_requested
47}
48
49/// How many shared router transports are still alive.
50#[cfg(any(test, feature = "test-utils"))]
51pub fn shared_transport_count() -> usize {
52    let mut wires = SHARED_WIRES.lock();
53    wires
54        .transports
55        .retain(|_, transport| transport.strong_count() > 0);
56    wires.transports.len()
57}
58
59/// The shared transport for the destination `api` dials, created at that
60/// destination's first stream (see the module docs and [`SHARED_WIRES`]).
61///
62/// A transport's ledger task is bound to the async runtime alive at its
63/// first use — the one-runtime-per-process reality of mobile, node, and
64/// agents. A process that tears its runtime down and starts another (some
65/// non-`nextest` test harnesses) would find the cached transport dead;
66/// `nextest`'s process-per-test model keeps tests clear of that.
67#[cfg(test)]
68pub(crate) fn shared_transport<C>(api: C) -> Arc<BidiTransport<BackendBinding>>
69where
70    C: XmtpMlsBidiStreams + ApiClientIdentity + Clone + Send + Sync + 'static,
71    C::SubscribeStream: 'static,
72{
73    let mutate = MutateLimits::from_limits(&api.bidi_limits());
74    shared_transport_within(api, mutate)
75}
76
77pub(in crate::subscriptions::incoming) fn shared_transport_within<C>(
78    api: C,
79    mutate: MutateLimits,
80) -> Arc<BidiTransport<BackendBinding>>
81where
82    C: XmtpMlsBidiStreams + ApiClientIdentity + Clone + Send + Sync + 'static,
83    C::SubscribeStream: 'static,
84{
85    // Chunk and pace below ApiClientWrapper with the owning client's snapshot.
86    // An API Arc can serve clients admitted with different limits.
87    let key = (api.host().to_owned(), api.api_client_identity(), mutate);
88    let mut wires = SHARED_WIRES.lock();
89    // A transport created while the app is backgrounded is born suspended —
90    // its first lease parks instead of dialing (see [`SharedWires`]).
91    let born_suspended = wires.suspend_requested;
92    wires
93        .transports
94        .retain(|_, transport| transport.strong_count() > 0);
95    if let Some(transport) = wires.transports.get(&key).and_then(Weak::upgrade) {
96        return transport;
97    }
98    tracing::info!(
99        host = %key.0,
100        suspended = born_suspended,
101        "bidi: initializing the shared transport for a destination"
102    );
103    let transport = Arc::new(BidiTransport::new_within(
104        move |initial| {
105            let api = api.clone();
106            async move {
107                BidiConnection::open(&api, initial)
108                    .await
109                    .map_err(OpenError::new)
110            }
111        },
112        born_suspended,
113        mutate,
114    ));
115    wires.transports.insert(key, Arc::downgrade(&transport));
116    transport
117}
118
119/// Suspend every shared wire and remember the intent for wires created
120/// later. Returns after every live wire acknowledges, with the first failure
121/// other than a closed wire.
122pub async fn suspend() -> Result<(), TransportError> {
123    // Enqueue each transport's command UNDER the same lock that flips the
124    // process flag, so a concurrent resume can't interleave its command send
125    // between our flag flip and ours — the actor then sees Suspend/Resume in
126    // the order the lock ordered the flags, not in scheduler order. The push is
127    // synchronous; only the reply await below runs outside the lock.
128    let (wires, pending): (Vec<_>, Vec<_>) = {
129        let mut shared = SHARED_WIRES.lock();
130        shared.suspend_requested = true;
131        shared
132            .transports
133            .iter()
134            .filter_map(|(host, weak)| {
135                weak.upgrade().map(|transport| {
136                    let pending = transport.enqueue_suspend();
137                    ((host.clone(), transport), pending)
138                })
139            })
140            .unzip()
141    };
142    settle_lifecycle(&wires, await_lifecycle_acks(pending).await)
143}
144
145/// Resume shared wires and clear the intent for wires created later.
146/// Reconnect and target processing continue in their tasks.
147pub fn resume() {
148    // Enqueue under the flag lock, same as [`suspend`]: the
149    // command order must match the flag order under all interleavings. The
150    // reply receivers are dropped — fire-and-forget — and an `Err` from a
151    // closed (tombstoned) transport is dropped with them: a destination
152    // permanently off the network is vacuously resumed.
153    let mut shared = SHARED_WIRES.lock();
154    shared.suspend_requested = false;
155    shared.transports.retain(|_, weak| {
156        if let Some(transport) = weak.upgrade() {
157            let _ = transport.enqueue_resume();
158            true
159        } else {
160            false
161        }
162    });
163}
164
165/// Await every enqueued lifecycle reply, mapping a dropped sender (a closed
166/// transport) to `Closed`. The command pushes already happened under the
167/// registry lock ([`suspend`]); only this await runs outside it,
168/// so a slow catch-up on one destination never delays another's command.
169async fn await_lifecycle_acks(
170    pending: Vec<Result<tokio::sync::oneshot::Receiver<()>, TransportError>>,
171) -> Vec<Result<(), TransportError>> {
172    futures::future::join_all(pending.into_iter().map(|reply| async move {
173        match reply {
174            Ok(rx) => rx.await.map_err(|_| TransportError::Closed),
175            Err(e) => Err(e),
176        }
177    }))
178    .await
179}
180
181/// Wait for every wire before reporting a lifecycle failure.
182fn settle_lifecycle(
183    wires: &[(WireKey, Arc<BidiTransport<BackendBinding>>)],
184    results: Vec<Result<(), TransportError>>,
185) -> Result<(), TransportError> {
186    let mut first_error = None;
187    for (_, result) in wires.iter().zip(results) {
188        match result {
189            Ok(()) => {}
190            Err(TransportError::Closed) => {}
191            Err(e) => first_error = first_error.or(Some(e)),
192        }
193    }
194    match first_error {
195        None => Ok(()),
196        Some(e) => Err(e),
197    }
198}
xmtp_mls/subscriptions/incoming/stream_api/shared_wires.rs

xmtp_mls/subscriptions/incoming/stream_api/
shared_wires.rs
