1use crate::{ApiClientWrapper, ApiError, Result, dyn_err};
3use futures::{StreamExt, TryStreamExt, stream};
4use prost::Message;
5use std::{
6 collections::{HashMap, HashSet, VecDeque},
7 future::Future,
8};
9use tonic::Code;
10use xmtp_common::{RetryableError, retry_async};
11use xmtp_configuration::*;
12use xmtp_proto::{
13 api::grpc_status,
14 api_client::XmtpBackendClient,
15 backend_v1 as wire,
16 types::{
17 CanonicalEnvelope, Cursor, IncomingBatchLimits, OrderedEnvelopeBatch, Topic, TopicCursor,
18 },
19};
20
21pub const MAX_PUBLISH_CHUNKS_IN_FLIGHT: usize = 4;
22pub const MAX_READ_CHUNKS_IN_FLIGHT: usize = 4;
23
24#[derive(Debug)]
26pub struct OrderedQueryPage {
27 pub batches: Vec<OrderedEnvelopeBatch>,
29 pub has_more: bool,
31}
32
33#[derive(Clone, Debug)]
34struct PublishEnvelope {
35 envelope: wire::ClientEnvelope,
36 canonical: CanonicalEnvelope,
37 topic: Topic,
38}
39
40#[derive(Clone, Debug)]
42pub struct PublishUnit {
43 envelopes: Vec<PublishEnvelope>,
44}
45impl PublishUnit {
46 pub fn new_within(
51 envelopes: Vec<wire::ClientEnvelope>,
52 limits: &LimitsConfiguration,
53 ) -> Result<Self> {
54 if envelopes.is_empty() {
55 return Err(ApiError::InvalidRequest("empty publish unit"));
56 }
57 let envelopes = envelopes
58 .into_iter()
59 .map(|envelope| {
60 let parsed = xmtp_mls_validation::parse_envelope(envelope)?;
61 if parsed.canonical.bytes.len() > limits.max_envelope_bytes {
62 return Err(ApiError::EnvelopeTooLarge);
63 }
64 Ok(PublishEnvelope {
65 envelope: parsed.envelope,
66 canonical: parsed.canonical,
67 topic: parsed.topic,
68 })
69 })
70 .collect::<Result<Vec<_>>>()?;
71 let unit = Self { envelopes };
72 let mut measure = PublishMeasure::new(limits);
73 measure.add(&unit);
74 if !measure.fits() {
75 return Err(ApiError::UnitTooLarge);
76 }
77 Ok(unit)
78 }
79
80 pub fn single_within(
81 envelope: wire::ClientEnvelope,
82 limits: &LimitsConfiguration,
83 ) -> Result<Self> {
84 Self::new_within(vec![envelope], limits)
85 }
86
87 pub fn new(envelopes: Vec<wire::ClientEnvelope>) -> Result<Self> {
90 Self::new_within(envelopes, &LimitsConfiguration::default())
91 }
92
93 pub fn single(envelope: wire::ClientEnvelope) -> Result<Self> {
94 Self::new(vec![envelope])
95 }
96}
97
98pub(crate) fn request(units: &[PublishUnit]) -> wire::PublishRequest {
99 wire::PublishRequest {
100 envelopes: units
101 .iter()
102 .flat_map(|unit| unit.envelopes.iter().map(|e| e.envelope.clone()))
103 .collect(),
104 }
105}
106struct PublishMeasure<'a> {
108 bytes: usize,
109 topics: HashSet<&'a Topic>,
110 limits: &'a LimitsConfiguration,
111}
112impl<'a> PublishMeasure<'a> {
113 fn new(limits: &'a LimitsConfiguration) -> Self {
114 Self {
115 bytes: 0,
116 topics: HashSet::new(),
117 limits,
118 }
119 }
120
121 fn add(&mut self, unit: &'a PublishUnit) {
122 for envelope in &unit.envelopes {
123 let len = envelope.canonical.bytes.len();
124 self.bytes += 1 + prost::length_delimiter_len(len) + len;
125 self.topics.insert(&envelope.topic);
126 }
127 }
128
129 fn fits(&self) -> bool {
130 self.bytes <= self.limits.max_request_bytes
131 && self.topics.len() <= self.limits.max_publish_topics
132 }
133}
134
135pub fn chunk_publish_within<'a>(
139 units: &'a [PublishUnit],
140 limits: &LimitsConfiguration,
141) -> Result<Vec<&'a [PublishUnit]>> {
142 let mut chunks = Vec::new();
143 let mut start = 0;
144 let mut measure = PublishMeasure::new(limits);
145 for (end, unit) in units.iter().enumerate() {
146 measure.add(unit);
147 if !measure.fits() {
148 if start == end {
149 return Err(ApiError::UnitTooLarge);
150 }
151 chunks.push(&units[start..end]);
152 start = end;
153 measure = PublishMeasure::new(limits);
154 measure.add(unit);
155 if !measure.fits() {
156 return Err(ApiError::UnitTooLarge);
157 }
158 }
159 }
160 if start < units.len() {
161 chunks.push(&units[start..]);
162 }
163 Ok(chunks)
164}
165
166pub fn chunk_publish(units: &[PublishUnit]) -> Result<Vec<&[PublishUnit]>> {
168 chunk_publish_within(units, &LimitsConfiguration::default())
169}
170
171const PUBLISH_ERROR_TYPE_URL: &str = "type.googleapis.com/xmtp.backend.v1.PublishError";
173
174pub(crate) fn publish_error_reason(
178 error: &(dyn std::error::Error + 'static),
179) -> Option<wire::publish_error::Reason> {
180 if crate::preflight::failure(error).is_some() {
181 return None;
182 }
183 let status = grpc_status(error).filter(|status| status.code() == Code::InvalidArgument)?;
184 let details = tonic_types::pb::Status::decode(status.details()).ok()?;
185 details
186 .details
187 .iter()
188 .filter(|detail| detail.type_url == PUBLISH_ERROR_TYPE_URL)
189 .find_map(|detail| wire::PublishError::decode(detail.value.as_slice()).ok())
190 .and_then(|error| wire::publish_error::Reason::try_from(error.reason).ok())
191}
192
193pub(crate) fn size_error(error: &(dyn std::error::Error + 'static)) -> bool {
194 if crate::preflight::failure(error).is_some() {
195 return false;
196 }
197 let Some(status) = grpc_status(error) else {
198 return false;
199 };
200 match status.code() {
201 Code::OutOfRange | Code::ResourceExhausted => true,
202 Code::InvalidArgument => {
203 publish_error_reason(error) == Some(wire::publish_error::Reason::TooLarge)
204 }
205 _ => false,
206 }
207}
208
209fn publish_size_error(error: &(dyn std::error::Error + 'static)) -> bool {
214 size_error(error) && grpc_status(error).is_none_or(|status| status.code() != Code::OutOfRange)
215}
216
217#[derive(Debug, thiserror::Error)]
219enum CallError<E> {
220 #[error(transparent)]
221 Resize(E),
222 #[error(transparent)]
223 Retry(E),
224}
225impl<E: RetryableError + 'static> RetryableError for CallError<E> {
226 fn is_retryable(&self) -> bool {
227 match self {
228 Self::Resize(_) => false,
229 Self::Retry(error) => {
230 crate::preflight::failure(error).is_none() && error.is_retryable()
231 }
232 }
233 }
234}
235
236impl<C> ApiClientWrapper<C> {
237 pub(crate) async fn retry_call<T, E, F, Fut>(
238 &self,
239 mut call: F,
240 resize: bool,
241 ) -> std::result::Result<T, E>
242 where
243 E: RetryableError + 'static,
244 F: FnMut() -> Fut,
245 Fut: Future<Output = std::result::Result<T, E>>,
246 {
247 retry_async!(
248 self.retry_strategy,
249 (async {
250 call().await.map_err(|error| {
251 if resize && size_error(&error) {
252 CallError::Resize(error)
253 } else {
254 CallError::Retry(error)
255 }
256 })
257 })
258 )
259 .map_err(|error| match error {
260 CallError::Resize(error) | CallError::Retry(error) => error,
261 })
262 }
263}
264
265impl<C: XmtpBackendClient> ApiClientWrapper<C> {
266 #[xmtp_common::rpc_span]
268 pub async fn publish_units(&self, units: Vec<PublishUnit>) -> Result<Vec<wire::EnvelopeMeta>> {
269 let chunks = chunk_publish_within(&units, self.limits())?
270 .into_iter()
271 .map(<[_]>::to_vec)
272 .collect::<Vec<_>>();
273 let mut responses: Vec<_> = stream::iter(chunks.into_iter().enumerate().map(
274 |(index, chunk)| async move {
275 Ok::<_, ApiError>((index, self.publish_chunk(&chunk).await?))
276 },
277 ))
278 .buffer_unordered(MAX_PUBLISH_CHUNKS_IN_FLIGHT)
279 .try_collect()
280 .await?;
281 responses.sort_by_key(|(index, _)| *index);
282 Ok(responses.into_iter().flat_map(|(_, metas)| metas).collect())
283 }
284
285 async fn publish_chunk(&self, units: &[PublishUnit]) -> Result<Vec<wire::EnvelopeMeta>> {
286 let mut pending = VecDeque::from([units]);
287 let mut metas = Vec::new();
288 while let Some(units) = pending.pop_front() {
289 let request = request(units);
290 match self
291 .retry_call(|| self.api_client.publish(request.clone()), units.len() > 1)
292 .await
293 {
294 Ok(response) => {
295 let expected: Vec<_> = units.iter().flat_map(|unit| &unit.envelopes).collect();
296 if response.envelope_metas.len() != expected.len() {
297 return Err(ApiError::InvalidResponse("publish metadata count"));
298 }
299 for (meta, envelope) in response.envelope_metas.iter().zip(expected) {
300 let (topic, _, _) =
301 xmtp_api_backend::envelope::metadata(meta, envelope.topic.kind())?;
302 if topic != envelope.topic {
303 return Err(ApiError::InvalidResponse("publish topic"));
304 }
305 }
306 metas.extend(response.envelope_metas);
307 }
308 Err(error) if publish_size_error(&error) && units.len() > 1 => {
309 let (left, right) = units.split_at(units.len() / 2);
310 pending.push_front(right);
311 pending.push_front(left);
312 }
313 Err(error) => return Err(dyn_err(error)),
314 }
315 }
316 Ok(metas)
317 }
318
319 pub async fn query_ordered_page(
323 &self,
324 cursors: TopicCursor,
325 limit: u32,
326 limits: IncomingBatchLimits,
327 ) -> Result<OrderedQueryPage> {
328 if limit == 0
329 || limit as usize > self.limits().max_query_limit
330 || limits.max_rows == 0
331 || limits.max_bytes == 0
332 {
333 return Err(ApiError::InvalidRequest("query limit"));
334 }
335 let topics: Vec<_> = cursors.into_iter().collect();
336 let mut pending: VecDeque<_> = topics
337 .chunks(self.limits().max_query_topics)
338 .map(|topics| (topics.to_vec(), limit))
339 .collect();
340 let mut page = OrderedQueryPage {
341 batches: Vec::new(),
342 has_more: false,
343 };
344 let mut remaining = limits;
345 while let Some((topics, requested_limit)) = pending.pop_front() {
346 if remaining.max_rows == 0 || remaining.max_bytes == 0 {
347 page.has_more = true;
348 break;
349 }
350 let requested_limit =
351 requested_limit.min(remaining.max_rows.min(u32::MAX as usize) as u32);
352 let request = wire::QueryRequest {
353 queries: topics
354 .iter()
355 .map(|(topic, cursor)| wire::TopicQuery {
356 topic: Some(wire::Topic {
357 topic: topic.cloned_vec(),
358 }),
359 cursor: Some((*cursor).into()),
360 })
361 .collect(),
362 limit: requested_limit,
363 };
364 let response = match self
365 .retry_call(
366 || self.api_client.query(request.clone()),
367 requested_limit > 1 || topics.len() > 1,
368 )
369 .await
370 {
371 Ok(response) => response,
372 Err(error) if size_error(&error) && requested_limit > 1 => {
373 pending.push_front((topics, (requested_limit / 2).max(1)));
374 continue;
375 }
376 Err(error) if size_error(&error) && topics.len() > 1 => {
377 let (left, right) = topics.split_at(topics.len() / 2);
378 pending.push_front((right.to_vec(), requested_limit));
379 pending.push_front((left.to_vec(), requested_limit));
380 continue;
381 }
382 Err(error) => return Err(dyn_err(error)),
383 };
384 let has_more = response
385 .continuation
386 .ok_or(ApiError::InvalidResponse("query continuation"))?
387 .has_more;
388 if has_more && response.envelopes.is_empty() {
389 return Err(ApiError::InvalidResponse("query has more without progress"));
390 }
391 if response.envelopes.len() > requested_limit as usize {
392 return Err(ApiError::InvalidResponse("query response limit"));
393 }
394 let mut positions = topics.iter().cloned().collect();
395 let batches = match xmtp_api_backend::envelope::ordered_batches(
396 &mut positions,
397 response.envelopes,
398 remaining,
399 ) {
400 Ok(batches) => batches,
401 Err(xmtp_api_backend::envelope::EnvelopeError::Capacity) if requested_limit > 1 => {
402 pending.push_front((topics, (requested_limit / 2).max(1)));
403 continue;
404 }
405 Err(xmtp_api_backend::envelope::EnvelopeError::Capacity)
406 if !page.batches.is_empty() =>
407 {
408 page.has_more = true;
409 break;
410 }
411 Err(error) => return Err(error.into()),
412 };
413 for batch in &batches {
414 remaining.max_rows -= batch.envelopes.len();
415 remaining.max_bytes -= batch
416 .envelopes
417 .iter()
418 .map(Message::encoded_len)
419 .sum::<usize>();
420 }
421 page.has_more |= has_more;
422 page.batches.extend(batches);
423 }
424 Ok(page)
425 }
426
427 pub async fn newest_topic_cursors(&self, topics: Vec<Topic>) -> Result<TopicCursor> {
431 let mut cursors: TopicCursor = topics
432 .iter()
433 .cloned()
434 .map(|topic| (topic, Cursor(0)))
435 .collect();
436 for result in self.newest(topics, false).await? {
437 let topic = Topic::parse(
438 &result
439 .topic
440 .ok_or(ApiError::InvalidResponse("newest topic"))?
441 .topic,
442 )?;
443 let meta = result
444 .meta
445 .ok_or(ApiError::InvalidResponse("newest metadata"))?;
446 let (_, cursor, _) = xmtp_api_backend::envelope::metadata(&meta, topic.kind())?;
447 cursors.insert(topic, cursor);
448 }
449 Ok(cursors)
450 }
451
452 pub async fn query_all(
454 &self,
455 cursors: TopicCursor,
456 limit: u32,
457 ) -> Result<Vec<wire::ServerEnvelope>> {
458 if limit == 0 || limit as usize > self.limits().max_query_limit {
459 return Err(ApiError::InvalidRequest("query limit"));
460 }
461 let topics: Vec<_> = cursors.into_iter().collect();
462 let results: Vec<Vec<_>> = stream::iter(
463 topics
464 .chunks(self.limits().max_query_topics)
465 .map(<[_]>::to_vec)
466 .collect::<Vec<_>>()
467 .into_iter()
468 .map(|chunk| self.query_chunk(chunk, limit)),
469 )
470 .buffer_unordered(MAX_READ_CHUNKS_IN_FLIGHT)
471 .try_collect()
472 .await?;
473 Ok(results.into_iter().flatten().collect())
474 }
475
476 async fn query_chunk(
477 &self,
478 topics: Vec<(Topic, Cursor)>,
479 limit: u32,
480 ) -> Result<Vec<wire::ServerEnvelope>> {
481 let mut output = Vec::new();
482 self.scan_chunk(topics, limit, |page| {
483 output.extend(page);
484 false
485 })
486 .await?;
487 Ok(output)
488 }
489
490 #[xmtp_common::rpc_span]
498 pub async fn settle_units(
499 &self,
500 units: &[PublishUnit],
501 ) -> Result<Vec<Option<wire::EnvelopeMeta>>> {
502 let envelopes: Vec<_> = units.iter().flat_map(|unit| &unit.envelopes).collect();
503 let mut pending: HashMap<_, Vec<_>> = HashMap::new();
504 for (index, envelope) in envelopes.iter().enumerate() {
505 pending
506 .entry(envelope.canonical.hash.to_vec())
507 .or_default()
508 .push(index);
509 }
510 let mut settled = vec![None; envelopes.len()];
511 let topics: HashSet<_> = envelopes.iter().map(|envelope| &envelope.topic).collect();
512 let topics: Vec<_> = topics
513 .into_iter()
514 .map(|topic| (topic.clone(), Cursor(0)))
515 .collect();
516 let limit = self.limits().max_query_limit as u32;
517 for chunk in topics.chunks(self.limits().max_query_topics) {
518 let mut invalid = false;
519 self.scan_chunk(chunk.to_vec(), limit, |page| {
520 invalid = page.into_iter().any(|stored| {
521 let Some(meta) = stored.meta else { return true };
522 let Ok(hash) = xmtp_api_backend::envelope::message_hash(&meta) else {
523 return true;
524 };
525 let Some(indices) = pending.remove(&hash) else {
526 return false;
527 };
528 let topic = &envelopes[indices[0]].topic;
530 let valid = xmtp_api_backend::envelope::metadata(&meta, topic.kind())
531 .is_ok_and(|(stored_topic, _, _)| stored_topic == *topic);
532 for index in indices {
533 settled[index] = Some(meta.clone());
534 }
535 !valid
536 });
537 invalid || pending.is_empty()
538 })
539 .await?;
540 if invalid {
541 return Err(ApiError::InvalidResponse("settled metadata"));
542 }
543 if pending.is_empty() {
544 break;
545 }
546 }
547 Ok(settled)
548 }
549
550 async fn scan_chunk(
553 &self,
554 topics: Vec<(Topic, Cursor)>,
555 limit: u32,
556 mut visit: impl FnMut(Vec<wire::ServerEnvelope>) -> bool,
557 ) -> Result<()> {
558 let mut pending = VecDeque::from([(topics, limit)]);
559 while let Some((mut topics, mut limit)) = pending.pop_front() {
560 loop {
561 let request = wire::QueryRequest {
562 queries: topics
563 .iter()
564 .map(|(topic, cursor)| wire::TopicQuery {
565 topic: Some(wire::Topic {
566 topic: topic.cloned_vec(),
567 }),
568 cursor: Some((*cursor).into()),
569 })
570 .collect(),
571 limit,
572 };
573 match self
574 .retry_call(
575 || self.api_client.query(request.clone()),
576 limit > 1 || topics.len() > 1,
577 )
578 .await
579 {
580 Ok(response) => {
581 let has_more = response
582 .continuation
583 .ok_or(ApiError::InvalidResponse("query continuation"))?
584 .has_more;
585 let mut cursors: HashMap<_, _> = topics
586 .iter_mut()
587 .map(|(topic, cursor)| (topic.clone(), cursor))
588 .collect();
589 let mut advanced = false;
590 for envelope in &response.envelopes {
591 let meta = envelope
592 .meta
593 .as_ref()
594 .ok_or(ApiError::InvalidResponse("query metadata"))?;
595 let topic = Topic::parse(
596 &meta
597 .topic
598 .as_ref()
599 .ok_or(ApiError::InvalidResponse("query topic"))?
600 .topic,
601 )?;
602 let cursor = cursors
603 .get_mut(&topic)
604 .ok_or(ApiError::InvalidResponse("unrequested query topic"))?;
605 let sequence = meta
606 .cursor
607 .as_ref()
608 .ok_or(ApiError::InvalidResponse("query cursor"))?
609 .sequence_id;
610 if sequence <= cursor.0 || sequence > i64::MAX as u64 {
611 return Err(ApiError::InvalidResponse(
612 "query cursor did not advance",
613 ));
614 }
615 **cursor = Cursor(sequence);
616 advanced = true;
617 }
618 if visit(response.envelopes) {
619 return Ok(());
620 }
621 if !has_more {
622 break;
623 }
624 if !advanced {
625 return Err(ApiError::InvalidResponse(
626 "query has more without progress",
627 ));
628 }
629 }
630 Err(error) if size_error(&error) && limit > 1 => {
631 limit = (limit / 2).max(1);
632 }
633 Err(error) if size_error(&error) && topics.len() > 1 => {
634 let right = topics.split_off(topics.len() / 2);
635 pending.push_front((right, limit));
636 }
637 Err(error) => return Err(dyn_err(error)),
638 }
639 }
640 }
641 Ok(())
642 }
643
644 pub(crate) async fn newest(
645 &self,
646 topics: Vec<Topic>,
647 full: bool,
648 ) -> Result<Vec<wire::query_newest_response::Result>> {
649 let topics: Vec<_> = topics
650 .into_iter()
651 .collect::<HashSet<_>>()
652 .into_iter()
653 .collect();
654 let cap = if full {
655 self.limits().max_newest_full_topics
656 } else {
657 self.limits().max_newest_metadata_topics
658 };
659 let results: Vec<Vec<_>> = stream::iter(
660 topics
661 .chunks(cap)
662 .map(<[_]>::to_vec)
663 .collect::<Vec<_>>()
664 .into_iter()
665 .map(|chunk| async move { self.newest_chunk(&chunk, full).await }),
666 )
667 .buffer_unordered(MAX_READ_CHUNKS_IN_FLIGHT)
668 .try_collect()
669 .await?;
670 Ok(results.into_iter().flatten().collect())
671 }
672
673 async fn newest_chunk(
674 &self,
675 topics: &[Topic],
676 full: bool,
677 ) -> Result<Vec<wire::query_newest_response::Result>> {
678 let mut pending = VecDeque::from([topics]);
679 let mut output = Vec::new();
680 while let Some(topics) = pending.pop_front() {
681 let request = wire::QueryNewestRequest {
682 topics: topics
683 .iter()
684 .map(|topic| wire::Topic {
685 topic: topic.cloned_vec(),
686 })
687 .collect(),
688 include_full_envelope: full,
689 };
690 match self
691 .retry_call(
692 || self.api_client.query_newest(request.clone()),
693 topics.len() > 1,
694 )
695 .await
696 {
697 Ok(response) => {
698 let mut seen = HashSet::new();
699 for result in &response.results {
700 let topic = Topic::parse(
701 &result
702 .topic
703 .as_ref()
704 .ok_or(ApiError::InvalidResponse("newest topic"))?
705 .topic,
706 )?;
707 if !topics.contains(&topic) || !seen.insert(topic.clone()) {
708 return Err(ApiError::InvalidResponse(
709 "unrequested or duplicate newest topic",
710 ));
711 }
712 let meta = result
713 .meta
714 .as_ref()
715 .ok_or(ApiError::InvalidResponse("newest metadata"))?;
716 let (meta_topic, _, _) =
717 xmtp_api_backend::envelope::metadata(meta, topic.kind())?;
718 if meta_topic != topic || (full && result.envelope.is_none()) {
719 return Err(ApiError::InvalidResponse("newest envelope"));
720 }
721 }
722 output.extend(response.results);
723 }
724 Err(error) if size_error(&error) && topics.len() > 1 => {
725 let (left, right) = topics.split_at(topics.len() / 2);
726 pending.push_front(right);
727 pending.push_front(left);
728 }
729 Err(error) => return Err(dyn_err(error)),
730 }
731 }
732 Ok(output)
733 }
734}
735
736#[cfg(test)]
737mod tests {
738 use super::*;
739 use xmtp_mls_validation::test_utils::{
740 GroupMessageKind, group_message_envelope, inline_welcome_envelope,
741 };
742
743 #[xmtp_common::test(unwrap_try = true)]
744 fn running_publish_measure_matches_encoded_mixed_request() {
745 let units = [
746 PublishUnit::single(inline_welcome_envelope([1; 32]))?,
747 PublishUnit::new(vec![
748 group_message_envelope([2; 16], GroupMessageKind::Proposal, [0; 127]),
749 group_message_envelope([2; 16], GroupMessageKind::Commit, [0; 16384]),
750 ])?,
751 PublishUnit::single(inline_welcome_envelope([1; 32]))?,
752 PublishUnit::single(inline_welcome_envelope([3; 32]))?,
753 ];
754 let limits = LimitsConfiguration::default();
755 let mut measure = PublishMeasure::new(&limits);
756 for (index, unit) in units.iter().enumerate() {
757 measure.add(unit);
758 assert_eq!(measure.bytes, request(&units[..=index]).encoded_len());
759 }
760 assert_eq!(measure.topics.len(), 3);
761 }
762}