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 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 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
139fn 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
222fn 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;