Skip to main content

xmtp_api_backend/
streams.rs

1use crate::{
2    BackendClient,
3    backend::SubscribeStatic,
4    envelope::{
5        decode_group_message, decode_welcome_message, ordered_batches, registration_targets,
6    },
7    queries::stream::try_extractor,
8};
9use futures::{StreamExt, stream};
10use std::collections::VecDeque;
11use xmtp_common::{
12    BoxDynStream,
13    time::{Duration, timeout},
14};
15use xmtp_configuration::BACKEND_DEFAULT_KEEPALIVE_INTERVAL_MS;
16use xmtp_proto::{
17    api::{ApiClientError, Client, QueryStreamExt},
18    api_client::{XmtpBackendClient, XmtpMlsStreams},
19    backend_v1 as wire,
20    types::{
21        Cursor, GroupId, GroupMessage, IncomingBatchLimits, IncomingEvent, IncomingSubscription,
22        InstallationId, Topic, TopicCursor, WelcomeMessage,
23    },
24};
25
26const SILENT_INTERVALS: u32 = 3;
27
28impl<C: Client> BackendClient<C> {
29    async fn newest_cursors(
30        &self,
31        topics: impl IntoIterator<Item = Topic>,
32    ) -> Result<TopicCursor, ApiClientError> {
33        let mut cursors: TopicCursor = topics.into_iter().map(|topic| (topic, Cursor(0))).collect();
34        let topics: Vec<_> = cursors.keys().cloned().collect();
35        for chunk in topics.chunks(self.limits().max_newest_metadata_topics) {
36            let response = self
37                .query_newest(wire::QueryNewestRequest {
38                    topics: chunk
39                        .iter()
40                        .map(|t| wire::Topic {
41                            topic: t.cloned_vec(),
42                        })
43                        .collect(),
44                    include_full_envelope: false,
45                })
46                .await?;
47            for result in response.results {
48                let topic =
49                    Topic::parse(&result.topic.ok_or_else(|| malformed("newest topic"))?.topic)?;
50                let cursor = result
51                    .meta
52                    .and_then(|m| m.cursor)
53                    .ok_or_else(|| malformed("newest cursor"))?;
54                *cursors
55                    .get_mut(&topic)
56                    .ok_or_else(|| malformed("unrequested newest topic"))? = cursor.into();
57            }
58        }
59        Ok(cursors)
60    }
61
62    /// Split static topic interest while preserving each registration's targets.
63    /// Any feed error ends all chunks; the caller reopens from durable receipt.
64    async fn static_events(
65        &self,
66        cursors: &TopicCursor,
67        limits: IncomingBatchLimits,
68    ) -> Result<BoxDynStream<'static, Result<IncomingEvent, ApiClientError>>, ApiClientError> {
69        if cursors.is_empty() {
70            return Ok(Box::pin(stream::pending()));
71        }
72        let topics: Vec<_> = cursors
73            .iter()
74            .map(|(topic, cursor)| (topic.clone(), *cursor))
75            .collect();
76        let mut streams = Vec::new();
77        for chunk in topics.chunks(self.limits().max_static_topics) {
78            let stream = SubscribeStatic(wire::SubscribeStaticRequest {
79                topics: chunk
80                    .iter()
81                    .map(|(topic, cursor)| wire::TopicQuery {
82                        topic: Some(wire::Topic {
83                            topic: topic.cloned_vec(),
84                        }),
85                        cursor: Some((*cursor).into()),
86                    })
87                    .collect(),
88            })
89            .subscribe(&self.client)
90            .await?;
91            streams.push(normalize_static_stream(
92                Box::pin(stream),
93                chunk.iter().cloned().collect(),
94                limits,
95            ));
96        }
97        // A wire error invalidates the complete subscription. Reopen all topics from durable cursors.
98        Ok(Box::pin(stream::unfold(
99            Some(stream::select_all(streams)),
100            |streams| async move {
101                let mut streams = streams?;
102                let item = streams.next().await?;
103                let ended = item.is_err() || matches!(item, Ok(IncomingEvent::Disconnected));
104                let remaining = if ended { None } else { Some(streams) };
105                Some((item, remaining))
106            },
107        )))
108    }
109
110    async fn static_envelopes(
111        &self,
112        cursors: &TopicCursor,
113    ) -> Result<
114        BoxDynStream<'static, Result<Vec<wire::ServerEnvelope>, ApiClientError>>,
115        ApiClientError,
116    > {
117        let events = self
118            .static_events(
119                cursors,
120                IncomingBatchLimits {
121                    max_rows: xmtp_configuration::BACKEND_DEFAULT_MAX_QUERY_LIMIT,
122                    max_bytes: xmtp_configuration::BACKEND_DEFAULT_MAX_REQUEST_BYTES,
123                },
124            )
125            .await?;
126        Ok(Box::pin(events.filter_map(|event| async move {
127            match event {
128                Ok(IncomingEvent::OrderedBatch(batch)) => Some(Ok(batch.envelopes)),
129                Ok(IncomingEvent::Registered { .. }) => None,
130                Ok(IncomingEvent::Disconnected) => None,
131                Err(error) => Some(Err(error)),
132            }
133        })))
134    }
135}
136
137/// Require `Started` before data and retain its fixed catch-up targets.
138/// Validate each complete frame before emitting batches; receipt stays external.
139fn normalize_static_stream(
140    stream: BoxDynStream<'static, Result<wire::SubscribeStaticResponse, ApiClientError>>,
141    starts: TopicCursor,
142    limits: IncomingBatchLimits,
143) -> BoxDynStream<'static, Result<IncomingEvent, ApiClientError>> {
144    let cursors = starts.clone();
145    Box::pin(stream::unfold(
146        (
147            stream,
148            Duration::from_millis(BACKEND_DEFAULT_KEEPALIVE_INTERVAL_MS),
149            starts,
150            cursors,
151            false,
152            VecDeque::new(),
153            false,
154        ),
155        move |(
156            mut stream,
157            mut interval,
158            starts,
159            mut cursors,
160            mut registered,
161            mut pending,
162            mut ended,
163        )| async move {
164            if ended {
165                return None;
166            }
167            let item = loop {
168                if let Some(event) = pending.pop_front() {
169                    break Ok(event);
170                }
171                let frame = match timeout(interval * SILENT_INTERVALS, stream.next()).await {
172                    Err(error) => break Err(ApiClientError::Expired(error)),
173                    Ok(None) => break Ok(IncomingEvent::Disconnected),
174                    Ok(Some(Err(error))) => break Err(error),
175                    Ok(Some(Ok(frame))) => frame,
176                };
177                match frame.response {
178                    Some(wire::subscribe_static_response::Response::Started(started))
179                        if !registered =>
180                    {
181                        let targets = match registration_targets(&starts, started.targets) {
182                            Ok(targets) => targets,
183                            Err(error) => break Err(error.into()),
184                        };
185                        registered = true;
186                        if started.keepalive_interval_ms != 0 {
187                            interval = Duration::from_millis(started.keepalive_interval_ms.into());
188                        }
189                        break Ok(IncomingEvent::Registered {
190                            starts: starts.clone(),
191                            targets,
192                        });
193                    }
194                    Some(wire::subscribe_static_response::Response::Keepalive(_)) if registered => {
195                    }
196                    Some(wire::subscribe_static_response::Response::Messages(messages))
197                        if registered =>
198                    {
199                        match ordered_batches(&mut cursors, messages.envelopes, limits) {
200                            Ok(batches) => {
201                                pending.extend(batches.into_iter().map(IncomingEvent::OrderedBatch))
202                            }
203                            Err(error) => break Err(error.into()),
204                        }
205                    }
206                    _ => break Err(malformed("static response order")),
207                }
208            };
209            ended = item.is_err() || matches!(item, Ok(IncomingEvent::Disconnected));
210            Some((
211                item,
212                (
213                    stream, interval, starts, cursors, registered, pending, ended,
214                ),
215            ))
216        },
217    ))
218}
219fn malformed(field: &str) -> ApiClientError {
220    ApiClientError::OtherUnretryable(format!("missing or invalid {field}").into())
221}
222
223#[xmtp_common::async_trait]
224impl<C: Client> XmtpMlsStreams for BackendClient<C> {
225    type Error = ApiClientError;
226    type GroupMessageStream = BoxDynStream<'static, Result<GroupMessage, ApiClientError>>;
227    type WelcomeMessageStream = BoxDynStream<'static, Result<WelcomeMessage, ApiClientError>>;
228    async fn subscribe_envelopes_with_cursors(
229        &self,
230        cursors: &TopicCursor,
231        limits: IncomingBatchLimits,
232    ) -> Result<IncomingSubscription<Self::Error>, Self::Error> {
233        Ok(IncomingSubscription::new(
234            self.static_events(cursors, limits).await?,
235            |_| {},
236        ))
237    }
238    async fn subscribe_group_messages(
239        &self,
240        groups: &[&GroupId],
241    ) -> Result<Self::GroupMessageStream, Self::Error> {
242        let cursors = self
243            .newest_cursors(
244                groups
245                    .iter()
246                    .map(|id| Topic::new_group_message(id.as_ref())),
247            )
248            .await?;
249        self.subscribe_group_messages_with_cursors(&cursors).await
250    }
251    async fn subscribe_group_messages_with_cursors(
252        &self,
253        cursors: &TopicCursor,
254    ) -> Result<Self::GroupMessageStream, Self::Error> {
255        Ok(Box::pin(try_extractor(
256            self.static_envelopes(cursors).await?,
257            |envelope| decode_group_message(envelope).map_err(Into::into),
258        )))
259    }
260    async fn subscribe_welcome_messages(
261        &self,
262        installations: &[&InstallationId],
263    ) -> Result<Self::WelcomeMessageStream, Self::Error> {
264        let cursors = self
265            .newest_cursors(
266                installations
267                    .iter()
268                    .map(|id| Topic::new_welcome_message(**id)),
269            )
270            .await?;
271        self.subscribe_welcome_messages_with_cursors(&cursors).await
272    }
273    async fn subscribe_welcome_messages_with_cursors(
274        &self,
275        cursors: &TopicCursor,
276    ) -> Result<Self::WelcomeMessageStream, Self::Error> {
277        Ok(Box::pin(try_extractor(
278            self.static_envelopes(cursors).await?,
279            |envelope| decode_welcome_message(envelope).map_err(Into::into),
280        )))
281    }
282}
283
284#[cfg(test)]
285mod tests;