xmtp_mls/subscriptions/
router_callbacks.rs

1//! Native callback streams over the backend subscription router.
2//!
3//! The receiver owns the shared wires and their suspension intent
4//! ([`shared_wires`]); this module keeps the app's suspend and resume calls.
5//! Retryable wire failures keep leases alive and reconnect with backoff.
6//! Each callback stream reports close on terminal failure or explicit close.
7
8use tokio::sync::oneshot;
9
10use xmtp_common::{MaybeSend, StreamHandle};
11use xmtp_db::consent_record::ConsentState;
12use xmtp_db::group::ConversationType;
13use xmtp_db::group_message::StoredGroupMessage;
14use xmtp_proto::api_client::{XmtpMlsBidiStreams, XmtpMlsStreams};
15use xmtp_proto::types::{GroupId, InstallationId};
16
17use xmtp_common::Event;
18use xmtp_macro::log_event;
19
20use super::{
21    Result, StreamKind, SubscribeError,
22    incoming::stream_api::shared_wires,
23    watchdog::{StreamCancel, close_reason},
24};
25use crate::context::XmtpSharedContext;
26use crate::groups::MlsGroup;
27
28pub use super::incoming::stream_api::shared_wires::ApiClientIdentity;
29#[cfg(any(test, feature = "test-utils"))]
30pub use super::incoming::stream_api::shared_wires::shared_transport_count;
31
32/// Suspend every shared wire and remember the intent for new wires.
33pub async fn suspend_bidi_streams() -> Result<()> {
34    shared_wires::suspend()
35        .await
36        .map_err(SubscribeError::Transport)
37}
38
39/// Resume shared wires. Reconnect and target processing continue in their tasks.
40pub async fn resume_bidi_streams() -> Result<()> {
41    shared_wires::resume();
42    Ok(())
43}
44
45/// Close telemetry and callback also run when the stream task is aborted.
46struct StreamClosedGuard<F: FnOnce()> {
47    kind: StreamKind,
48    installation: InstallationId,
49    on_close: Option<F>,
50}
51
52impl<F: FnOnce()> Drop for StreamClosedGuard<F> {
53    fn drop(&mut self) {
54        log_event!(Event::StreamClosed, self.installation, kind = ?self.kind);
55        if let Some(on_close) = self.on_close.take() {
56            on_close();
57        }
58    }
59}
60
61pub(crate) struct StreamOrigin {
62    pub(crate) kind: StreamKind,
63    pub(crate) installation: InstallationId,
64    pub(crate) cancel: StreamCancel,
65}
66
67impl StreamOrigin {
68    fn new<Context: XmtpSharedContext>(kind: StreamKind, context: &Context) -> Self {
69        Self {
70            kind,
71            installation: context.installation_id(),
72            cancel: StreamCancel::new(context),
73        }
74    }
75}
76
77/// Deliver local stream items. The next poll acknowledges the preceding callback return.
78pub(crate) fn pump_stream<T, S, St>(
79    origin: StreamOrigin,
80    subscribe: S,
81    mut callback: impl FnMut(Result<T>) + MaybeSend + 'static,
82    on_close: impl FnOnce() + MaybeSend + 'static,
83) -> impl StreamHandle<StreamOutput = Result<()>>
84where
85    T: MaybeSend + 'static,
86    St: futures::Stream<Item = Result<T>> + MaybeSend + Unpin + 'static,
87    S: Future<Output = Result<St>> + MaybeSend + 'static,
88{
89    use futures::StreamExt;
90    let (tx, rx) = oneshot::channel();
91    let task = async move {
92        let StreamOrigin {
93            kind,
94            installation,
95            cancel,
96        } = origin;
97        log_event!(Event::StreamOpened, installation, kind = ?kind);
98        let _closed = StreamClosedGuard {
99            kind,
100            installation,
101            on_close: Some(on_close),
102        };
103        let mut stream = tokio::select! {
104            biased;
105            _ = cancel.cancelled() => return close_reason(&cancel, true, &mut callback),
106            result = subscribe => match result {
107                Ok(stream) => stream,
108                Err(error) => {
109                    close_reason(&cancel, true, &mut callback)?;
110                    callback(Err(error));
111                    return Ok(());
112                }
113            },
114        };
115        let _ = tx.send(());
116        loop {
117            tokio::select! {
118                biased;
119                _ = cancel.cancelled() => return close_reason(&cancel, true, &mut callback),
120                next = stream.next() => match next {
121                    Some(item) => {
122                        // The connection can become blocked while the local read is in progress.
123                        // Report it before another item reaches the callback.
124                        close_reason(&cancel, true, &mut callback)?;
125                        let terminal = item.is_err();
126                        callback(item);
127                        if terminal {
128                            return Ok(());
129                        }
130                    },
131                    None => return close_reason(&cancel, true, &mut callback),
132                }
133            }
134        }
135    };
136    xmtp_common::spawn(Some(rx), xmtp_common::bind_task_hub(task))
137}
138
139/// `keep_alive` lives until the subscription opens.
140pub(crate) fn stream_all_messages_with_callback_dispatch<C>(
141    context: C,
142    keep_alive: impl MaybeSend + 'static,
143    conversation_type: Option<ConversationType>,
144    consent_states: Option<Vec<ConsentState>>,
145    callback: impl FnMut(Result<StoredGroupMessage>) + MaybeSend + 'static,
146    on_close: impl FnOnce() + MaybeSend + 'static,
147) -> impl StreamHandle<StreamOutput = Result<()>>
148where
149    C: XmtpSharedContext + 'static,
150    C::ApiClient:
151        XmtpMlsBidiStreams + XmtpMlsStreams + ApiClientIdentity + Clone + Send + Sync + 'static,
152    <C::ApiClient as XmtpMlsBidiStreams>::SubscribeStream: 'static,
153{
154    let origin = StreamOrigin::new(StreamKind::All, &context);
155    let subscribe = async move {
156        let _keep_alive = keep_alive;
157        super::stream_all::StreamAllMessages::new_owned(context, conversation_type, consent_states)
158            .await
159    };
160    pump_stream(origin, subscribe, callback, on_close)
161}
162
163/// `keep_alive` lives until the subscription opens.
164pub(crate) fn stream_conversations_with_callback_dispatch<C>(
165    context: C,
166    keep_alive: impl MaybeSend + 'static,
167    conversation_type: Option<ConversationType>,
168    include_duplicate_dms: bool,
169    callback: impl FnMut(Result<MlsGroup<C>>) + MaybeSend + 'static,
170    on_close: impl FnOnce() + MaybeSend + 'static,
171) -> impl StreamHandle<StreamOutput = Result<()>>
172where
173    C: XmtpSharedContext + 'static,
174    C::ApiClient:
175        XmtpMlsBidiStreams + XmtpMlsStreams + ApiClientIdentity + Clone + Send + Sync + 'static,
176    <C::ApiClient as XmtpMlsBidiStreams>::SubscribeStream: 'static,
177{
178    let origin = StreamOrigin::new(StreamKind::Conversations, &context);
179    let subscribe = async move {
180        let _keep_alive = keep_alive;
181        super::stream_conversations::StreamConversations::new_owned(
182            context,
183            conversation_type,
184            include_duplicate_dms,
185            None,
186        )
187        .await
188    };
189    pump_stream(origin, subscribe, callback, on_close)
190}
191
192pub fn stream_conversation_messages_with_callback_dispatch<C>(
193    context: C,
194    group_id: GroupId,
195    callback: impl FnMut(Result<StoredGroupMessage>) + MaybeSend + 'static,
196    on_close: impl FnOnce() + MaybeSend + 'static,
197) -> impl StreamHandle<StreamOutput = Result<()>>
198where
199    C: XmtpSharedContext + 'static,
200    C::ApiClient:
201        XmtpMlsBidiStreams + XmtpMlsStreams + ApiClientIdentity + Clone + Send + Sync + 'static,
202    <C::ApiClient as XmtpMlsBidiStreams>::SubscribeStream: 'static,
203{
204    let origin = StreamOrigin::new(StreamKind::Messages, &context);
205    let subscribe = async move {
206        super::stream_messages::StreamGroupMessages::new_owned(context, vec![group_id]).await
207    };
208    pump_stream(origin, subscribe, callback, on_close)
209}
xmtp_mls/subscriptions/router_callbacks.rs

xmtp_mls/subscriptions/
router_callbacks.rs
