use crate::fetch::ObjectEntry;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use sha2::{Digest as _, Sha256};
use spate_core::coordination::{CoordinationError, CoordinationErrorKind, SplitId};
use std::collections::VecDeque;
pub const DESCRIPTOR_VERSION: u32 = 1;
pub(crate) const PACKING_VERSION: u32 = 1;
pub(crate) const PACKING_LOOKBACK: usize = 10;
pub(crate) const OPEN_COST_DIVISOR: u64 = 16;
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DescriptorObject {
pub key: String,
pub size: u64,
pub etag: Option<String>,
pub last_modified_ms: i64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SplitDescriptor {
pub(crate) v: u32,
pub objects: Vec<DescriptorObject>,
}
#[derive(Deserialize)]
struct VersionProbe {
v: u32,
}
impl SplitDescriptor {
#[must_use]
pub fn new(objects: Vec<DescriptorObject>) -> SplitDescriptor {
SplitDescriptor {
v: DESCRIPTOR_VERSION,
objects,
}
}
#[must_use]
pub fn version(&self) -> u32 {
self.v
}
pub(crate) fn to_entries(&self) -> Vec<ObjectEntry> {
self.objects
.iter()
.map(|o| ObjectEntry {
key: o.key.clone(),
size: o.size,
etag: o.etag.clone(),
last_modified_ms: o.last_modified_ms,
})
.collect()
}
pub(crate) fn from_entries(entries: &[ObjectEntry]) -> SplitDescriptor {
SplitDescriptor::new(
entries
.iter()
.map(|e| DescriptorObject {
key: e.key.clone(),
size: e.size,
etag: e.etag.clone(),
last_modified_ms: e.last_modified_ms,
})
.collect(),
)
}
pub fn encode(&self) -> Result<Vec<u8>, CoordinationError> {
if self.v != DESCRIPTOR_VERSION {
return Err(CoordinationError::new(
CoordinationErrorKind::Fatal,
format!(
"descriptor version {} is not the supported {DESCRIPTOR_VERSION}; \
construct via SplitDescriptor::new",
self.v
),
));
}
Ok(serde_json::to_vec(self).expect("descriptor serialization is infallible: no non-string map keys, no fallible Serialize impls"))
}
pub fn decode(bytes: &[u8]) -> Result<SplitDescriptor, CoordinationError> {
let fatal = |reason: String| CoordinationError::new(CoordinationErrorKind::Fatal, reason);
let probe: VersionProbe = serde_json::from_slice(bytes)
.map_err(|e| fatal(format!("split descriptor is not valid JSON: {e}")))?;
if probe.v != DESCRIPTOR_VERSION {
return Err(fatal(format!(
"split descriptor version {} is not this release's version \
{DESCRIPTOR_VERSION}; the split was planned by an incompatible release",
probe.v
)));
}
serde_json::from_slice(bytes)
.map_err(|e| fatal(format!("split descriptor failed to decode: {e}")))
}
}
pub fn split_id_for<'a, I>(members: I) -> Result<SplitId, CoordinationError>
where
I: IntoIterator<Item = (&'a str, Option<&'a str>)>,
{
split_id_with_version(members, PACKING_VERSION)
}
fn split_id_with_version<'a, I>(members: I, version: u32) -> Result<SplitId, CoordinationError>
where
I: IntoIterator<Item = (&'a str, Option<&'a str>)>,
{
let mut members: Vec<(&str, Option<&str>)> = members.into_iter().collect();
if members.is_empty() {
return Err(CoordinationError::new(
CoordinationErrorKind::Fatal,
"cannot mint a split id for an empty member set",
));
}
members.sort_unstable_by_key(|(key, _)| *key);
let mut hasher = Sha256::new();
hasher.update(b"spate-s3-split\n");
hasher.update(version.to_le_bytes());
for (key, etag) in members {
hasher.update(u32::try_from(key.len()).unwrap_or(u32::MAX).to_le_bytes());
hasher.update(key.as_bytes());
match etag {
Some(etag) => {
hasher.update([0x01]);
hasher.update(u32::try_from(etag.len()).unwrap_or(u32::MAX).to_le_bytes());
hasher.update(etag.as_bytes());
}
None => hasher.update([0x00]),
}
}
let digest = hasher.finalize();
SplitId::new(format!("s3-{}", URL_SAFE_NO_PAD.encode(&digest[..16])))
}
pub(crate) fn pack(entries: Vec<ObjectEntry>, target_bytes: u64) -> Vec<Vec<ObjectEntry>> {
debug_assert!(
target_bytes > 0,
"config validation rejects a zero split target"
);
struct Bin {
members: Vec<ObjectEntry>,
cost: u64,
}
let floor = (target_bytes / OPEN_COST_DIVISOR).max(1);
let mut bins: Vec<Bin> = Vec::new();
let mut open: VecDeque<usize> = VecDeque::new();
for entry in entries {
let cost = entry.size.max(floor);
let idx = match open
.iter()
.position(|&i| bins[i].cost.saturating_add(cost) <= target_bytes)
{
Some(pos) => open[pos],
None => {
if open.len() == PACKING_LOOKBACK {
open.pop_front();
}
bins.push(Bin {
members: Vec::new(),
cost: 0,
});
let idx = bins.len() - 1;
open.push_back(idx);
idx
}
};
bins[idx].members.push(entry);
bins[idx].cost = bins[idx].cost.saturating_add(cost);
if bins[idx].cost >= target_bytes
&& let Some(pos) = open.iter().position(|&i| i == idx)
{
open.remove(pos);
}
}
bins.into_iter().map(|b| b.members).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
fn entry(key: &str, size: u64) -> ObjectEntry {
ObjectEntry {
key: key.to_string(),
size,
etag: Some(format!("\"etag-{key}\"")),
last_modified_ms: 1_760_000_000_000,
}
}
const MB: u64 = 1024 * 1024;
#[test]
fn digest_id_is_pinned() {
let id = split_id_for([
("exports/2026/part-000.ndjson", Some("\"9b2cf5\"")),
("exports/2026/part-001.ndjson", None),
])
.unwrap();
assert_eq!(id.as_str(), "s3-_rD2rPZklAFVV4pYEYaWxg");
}
#[test]
fn digest_is_order_insensitive_but_content_sensitive() {
let forward = split_id_for([("a", Some("1")), ("b", Some("2"))]).unwrap();
let reversed = split_id_for([("b", Some("2")), ("a", Some("1"))]).unwrap();
assert_eq!(forward, reversed);
let other_key = split_id_for([("a", Some("1")), ("c", Some("2"))]).unwrap();
assert_ne!(forward, other_key);
}
#[test]
fn digest_changes_when_an_etag_changes() {
let before = split_id_for([("a", Some("v1")), ("b", Some("x"))]).unwrap();
let overwritten = split_id_for([("a", Some("v2")), ("b", Some("x"))]).unwrap();
let dropped = split_id_for([("a", None), ("b", Some("x"))]).unwrap();
assert_ne!(before, overwritten);
assert_ne!(before, dropped);
}
#[test]
fn digest_folds_in_the_packing_version() {
let v1 = split_id_with_version([("a", Some("1"))], 1).unwrap();
let v2 = split_id_with_version([("a", Some("1"))], 2).unwrap();
assert_ne!(v1, v2);
}
#[test]
fn empty_member_set_is_rejected() {
assert!(split_id_for(std::iter::empty()).is_err());
}
#[test]
fn ambiguous_concatenations_do_not_collide() {
let a = split_id_for([("ab", Some("c"))]).unwrap();
let b = split_id_for([("a", Some("bc"))]).unwrap();
let c = split_id_for([("abc", None)]).unwrap();
assert_ne!(a, b);
assert_ne!(a, c);
assert_ne!(b, c);
}
#[test]
fn descriptor_round_trips() {
let desc = SplitDescriptor::from_entries(&[
entry("exports/part-000.ndjson.gz", 52 * MB),
ObjectEntry {
key: "exports/part-001.ndjson.gz".to_string(),
size: 9,
etag: None,
last_modified_ms: 1,
},
]);
let decoded = SplitDescriptor::decode(&desc.encode().unwrap()).unwrap();
assert_eq!(decoded, desc);
assert_eq!(decoded.v, DESCRIPTOR_VERSION);
}
#[test]
fn descriptor_encoding_is_pinned() {
let desc = SplitDescriptor::from_entries(&[entry("k", 5)]);
assert_eq!(
String::from_utf8(desc.encode().unwrap()).unwrap(),
r#"{"v":1,"objects":[{"key":"k","size":5,"etag":"\"etag-k\"","last_modified_ms":1760000000000}]}"#,
);
}
#[test]
fn encode_refuses_a_descriptor_not_built_by_new() {
let rogue = SplitDescriptor {
v: 0,
objects: vec![],
};
let err = rogue.encode().unwrap_err();
assert_eq!(err.kind, CoordinationErrorKind::Fatal);
assert!(
err.reason.contains("SplitDescriptor::new"),
"reason: {}",
err.reason
);
assert_eq!(SplitDescriptor::new(vec![]).version(), DESCRIPTOR_VERSION);
}
#[test]
fn unknown_descriptor_version_is_rejected_actionably() {
let err = SplitDescriptor::decode(br#"{"v":999,"objects":[]}"#).unwrap_err();
assert_eq!(err.kind, CoordinationErrorKind::Fatal);
assert!(err.reason.contains("version 999"), "reason: {}", err.reason);
assert!(
err.reason.contains("incompatible release"),
"reason: {}",
err.reason
);
let garbage = SplitDescriptor::decode(b"not json").unwrap_err();
assert_eq!(garbage.kind, CoordinationErrorKind::Fatal);
}
#[test]
fn packing_is_deterministic_for_a_fixed_listing() {
let listing: Vec<ObjectEntry> = (0..200)
.map(|i| entry(&format!("k{i:04}"), (i % 40) * MB))
.collect();
let a = pack(listing.clone(), 64 * MB);
let b = pack(listing, 64 * MB);
assert_eq!(a, b);
}
#[test]
fn tiny_objects_coalesce_under_the_open_cost_floor() {
let listing: Vec<ObjectEntry> = (0..64).map(|i| entry(&format!("k{i:02}"), 1)).collect();
let bins = pack(listing, 64 * MB);
assert_eq!(bins.len(), 4);
assert!(bins.iter().all(|b| b.len() == 16));
}
#[test]
fn adversarial_listing_sizes_do_not_overflow_the_fit_test() {
let listing = vec![
entry("a", 10 * MB),
entry("huge", u64::MAX),
entry("z", 10 * MB),
];
let bins = pack(listing, 64 * MB);
let huge_bin = bins
.iter()
.find(|b| b.iter().any(|e| e.key == "huge"))
.unwrap();
assert_eq!(huge_bin.len(), 1, "the oversized object lands alone");
let total: usize = bins.iter().map(Vec::len).sum();
assert_eq!(total, 3, "nothing lost, nothing duplicated");
}
#[test]
fn oversized_object_gets_its_own_split() {
let listing = vec![
entry("a", 10 * MB),
entry("huge", 500 * MB),
entry("b", 10 * MB),
];
let bins = pack(listing, 64 * MB);
let huge_bin = bins
.iter()
.find(|b| b.iter().any(|e| e.key == "huge"))
.unwrap();
assert_eq!(
huge_bin.len(),
1,
"an oversized object never shares a split"
);
}
#[test]
fn packing_preserves_listing_order_within_and_across_bins() {
let listing: Vec<ObjectEntry> = (0..50)
.map(|i| {
entry(
&format!("k{i:02}"),
if i % 7 == 0 { 60 * MB } else { 3 * MB },
)
})
.collect();
let bins = pack(listing, 64 * MB);
for bin in &bins {
let keys: Vec<&str> = bin.iter().map(|e| e.key.as_str()).collect();
let mut sorted = keys.clone();
sorted.sort_unstable();
assert_eq!(keys, sorted, "members stay in listing order within a bin");
}
let firsts: Vec<&str> = bins.iter().map(|b| b[0].key.as_str()).collect();
let mut sorted = firsts.clone();
sorted.sort_unstable();
assert_eq!(
firsts, sorted,
"bins emerge in listing order of their first member"
);
}
#[test]
fn lookback_bounds_how_long_a_bin_stays_open() {
let mut listing = vec![entry("a-half-full", 40 * MB)];
for i in 0..PACKING_LOOKBACK + 2 {
listing.push(entry(&format!("b{i:02}"), 60 * MB));
}
listing.sort_unstable_by(|a, b| a.key.cmp(&b.key));
let bins = pack(listing, 64 * MB);
assert_eq!(bins.len(), PACKING_LOOKBACK + 3);
assert!(bins.iter().all(|b| b.len() == 1));
}
proptest! {
#[test]
fn prop_packing_partitions_the_listing_exactly(
sizes in proptest::collection::vec(0u64..300 * MB, 0..120),
target_mb in 1u64..129,
) {
let listing: Vec<ObjectEntry> = sizes
.iter()
.enumerate()
.map(|(i, &s)| entry(&format!("k{i:04}"), s))
.collect();
let bins = pack(listing.clone(), target_mb * MB);
let repacked: Vec<ObjectEntry> = {
let mut all: Vec<ObjectEntry> = bins.iter().flatten().cloned().collect();
all.sort_unstable_by(|a, b| a.key.cmp(&b.key));
all
};
prop_assert_eq!(repacked, listing);
prop_assert!(bins.iter().all(|b| !b.is_empty()));
}
#[test]
fn prop_member_count_is_bounded_by_the_floor(
sizes in proptest::collection::vec(0u64..300 * MB, 0..120),
) {
let target = 64 * MB;
let listing: Vec<ObjectEntry> = sizes
.iter()
.enumerate()
.map(|(i, &s)| entry(&format!("k{i:04}"), s))
.collect();
let bins = pack(listing, target);
prop_assert!(bins.iter().all(|b| b.len() <= 16));
}
#[test]
fn prop_split_ids_are_valid_for_arbitrary_keys(
keys in proptest::collection::btree_set("[ -~]{1,64}", 1..8),
etag in proptest::option::of("[ -~]{1,16}"),
) {
let members: Vec<(&str, Option<&str>)> =
keys.iter().map(|k| (k.as_str(), etag.as_deref())).collect();
let id = split_id_for(members).unwrap();
prop_assert!(id.as_str().starts_with("s3-"));
prop_assert_eq!(id.as_str().len(), 25);
}
#[test]
fn prop_packing_and_ids_are_deterministic(
sizes in proptest::collection::vec(0u64..200 * MB, 1..60),
) {
let listing: Vec<ObjectEntry> = sizes
.iter()
.enumerate()
.map(|(i, &s)| entry(&format!("k{i:04}"), s))
.collect();
let ids = |bins: &[Vec<ObjectEntry>]| -> Vec<SplitId> {
bins.iter()
.map(|b| {
split_id_for(b.iter().map(|e| (e.key.as_str(), e.etag.as_deref())))
.unwrap()
})
.collect()
};
let a = pack(listing.clone(), 32 * MB);
let b = pack(listing, 32 * MB);
prop_assert_eq!(ids(&a), ids(&b));
}
}
}