Skip to main content

xmtp_api_backend/
streams.rs

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