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
23pub(crate) fn upload_one(
24 transport: &mut dyn Transport,
25 id: &str,
26 bytes: &[u8],
27 media_type: Option<&str>,
28 fixed_now: Option<i64>,
29) -> Result<(), TransportError> {
30 match transport.blob_upload_grant(id, bytes.len() as u64, media_type)? {
31 BlobUploadGrant::Present => return Ok(()),
32 BlobUploadGrant::Url {
33 url,
34 url_expires_at_ms,
35 } => {
36 if url_expires_at_ms.is_none_or(|expiry| expiry > now_ms(fixed_now))
37 && transport.blob_put_url(&url, bytes, media_type).is_ok()
38 {
39 return Ok(());
40 }
41 }
42 BlobUploadGrant::None => {}
43 }
44 transport.blob_upload(id, bytes, media_type)
45}
46
47pub struct PreparedSyncRound {
50 pub(crate) id: uuid::Uuid,
51 pub(crate) started_at_ms: i64,
52 pub(crate) meta: RequestMeta,
53 pub(crate) message: Message,
54 pub(crate) realtime: bool,
55 pub(crate) uploads: std::collections::VecDeque<PendingUpload>,
56 pub(crate) scopes: BTreeMap<String, String>,
57 #[cfg(feature = "bench-internals")]
58 pub(crate) benchmark_phases: crate::bench::Recorder,
59 pub(crate) progress: ProgressObserver,
60 pub(crate) fixed_now: Option<i64>,
61 pub(crate) round_deadline_at: Option<std::time::Instant>,
66}
67
68pub struct CompletedSyncRound {
70 pub(crate) prepared: PreparedSyncRound,
71 pub(crate) exchange: ExchangeResult,
72}
73
74pub(crate) enum ExchangeResult {
75 Upload {
76 id: String,
77 result: Result<(), TransportError>,
78 },
79 Reply {
80 response: Result<Message, TransportError>,
81 transport_failed: bool,
82 downloads: DownloadResults,
83 },
84}
85
86pub enum AppliedSyncRound {
89 Continue(PreparedSyncRound),
90 Complete {
91 outcome: SyncOutcome,
92 controls: Vec<String>,
93 more: bool,
95 bootstrap_advanced: bool,
97 },
98}
99
100impl PreparedSyncRound {
101 fn scope_round_deadline(
106 &mut self,
107 transport: &mut dyn Transport,
108 ) -> Result<(), TransportError> {
109 if self.round_deadline_at.is_none() {
110 if let Some(budget) = transport.round_deadline() {
111 #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
115 {
116 let _ = budget;
117 return Err(TransportError::new(
118 "sync.invalid_request",
119 "transport round deadline is unsupported on wasm32-unknown-unknown",
120 ));
121 }
122 #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
123 match std::time::Instant::now().checked_add(budget) {
124 Some(deadline) => self.round_deadline_at = Some(deadline),
125 None => {
127 return Err(TransportError::new(
128 "sync.invalid_request",
129 "transport round deadline is out of range",
130 ))
131 }
132 }
133 }
134 }
135 transport.set_round_deadline(self.round_deadline_at);
136 Ok(())
137 }
138
139 pub fn exchange(mut self, transport: &mut dyn Transport) -> CompletedSyncRound {
142 if let Some(upload) = self.uploads.pop_front() {
143 let result = match upload.bytes {
147 Ok(bytes) => self.scope_round_deadline(transport).and_then(|()| {
148 upload_one(
149 transport,
150 &upload.id,
151 &bytes,
152 upload.media_type.as_deref(),
153 self.fixed_now,
154 )
155 }),
156 Err(error) => Err(error),
157 };
158 transport.set_round_deadline(None);
159 return CompletedSyncRound {
160 prepared: self,
161 exchange: ExchangeResult::Upload {
162 id: upload.id,
163 result,
164 },
165 };
166 }
167 let mut transport_failed = false;
168 let response = (|| {
169 #[cfg(feature = "bench-internals")]
170 let encode_phase = self
171 .benchmark_phases
172 .start(crate::bench::Phase::RequestEncode);
173 let request = encode_message(&self.message);
174 #[cfg(feature = "bench-internals")]
175 drop(encode_phase);
176 self.scope_round_deadline(transport)?;
180 let round = if self.realtime {
181 transport.realtime_sync(&request)
182 } else {
183 transport.sync(&request)
184 };
185 transport_failed = round.is_err();
186 let bytes = round?;
187 #[cfg(feature = "bench-internals")]
188 let _decode_phase = self
189 .benchmark_phases
190 .start(crate::bench::Phase::ResponseDecode);
191 let response = decode_message(&bytes)
192 .map_err(|e| TransportError::new(e.code.as_str(), e.detail))?;
193 if response.msg_kind != MsgKind::Response {
194 return Err(TransportError::new(
195 "sync.invalid_request",
196 "expected a response message",
197 ));
198 }
199 Ok(response)
200 })();
201 let mut downloads = DownloadResults::default();
202 if let Ok(response) = &response {
203 if matches!(response.frames.first(), Some(Frame::RespHeader { required_schema_version: None, reset_required: Some(false), log_epoch: Some(epoch), .. }) if self.message.frames.first().is_some_and(|f| matches!(f, Frame::ReqHeader { log_epoch: Some(request_epoch), .. } if request_epoch == epoch)))
205 {
206 let mut scope: Option<(&str, &str)> = None;
207 let batch_started_at = now_ms(self.fixed_now);
211 let mut planned: Vec<PlannedFetch> = Vec::new();
212 for frame in &response.frames {
213 match frame {
214 Frame::SubStart {
215 id,
216 status: SubStatus::Active,
217 ..
218 } => {
219 scope = self
220 .scopes
221 .get(id)
222 .map(|scopes| (id.as_str(), scopes.as_str()));
223 }
224 Frame::SubStart { .. } | Frame::SubEnd { .. } => scope = None,
225 Frame::Error { .. } => break,
226 Frame::SegmentRef {
227 segment_id,
228 table,
229 media_type,
230 url,
231 url_expires_at_ms,
232 byte_length,
233 row_count,
234 ..
235 } => {
236 let Some((id, scopes)) = scope else { continue };
237 let advertised = match media_type {
238 MediaType::Rows => self.meta.accept & 3 != 0,
239 MediaType::Sqlite => self.meta.accept & 4 != 0,
240 };
241 if !advertised || (url.is_some() && self.meta.accept & 8 == 0) {
242 break;
243 }
244 let expired = url.is_some()
245 && url_expires_at_ms
246 .is_some_and(|expiry| expiry <= batch_started_at);
247 let (key, fetch) = match url {
248 Some(url) => (
249 DownloadKey::Url(url.clone()),
250 if expired {
251 None
252 } else {
253 Some(SegmentFetch::Url(url.clone()))
254 },
255 ),
256 None => (
257 DownloadKey::Segment(
258 segment_id.clone(),
259 table.clone(),
260 scopes.to_owned(),
261 ),
262 Some(SegmentFetch::Direct(SegmentRequest {
263 segment_id: segment_id.clone(),
264 table: table.clone(),
265 requested_scopes_json: scopes.to_owned(),
266 })),
267 ),
268 };
269 planned.push(PlannedFetch {
270 key,
271 subscription_id: id.to_owned(),
272 table: table.clone(),
273 segment_id: segment_id.clone(),
274 bytes_total: u64::try_from(*byte_length).ok(),
275 rows_total: u64::try_from(*row_count).ok(),
276 fetch,
277 });
278 if expired {
279 break;
280 }
281 }
282 _ => {}
283 }
284 }
285 if !planned.is_empty() {
286 let progress = self.progress.clone();
287 let mut on_progress = |index: usize, bytes: u64| {
288 let Some(planned_fetch) = planned.get(index) else {
289 return;
290 };
291 progress.update(|p| {
292 p.phase = ProgressPhase::Download;
293 p.subscription_id = Some(planned_fetch.subscription_id.clone());
294 p.table = Some(planned_fetch.table.clone());
295 p.segment_id = Some(planned_fetch.segment_id.clone());
296 p.bytes_received = bytes;
297 p.bytes_total = planned_fetch.bytes_total;
298 p.rows_processed = 0;
299 p.rows_total = planned_fetch.rows_total;
300 });
301 };
302 on_progress(0, 0);
304 let fetches: Vec<SegmentFetch> = planned
305 .iter()
306 .filter_map(|planned_fetch| planned_fetch.fetch.clone())
307 .collect();
308 let mut results = transport
309 .download_segments(&fetches, &mut on_progress)
310 .into_iter();
311 for planned_fetch in planned {
312 let result = match planned_fetch.fetch {
313 None => Err(TransportError::new(
314 "sync.segment_expired",
315 "signed segment URL expired before fetch",
316 )),
317 Some(_) => match results.next() {
320 Some(result) => result,
321 None => break,
322 },
323 };
324 let failed = result.is_err();
325 downloads
326 .entries
327 .push_back((planned_fetch.key, batch_started_at, result));
328 if failed {
329 break;
330 }
331 }
332 }
333 }
334 }
335 transport.set_round_deadline(None);
336 CompletedSyncRound {
337 prepared: self,
338 exchange: ExchangeResult::Reply {
339 response,
340 downloads,
341 transport_failed,
342 },
343 }
344 }
345}
346
347#[derive(PartialEq)]
348enum DownloadKey {
349 Url(String),
350 Segment(String, String, String),
351}
352
353struct PlannedFetch {
356 key: DownloadKey,
357 subscription_id: String,
358 table: String,
359 segment_id: String,
360 bytes_total: Option<u64>,
361 rows_total: Option<u64>,
362 fetch: Option<SegmentFetch>,
367}
368
369#[derive(Default)]
370pub(crate) struct DownloadResults {
371 entries: std::collections::VecDeque<(DownloadKey, i64, Result<Vec<u8>, TransportError>)>,
372 pub controls: Vec<String>,
373}
374
375impl DownloadResults {
376 fn take(&mut self, key: DownloadKey) -> Result<Vec<u8>, TransportError> {
377 let position = self
378 .entries
379 .iter()
380 .position(|(candidate, ..)| *candidate == key);
381 position
382 .and_then(|i| self.entries.remove(i))
383 .map(|(_, _, bytes)| bytes)
384 .unwrap_or_else(|| {
385 Err(TransportError::new(
386 "client.round_download_missing",
387 "round apply requested an uncaptured segment",
388 ))
389 })
390 }
391}
392
393impl Transport for DownloadResults {
394 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
395 Err(TransportError::new(
396 "client.round_network_on_owner",
397 "network requests must execute outside the client owner",
398 ))
399 }
400 fn realtime_sync(&mut self, request: &[u8]) -> Result<Vec<u8>, TransportError> {
401 self.sync(request)
402 }
403 fn download_segment(
404 &mut self,
405 request: &SegmentRequest,
406 _: &mut dyn FnMut(u64),
407 ) -> Result<Vec<u8>, TransportError> {
408 self.take(DownloadKey::Segment(
409 request.segment_id.clone(),
410 request.table.clone(),
411 request.requested_scopes_json.clone(),
412 ))
413 }
414 fn fetch_url(&mut self, url: &str, _: &mut dyn FnMut(u64)) -> Result<Vec<u8>, TransportError> {
415 self.take(DownloadKey::Url(url.to_owned()))
416 }
417 fn segment_fetch_started_at(&self, url: &str) -> Option<i64> {
418 self.entries
419 .iter()
420 .find_map(|(key, time, _)| (*key == DownloadKey::Url(url.to_owned())).then_some(*time))
421 }
422 fn realtime_connect(&mut self) -> Result<(), TransportError> {
423 Err(TransportError::new(
424 "client.round_network_on_owner",
425 "network connections must execute outside the client owner",
426 ))
427 }
428 fn realtime_send(&mut self, text: &str) -> Result<(), TransportError> {
429 self.controls.push(text.to_owned());
430 Ok(())
431 }
432 fn realtime_close(&mut self) -> Result<(), TransportError> {
433 Ok(())
434 }
435}
436
437#[cfg(test)]
438mod tests {
439 use super::*;
440 use crate::{BlobUploadGrant, SegmentRequest, TransportError};
441
442 struct SpentDeadline;
446
447 impl Transport for SpentDeadline {
448 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
449 unreachable!()
450 }
451 fn realtime_sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
452 unreachable!()
453 }
454 fn download_segment(
455 &mut self,
456 _: &SegmentRequest,
457 _: &mut dyn FnMut(u64),
458 ) -> Result<Vec<u8>, TransportError> {
459 unreachable!()
460 }
461 fn blob_upload_grant(
462 &mut self,
463 _: &str,
464 _: u64,
465 _: Option<&str>,
466 ) -> Result<BlobUploadGrant, TransportError> {
467 Ok(BlobUploadGrant::Url {
468 url: "https://storage.example/signed".into(),
469 url_expires_at_ms: None,
470 })
471 }
472 fn blob_put_url(
473 &mut self,
474 _: &str,
475 _: &[u8],
476 _: Option<&str>,
477 ) -> Result<(), TransportError> {
478 Err(TransportError::new(
479 "transport.timeout",
480 "signed upload timed out",
481 ))
482 }
483 fn blob_upload(
484 &mut self,
485 _: &str,
486 _: &[u8],
487 _: Option<&str>,
488 ) -> Result<(), TransportError> {
489 Err(TransportError::new(
490 "transport.timeout",
491 "direct upload timed out",
492 ))
493 }
494 fn realtime_connect(&mut self) -> Result<(), TransportError> {
495 Ok(())
496 }
497 fn realtime_send(&mut self, _: &str) -> Result<(), TransportError> {
498 Ok(())
499 }
500 fn realtime_close(&mut self) -> Result<(), TransportError> {
501 Ok(())
502 }
503 }
504
505 #[test]
506 fn a_spent_round_deadline_survives_the_signed_upload_fallback() {
507 let mut transport = SpentDeadline;
508 let error = upload_one(&mut transport, "sha256:test", &[0u8; 4], None, None).unwrap_err();
509 assert_eq!(error.code, "transport.timeout");
510 }
511
512 #[derive(Default)]
514 struct CountingFetches {
515 calls: Vec<String>,
516 }
517
518 impl Transport for CountingFetches {
519 fn sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
520 unreachable!()
521 }
522 fn realtime_sync(&mut self, _: &[u8]) -> Result<Vec<u8>, TransportError> {
523 unreachable!()
524 }
525 fn download_segment(
526 &mut self,
527 request: &SegmentRequest,
528 on_progress: &mut dyn FnMut(u64),
529 ) -> Result<Vec<u8>, TransportError> {
530 self.calls.push(request.segment_id.clone());
531 on_progress(4);
532 if request.segment_id == "sha256:b" {
533 Err(TransportError::new(
534 "transport.timeout",
535 "second fetch failed",
536 ))
537 } else {
538 Ok(request.segment_id.clone().into_bytes())
539 }
540 }
541 fn realtime_connect(&mut self) -> Result<(), TransportError> {
542 Ok(())
543 }
544 fn realtime_send(&mut self, _: &str) -> Result<(), TransportError> {
545 Ok(())
546 }
547 fn realtime_close(&mut self) -> Result<(), TransportError> {
548 Ok(())
549 }
550 }
551
552 #[test]
553 fn the_default_batch_runs_the_single_fetches_in_order_and_stops_at_the_first_error() {
554 let mut transport = CountingFetches::default();
555 let fetches: Vec<SegmentFetch> = ["a", "b", "c"]
556 .into_iter()
557 .map(|id| {
558 SegmentFetch::Direct(SegmentRequest {
559 segment_id: format!("sha256:{id}"),
560 table: "tasks".to_owned(),
561 requested_scopes_json: "{}".to_owned(),
562 })
563 })
564 .collect();
565 let mut reported = Vec::new();
566 let results = transport.download_segments(&fetches, &mut |index, bytes| {
567 reported.push((index, bytes));
568 });
569 assert_eq!(
570 transport.calls,
571 vec!["sha256:a", "sha256:b"],
572 "the third fetch was never attempted"
573 );
574 assert_eq!(results.len(), 2, "one result per attempted fetch");
575 assert_eq!(results[0].as_ref().unwrap(), b"sha256:a");
576 assert_eq!(results[1].as_ref().unwrap_err().code, "transport.timeout");
577 assert_eq!(reported, vec![(0, 4), (1, 4)]);
578 }
579}