use std::fmt::Write as _;
use arrow_array::RecordBatch;
use arrow_schema::Schema;
use async_trait::async_trait;
use crate::{
error::{BifrostError, Result},
state::SyncState,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SnapshotWindow {
pub from: Option<i64>,
pub to: i64,
}
impl SnapshotWindow {
pub const fn new(from: Option<i64>, to: i64) -> Self {
Self { from, to }
}
}
#[derive(Debug, Clone, Copy)]
pub struct TableWindow<'a> {
pub table: &'a str,
pub snapshots: SnapshotWindow,
pub upserts: &'a [RecordBatch],
pub deletes: &'a [RecordBatch],
pub needs_full_reload: bool,
pub sequence: Option<i64>,
pub key_columns: &'a [String],
}
impl TableWindow<'_> {
pub fn upsert_rows(&self) -> usize {
self.upserts.iter().map(RecordBatch::num_rows).sum()
}
pub fn delete_rows(&self) -> usize {
self.deletes.iter().map(RecordBatch::num_rows).sum()
}
pub fn is_empty(&self) -> bool {
self.upsert_rows() == 0 && self.delete_rows() == 0
}
pub fn schema(&self) -> Option<arrow_schema::SchemaRef> {
self.upserts
.first()
.or_else(|| self.deletes.first())
.map(RecordBatch::schema)
}
}
#[derive(Debug)]
pub struct AckToken(());
impl AckToken {
#[must_use]
pub fn ack(self, to_snapshot: i64, labels: Vec<String>) -> LoadAck {
LoadAck {
to_snapshot,
labels,
rows_loaded: 0,
rows_deleted: 0,
rows_rejected_duplicate: 0,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LoadAck {
to_snapshot: i64,
pub labels: Vec<String>,
pub rows_loaded: usize,
pub rows_deleted: usize,
pub rows_rejected_duplicate: usize,
}
impl LoadAck {
pub const fn to_snapshot(&self) -> i64 {
self.to_snapshot
}
#[must_use]
pub fn with_counts(mut self, loaded: usize, deleted: usize, duplicate: usize) -> Self {
self.rows_loaded = loaded;
self.rows_deleted = deleted;
self.rows_rejected_duplicate = duplicate;
self
}
}
#[async_trait]
pub trait TableSink: Send {
async fn probe(&mut self, schema: &Schema) -> Result<()>;
async fn apply(&mut self, window: &TableWindow<'_>, token: AckToken) -> Result<LoadAck>;
}
pub async fn deliver_window<S: TableSink + ?Sized>(
sink: &mut S,
state: &mut SyncState,
window: &TableWindow<'_>,
) -> Result<LoadAck> {
refuse_full_reload(window.table, window.needs_full_reload)?;
let ack = sink.apply(window, AckToken(())).await?;
if ack.to_snapshot != window.snapshots.to {
return Err(BifrostError::Other(format!(
"{}: sink acked snapshot {} but the window ended at {} — refusing to \
move the cursor",
window.table, ack.to_snapshot, window.snapshots.to,
)));
}
let mut next = state.get(window.table);
next.snapshot_id = Some(ack.to_snapshot);
state.advance_staged(window.table, next).await?;
state.flush()?;
tracing::info!(
target: "bifrost::sink",
event = "table_window_delivered",
table = %window.table,
from = ?window.snapshots.from,
to = ack.to_snapshot,
labels = ack.labels.len(),
rows = ack.rows_loaded,
deleted = ack.rows_deleted,
duplicate = ack.rows_rejected_duplicate,
"window loaded, then cursor advanced",
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::cursor_after_ack",
true,
&format!(
"{} advanced to {} after {} label(s)",
window.table,
ack.to_snapshot,
ack.labels.len()
),
);
Ok(ack)
}
pub fn refuse_full_reload(table: &str, needs_full_reload: bool) -> Result<()> {
if !needs_full_reload {
return Ok(());
}
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::full_reload_refused",
false,
&format!("{table}: window needs a full reload — refused, cursor held"),
);
tracing::error!(
target: "bifrost::sink",
event = "full_reload_refused",
table = %table,
"window needs a full reload — refusing (cursor not advanced)",
);
Err(BifrostError::Other(format!(
"{table}: this window needs a FULL RELOAD (an overwrite/replace commit or \
a position-delete file defeats the incremental delta). Refusing: a silent \
fallback would rewrite the whole table. Run an explicit one-shot reload, \
then resume."
)))
}
pub const MAX_LABEL_LEN: usize = 128;
pub fn stream_load_label(
db: &str,
table: &str,
from: Option<i64>,
to: i64,
chunk: usize,
) -> String {
let mut tail = String::new();
match from {
Some(f) => write!(tail, "-{f}-{to}-{chunk}").expect("string write"),
None => write!(tail, "-init-{to}-{chunk}").expect("string write"),
}
let full_name = format!("{}:{}", label_component(db), label_component(table));
let budget = MAX_LABEL_LEN.saturating_sub("skade-".len() + tail.len());
let name = if full_name.len() <= budget {
full_name
} else {
let digest = format!("_{:016x}", fnv1a64(full_name.as_bytes()));
let keep = budget.saturating_sub(digest.len());
let mut short: String = full_name.chars().take(keep).collect();
short.push_str(&digest);
short
};
format!("skade-{name}{tail}")
}
fn label_component(s: &str) -> String {
if s.is_empty() {
return "_e".into();
}
let mut changed = false;
let sanitized: String = s
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '_' {
c
} else {
changed = true;
'_'
}
})
.collect();
if changed {
format!("{sanitized}_{:08x}", fnv1a64(s.as_bytes()) as u32)
} else {
sanitized
}
}
pub fn is_valid_doris_label(label: &str) -> bool {
!label.is_empty()
&& label.chars().count() <= MAX_LABEL_LEN
&& label.chars().all(|c| {
c == '-' || c == '_' || c == ':' || c.is_ascii_alphanumeric() || c.is_alphabetic()
})
}
fn fnv1a64(bytes: &[u8]) -> u64 {
let mut h: u64 = 0xcbf2_9ce4_8422_2325;
for b in bytes {
h ^= u64::from(*b);
h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
h
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use arrow_array::{Int64Array, StringArray};
use arrow_schema::{DataType, Field};
use std::sync::Arc as StdArc;
fn batch(ids: Vec<i64>) -> RecordBatch {
let n = ids.len();
RecordBatch::try_new(
StdArc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
StdArc::new(Int64Array::from(ids)),
StdArc::new(StringArray::from(
(0..n).map(|i| format!("n{i}")).collect::<Vec<_>>(),
)),
],
)
.unwrap()
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Event {
Applied(Option<i64>),
}
struct SpySink {
log: Arc<Mutex<Vec<Event>>>,
state_path: std::path::PathBuf,
fail: bool,
ack_snapshot: Option<i64>,
}
#[async_trait]
impl TableSink for SpySink {
async fn probe(&mut self, _schema: &Schema) -> Result<()> {
Ok(())
}
async fn apply(&mut self, window: &TableWindow<'_>, token: AckToken) -> Result<LoadAck> {
let seen = std::fs::read_to_string(&self.state_path)
.ok()
.and_then(|s| serde_json::from_str::<serde_json::Value>(&s).ok())
.and_then(|v| {
v.get(window.table)
.and_then(|t| t.get("snapshot_id"))
.and_then(serde_json::Value::as_i64)
});
self.log.lock().unwrap().push(Event::Applied(seen));
if self.fail {
return Err(BifrostError::Other("sink refused".into()));
}
let to = self.ack_snapshot.unwrap_or(window.snapshots.to);
Ok(token.ack(to, vec!["l0".into()]).with_counts(
window.upsert_rows(),
window.delete_rows(),
0,
))
}
}
fn spy(dir: &std::path::Path, fail: bool) -> (SpySink, Arc<Mutex<Vec<Event>>>) {
let log = Arc::new(Mutex::new(Vec::new()));
(
SpySink {
log: Arc::clone(&log),
state_path: dir.join("state.json"),
fail,
ack_snapshot: None,
},
log,
)
}
#[tokio::test]
async fn cursor_is_committed_strictly_after_the_load() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("state.json");
let mut state = SyncState::open(&path).unwrap();
let (mut sink, log) = spy(dir.path(), false);
let ups = [batch(vec![1, 2, 3])];
let w = TableWindow {
table: "bench.doris2g",
snapshots: SnapshotWindow::new(Some(1042), 1043),
upserts: &ups,
deletes: &[],
needs_full_reload: false,
sequence: None,
key_columns: &[],
};
let ack = deliver_window(&mut sink, &mut state, &w).await.unwrap();
assert_eq!(ack.to_snapshot(), 1043);
assert_eq!(ack.rows_loaded, 3);
assert_eq!(
log.lock().unwrap().as_slice(),
&[Event::Applied(None)],
"the load must run while the cursor is still unset"
);
assert_eq!(state.get("bench.doris2g").snapshot_id, Some(1043));
let w2 = TableWindow {
table: "bench.doris2g",
snapshots: SnapshotWindow::new(Some(1043), 1044),
upserts: &ups,
deletes: &[],
needs_full_reload: false,
sequence: None,
key_columns: &[],
};
deliver_window(&mut sink, &mut state, &w2).await.unwrap();
assert_eq!(
log.lock().unwrap()[1],
Event::Applied(Some(1043)),
"the load must run BEFORE the cursor moves to 1044"
);
assert_eq!(state.get("bench.doris2g").snapshot_id, Some(1044));
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::cursor_strictly_after_load",
state.get("bench.doris2g").snapshot_id == Some(1044),
"apply observed the old cursor both rounds",
);
}
#[tokio::test]
async fn a_failed_load_does_not_advance_the_cursor() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("state.json");
let mut state = SyncState::open(&path).unwrap();
let (mut sink, _log) = spy(dir.path(), true);
let ups = [batch(vec![1])];
let w = TableWindow {
table: "t",
snapshots: SnapshotWindow::new(Some(7), 8),
upserts: &ups,
deletes: &[],
needs_full_reload: false,
sequence: None,
key_columns: &[],
};
assert!(deliver_window(&mut sink, &mut state, &w).await.is_err());
assert_eq!(
state.get("t").snapshot_id,
None,
"a refused load must not move the cursor"
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::failed_load_holds_cursor",
state.get("t").snapshot_id.is_none(),
"cursor unchanged after a failed apply",
);
}
#[tokio::test]
async fn a_mismatched_ack_is_refused() {
let dir = tempfile::tempdir().unwrap();
let mut state = SyncState::open(dir.path().join("state.json")).unwrap();
let (mut sink, _log) = spy(dir.path(), false);
sink.ack_snapshot = Some(9999);
let ups = [batch(vec![1])];
let w = TableWindow {
table: "t",
snapshots: SnapshotWindow::new(Some(7), 8),
upserts: &ups,
deletes: &[],
needs_full_reload: false,
sequence: None,
key_columns: &[],
};
let err = deliver_window(&mut sink, &mut state, &w)
.await
.expect_err("mismatched ack must be refused");
assert!(
err.to_string().contains("refusing to move the cursor"),
"{err}"
);
assert_eq!(state.get("t").snapshot_id, None);
}
#[tokio::test]
async fn a_full_reload_window_is_refused_and_never_loaded() {
let dir = tempfile::tempdir().unwrap();
let mut state = SyncState::open(dir.path().join("state.json")).unwrap();
let (mut sink, log) = spy(dir.path(), false);
let ups = [batch(vec![1, 2])];
let w = TableWindow {
table: "bench.doris2g",
snapshots: SnapshotWindow::new(Some(1042), 1043),
upserts: &ups,
deletes: &[],
needs_full_reload: true,
sequence: None,
key_columns: &[],
};
let err = deliver_window(&mut sink, &mut state, &w)
.await
.expect_err("a full-reload window must be refused");
assert!(err.to_string().contains("FULL RELOAD"), "{err}");
assert!(
log.lock().unwrap().is_empty(),
"the sink must not be called at all"
);
assert_eq!(
state.get("bench.doris2g").snapshot_id,
None,
"cursor must not move on a refused window"
);
assert!(refuse_full_reload("t", false).is_ok());
assert!(refuse_full_reload("t", true).is_err());
}
#[test]
fn label_is_derived_and_stable() {
let a = stream_load_label("fake", "personer2g", Some(1042), 1043, 0);
let b = stream_load_label("fake", "personer2g", Some(1042), 1043, 0);
assert_eq!(a, b);
assert_eq!(a, "skade-fake:personer2g-1042-1043-0");
assert_eq!(
stream_load_label("fake", "personer2g", None, 1043, 0),
"skade-fake:personer2g-init-1043-0"
);
}
#[test]
fn labels_do_not_collide() {
let mut seen = HashSet::new();
let names = [
("fake", "personer2g"),
("fake", "personer2g-1"),
("fake-1", "personer2g"),
("fake", "personer.2g"),
("fake", "personer_2g"),
("f", "a"),
];
let mut n = 0;
for (db, t) in names {
for from in [None, Some(0), Some(1), Some(2), Some(1042)] {
for to in [1_i64, 2, 3, 1043] {
for chunk in 0..4 {
let l = stream_load_label(db, t, from, to, chunk);
assert!(
l.len() <= MAX_LABEL_LEN,
"label over {MAX_LABEL_LEN} chars: {l}"
);
assert!(
seen.insert(l.clone()),
"label collision on ({db}, {t}, {from:?}, {to}, {chunk}): {l}"
);
n += 1;
}
}
}
}
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::labels_collision_free",
seen.len() == n,
&format!("{n} distinct windows ⇒ {} distinct labels", seen.len()),
);
}
#[test]
fn every_label_satisfies_the_doris_regex() {
let nasty = [
("fake", "personer2g"),
("fake", "personer.2g"),
("my db", "tab/le"),
("db", "t\u{e4}bell"),
("db", "\u{4eba}"),
("db", &*"x".repeat(400)),
("", ""),
];
let mut n = 0;
for (db, t) in nasty {
for from in [None, Some(1_042)] {
for chunk in [0usize, 17] {
let l = stream_load_label(db, t, from, 1_043, chunk);
assert!(
is_valid_doris_label(&l),
"label rejected by Doris's regex: {l:?} (from {db:?}/{t:?})"
);
assert!(!l.contains('.'), "a dot is invalid in a Doris label: {l}");
assert!(l.len() <= MAX_LABEL_LEN, "over 128 chars: {l}");
n += 1;
}
}
}
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"table_sink::label_matches_doris_regex",
true,
&format!("{n} labels, all matching ^[-_A-Za-z0-9:\\p{{L}}]{{1,128}}$"),
);
}
#[test]
fn truncated_labels_stay_unique_and_in_bounds() {
let long_a = "a".repeat(400);
let long_b = format!("{}b", "a".repeat(399));
let la = stream_load_label("db", &long_a, Some(1), 2, 0);
let lb = stream_load_label("db", &long_b, Some(1), 2, 0);
assert!(la.len() <= MAX_LABEL_LEN && lb.len() <= MAX_LABEL_LEN);
assert_ne!(la, lb, "truncation must not reintroduce a collision");
assert!(la.ends_with("-1-2-0"), "{la}");
}
}