1use crate::fetch::ObjectEntry;
40use base64::Engine as _;
41use base64::engine::general_purpose::URL_SAFE_NO_PAD;
42use serde::{Deserialize, Serialize};
43use sha2::{Digest as _, Sha256};
44use spate_core::coordination::{CoordinationError, CoordinationErrorKind, SplitId};
45use std::collections::VecDeque;
46
47pub const DESCRIPTOR_VERSION: u32 = 1;
51
52pub(crate) const PACKING_VERSION: u32 = 1;
56
57pub(crate) const PACKING_LOOKBACK: usize = 10;
60
61pub(crate) const OPEN_COST_DIVISOR: u64 = 16;
64
65#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
70pub struct DescriptorObject {
71 pub key: String,
73 pub size: u64,
75 pub etag: Option<String>,
79 pub last_modified_ms: i64,
81}
82
83#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
95pub struct SplitDescriptor {
96 pub(crate) v: u32,
100 pub objects: Vec<DescriptorObject>,
102}
103
104#[derive(Deserialize)]
107struct VersionProbe {
108 v: u32,
109}
110
111impl SplitDescriptor {
112 #[must_use]
117 pub fn new(objects: Vec<DescriptorObject>) -> SplitDescriptor {
118 SplitDescriptor {
119 v: DESCRIPTOR_VERSION,
120 objects,
121 }
122 }
123
124 #[must_use]
127 pub fn version(&self) -> u32 {
128 self.v
129 }
130
131 pub(crate) fn to_entries(&self) -> Vec<ObjectEntry> {
134 self.objects
135 .iter()
136 .map(|o| ObjectEntry {
137 key: o.key.clone(),
138 size: o.size,
139 etag: o.etag.clone(),
140 last_modified_ms: o.last_modified_ms,
141 })
142 .collect()
143 }
144
145 pub(crate) fn from_entries(entries: &[ObjectEntry]) -> SplitDescriptor {
147 SplitDescriptor::new(
148 entries
149 .iter()
150 .map(|e| DescriptorObject {
151 key: e.key.clone(),
152 size: e.size,
153 etag: e.etag.clone(),
154 last_modified_ms: e.last_modified_ms,
155 })
156 .collect(),
157 )
158 }
159
160 pub fn encode(&self) -> Result<Vec<u8>, CoordinationError> {
168 if self.v != DESCRIPTOR_VERSION {
169 return Err(CoordinationError::new(
170 CoordinationErrorKind::Fatal,
171 format!(
172 "descriptor version {} is not the supported {DESCRIPTOR_VERSION}; \
173 construct via SplitDescriptor::new",
174 self.v
175 ),
176 ));
177 }
178 Ok(serde_json::to_vec(self).expect("descriptor serialization is infallible: no non-string map keys, no fallible Serialize impls"))
179 }
180
181 pub fn decode(bytes: &[u8]) -> Result<SplitDescriptor, CoordinationError> {
189 let fatal = |reason: String| CoordinationError::new(CoordinationErrorKind::Fatal, reason);
190 let probe: VersionProbe = serde_json::from_slice(bytes)
191 .map_err(|e| fatal(format!("split descriptor is not valid JSON: {e}")))?;
192 if probe.v != DESCRIPTOR_VERSION {
193 return Err(fatal(format!(
194 "split descriptor version {} is not this release's version \
195 {DESCRIPTOR_VERSION}; the split was planned by an incompatible release",
196 probe.v
197 )));
198 }
199 serde_json::from_slice(bytes)
200 .map_err(|e| fatal(format!("split descriptor failed to decode: {e}")))
201 }
202}
203
204pub fn split_id_for<'a, I>(members: I) -> Result<SplitId, CoordinationError>
249where
250 I: IntoIterator<Item = (&'a str, Option<&'a str>)>,
251{
252 split_id_with_version(members, PACKING_VERSION)
253}
254
255fn split_id_with_version<'a, I>(members: I, version: u32) -> Result<SplitId, CoordinationError>
258where
259 I: IntoIterator<Item = (&'a str, Option<&'a str>)>,
260{
261 let mut members: Vec<(&str, Option<&str>)> = members.into_iter().collect();
262 if members.is_empty() {
263 return Err(CoordinationError::new(
264 CoordinationErrorKind::Fatal,
265 "cannot mint a split id for an empty member set",
266 ));
267 }
268 members.sort_unstable_by_key(|(key, _)| *key);
269
270 let mut hasher = Sha256::new();
271 hasher.update(b"spate-s3-split\n");
272 hasher.update(version.to_le_bytes());
273 for (key, etag) in members {
274 hasher.update(u32::try_from(key.len()).unwrap_or(u32::MAX).to_le_bytes());
275 hasher.update(key.as_bytes());
276 match etag {
277 Some(etag) => {
278 hasher.update([0x01]);
279 hasher.update(u32::try_from(etag.len()).unwrap_or(u32::MAX).to_le_bytes());
280 hasher.update(etag.as_bytes());
281 }
282 None => hasher.update([0x00]),
283 }
284 }
285 let digest = hasher.finalize();
286 SplitId::new(format!("s3-{}", URL_SAFE_NO_PAD.encode(&digest[..16])))
287}
288
289pub(crate) fn pack(entries: Vec<ObjectEntry>, target_bytes: u64) -> Vec<Vec<ObjectEntry>> {
298 debug_assert!(
299 target_bytes > 0,
300 "config validation rejects a zero split target"
301 );
302 struct Bin {
303 members: Vec<ObjectEntry>,
304 cost: u64,
305 }
306 let floor = (target_bytes / OPEN_COST_DIVISOR).max(1);
307 let mut bins: Vec<Bin> = Vec::new();
308 let mut open: VecDeque<usize> = VecDeque::new();
310 for entry in entries {
311 let cost = entry.size.max(floor);
312 let idx = match open
315 .iter()
316 .position(|&i| bins[i].cost.saturating_add(cost) <= target_bytes)
317 {
318 Some(pos) => open[pos],
319 None => {
320 if open.len() == PACKING_LOOKBACK {
321 open.pop_front();
322 }
323 bins.push(Bin {
324 members: Vec::new(),
325 cost: 0,
326 });
327 let idx = bins.len() - 1;
328 open.push_back(idx);
329 idx
330 }
331 };
332 bins[idx].members.push(entry);
333 bins[idx].cost = bins[idx].cost.saturating_add(cost);
334 if bins[idx].cost >= target_bytes
335 && let Some(pos) = open.iter().position(|&i| i == idx)
336 {
337 open.remove(pos);
338 }
339 }
340 bins.into_iter().map(|b| b.members).collect()
341}
342
343#[cfg(test)]
344mod tests {
345 use super::*;
346 use proptest::prelude::*;
347
348 fn entry(key: &str, size: u64) -> ObjectEntry {
349 ObjectEntry {
350 key: key.to_string(),
351 size,
352 etag: Some(format!("\"etag-{key}\"")),
353 last_modified_ms: 1_760_000_000_000,
354 }
355 }
356
357 const MB: u64 = 1024 * 1024;
358
359 #[test]
362 fn digest_id_is_pinned() {
363 let id = split_id_for([
367 ("exports/2026/part-000.ndjson", Some("\"9b2cf5\"")),
368 ("exports/2026/part-001.ndjson", None),
369 ])
370 .unwrap();
371 assert_eq!(id.as_str(), "s3-_rD2rPZklAFVV4pYEYaWxg");
372 }
373
374 #[test]
375 fn digest_is_order_insensitive_but_content_sensitive() {
376 let forward = split_id_for([("a", Some("1")), ("b", Some("2"))]).unwrap();
377 let reversed = split_id_for([("b", Some("2")), ("a", Some("1"))]).unwrap();
378 assert_eq!(forward, reversed);
379
380 let other_key = split_id_for([("a", Some("1")), ("c", Some("2"))]).unwrap();
381 assert_ne!(forward, other_key);
382 }
383
384 #[test]
385 fn digest_changes_when_an_etag_changes() {
386 let before = split_id_for([("a", Some("v1")), ("b", Some("x"))]).unwrap();
387 let overwritten = split_id_for([("a", Some("v2")), ("b", Some("x"))]).unwrap();
388 let dropped = split_id_for([("a", None), ("b", Some("x"))]).unwrap();
389 assert_ne!(before, overwritten);
390 assert_ne!(before, dropped);
391 }
392
393 #[test]
394 fn digest_folds_in_the_packing_version() {
395 let v1 = split_id_with_version([("a", Some("1"))], 1).unwrap();
396 let v2 = split_id_with_version([("a", Some("1"))], 2).unwrap();
397 assert_ne!(v1, v2);
398 }
399
400 #[test]
401 fn empty_member_set_is_rejected() {
402 assert!(split_id_for(std::iter::empty()).is_err());
403 }
404
405 #[test]
406 fn ambiguous_concatenations_do_not_collide() {
407 let a = split_id_for([("ab", Some("c"))]).unwrap();
410 let b = split_id_for([("a", Some("bc"))]).unwrap();
411 let c = split_id_for([("abc", None)]).unwrap();
412 assert_ne!(a, b);
413 assert_ne!(a, c);
414 assert_ne!(b, c);
415 }
416
417 #[test]
420 fn descriptor_round_trips() {
421 let desc = SplitDescriptor::from_entries(&[
422 entry("exports/part-000.ndjson.gz", 52 * MB),
423 ObjectEntry {
424 key: "exports/part-001.ndjson.gz".to_string(),
425 size: 9,
426 etag: None,
427 last_modified_ms: 1,
428 },
429 ]);
430 let decoded = SplitDescriptor::decode(&desc.encode().unwrap()).unwrap();
431 assert_eq!(decoded, desc);
432 assert_eq!(decoded.v, DESCRIPTOR_VERSION);
433 }
434
435 #[test]
436 fn descriptor_encoding_is_pinned() {
437 let desc = SplitDescriptor::from_entries(&[entry("k", 5)]);
441 assert_eq!(
442 String::from_utf8(desc.encode().unwrap()).unwrap(),
443 r#"{"v":1,"objects":[{"key":"k","size":5,"etag":"\"etag-k\"","last_modified_ms":1760000000000}]}"#,
444 );
445 }
446
447 #[test]
448 fn encode_refuses_a_descriptor_not_built_by_new() {
449 let rogue = SplitDescriptor {
452 v: 0,
453 objects: vec![],
454 };
455 let err = rogue.encode().unwrap_err();
456 assert_eq!(err.kind, CoordinationErrorKind::Fatal);
457 assert!(
458 err.reason.contains("SplitDescriptor::new"),
459 "reason: {}",
460 err.reason
461 );
462 assert_eq!(SplitDescriptor::new(vec![]).version(), DESCRIPTOR_VERSION);
463 }
464
465 #[test]
466 fn unknown_descriptor_version_is_rejected_actionably() {
467 let err = SplitDescriptor::decode(br#"{"v":999,"objects":[]}"#).unwrap_err();
468 assert_eq!(err.kind, CoordinationErrorKind::Fatal);
469 assert!(err.reason.contains("version 999"), "reason: {}", err.reason);
470 assert!(
471 err.reason.contains("incompatible release"),
472 "reason: {}",
473 err.reason
474 );
475
476 let garbage = SplitDescriptor::decode(b"not json").unwrap_err();
477 assert_eq!(garbage.kind, CoordinationErrorKind::Fatal);
478 }
479
480 #[test]
483 fn packing_is_deterministic_for_a_fixed_listing() {
484 let listing: Vec<ObjectEntry> = (0..200)
485 .map(|i| entry(&format!("k{i:04}"), (i % 40) * MB))
486 .collect();
487 let a = pack(listing.clone(), 64 * MB);
488 let b = pack(listing, 64 * MB);
489 assert_eq!(a, b);
490 }
491
492 #[test]
493 fn tiny_objects_coalesce_under_the_open_cost_floor() {
494 let listing: Vec<ObjectEntry> = (0..64).map(|i| entry(&format!("k{i:02}"), 1)).collect();
497 let bins = pack(listing, 64 * MB);
498 assert_eq!(bins.len(), 4);
499 assert!(bins.iter().all(|b| b.len() == 16));
500 }
501
502 #[test]
503 fn adversarial_listing_sizes_do_not_overflow_the_fit_test() {
504 let listing = vec![
507 entry("a", 10 * MB),
508 entry("huge", u64::MAX),
509 entry("z", 10 * MB),
510 ];
511 let bins = pack(listing, 64 * MB);
512 let huge_bin = bins
513 .iter()
514 .find(|b| b.iter().any(|e| e.key == "huge"))
515 .unwrap();
516 assert_eq!(huge_bin.len(), 1, "the oversized object lands alone");
517 let total: usize = bins.iter().map(Vec::len).sum();
518 assert_eq!(total, 3, "nothing lost, nothing duplicated");
519 }
520
521 #[test]
522 fn oversized_object_gets_its_own_split() {
523 let listing = vec![
524 entry("a", 10 * MB),
525 entry("huge", 500 * MB),
526 entry("b", 10 * MB),
527 ];
528 let bins = pack(listing, 64 * MB);
529 let huge_bin = bins
530 .iter()
531 .find(|b| b.iter().any(|e| e.key == "huge"))
532 .unwrap();
533 assert_eq!(
534 huge_bin.len(),
535 1,
536 "an oversized object never shares a split"
537 );
538 }
539
540 #[test]
541 fn packing_preserves_listing_order_within_and_across_bins() {
542 let listing: Vec<ObjectEntry> = (0..50)
543 .map(|i| {
544 entry(
545 &format!("k{i:02}"),
546 if i % 7 == 0 { 60 * MB } else { 3 * MB },
547 )
548 })
549 .collect();
550 let bins = pack(listing, 64 * MB);
551 for bin in &bins {
552 let keys: Vec<&str> = bin.iter().map(|e| e.key.as_str()).collect();
553 let mut sorted = keys.clone();
554 sorted.sort_unstable();
555 assert_eq!(keys, sorted, "members stay in listing order within a bin");
556 }
557 let firsts: Vec<&str> = bins.iter().map(|b| b[0].key.as_str()).collect();
558 let mut sorted = firsts.clone();
559 sorted.sort_unstable();
560 assert_eq!(
561 firsts, sorted,
562 "bins emerge in listing order of their first member"
563 );
564 }
565
566 #[test]
567 fn lookback_bounds_how_long_a_bin_stays_open() {
568 let mut listing = vec![entry("a-half-full", 40 * MB)];
571 for i in 0..PACKING_LOOKBACK + 2 {
574 listing.push(entry(&format!("b{i:02}"), 60 * MB));
575 }
576 listing.sort_unstable_by(|a, b| a.key.cmp(&b.key));
577 let bins = pack(listing, 64 * MB);
578 assert_eq!(bins.len(), PACKING_LOOKBACK + 3);
581 assert!(bins.iter().all(|b| b.len() == 1));
582 }
583
584 proptest! {
585 #[test]
586 fn prop_packing_partitions_the_listing_exactly(
587 sizes in proptest::collection::vec(0u64..300 * MB, 0..120),
588 target_mb in 1u64..129,
589 ) {
590 let listing: Vec<ObjectEntry> = sizes
591 .iter()
592 .enumerate()
593 .map(|(i, &s)| entry(&format!("k{i:04}"), s))
594 .collect();
595 let bins = pack(listing.clone(), target_mb * MB);
596 let repacked: Vec<ObjectEntry> = {
597 let mut all: Vec<ObjectEntry> = bins.iter().flatten().cloned().collect();
598 all.sort_unstable_by(|a, b| a.key.cmp(&b.key));
599 all
600 };
601 prop_assert_eq!(repacked, listing);
603 prop_assert!(bins.iter().all(|b| !b.is_empty()));
604 }
605
606 #[test]
607 fn prop_member_count_is_bounded_by_the_floor(
608 sizes in proptest::collection::vec(0u64..300 * MB, 0..120),
609 ) {
610 let target = 64 * MB;
611 let listing: Vec<ObjectEntry> = sizes
612 .iter()
613 .enumerate()
614 .map(|(i, &s)| entry(&format!("k{i:04}"), s))
615 .collect();
616 let bins = pack(listing, target);
617 prop_assert!(bins.iter().all(|b| b.len() <= 16));
620 }
621
622 #[test]
623 fn prop_split_ids_are_valid_for_arbitrary_keys(
624 keys in proptest::collection::btree_set("[ -~]{1,64}", 1..8),
625 etag in proptest::option::of("[ -~]{1,16}"),
626 ) {
627 let members: Vec<(&str, Option<&str>)> =
630 keys.iter().map(|k| (k.as_str(), etag.as_deref())).collect();
631 let id = split_id_for(members).unwrap();
632 prop_assert!(id.as_str().starts_with("s3-"));
633 prop_assert_eq!(id.as_str().len(), 25);
634 }
635
636 #[test]
637 fn prop_packing_and_ids_are_deterministic(
638 sizes in proptest::collection::vec(0u64..200 * MB, 1..60),
639 ) {
640 let listing: Vec<ObjectEntry> = sizes
641 .iter()
642 .enumerate()
643 .map(|(i, &s)| entry(&format!("k{i:04}"), s))
644 .collect();
645 let ids = |bins: &[Vec<ObjectEntry>]| -> Vec<SplitId> {
646 bins.iter()
647 .map(|b| {
648 split_id_for(b.iter().map(|e| (e.key.as_str(), e.etag.as_deref())))
649 .unwrap()
650 })
651 .collect()
652 };
653 let a = pack(listing.clone(), 32 * MB);
654 let b = pack(listing, 32 * MB);
655 prop_assert_eq!(ids(&a), ids(&b));
656 }
657 }
658}