use crate::config::Compression;
use crate::fetch::ObjectEntry;
use crate::framer::Codec;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use sha2::{Digest as _, Sha256};
pub const DESCRIPTOR_VERSION: u32 = 2;
pub const MAX_UNIT_MEMBERS: usize = 4096;
pub(crate) const PACK_MEMBERS: usize = 16;
pub(crate) const PACK_KEY_MAX_BYTES: usize = 1024;
pub(crate) const PACK_ETAG_MAX_BYTES: usize = 256;
pub(crate) const PACKED_BODY_MAX_BYTES: usize = 48 * 1024;
const _: () = assert!(PACK_MEMBERS <= MAX_UNIT_MEMBERS);
const _: () = assert!(PACKED_BODY_MAX_BYTES + 1024 + 25 <= 1_000_000 - 36);
const _: () =
assert!(PACKED_BODY_MAX_BYTES + 1024 + 25 <= 1_048_576 - "spate.unit".len() - "String".len());
pub(crate) const PACKING_VERSION: u32 = 2;
pub(crate) const MAX_SUBDIVIDABLE_BYTES: u64 = 50_000 << 30;
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct SplitError {
pub reason: String,
}
impl SplitError {
fn new(reason: impl Into<String>) -> SplitError {
SplitError {
reason: reason.into(),
}
}
}
impl std::fmt::Display for SplitError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.reason)
}
}
impl std::error::Error for SplitError {}
#[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, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SplitRange {
pub start: u64,
pub end: u64,
pub delimiter: u8,
}
impl SplitRange {
#[must_use]
pub fn new(start: u64, end: u64, delimiter: u8) -> SplitRange {
SplitRange {
start,
end,
delimiter,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SplitDescriptor {
pub(crate) v: u32,
pub objects: Vec<DescriptorObject>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub range: Option<SplitRange>,
}
#[derive(Deserialize)]
struct VersionProbe {
v: u32,
}
impl SplitDescriptor {
#[must_use]
pub fn new(objects: Vec<DescriptorObject>) -> SplitDescriptor {
SplitDescriptor {
v: DESCRIPTOR_VERSION,
objects,
range: None,
}
}
#[must_use]
pub fn with_range(object: DescriptorObject, range: SplitRange) -> SplitDescriptor {
SplitDescriptor {
v: DESCRIPTOR_VERSION,
objects: vec![object],
range: Some(range),
}
}
#[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()
}
#[cfg(test)]
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>, SplitError> {
if self.v != DESCRIPTOR_VERSION {
return Err(SplitError::new(format!(
"descriptor version {} is not the supported {DESCRIPTOR_VERSION}; \
construct via SplitDescriptor::new",
self.v
)));
}
self.validate_range()?;
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, SplitError> {
let fatal = SplitError::new;
let probe: VersionProbe = serde_json::from_slice(bytes)
.map_err(|e| fatal(format!("unit descriptor is not valid JSON: {e}")))?;
if probe.v != DESCRIPTOR_VERSION {
return Err(fatal(format!(
"unit descriptor version {} is not this release's version \
{DESCRIPTOR_VERSION}; the unit was written by an incompatible release",
probe.v
)));
}
let descriptor: SplitDescriptor = serde_json::from_slice(bytes)
.map_err(|e| fatal(format!("unit descriptor failed to decode: {e}")))?;
descriptor.validate_range()?;
Ok(descriptor)
}
fn validate_range(&self) -> Result<(), SplitError> {
let Some(range) = self.range else {
return Ok(());
};
let fatal = SplitError::new;
let [object] = self.objects.as_slice() else {
return Err(fatal(format!(
"a ranged unit descriptor has exactly one member, this one has {}",
self.objects.len()
)));
};
if object.etag.is_none() {
return Err(fatal(format!(
"ranged unit descriptor over \"{}\" has no ETag to pin its reads to",
object.key
)));
}
if range.start >= range.end || range.end > object.size {
return Err(fatal(format!(
"ranged unit descriptor over \"{}\" has range [{}, {}), which is empty or \
extends past the object's {} bytes",
object.key, range.start, range.end, object.size
)));
}
Ok(())
}
}
pub fn split_id_for<'a, I>(members: I) -> Result<String, SplitError>
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<String, SplitError>
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(SplitError::new(
"cannot mint a unit 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 {
digest_member(&mut hasher, key, etag);
}
id_from_digest(hasher)
}
pub fn split_id_for_range(key: &str, etag: &str, range: SplitRange) -> Result<String, SplitError> {
if range.start >= range.end {
return Err(SplitError::new(format!(
"cannot mint a unit id for the empty range [{}, {})",
range.start, range.end
)));
}
let mut hasher = Sha256::new();
hasher.update(b"spate-s3-range\n");
hasher.update(PACKING_VERSION.to_le_bytes());
digest_member(&mut hasher, key, Some(etag));
hasher.update(range.start.to_le_bytes());
hasher.update(range.end.to_le_bytes());
hasher.update([range.delimiter]);
id_from_digest(hasher)
}
fn digest_member(hasher: &mut Sha256, key: &str, etag: Option<&str>) {
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]),
}
}
fn id_from_digest(hasher: Sha256) -> Result<String, SplitError> {
let digest = hasher.finalize();
Ok(format!("s3-{}", URL_SAFE_NO_PAD.encode(&digest[..16])))
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct Packing {
pub(crate) target_bytes: u64,
pub(crate) compression: Compression,
pub(crate) delimiter: Option<u8>,
}
impl Packing {
pub(crate) fn delimiter_for(&self, entry: &ObjectEntry) -> Option<u8> {
let delimiter = self.delimiter?;
let cut = entry.size > self.target_bytes
&& entry.size <= MAX_SUBDIVIDABLE_BYTES
&& entry.etag.is_some()
&& Codec::resolve(self.compression, &entry.key) == Codec::Plain;
cut.then_some(delimiter)
}
}
pub(crate) fn tile(size: u64, target: u64, delimiter: u8) -> impl Iterator<Item = SplitRange> {
let n = size.div_ceil(target);
let boundary = move |i: u64| (u128::from(i) * u128::from(size) / u128::from(n)) as u64;
(0..n).map(move |i| SplitRange::new(boundary(i), boundary(i + 1), delimiter))
}
#[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;
fn whole(target: u64) -> Packing {
Packing {
target_bytes: target,
compression: Compression::Auto,
delimiter: None,
}
}
fn cut(target: u64) -> Packing {
Packing {
delimiter: Some(b'\n'),
..whole(target)
}
}
#[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, "s3-xje_MfL-wSNqMCft5ellCg");
}
#[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() {
let err = split_id_for(std::iter::empty()).unwrap_err();
assert_eq!(err.to_string(), err.reason);
assert!(err.reason.contains("empty member set"), "{}", err.reason);
}
#[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":2,"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![],
range: None,
};
let err = rogue.encode().unwrap_err();
assert!(
err.reason.contains("SplitDescriptor::new"),
"reason: {}",
err.reason
);
assert_eq!(SplitDescriptor::new(vec![]).version(), DESCRIPTOR_VERSION);
}
#[test]
fn a_version_1_descriptor_is_refused() {
let err = SplitDescriptor::decode(
br#"{"v":1,"objects":[{"key":"k","size":5,"etag":null,"last_modified_ms":1}]}"#,
)
.unwrap_err();
assert!(err.reason.contains("version 1"), "reason: {}", err.reason);
}
#[test]
fn unknown_descriptor_version_is_rejected_actionably() {
let err = SplitDescriptor::decode(br#"{"v":999,"objects":[]}"#).unwrap_err();
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!(
garbage.reason.contains("not valid JSON"),
"{}",
garbage.reason
);
}
fn ranged_object(size: u64, etag: Option<&str>) -> DescriptorObject {
DescriptorObject {
key: "big.ndjson".to_string(),
size,
etag: etag.map(str::to_owned),
last_modified_ms: 7,
}
}
#[test]
fn ranged_descriptor_round_trips_and_its_encoding_is_pinned() {
let desc = SplitDescriptor::with_range(
ranged_object(100, Some("\"e\"")),
SplitRange::new(40, 80, b'\n'),
);
let encoded = desc.encode().unwrap();
assert_eq!(
String::from_utf8(encoded.clone()).unwrap(),
r#"{"v":2,"objects":[{"key":"big.ndjson","size":100,"etag":"\"e\"","last_modified_ms":7}],"range":{"start":40,"end":80,"delimiter":10}}"#,
);
assert_eq!(SplitDescriptor::decode(&encoded).unwrap(), desc);
}
#[test]
fn invalid_ranged_descriptors_are_refused_on_both_sides() {
let pinned = || ranged_object(100, Some("\"e\""));
let two_members = SplitDescriptor {
objects: vec![pinned(), pinned()],
..SplitDescriptor::with_range(pinned(), SplitRange::new(0, 10, b'\n'))
};
let no_members = SplitDescriptor {
objects: vec![],
..SplitDescriptor::with_range(pinned(), SplitRange::new(0, 10, b'\n'))
};
let cases = [
("member", two_members),
("member", no_members),
(
"ETag",
SplitDescriptor::with_range(
ranged_object(100, None),
SplitRange::new(0, 10, b'\n'),
),
),
(
"empty",
SplitDescriptor::with_range(pinned(), SplitRange::new(10, 10, b'\n')),
),
(
"empty",
SplitDescriptor::with_range(pinned(), SplitRange::new(11, 10, b'\n')),
),
(
"past the object",
SplitDescriptor::with_range(pinned(), SplitRange::new(0, 101, b'\n')),
),
];
for (needle, desc) in cases {
let err = desc.encode().unwrap_err();
assert!(err.reason.contains(needle), "encode: {}", err.reason);
let raw = serde_json::to_vec(&desc).unwrap();
let err = SplitDescriptor::decode(&raw).unwrap_err();
assert!(err.reason.contains(needle), "decode: {}", err.reason);
}
let whole = SplitDescriptor::with_range(pinned(), SplitRange::new(0, 100, b'\n'));
SplitDescriptor::decode(&whole.encode().unwrap()).unwrap();
}
#[test]
fn ranged_id_is_pinned() {
let id = split_id_for_range(
"exports/2026/part-000.ndjson",
"\"9b2cf5\"",
SplitRange::new(0, 64 * MB, b'\n'),
)
.unwrap();
assert_eq!(id, "s3-fA0wNRW2MyDxquwFvNy0JA");
}
#[test]
fn ranged_ids_digest_every_input_and_never_equal_member_set_ids() {
let base = split_id_for_range("k", "e", SplitRange::new(0, 10, b'\n')).unwrap();
for other in [
split_id_for_range("k2", "e", SplitRange::new(0, 10, b'\n')),
split_id_for_range("k", "e2", SplitRange::new(0, 10, b'\n')),
split_id_for_range("k", "e", SplitRange::new(1, 10, b'\n')),
split_id_for_range("k", "e", SplitRange::new(0, 11, b'\n')),
split_id_for_range("k", "e", SplitRange::new(0, 10, b';')),
] {
assert_ne!(other.unwrap(), base);
}
assert_ne!(split_id_for([("k", Some("e"))]).unwrap(), base);
assert!(split_id_for_range("k", "e", SplitRange::new(10, 10, b'\n')).is_err());
}
#[test]
fn an_object_above_the_subdivision_limit_stays_whole() {
let packing = cut(1 << 30);
assert_eq!(packing.delimiter_for(&entry("6tib", 6 << 40)), Some(b'\n'));
let at = entry("at", MAX_SUBDIVIDABLE_BYTES);
assert_eq!(packing.delimiter_for(&at), Some(b'\n'));
let over = entry("over", MAX_SUBDIVIDABLE_BYTES + 1);
assert_eq!(packing.delimiter_for(&over), None);
}
#[test]
fn objects_that_cannot_be_cut_stay_whole() {
let target = 64 * MB;
let plain = entry("plain.ndjson", 3 * target);
let gzip_config = Packing {
compression: Compression::Gzip,
..cut(target)
};
let no_etag = ObjectEntry {
etag: None,
..plain.clone()
};
for (packing, object) in [
(cut(target), entry("gz.ndjson.gz", 3 * target)),
(cut(target), entry("zst.ndjson.zst", 3 * target)),
(gzip_config, plain.clone()),
(cut(target), no_etag),
(whole(target), plain.clone()),
(cut(target), entry("at.ndjson", target)),
] {
assert_eq!(
packing.delimiter_for(&object),
None,
"{packing:?} {object:?}"
);
}
assert_eq!(
cut(target).delimiter_for(&entry("over.ndjson", target + 1)),
Some(b'\n'),
"one byte above the target is cut"
);
}
#[test]
fn ranges_tile_the_object_within_half_to_one_target() {
for target in 1..=64u64 {
for size in target + 1..=8 * target {
assert_tiles(size, target);
}
}
assert_tiles(MAX_SUBDIVIDABLE_BYTES, MB);
assert_tiles(MAX_SUBDIVIDABLE_BYTES - 1, MB);
}
fn assert_tiles(size: u64, target: u64) {
let mut count = 0;
let mut next = 0;
for range in tile(size, target, b'\n') {
count += 1;
assert_eq!(range.start, next, "size {size} target {target}");
let len = range.end - range.start;
assert!(
target / 2 <= len && len <= target,
"size {size} target {target}: range {range:?} is {len} bytes"
);
next = range.end;
}
assert_eq!(count, size.div_ceil(target), "size {size} target {target}");
assert_eq!(next, size, "size {size} target {target}");
}
proptest! {
#[test]
fn prop_ranges_tile_large_objects(
target in 1u64..=1 << 32,
n in 2u64..=1 << 12,
extra in any::<u64>(),
) {
assert_tiles((n - 1) * target + 1 + extra % target, target);
assert_tiles(n * target, target);
}
#[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.starts_with("s3-"));
prop_assert_eq!(id.len(), 25);
}
}
}