1use std::collections::BTreeMap;
3
4use ssp2::model::{Frame, MediaType, Message, MsgKind, SubStatus};
5use ssp2::{decode_message, encode_message};
6
7use crate::client::RequestMeta;
8use crate::{
9 BlobUploadGrant, ProgressObserver, ProgressPhase, SegmentFetch, SegmentRequest, SyncOutcome,
10 Transport, TransportError,
11};
12
13pub(crate) fn now_ms(fixed: Option<i64>) -> i64 {
14 fixed.unwrap_or_else(crate::wall_clock_ms)
15}
16
17pub(crate) struct PendingUpload {
18 pub id: String,
19 pub bytes: Result<Vec<u8>, TransportError>,
20 pub media_type: Option<String>,
21}
22
23#[derive(Debug)]
31pub enum RoundOperation<'a> {
32 UploadGrant {
34 blob_id: &'a str,
35 byte_length: u64,
36 media_type: Option<&'a str>,
37 },
38 UploadPut {
41 blob_id: &'a str,
42 url: &'a str,
43 bytes: &'a [u8],
44 media_type: Option<&'a str>,
45 },
46 Upload {
48 blob_id: &'a str,
49 bytes: &'a [u8],
50 media_type: Option<&'a str>,
51 },
52 Sync { request: &'a [u8], realtime: bool },
55 Downloads(&'a [SegmentFetch]),
59}
60
61#[derive(Debug)]
63pub enum RoundReply {
64 UploadGrant(Result<BlobUploadGrant, TransportError>),
66 Upload(Result<(), TransportError>),
69 Sync(Result<Vec<u8>, TransportError>),
71 Downloads(Vec<Result<Vec<u8>, TransportError>>),
74}
75
76pub(crate) enum RoundPhase {
80 UploadGrant,
81 UploadPut {
82 url: String,
83 },
84 UploadDirect,
85 Sync,
86 Downloads {
87 planned: Vec<PlannedFetch>,
88 fetches: Vec<SegmentFetch>,
89 response: Message,
90 batch_started_at: i64,
91 },
92}
93
94pub struct PreparedSyncRound {
97 pub(crate) id: uuid::Uuid,
98 pub(crate) started_at_ms: i64,
99 pub(crate) meta: RequestMeta,
100 pub(crate) message: Message,
101 pub(crate) realtime: bool,
102 pub(crate) uploads: std::collections::VecDeque<PendingUpload>,
103 pub(crate) scopes: BTreeMap<String, String>,
104 #[cfg(feature = "bench-internals")]
105 pub(crate) benchmark_phases: crate::bench::Recorder,
106 pub(crate) progress: ProgressObserver,
107 pub(crate) fixed_now: Option<i64>,
108 pub(crate) round_deadline_at: Option<std::time::Instant>,
113 pub(crate) encoded: Option<Vec<u8>>,
115 pub(crate) phase: RoundPhase,
116 pub(crate) continuation: Option<RoundPhase>,
119 pub(crate) completed: Option<ExchangeResult>,
122}
123
124pub struct CompletedSyncRound {
126 pub(crate) prepared: PreparedSyncRound,
127 pub(crate) exchange: ExchangeResult,
128}
129
130pub(crate) enum ExchangeResult {
131 Upload {
132 id: String,
133 result: Result<(), TransportError>,
134 },
135 Reply {
136 response: Result<Message, TransportError>,
137 transport_failed: bool,
138 downloads: DownloadResults,
139 },
140}
141
142#[allow(clippy::large_enum_variant)]
147pub enum AppliedSyncRound {
148 Continue(PreparedSyncRound),
149 Complete {
150 outcome: SyncOutcome,
151 controls: Vec<String>,
152 more: bool,
154 bootstrap_advanced: bool,
156 },
157}
158
159impl PreparedSyncRound {
160 fn scope_round_deadline(
165 &mut self,
166 transport: &mut dyn Transport,
167 ) -> Result<(), TransportError> {
168 if self.round_deadline_at.is_none() {
169 if let Some(budget) = transport.round_deadline() {
170 #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
174 {
175 let _ = budget;
176 return Err(TransportError::new(
177 "sync.invalid_request",
178 "transport round deadline is unsupported on wasm32-unknown-unknown",
179 ));
180 }
181 #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
182 match std::time::Instant::now().checked_add(budget) {
183 Some(deadline) => self.round_deadline_at = Some(deadline),
184 None => {
186 return Err(TransportError::new(
187 "sync.invalid_request",
188 "transport round deadline is out of range",
189 ))
190 }
191 }
192 }
193 }
194 transport.set_round_deadline(self.round_deadline_at);
195 Ok(())
196 }
197
198 pub fn exchange(mut self, transport: &mut dyn Transport) -> CompletedSyncRound {
205 let mut scoped = false;
206 while self.completed.is_none() {
207 let reply = {
208 let Some(operation) = self.next_operation_scoped(transport, &mut scoped) else {
209 break;
210 };
211 match operation {
212 RoundOperation::UploadGrant {
213 blob_id,
214 byte_length,
215 media_type,
216 } => RoundReply::UploadGrant(transport.blob_upload_grant(
217 blob_id,
218 byte_length,
219 media_type,
220 )),
221 RoundOperation::UploadPut {
222 blob_id,
223 url,
224 bytes,
225 media_type,
226 } => {
227 RoundReply::Upload(transport.blob_put_url(blob_id, url, bytes, media_type))
228 }
229 RoundOperation::Upload {
230 blob_id,
231 bytes,
232 media_type,
233 } => RoundReply::Upload(transport.blob_upload(blob_id, bytes, media_type)),
234 RoundOperation::Sync { request, realtime } => RoundReply::Sync(if realtime {
235 transport.realtime_sync(request)
236 } else {
237 transport.sync(request)
238 }),
239 RoundOperation::Downloads(fetches) => {
240 RoundReply::Downloads(transport.download_segments(fetches, &mut |_, _| {}))
241 }
242 }
243 };
244 self.submit(reply)
245 .expect("the exchange answers its own operation");
246 }
247 transport.set_round_deadline(None);
248 self.finish().expect("the exchange completed the round")
249 }
250
251 pub fn request_bytes(&self) -> Vec<u8> {
255 self.encoded
256 .clone()
257 .unwrap_or_else(|| encode_message(&self.message))
258 }
259
260 pub fn next_operation(&mut self) -> Option<RoundOperation<'_>> {
269 let mut scoped = false;
270 self.next_operation_impl(None, &mut scoped)
271 }
272
273 pub fn submit(&mut self, reply: RoundReply) -> Result<(), TransportError> {
277 if self.completed.is_some() {
278 return Err(TransportError::new(
279 "client.invalid_round_reply",
280 "the exchange already has its reply",
281 ));
282 }
283 let phase = std::mem::replace(&mut self.phase, RoundPhase::Sync);
284 match (phase, reply) {
285 (RoundPhase::UploadGrant, RoundReply::UploadGrant(Err(error))) => {
286 let upload = self
287 .uploads
288 .pop_front()
289 .expect("an upload phase owns an upload");
290 self.completed = Some(ExchangeResult::Upload {
291 id: upload.id,
292 result: Err(error),
293 });
294 }
295 (RoundPhase::UploadGrant, RoundReply::UploadGrant(Ok(grant))) => match grant {
296 BlobUploadGrant::Present => self.upload_complete(),
297 BlobUploadGrant::Url {
298 url,
299 url_expires_at_ms,
300 } => {
301 if url_expires_at_ms.is_none_or(|expiry| expiry > now_ms(self.fixed_now)) {
304 self.phase = RoundPhase::UploadPut { url };
305 } else {
306 self.phase = RoundPhase::UploadDirect;
307 }
308 }
309 BlobUploadGrant::None => self.phase = RoundPhase::UploadDirect,
310 },
311 (RoundPhase::UploadPut { .. }, RoundReply::Upload(Ok(()))) => self.upload_complete(),
312 (RoundPhase::UploadPut { .. }, RoundReply::Upload(Err(_))) => {
315 self.phase = RoundPhase::UploadDirect;
316 }
317 (RoundPhase::UploadDirect, RoundReply::Upload(Ok(()))) => self.upload_complete(),
318 (RoundPhase::UploadDirect, RoundReply::Upload(Err(error))) => {
319 let upload = self
320 .uploads
321 .pop_front()
322 .expect("an upload phase owns an upload");
323 self.completed = Some(ExchangeResult::Upload {
324 id: upload.id,
325 result: Err(error),
326 });
327 }
328 (RoundPhase::Sync, RoundReply::Sync(reply)) => {
329 let transport_failed = reply.is_err();
330 let response = reply.and_then(|bytes| {
331 #[cfg(feature = "bench-internals")]
332 let _decode_phase = self
333 .benchmark_phases
334 .start(crate::bench::Phase::ResponseDecode);
335 let response = decode_message(&bytes)
336 .map_err(|e| TransportError::new(e.code.as_str(), e.detail))?;
337 if response.msg_kind != MsgKind::Response {
338 return Err(TransportError::new(
339 "sync.invalid_request",
340 "expected a response message",
341 ));
342 }
343 Ok(response)
344 });
345 let Ok(response) = response else {
346 self.completed = Some(ExchangeResult::Reply {
347 response,
348 transport_failed,
349 downloads: DownloadResults::default(),
350 });
351 return Ok(());
352 };
353 let mut planned: Vec<PlannedFetch> = Vec::new();
354 if matches!(response.frames.first(), Some(Frame::RespHeader { required_schema_version: None, reset_required: false, log_epoch: epoch, .. }) if self.message.frames.first().is_some_and(|f| matches!(f, Frame::ReqHeader { log_epoch: Some(request_epoch), .. } if request_epoch == epoch)))
357 {
358 let mut scope: Option<(&str, &str)> = None;
359 let batch_started_at = now_ms(self.fixed_now);
363 for frame in &response.frames {
364 match frame {
365 Frame::SubStart {
366 id,
367 status: SubStatus::Active,
368 ..
369 } => {
370 scope = self
371 .scopes
372 .get(id)
373 .map(|scopes| (id.as_str(), scopes.as_str()));
374 }
375 Frame::SubStart { .. } | Frame::SubEnd { .. } => scope = None,
376 Frame::Error { .. } => break,
377 Frame::SegmentRef {
378 segment_id,
379 table,
380 media_type,
381 url,
382 url_expires_at_ms,
383 byte_length,
384 row_count,
385 ..
386 } => {
387 let Some((id, scopes)) = scope else { continue };
388 let advertised = match media_type {
389 MediaType::Rows => self.meta.accept & 3 != 0,
390 MediaType::Sqlite => self.meta.accept & 4 != 0,
391 };
392 if !advertised || (url.is_some() && self.meta.accept & 8 == 0) {
393 break;
394 }
395 let cap = u64::try_from(*byte_length).unwrap_or(0);
397 let expired = url.is_some()
398 && url_expires_at_ms
399 .is_some_and(|expiry| expiry <= batch_started_at);
400 let (key, fetch) = match url {
401 Some(url) => (
402 DownloadKey::Url(url.clone()),
403 if expired {
404 None
405 } else {
406 Some(SegmentFetch::Url {
407 url: url.clone(),
408 byte_length: cap,
409 })
410 },
411 ),
412 None => (
413 DownloadKey::Segment(
414 segment_id.clone(),
415 table.clone(),
416 scopes.to_owned(),
417 ),
418 Some(SegmentFetch::Direct(SegmentRequest {
419 segment_id: segment_id.clone(),
420 table: table.clone(),
421 requested_scopes_json: scopes.to_owned(),
422 byte_length: cap,
423 })),
424 ),
425 };
426 planned.push(PlannedFetch {
427 key,
428 unit_id: id.to_owned(),
429 table: table.clone(),
430 segment_id: segment_id.clone(),
431 bytes_total: u64::try_from(*byte_length).ok(),
432 rows_total: u64::try_from(*row_count).ok(),
433 fetch,
434 });
435 if expired {
436 break;
437 }
438 }
439 _ => {}
440 }
441 }
442 if !planned.is_empty() {
443 let progress = self.progress.clone();
444 let on_progress = |index: usize, bytes: u64| {
445 let Some(planned_fetch) = planned.get(index) else {
446 return;
447 };
448 progress.update(|p| {
449 p.phase = ProgressPhase::Download;
450 p.unit_id = Some(planned_fetch.unit_id.clone());
451 p.table = Some(planned_fetch.table.clone());
452 p.segment_id = Some(planned_fetch.segment_id.clone());
453 p.bytes_received = bytes;
454 p.bytes_total = planned_fetch.bytes_total;
455 p.rows_processed = 0;
456 p.rows_total = planned_fetch.rows_total;
457 });
458 };
459 on_progress(0, 0);
461 let fetches: Vec<SegmentFetch> = planned
462 .iter()
463 .filter_map(|planned_fetch| planned_fetch.fetch.clone())
464 .collect();
465 self.phase = RoundPhase::Downloads {
466 planned,
467 fetches,
468 response,
469 batch_started_at,
470 };
471 return Ok(());
472 }
473 }
474 self.completed = Some(ExchangeResult::Reply {
475 response: Ok(response),
476 transport_failed,
477 downloads: DownloadResults::default(),
478 });
479 }
480 (
481 RoundPhase::Downloads {
482 planned,
483 batch_started_at,
484 response,
485 ..
486 },
487 RoundReply::Downloads(results),
488 ) => {
489 let mut downloads = DownloadResults::default();
490 let progress = self.progress.clone();
491 let mut results = results.into_iter();
492 for planned_fetch in planned {
493 let result = match planned_fetch.fetch {
494 None => Err(TransportError::new(
495 "sync.segment_expired",
496 "signed segment URL expired before fetch",
497 )),
498 Some(_) => match results.next() {
501 Some(result) => result,
502 None => break,
503 },
504 };
505 if let Ok(bytes) = &result {
506 progress.update(|p| {
507 p.phase = ProgressPhase::Download;
508 p.unit_id = Some(planned_fetch.unit_id.clone());
509 p.table = Some(planned_fetch.table.clone());
510 p.segment_id = Some(planned_fetch.segment_id.clone());
511 p.bytes_received = bytes.len() as u64;
512 p.bytes_total = planned_fetch.bytes_total;
513 p.rows_processed = 0;
514 p.rows_total = planned_fetch.rows_total;
515 });
516 }
517 let failed = result.is_err();
518 downloads
519 .entries
520 .push_back((planned_fetch.key, batch_started_at, result));
521 if failed {
522 break;
523 }
524 }
525 self.completed = Some(ExchangeResult::Reply {
526 response: Ok(response),
527 transport_failed: false,
528 downloads,
529 });
530 }
531 (phase, _) => {
532 self.phase = phase;
533 return Err(TransportError::new(
534 "client.invalid_round_reply",
535 "the reply does not answer the pending operation",
536 ));
537 }
538 }
539 Ok(())
540 }
541
542 pub fn finish(mut self) -> Option<CompletedSyncRound> {
545 let exchange = self.completed.take()?;
546 Some(CompletedSyncRound {
547 prepared: self,
548 exchange,
549 })
550 }
551
552 fn upload_complete(&mut self) {
553 let upload = self
554 .uploads
555 .pop_front()
556 .expect("an upload phase owns an upload");
557 self.completed = Some(ExchangeResult::Upload {
558 id: upload.id,
559 result: Ok(()),
560 });
561 self.continuation = Some(if self.uploads.is_empty() {
562 RoundPhase::Sync
563 } else {
564 RoundPhase::UploadGrant
565 });
566 }
567
568 fn ensure_encoded(&mut self) {
569 if self.encoded.is_none() {
570 #[cfg(feature = "bench-internals")]
571 let _phase = self
572 .benchmark_phases
573 .start(crate::bench::Phase::RequestEncode);
574 self.encoded = Some(encode_message(&self.message));
575 }
576 }
577
578 fn next_operation_scoped<'a>(
579 &'a mut self,
580 transport: &mut dyn Transport,
581 scoped: &mut bool,
582 ) -> Option<RoundOperation<'a>> {
583 self.next_operation_impl(Some(transport), scoped)
584 }
585
586 fn next_operation_impl<'a>(
587 &'a mut self,
588 mut scope: Option<&mut dyn Transport>,
589 scoped: &mut bool,
590 ) -> Option<RoundOperation<'a>> {
591 if self.completed.is_some() {
592 return None;
593 }
594 if let Some(next) = self.continuation.take() {
596 self.phase = next;
597 }
598 if matches!(self.phase, RoundPhase::UploadGrant) && self.uploads.is_empty() {
599 self.phase = RoundPhase::Sync;
600 }
601 if matches!(self.phase, RoundPhase::UploadGrant)
605 && self
606 .uploads
607 .front()
608 .is_some_and(|upload| upload.bytes.is_err())
609 {
610 let upload = self
611 .uploads
612 .pop_front()
613 .expect("the failed upload is staged");
614 self.completed = Some(ExchangeResult::Upload {
615 id: upload.id,
616 result: Err(upload.bytes.unwrap_err()),
617 });
618 return None;
619 }
620 if matches!(self.phase, RoundPhase::Sync) {
621 self.ensure_encoded();
622 }
623 let needs_network = !matches!(self.phase, RoundPhase::Downloads { .. });
624 if needs_network && !*scoped {
628 *scoped = true;
629 if let Some(transport) = scope.as_mut() {
630 if let Err(error) = self.scope_round_deadline(&mut **transport) {
631 if matches!(
632 self.phase,
633 RoundPhase::UploadGrant
634 | RoundPhase::UploadPut { .. }
635 | RoundPhase::UploadDirect
636 ) {
637 let upload = self
638 .uploads
639 .pop_front()
640 .expect("an upload phase owns an upload");
641 self.completed = Some(ExchangeResult::Upload {
642 id: upload.id,
643 result: Err(error),
644 });
645 } else {
646 self.completed = Some(ExchangeResult::Reply {
647 response: Err(error),
648 transport_failed: false,
649 downloads: DownloadResults::default(),
650 });
651 }
652 return None;
653 }
654 }
655 }
656 match &self.phase {
657 RoundPhase::UploadGrant => {
658 let upload = self
659 .uploads
660 .front()
661 .expect("an upload phase owns an upload");
662 let bytes = upload
663 .bytes
664 .as_ref()
665 .expect("a staged body that failed to read has no operation");
666 Some(RoundOperation::UploadGrant {
667 blob_id: &upload.id,
668 byte_length: bytes.len() as u64,
669 media_type: upload.media_type.as_deref(),
670 })
671 }
672 RoundPhase::UploadPut { url, .. } => {
673 let upload = self
674 .uploads
675 .front()
676 .expect("an upload phase owns an upload");
677 let bytes = upload
678 .bytes
679 .as_ref()
680 .expect("a staged body that failed to read has no operation");
681 Some(RoundOperation::UploadPut {
682 blob_id: &upload.id,
683 url,
684 bytes,
685 media_type: upload.media_type.as_deref(),
686 })
687 }
688 RoundPhase::UploadDirect => {
689 let upload = self
690 .uploads
691 .front()
692 .expect("an upload phase owns an upload");
693 let bytes = upload
694 .bytes
695 .as_ref()
696 .expect("a staged body that failed to read has no operation");
697 Some(RoundOperation::Upload {
698 blob_id: &upload.id,
699 bytes,
700 media_type: upload.media_type.as_deref(),
701 })
702 }
703 RoundPhase::Sync => Some(RoundOperation::Sync {
704 request: self.encoded.as_deref().expect("the request is encoded"),
705 realtime: self.realtime,
706 }),
707 RoundPhase::Downloads { fetches, .. } => Some(RoundOperation::Downloads(fetches)),
708 }
709 }
710}
711
712#[derive(PartialEq)]
713enum DownloadKey {
714 Url(String),
715 Segment(String, String, String),
716}
717
718pub(crate) struct PlannedFetch {
721 key: DownloadKey,
722 unit_id: String,
723 table: String,
724 segment_id: String,
725 bytes_total: Option<u64>,
726 rows_total: Option<u64>,
727 fetch: Option<SegmentFetch>,
732}
733
734#[derive(Default)]
735pub(crate) struct DownloadResults {
736 entries: std::collections::VecDeque<(DownloadKey, i64, Result<Vec<u8>, TransportError>)>,
737 pub controls: Vec<String>,
738}
739
740impl DownloadResults {
741 fn take(&mut self, key: DownloadKey) -> Result<Vec<u8>, TransportError> {
742 let position = self
743 .entries
744 .iter()
745 .position(|(candidate, ..)| *candidate == key);
746 position
747 .and_then(|i| self.entries.remove(i))
748 .map(|(_, _, bytes)| bytes)
749 .unwrap_or_else(|| {
750 Err(TransportError::new(
751 "client.round_download_missing",
752 "round apply requested an uncaptured segment",
753 ))
754 })
755 }
756}
757
758impl Transport for DownloadResults {
759 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
760 Err(TransportError::new(
761 "client.round_network_on_owner",
762 "network requests must execute outside the client owner",
763 ))
764 }
765 fn realtime_sync(&mut self, request: &[u8]) -> Result<Vec<u8>, TransportError> {
766 self.sync(request)
767 }
768 fn download_segment(
769 &mut self,
770 request: &SegmentRequest,
771 _: &mut dyn FnMut(u64),
772 ) -> Result<Vec<u8>, TransportError> {
773 self.take(DownloadKey::Segment(
774 request.segment_id.clone(),
775 request.table.clone(),
776 request.requested_scopes_json.clone(),
777 ))
778 }
779 fn fetch_url(
780 &mut self,
781 url: &str,
782 _: u64,
783 _: &mut dyn FnMut(u64),
784 ) -> Result<Vec<u8>, TransportError> {
785 self.take(DownloadKey::Url(url.to_owned()))
786 }
787 fn segment_fetch_started_at(&self, url: &str) -> Option<i64> {
788 self.entries
789 .iter()
790 .find_map(|(key, time, _)| (*key == DownloadKey::Url(url.to_owned())).then_some(*time))
791 }
792 fn realtime_connect(&mut self) -> Result<(), TransportError> {
793 Err(TransportError::new(
794 "client.round_network_on_owner",
795 "network connections must execute outside the client owner",
796 ))
797 }
798 fn realtime_send(&mut self, text: &str) -> Result<(), TransportError> {
799 self.controls.push(text.to_owned());
800 Ok(())
801 }
802 fn realtime_close(&mut self) -> Result<(), TransportError> {
803 Ok(())
804 }
805}
806
807#[cfg(test)]
808pub(crate) mod test_support {
809 use super::*;
810
811 pub(crate) fn drive_round(
814 mut prepared: PreparedSyncRound,
815 transport: &mut dyn Transport,
816 ) -> CompletedSyncRound {
817 while let Some(operation) = prepared.next_operation() {
818 let reply = match operation {
819 RoundOperation::UploadGrant {
820 blob_id,
821 byte_length,
822 media_type,
823 } => RoundReply::UploadGrant(transport.blob_upload_grant(
824 blob_id,
825 byte_length,
826 media_type,
827 )),
828 RoundOperation::UploadPut {
829 blob_id,
830 url,
831 bytes,
832 media_type,
833 } => RoundReply::Upload(transport.blob_put_url(blob_id, url, bytes, media_type)),
834 RoundOperation::Upload {
835 blob_id,
836 bytes,
837 media_type,
838 } => RoundReply::Upload(transport.blob_upload(blob_id, bytes, media_type)),
839 RoundOperation::Sync { request, realtime } => RoundReply::Sync(if realtime {
840 transport.realtime_sync(request)
841 } else {
842 transport.sync(request)
843 }),
844 RoundOperation::Downloads(fetches) => {
845 RoundReply::Downloads(transport.download_segments(fetches, &mut |_, _| {}))
846 }
847 };
848 prepared
849 .submit(reply)
850 .expect("the driver answers the pending operation");
851 }
852 prepared.finish().expect("the round has no work left")
853 }
854}
855
856#[cfg(test)]
857mod tests {
858 use super::*;
859 use crate::{BlobUploadGrant, SegmentRequest, SyncClient, TransportError};
860
861 struct SpentDeadline;
865
866 impl Transport for SpentDeadline {
867 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
868 unreachable!()
869 }
870 fn realtime_sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
871 unreachable!()
872 }
873 fn download_segment(
874 &mut self,
875 _: &SegmentRequest,
876 _: &mut dyn FnMut(u64),
877 ) -> Result<Vec<u8>, TransportError> {
878 unreachable!()
879 }
880 fn blob_upload_grant(
881 &mut self,
882 _: &str,
883 _: u64,
884 _: Option<&str>,
885 ) -> Result<BlobUploadGrant, TransportError> {
886 Ok(BlobUploadGrant::Url {
887 url: "https://storage.example/signed".into(),
888 url_expires_at_ms: None,
889 })
890 }
891 fn blob_put_url(
892 &mut self,
893 _: &str,
894 _: &str,
895 _: &[u8],
896 _: Option<&str>,
897 ) -> Result<(), TransportError> {
898 Err(TransportError::new(
899 "transport.timeout",
900 "signed upload timed out",
901 ))
902 }
903 fn blob_upload(
904 &mut self,
905 _: &str,
906 _: &[u8],
907 _: Option<&str>,
908 ) -> Result<(), TransportError> {
909 Err(TransportError::new(
910 "transport.timeout",
911 "direct upload timed out",
912 ))
913 }
914 fn realtime_connect(&mut self) -> Result<(), TransportError> {
915 Ok(())
916 }
917 fn realtime_send(&mut self, _: &str) -> Result<(), TransportError> {
918 Ok(())
919 }
920 fn realtime_close(&mut self) -> Result<(), TransportError> {
921 Ok(())
922 }
923 }
924
925 #[test]
926 fn a_failed_presigned_put_falls_back_to_the_direct_endpoint() {
927 let schema = serde_json::json!({"version":1,"tables":[{"name":"attachments","primaryKey":"id","columns":[
928 {"name":"id","type":"string","nullable":false},
929 {"name":"file","type":"blob_ref","nullable":true}],"scopes":[]}]});
930 let mut client =
931 SyncClient::new("spent-deadline".into(), &schema, Default::default()).unwrap();
932 client
933 .set_meta(crate::client::LOG_EPOCH_KEY, "epoch")
934 .unwrap();
935 client
936 .upload_blob(b"body", Some("text/plain".into()), None)
937 .unwrap();
938 let mut prepared = client.prepare_sync_round(false).unwrap();
939 let mut transport = SpentDeadline;
940 let mut steps = Vec::new();
941 while let Some(operation) = prepared.next_operation() {
942 let reply = match operation {
943 RoundOperation::UploadGrant {
944 blob_id,
945 byte_length,
946 media_type,
947 } => {
948 steps.push("grant");
949 RoundReply::UploadGrant(transport.blob_upload_grant(
950 blob_id,
951 byte_length,
952 media_type,
953 ))
954 }
955 RoundOperation::UploadPut {
956 blob_id,
957 url,
958 bytes,
959 media_type,
960 } => {
961 steps.push("put");
962 RoundReply::Upload(transport.blob_put_url(blob_id, url, bytes, media_type))
963 }
964 RoundOperation::Upload {
965 blob_id,
966 bytes,
967 media_type,
968 } => {
969 steps.push("direct");
970 RoundReply::Upload(transport.blob_upload(blob_id, bytes, media_type))
971 }
972 _ => panic!("only upload operations precede the round message"),
973 };
974 prepared.submit(reply).unwrap();
975 }
976 assert_eq!(steps, ["grant", "put", "direct"]);
977 let outcome = match client.apply_sync_round(prepared.finish().unwrap()) {
978 AppliedSyncRound::Complete { outcome, .. } => outcome,
979 AppliedSyncRound::Continue(_) => panic!("an upload failure does not continue"),
980 };
981 match outcome {
982 SyncOutcome::Failed { error_code, .. } => {
983 assert_eq!(error_code, "transport.timeout")
984 }
985 other => panic!("expected a typed upload failure, got {other:?}"),
986 }
987 }
988
989 struct OversizedSegment {
992 unit_id: String,
993 segment_id: String,
994 body: Vec<u8>,
995 }
996
997 impl Transport for OversizedSegment {
998 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
999 Ok(encode_message(&Message {
1000 wire_version: ssp2::decode::WIRE_VERSION,
1001 msg_kind: MsgKind::Response,
1002 frames: vec![
1003 Frame::RespHeader {
1004 required_schema_version: None,
1005 latest_schema_version: None,
1006 log_epoch: "epoch".into(),
1007 reset_required: false,
1008 },
1009 Frame::SubStart {
1010 id: self.unit_id.clone(),
1011 status: SubStatus::Active,
1012 reason_code: String::new(),
1013 bootstrap: true,
1014 fenced: Vec::new(),
1015 },
1016 Frame::SegmentRef {
1017 segment_id: self.segment_id.clone(),
1018 media_type: MediaType::Rows,
1019 table: "tasks".into(),
1020 byte_length: 4,
1021 row_count: 1,
1022 as_of_commit_seq: 1,
1023 scope_digest: "digest".into(),
1024 row_cursor: None,
1025 next_row_cursor: None,
1026 url: None,
1027 url_expires_at_ms: None,
1028 },
1029 Frame::SubEnd {
1030 next_cursor: 1,
1031 bootstrap_state: None,
1032 },
1033 ],
1034 }))
1035 }
1036 fn realtime_sync(&mut self, request: &[u8]) -> Result<Vec<u8>, TransportError> {
1037 self.sync(request)
1038 }
1039 fn download_segment(
1040 &mut self,
1041 _: &SegmentRequest,
1042 _: &mut dyn FnMut(u64),
1043 ) -> Result<Vec<u8>, TransportError> {
1044 Ok(self.body.clone())
1045 }
1046 fn realtime_connect(&mut self) -> Result<(), TransportError> {
1047 Ok(())
1048 }
1049 fn realtime_send(&mut self, _: &str) -> Result<(), TransportError> {
1050 Ok(())
1051 }
1052 fn realtime_close(&mut self) -> Result<(), TransportError> {
1053 Ok(())
1054 }
1055 }
1056
1057 #[test]
1058 fn a_segment_body_longer_than_its_byte_length_fails_before_the_content_check() {
1059 use sha2::Digest;
1060 let schema = serde_json::json!({"version":1,"tables":[{"name":"tasks","primaryKey":"id","columns":[
1061 {"name":"id","type":"string","nullable":false},
1062 {"name":"project_id","type":"string","nullable":false}],
1063 "scopes":[{"pattern":"project:{project_id}"}]}]});
1064 let mut client = SyncClient::new("oversized".into(), &schema, Default::default()).unwrap();
1065 client
1066 .set_meta(crate::client::LOG_EPOCH_KEY, "epoch")
1067 .unwrap();
1068 let unit_id = client
1069 .subscribe(crate::UnitScopes::from([(
1070 "project_id".into(),
1071 "p1".into(),
1072 )]))
1073 .unwrap();
1074 let body = b"twelve bytes".to_vec();
1077 let segment_id = format!(
1078 "sha256:{}",
1079 sha2::Sha256::digest(&body)
1080 .iter()
1081 .map(|byte| format!("{byte:02x}"))
1082 .collect::<String>()
1083 );
1084 let mut transport = OversizedSegment {
1085 unit_id,
1086 segment_id,
1087 body,
1088 };
1089 match client.sync(&mut transport) {
1090 SyncOutcome::Failed {
1091 error_code,
1092 message,
1093 ..
1094 } => {
1095 assert_eq!(error_code, "sync.invalid_request");
1096 assert_eq!(message, "segment exceeds its byteLength");
1097 }
1098 other => panic!("expected the byteLength cap to fail the round, got {other:?}"),
1099 }
1100 }
1101
1102 #[derive(Default)]
1104 struct CountingFetches {
1105 calls: Vec<String>,
1106 }
1107
1108 impl Transport for CountingFetches {
1109 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
1110 unreachable!()
1111 }
1112 fn realtime_sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
1113 unreachable!()
1114 }
1115 fn download_segment(
1116 &mut self,
1117 request: &SegmentRequest,
1118 on_progress: &mut dyn FnMut(u64),
1119 ) -> Result<Vec<u8>, TransportError> {
1120 self.calls.push(request.segment_id.clone());
1121 on_progress(4);
1122 if request.segment_id == "sha256:b" {
1123 Err(TransportError::new(
1124 "transport.timeout",
1125 "second fetch failed",
1126 ))
1127 } else {
1128 Ok(request.segment_id.clone().into_bytes())
1129 }
1130 }
1131 fn realtime_connect(&mut self) -> Result<(), TransportError> {
1132 Ok(())
1133 }
1134 fn realtime_send(&mut self, _: &str) -> Result<(), TransportError> {
1135 Ok(())
1136 }
1137 fn realtime_close(&mut self) -> Result<(), TransportError> {
1138 Ok(())
1139 }
1140 }
1141
1142 #[test]
1143 fn the_default_batch_runs_the_single_fetches_in_order_and_stops_at_the_first_error() {
1144 let mut transport = CountingFetches::default();
1145 let fetches: Vec<SegmentFetch> = ["a", "b", "c"]
1146 .into_iter()
1147 .map(|id| {
1148 SegmentFetch::Direct(SegmentRequest {
1149 segment_id: format!("sha256:{id}"),
1150 table: "tasks".to_owned(),
1151 requested_scopes_json: "{}".to_owned(),
1152 byte_length: 4,
1153 })
1154 })
1155 .collect();
1156 let mut reported = Vec::new();
1157 let results = transport.download_segments(&fetches, &mut |index, bytes| {
1158 reported.push((index, bytes));
1159 });
1160 assert_eq!(
1161 transport.calls,
1162 vec!["sha256:a", "sha256:b"],
1163 "the third fetch was never attempted"
1164 );
1165 assert_eq!(results.len(), 2, "one result per attempted fetch");
1166 assert_eq!(results[0].as_ref().unwrap(), b"sha256:a");
1167 assert_eq!(results[1].as_ref().unwrap_err().code, "transport.timeout");
1168 assert_eq!(reported, vec![(0, 4), (1, 4)]);
1169 }
1170}