#![allow(async_fn_in_trait)]
#![deny(missing_docs)]
pub use weir_core::Payload;
pub trait SinkError: Send + Sync + std::error::Error + 'static {
fn is_transient(&self) -> bool;
fn retry_after(&self) -> Option<std::time::Duration> {
None
}
}
impl SinkError for std::convert::Infallible {
fn is_transient(&self) -> bool {
match *self {}
}
}
#[derive(Debug, Clone)]
pub struct BasicSinkError {
message: String,
transient: bool,
}
impl BasicSinkError {
pub fn transient(message: impl Into<String>) -> Self {
Self {
message: message.into(),
transient: true,
}
}
pub fn permanent(message: impl Into<String>) -> Self {
Self {
message: message.into(),
transient: false,
}
}
}
impl std::fmt::Display for BasicSinkError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.message)
}
}
impl std::error::Error for BasicSinkError {}
impl SinkError for BasicSinkError {
fn is_transient(&self) -> bool {
self.transient
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct CommitResult {
pub committed: Vec<Payload>,
pub dead_lettered: Vec<(Payload, String)>,
}
impl CommitResult {
#[must_use]
pub fn new(committed: Vec<Payload>, dead_lettered: Vec<(Payload, String)>) -> Self {
Self {
committed,
dead_lettered,
}
}
}
#[derive(Clone, Debug)]
#[non_exhaustive]
pub enum SinkHealth {
Healthy,
Degraded(String),
Down(String),
}
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct DedupToken([u8; 32]);
impl DedupToken {
pub fn for_payloads<'a>(payloads: impl IntoIterator<Item = &'a Payload>) -> Self {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
for p in payloads {
hasher.update((p.len() as u64).to_le_bytes());
hasher.update(p.as_ref());
}
Self(hasher.finalize().into())
}
pub fn from_bytes(bytes: [u8; 32]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; 32] {
&self.0
}
pub fn to_hex(&self) -> String {
use std::fmt::Write as _;
let mut out = String::with_capacity(64);
for b in self.0 {
let _ = write!(out, "{b:02x}");
}
out
}
pub fn to_prefixed_hex(&self) -> String {
format!("sha256:{}", self.to_hex())
}
}
impl std::fmt::Debug for DedupToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "DedupToken({})", self.to_hex())
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct RecordId([u8; 32]);
impl RecordId {
pub fn for_record(segment: &str, index: u64, payload: &Payload) -> Self {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update((segment.len() as u64).to_le_bytes());
hasher.update(segment.as_bytes());
hasher.update(index.to_le_bytes());
hasher.update((payload.len() as u64).to_le_bytes());
hasher.update(payload.as_ref());
Self(hasher.finalize().into())
}
pub fn from_bytes(bytes: [u8; 32]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; 32] {
&self.0
}
pub fn to_hex(&self) -> String {
use std::fmt::Write as _;
let mut out = String::with_capacity(64);
for b in self.0 {
let _ = write!(out, "{b:02x}");
}
out
}
pub fn to_prefixed_hex(&self) -> String {
format!("sha256:{}", self.to_hex())
}
}
impl std::fmt::Debug for RecordId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "RecordId({})", self.to_hex())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SinkBatch {
records: Vec<Payload>,
dedup_token: DedupToken,
record_ids: Option<Vec<RecordId>>,
segment_created_at: Option<i64>,
}
impl SinkBatch {
pub fn new(records: Vec<Payload>, dedup_token: DedupToken) -> Self {
Self {
records,
dedup_token,
record_ids: None,
segment_created_at: None,
}
}
pub fn with_record_ids(
records: Vec<Payload>,
dedup_token: DedupToken,
record_ids: Vec<RecordId>,
) -> Self {
assert_eq!(
records.len(),
record_ids.len(),
"SinkBatch::with_record_ids: {} records but {} ids",
records.len(),
record_ids.len()
);
Self {
records,
dedup_token,
record_ids: Some(record_ids),
segment_created_at: None,
}
}
#[must_use]
pub fn with_segment_context(
records: Vec<Payload>,
dedup_token: DedupToken,
record_ids: Vec<RecordId>,
segment_created_at: i64,
) -> Self {
assert_eq!(
records.len(),
record_ids.len(),
"with_segment_context: record_ids must be parallel to records"
);
Self {
records,
dedup_token,
record_ids: Some(record_ids),
segment_created_at: Some(segment_created_at),
}
}
#[must_use]
pub fn segment_created_at(&self) -> Option<i64> {
self.segment_created_at
}
pub fn record_ids(&self) -> Option<&[RecordId]> {
self.record_ids.as_deref()
}
pub fn into_records_with_ids(self) -> (Vec<Payload>, Option<Vec<RecordId>>) {
(self.records, self.record_ids)
}
pub fn records(&self) -> &[Payload] {
&self.records
}
pub fn into_records(self) -> Vec<Payload> {
self.records
}
pub fn into_parts(self) -> (Vec<Payload>, DedupToken) {
(self.records, self.dedup_token)
}
pub fn dedup_token(&self) -> &DedupToken {
&self.dedup_token
}
pub fn len(&self) -> usize {
self.records.len()
}
pub fn is_empty(&self) -> bool {
self.records.is_empty()
}
}
impl From<Vec<Payload>> for SinkBatch {
fn from(records: Vec<Payload>) -> Self {
let dedup_token = DedupToken::for_payloads(records.iter());
Self {
records,
dedup_token,
record_ids: None,
segment_created_at: None,
}
}
}
pub trait Sink: Send + Sync + 'static {
type Error: SinkError;
async fn commit(&self, batch: SinkBatch) -> Result<CommitResult, Self::Error>;
fn max_batch_size(&self) -> usize {
1000
}
async fn health(&self) -> SinkHealth;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn identical_payloads_at_different_coordinates_get_different_ids() {
let p = Payload::from(b"heartbeat".as_ref());
let a = RecordId::for_record("seg-000042.wab", 7, &p);
let b = RecordId::for_record("seg-000042.wab", 8, &p);
let c = RecordId::for_record("seg-000043.wab", 7, &p);
assert_ne!(a, b, "same segment, different index must differ");
assert_ne!(a, c, "same index, different segment must differ");
assert_ne!(b, c);
}
#[test]
fn ids_do_not_move_when_the_records_are_re_split_into_other_batches() {
let segment = "shard_00/seg_00000001.wab.sealed";
let payloads: Vec<Payload> = (1..=10)
.map(|i| Payload::copy_from_slice(format!("record-{i:02}").as_bytes()))
.collect();
let ids: Vec<RecordId> = payloads
.iter()
.enumerate()
.map(|(i, p)| RecordId::for_record(segment, i as u64 + 1, p))
.collect();
let at = |n: usize| -> Vec<RecordId> { ids.chunks(n).flat_map(|c| c.to_vec()).collect() };
assert_eq!(at(3).len(), 10);
assert_eq!(at(3), at(7), "a re-split must not change any id");
let batch_relative = |n: usize| -> Vec<RecordId> {
payloads
.chunks(n)
.flat_map(|c| {
c.iter()
.enumerate()
.map(|(i, p)| RecordId::for_record(segment, i as u64 + 1, p))
.collect::<Vec<_>>()
})
.collect()
};
assert_ne!(
batch_relative(3),
batch_relative(7),
"precondition: a batch-relative index is what would break this"
);
}
#[test]
fn the_same_record_reread_gets_the_same_id() {
let p = Payload::from(b"heartbeat".as_ref());
assert_eq!(
RecordId::for_record("seg-000042.wab", 7, &p),
RecordId::for_record("seg-000042.wab", 7, &p)
);
}
#[test]
fn field_boundaries_cannot_be_confused() {
let p = Payload::from(b"x".as_ref());
assert_ne!(
RecordId::for_record("ab", 1, &p),
RecordId::for_record("a", 1, &p),
"a shorter segment name must not collide with a longer one"
);
}
#[test]
fn payload_bytes_still_participate() {
let a = RecordId::for_record("seg.wab", 0, &Payload::from(b"one".as_ref()));
let b = RecordId::for_record("seg.wab", 0, &Payload::from(b"two".as_ref()));
assert_ne!(a, b);
}
#[test]
fn record_id_hex_forms_are_well_shaped() {
let id = RecordId::for_record("seg.wab", 0, &Payload::from(b"x".as_ref()));
assert_eq!(id.to_hex().len(), 64);
assert!(id.to_hex().chars().all(|c| c.is_ascii_hexdigit()));
assert_eq!(id.to_prefixed_hex(), format!("sha256:{}", id.to_hex()));
assert!(format!("{id:?}").contains(&id.to_hex()));
}
#[test]
fn with_record_ids_round_trips() {
let records = vec![Payload::from(b"a".as_ref()), Payload::from(b"b".as_ref())];
let ids: Vec<_> = records
.iter()
.enumerate()
.map(|(i, p)| RecordId::for_record("seg.wab", i as u64, p))
.collect();
let batch = SinkBatch::with_record_ids(
records.clone(),
DedupToken::for_payloads(records.iter()),
ids.clone(),
);
assert_eq!(batch.record_ids(), Some(&ids[..]));
let (r, got) = batch.into_records_with_ids();
assert_eq!(r, records);
assert_eq!(got, Some(ids));
}
#[test]
fn plain_batches_carry_no_record_ids() {
let records = vec![Payload::from(b"a".as_ref())];
let batch = SinkBatch::new(records.clone(), DedupToken::for_payloads(records.iter()));
assert_eq!(batch.record_ids(), None);
assert_eq!(SinkBatch::from(records).record_ids(), None);
}
#[test]
#[should_panic(expected = "2 records but 1 ids")]
fn mismatched_id_count_panics_rather_than_mispairing() {
let records = vec![Payload::from(b"a".as_ref()), Payload::from(b"b".as_ref())];
let token = DedupToken::for_payloads(records.iter());
SinkBatch::with_record_ids(records, token, vec![RecordId::from_bytes([0u8; 32])]);
}
fn block_on<F: std::future::Future>(fut: F) -> F::Output {
use std::task::{Context, Poll};
let mut fut = std::pin::pin!(fut);
let waker = std::task::Waker::noop();
let mut cx = Context::from_waker(waker);
loop {
if let Poll::Ready(v) = fut.as_mut().poll(&mut cx) {
return v;
}
}
}
#[test]
fn a_custom_sink_can_be_driven_and_unit_tested() {
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug)]
struct Never;
impl std::fmt::Display for Never {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "never")
}
}
impl std::error::Error for Never {}
impl SinkError for Never {
fn is_transient(&self) -> bool {
true
}
}
struct CountingSink {
committed: AtomicUsize,
}
impl Sink for CountingSink {
type Error = Never;
async fn commit(&self, batch: SinkBatch) -> Result<CommitResult, Never> {
let batch = batch.into_records();
let (mut ok, mut dead) = (Vec::new(), Vec::new());
for r in batch {
if r.as_ref() == b"reject" {
dead.push((r, "rejected by CountingSink".to_string()));
} else {
ok.push(r);
}
}
self.committed.fetch_add(ok.len(), Ordering::Relaxed);
Ok(CommitResult::new(ok, dead))
}
async fn health(&self) -> SinkHealth {
SinkHealth::Healthy
}
}
let sink = CountingSink {
committed: AtomicUsize::new(0),
};
let records = vec![
Payload::copy_from_slice(b"keep-1"),
Payload::copy_from_slice(b"reject"),
Payload::copy_from_slice(b"keep-2"),
];
let batch = SinkBatch::new(records.clone(), DedupToken::for_payloads(records.iter()));
let result = block_on(sink.commit(batch)).unwrap();
assert_eq!(result.committed.len(), 2);
assert_eq!(result.dead_lettered.len(), 1);
assert_eq!(&result.dead_lettered[0].0[..], b"reject");
assert_eq!(sink.committed.load(Ordering::Relaxed), 2);
assert!(matches!(block_on(sink.health()), SinkHealth::Healthy));
}
#[test]
fn basic_sink_error_classifies_and_displays() {
let t = BasicSinkError::transient("503 from upstream");
assert!(t.is_transient());
assert_eq!(t.to_string(), "503 from upstream");
let p = BasicSinkError::permanent("400 bad request");
assert!(!p.is_transient());
let _: &dyn SinkError = &p;
let _: &dyn std::error::Error = &p;
}
#[test]
fn commit_result_new_keeps_both_partitions() {
let r = CommitResult::new(
vec![Payload::from(b"a".as_ref())],
vec![(Payload::from(b"b".as_ref()), "rejected".to_string())],
);
assert_eq!(r.committed.len(), 1);
assert_eq!(r.dead_lettered.len(), 1);
assert_eq!(&r.committed[0][..], b"a");
assert_eq!(&r.dead_lettered[0].0[..], b"b");
assert_eq!(r.dead_lettered[0].1, "rejected");
}
fn p(bytes: &[u8]) -> Payload {
Payload::from(bytes)
}
const KNOWN_ANSWER_1_X: &str =
"5bc1fe58cc34db881c67b2acd898651311f6dfc576285c906c8f97e049b15342";
#[test]
fn dedup_token_is_deterministic() {
let batch = [p(b"a"), p(b"b")];
assert_eq!(
DedupToken::for_payloads(&batch).to_hex(),
DedupToken::for_payloads(&batch).to_hex()
);
}
#[test]
fn dedup_token_changes_on_reorder() {
let a = [p(b"a"), p(b"b")];
let b = [p(b"b"), p(b"a")];
assert_ne!(
DedupToken::for_payloads(&a).to_hex(),
DedupToken::for_payloads(&b).to_hex()
);
}
#[test]
fn dedup_token_distinguishes_different_batch_boundaries() {
let a = [p(b"ab"), p(b"c")];
let b = [p(b"a"), p(b"bc")];
assert_ne!(
DedupToken::for_payloads(&a).to_hex(),
DedupToken::for_payloads(&b).to_hex()
);
}
#[test]
fn dedup_token_empty_batch_is_the_empty_sha256() {
assert_eq!(
DedupToken::for_payloads(&[] as &[Payload]).to_hex(),
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
);
}
#[test]
fn dedup_token_hex_forms_agree() {
let t = DedupToken::for_payloads(&[p(b"x")]);
assert_eq!(t.to_prefixed_hex(), format!("sha256:{}", t.to_hex()));
assert_eq!(t.to_hex().len(), 64);
assert!(
t.to_hex()
.chars()
.all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
);
}
#[test]
fn dedup_token_round_trips_through_bytes() {
let t = DedupToken::for_payloads(&[p(b"x")]);
assert_eq!(DedupToken::from_bytes(*t.as_bytes()).to_hex(), t.to_hex());
}
#[test]
fn dedup_token_debug_is_hex_not_a_byte_array() {
let t = DedupToken::for_payloads(&[p(b"x")]);
let d = format!("{t:?}");
assert!(
d.contains(&t.to_hex()),
"Debug should show the hex digest: {d}"
);
assert!(
!d.contains('['),
"Debug should not dump the byte array: {d}"
);
}
#[test]
fn dedup_token_matches_the_weir_1_x_digest() {
let batch = [p(b"alpha"), p(b"beta"), p(b"gamma")];
assert_eq!(
DedupToken::for_payloads(&batch).to_hex(),
KNOWN_ANSWER_1_X,
"the dedup digest changed — see the doc comment"
);
}
#[test]
fn sink_batch_exposes_records_and_token() {
let token = DedupToken::for_payloads(&[p(b"a")]);
let batch = SinkBatch::new(vec![p(b"a"), p(b"b")], token);
assert_eq!(batch.len(), 2);
assert!(!batch.is_empty());
assert_eq!(batch.records(), &[p(b"a"), p(b"b")][..]);
assert_eq!(batch.dedup_token().to_hex(), token.to_hex());
}
#[test]
fn sink_batch_into_records_is_the_migration_path() {
let batch = SinkBatch::new(vec![p(b"a")], DedupToken::for_payloads(&[p(b"a")]));
let records: Vec<Payload> = batch.into_records();
assert_eq!(records, vec![p(b"a")]);
}
#[test]
fn sink_batch_into_parts_yields_both() {
let token = DedupToken::for_payloads(&[p(b"a")]);
let (records, t) = SinkBatch::new(vec![p(b"a")], token).into_parts();
assert_eq!(records, vec![p(b"a")]);
assert_eq!(t.to_hex(), token.to_hex());
}
#[test]
fn sink_batch_empty_is_empty() {
let batch = SinkBatch::new(vec![], DedupToken::for_payloads(&[] as &[Payload]));
assert!(batch.is_empty());
assert_eq!(batch.len(), 0);
}
#[test]
fn segment_created_at_is_none_from_the_plain_constructors() {
let token = DedupToken::for_payloads(&[p(b"a")]);
assert_eq!(
SinkBatch::new(vec![p(b"a")], token).segment_created_at(),
None
);
assert_eq!(SinkBatch::from(vec![p(b"a")]).segment_created_at(), None);
assert_eq!(
SinkBatch::with_record_ids(
vec![p(b"a")],
token,
vec![RecordId::for_record(
"shard_00/seg_00000001.wab",
1,
&p(b"a")
)],
)
.segment_created_at(),
None
);
}
#[test]
fn segment_created_at_round_trips_through_with_segment_context() {
let payload = p(b"a");
let batch = SinkBatch::with_segment_context(
vec![payload.clone()],
DedupToken::for_payloads(std::slice::from_ref(&payload)),
vec![RecordId::for_record(
"shard_00/seg_00000001.wab",
1,
&payload,
)],
1_788_704_730_000_000_000,
);
assert_eq!(batch.segment_created_at(), Some(1_788_704_730_000_000_000));
assert_eq!(batch.record_ids().map(<[RecordId]>::len), Some(1));
}
#[test]
#[should_panic(expected = "with_segment_context: record_ids must be parallel")]
fn with_segment_context_keeps_the_length_mismatch_assertion() {
let a = p(b"a");
let _ = SinkBatch::with_segment_context(
vec![a.clone(), p(b"b")],
DedupToken::for_payloads(std::slice::from_ref(&a)),
vec![RecordId::for_record("s", 1, &a)],
0,
);
}
}