crates/mimir-ingest/tests/ingest/ring_split.rs
//! The dependency ring in its own tables, the repo tables partitioned by
//! repo: a new lake is born that way; a lake written before is migrated by
//! `EventLog::split_ring` (copy + verify + switch, the old rows kept) and
//! brought to the partitioned shape by the next compaction (which drops the
//! ring rows from the repo tables, verified by count). A filtered scan of
//! one repo then opens that repo's files only.
use std::path::Path;
use mimir_ingest::{CompactOptions, Consumer, EventLog, FoldScope, Kind, RepoSpec, SplitOptions};
use skade::ScanFilter;
use skade::arrow_array::RecordBatch;
use crate::common::{FakeDumper, Scratch, open_log, specs};
fn repo_a(src: &Scratch) {
src.write("crates/a/src/lib.rs", "fn f (x: i32) | one\nfn g () | two\ncall g -> f\n");
src.commit("a one");
src.write("crates/a/src/lib.rs", "fn f (x: i32) | uno\nfn g () | two\ncall g -> f\ncall f -> g\n");
src.commit("a two");
}
fn repo_b(src: &Scratch) {
src.write("crates/b/src/lib.rs", "fn h () | three\ncall h -> h\n");
src.commit("b one");
}
fn spec(tmp: &Path, name: &str, src: &Scratch) -> RepoSpec {
RepoSpec {
name: name.into(),
source: src.root.to_string_lossy().into_owned(),
branch: "main".into(),
worktree: tmp.join(format!("clone-{name}")),
initial_offset: None,
}
}
fn rows_of(b: &RecordBatch) -> Vec<Vec<u8>> {
use arrow_row::{RowConverter, SortField};
let fields: Vec<SortField> = b.schema().fields().iter().map(|f| SortField::new(f.data_type().clone())).collect();
let conv = RowConverter::new(fields).unwrap();
let rows = conv.convert_columns(b.columns()).unwrap();
(0..b.num_rows()).map(|i| rows.row(i).as_ref().to_vec()).collect()
}
/// Every event row of the source stream, repo tables AND ring tables, sorted.
async fn all_source_rows(log: &EventLog) -> Vec<Vec<u8>> {
let mut out = Vec::new();
for kind in Kind::SOURCE {
for b in log.events(kind).await.unwrap() {
out.extend(rows_of(&b));
}
}
out.sort();
out
}
fn repo_of(b: &RecordBatch, i: usize) -> String {
use skade::arrow_array::{Array, StringArray};
b.column_by_name("repo").unwrap().as_any().downcast_ref::<StringArray>().unwrap().value(i).to_string()
}
/// (rows with repo == @universe, rows with another repo) of a table.
async fn ring_and_repo_rows(log: &EventLog, table: &str) -> (usize, usize) {
let t = log.warehouse().table(&format!("mimir_v2.{table}")).await.unwrap();
let (mut ring, mut repo) = (0, 0);
for b in skade::read_all(t.inner()).await.unwrap() {
for i in 0..b.num_rows() {
if repo_of(&b, i) == "@universe" { ring += 1 } else { repo += 1 }
}
}
(ring, repo)
}
async fn plan_files(log: &EventLog, table: &str, repo: Option<&str>) -> usize {
let t = log.warehouse().table(&format!("mimir_v2.{table}")).await.unwrap();
let f = repo.map(|r| ScanFilter::eq("repo", r));
skade::plan_stats(t.inner(), f.as_ref()).await.unwrap().data_files as usize
}
async fn partitioned_by_repo(log: &EventLog, table: &str) -> bool {
let t = log.warehouse().table(&format!("mimir_v2.{table}")).await.unwrap();
let md = t.inner().metadata();
let names: Vec<&str> = md.default_partition_spec().fields().iter().filter_map(|f| md.current_schema().name_by_field_id(f.source_id)).collect();
names == ["repo"]
}
async fn consume_two(tmp: &Path, log: EventLog) -> EventLog {
let a = Scratch::new(tmp.join("src-a"));
repo_a(&a);
let b = Scratch::new(tmp.join("src-b"));
repo_b(&b);
let mut c = Consumer::boot(log, FakeDumper::default()).await.unwrap();
c.run(&spec(tmp, "ra", &a)).await.unwrap();
c.run(&spec(tmp, "rb", &b)).await.unwrap();
let (log, _) = c.into_parts();
log
}
fn ring_ids(rows: &mimir_ingest::FoldedRows) -> Vec<String> {
use skade::arrow_array::{Array, StringArray};
let mut v: Vec<String> = rows
.open
.iter()
.map(|&(b, r)| {
let batch = &rows.batches[b as usize];
let id = batch.column_by_name("moniker").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
format!("{}\u{1f}{}", repo_of(batch, r as usize), id.value(r as usize))
})
.filter(|s| s.starts_with("@universe"))
.collect();
v.sort();
v
}
#[tokio::test(flavor = "current_thread")]
async fn a_new_lake_keeps_the_ring_in_its_own_tables_and_partitions_by_repo() {
let tmp = tempfile::tempdir().unwrap();
let log = open_log(&tmp.path().join("lake")).await;
assert!(log.ring_split(), "a new lake is born split");
let log = consume_two(tmp.path(), log).await;
// the ring (dep@1.0.0) is in the ring tables only
for (kind, table) in [(Kind::Nodes, "code_node_events"), (Kind::Edges, "code_edge_events"), (Kind::Spans, "code_span_events")] {
let (ring, repo) = ring_and_repo_rows(&log, table).await;
assert_eq!(ring, 0, "{table} holds no ring rows");
assert!(repo > 0, "{table} holds the repos' rows");
let (ring, repo) = ring_and_repo_rows(&log, EventLog::ring_table_name(kind).unwrap()).await;
assert_eq!(repo, 0, "the ring table of {kind:?} holds ring rows only");
assert!(ring > 0 || kind != Kind::Nodes, "the ring table of {kind:?} holds the dep's rows");
assert!(partitioned_by_repo(&log, table).await, "{table} is partitioned by repo");
}
// the fold sees the ring when asked (the dep's helper is a universe node)
let ring = log.fold_rows(Kind::Nodes, &FoldScope { repos: Default::default(), universe: true }, &[]).await.unwrap();
assert_eq!(ring_ids(&ring), vec!["@universe\u{1f}dep@1.0.0 helper".to_string()]);
// one repo's scan opens that repo's files only
let all = plan_files(&log, "code_node_events", None).await;
let ra = plan_files(&log, "code_node_events", Some("ra")).await;
let rb = plan_files(&log, "code_node_events", Some("rb")).await;
assert!(ra >= 2 && rb >= 1 && ra + rb == all, "all {all}, ra {ra}, rb {rb}");
// the repo's state is whole
let ra_fold = log.fold(Some("ra"), None).await.unwrap();
assert_eq!(ra_fold.repo(Kind::Edges, "ra").unwrap().len(), 2);
}
#[tokio::test(flavor = "current_thread")]
async fn a_legacy_lake_is_split_then_partitioned_by_the_compaction() {
let tmp = tempfile::tempdir().unwrap();
let lake = tmp.path().join("lake");
let wh = skade::open(&lake).await.unwrap();
let log = EventLog::open_legacy_for_test(wh, specs()).await.unwrap();
assert!(!log.ring_split());
let log = consume_two(tmp.path(), log).await;
// the legacy shape: ring rows inside the repo tables, no partition
let (ring, _) = ring_and_repo_rows(&log, "code_node_events").await;
assert!(ring > 0);
assert!(!partitioned_by_repo(&log, "code_node_events").await);
let rows_before = all_source_rows(&log).await;
let fold_before = log.fold(None, None).await.unwrap();
let ring_before = ring_ids(&log.fold_rows(Kind::Nodes, &FoldScope { repos: Default::default(), universe: true }, &[]).await.unwrap());
assert!(!ring_before.is_empty());
let stamp_before = log.stamp().await.unwrap();
// a crash after the copy: nothing switched, the next run completes
let crash = SplitOptions { fail_after_copy: true, ..Default::default() };
let err = log.split_ring(&crash).await.unwrap_err();
assert!(format!("{err:#}").contains("fail_after_copy"), "{err:#}");
assert!(!log.ring_split(), "not switched before the verify");
// the next run finds the copy complete, verifies and switches
let report = log.split_ring(&SplitOptions::default()).await.unwrap();
assert!(log.ring_split(), "{report}");
assert!(report.copied == 0 && report.verified > 0 && report.switched, "{report}");
// reads: the same ring, the same fold; the old rows are still in the repo tables
assert_eq!(ring_ids(&log.fold_rows(Kind::Nodes, &FoldScope { repos: Default::default(), universe: true }, &[]).await.unwrap()), ring_before);
assert_eq!(log.fold(None, None).await.unwrap(), fold_before);
assert_eq!(log.stamp().await.unwrap(), stamp_before);
let (still, _) = ring_and_repo_rows(&log, "code_node_events").await;
assert_eq!(still, ring, "kept until the compaction drops them");
// again: nothing to do
let again = log.split_ring(&SplitOptions::default()).await.unwrap();
assert_eq!(again.copied, 0, "{again}");
// the compaction partitions the repo tables and drops their ring rows
let opts = CompactOptions { grace: std::time::Duration::ZERO, ..Default::default() };
log.compact(&opts).await.unwrap();
for table in ["code_node_events", "code_edge_events", "code_span_events"] {
let (ring, repo) = ring_and_repo_rows(&log, table).await;
assert_eq!(ring, 0, "{table}");
assert!(repo > 0, "{table}");
assert!(partitioned_by_repo(&log, table).await, "{table}");
}
assert_eq!(all_source_rows(&log).await, rows_before, "every row is still there, in one of the two tables");
assert_eq!(log.fold(None, None).await.unwrap(), fold_before);
assert_eq!(ring_ids(&log.fold_rows(Kind::Nodes, &FoldScope { repos: Default::default(), universe: true }, &[]).await.unwrap()), ring_before);
let all = plan_files(&log, "code_edge_events", None).await;
let ra = plan_files(&log, "code_edge_events", Some("ra")).await;
assert!(ra < all && ra >= 1, "all {all}, ra {ra}");
// the log goes on: a release of ra lands in ra's partition, ring rows in the ring tables
let a = Scratch { root: tmp.path().join("src-a") };
a.write("crates/a/src/lib.rs", "fn f (x: i64) | dos\nfn g () | two\ncall f -> g\n");
a.commit("a three");
let mut c = Consumer::boot(log, FakeDumper::default()).await.unwrap();
c.run(&spec(tmp.path(), "ra", &a)).await.unwrap();
let (log, _) = c.into_parts();
let (ring, _) = ring_and_repo_rows(&log, "code_node_events").await;
assert_eq!(ring, 0);
assert_eq!(log.fold(Some("ra"), None).await.unwrap().repo(Kind::Edges, "ra").unwrap().len(), 1);
}