use std::collections::HashMap;
use nodedb_types::{CollectionKey, DatabaseId};
use crate::record::{RecordType, WalRecord, WriteAbortedPayload};
use crate::replay::aborted::AbortedWrites;
use crate::tombstone::CollectionTombstonePayload;
#[derive(Debug, Clone)]
struct Tombstone {
bare_name: String,
purge_lsn: u64,
}
#[derive(Debug, Default, Clone)]
pub struct TombstoneSet {
entries: HashMap<(u64, u64, String), Tombstone>,
}
impl TombstoneSet {
pub fn new() -> Self {
Self::default()
}
pub fn for_database(&self, database_id: u64) -> DatabaseTombstones<'_> {
DatabaseTombstones {
set: self,
database_id,
}
}
pub fn insert(&mut self, key: CollectionKey<'_>, tenant_id: u64, purge_lsn: u64) {
let storage = key.qualified();
self.merge(
(
key.database_id().as_u64(),
tenant_id,
storage.as_str().to_string(),
),
Tombstone {
bare_name: key.name().to_string(),
purge_lsn,
},
);
}
fn merge(&mut self, slot: (u64, u64, String), tombstone: Tombstone) {
self.entries
.entry(slot)
.and_modify(|existing| {
if tombstone.purge_lsn > existing.purge_lsn {
existing.purge_lsn = tombstone.purge_lsn;
}
})
.or_insert(tombstone);
}
pub fn is_tombstoned(
&self,
database_id: u64,
tenant_id: u64,
storage_name: &str,
lsn: u64,
) -> bool {
self.entries
.get(&(database_id, tenant_id, storage_name.to_string()))
.is_some_and(|tombstone| lsn < tombstone.purge_lsn)
}
pub fn is_key_tombstoned(&self, key: CollectionKey<'_>, tenant_id: u64, lsn: u64) -> bool {
self.is_tombstoned(
key.database_id().as_u64(),
tenant_id,
key.qualified().as_str(),
lsn,
)
}
pub fn purge_lsn(&self, key: CollectionKey<'_>, tenant_id: u64) -> Option<u64> {
let storage = key.qualified();
self.entries
.get(&(
key.database_id().as_u64(),
tenant_id,
storage.as_str().to_string(),
))
.map(|tombstone| tombstone.purge_lsn)
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (u64, u64, &str, u64)> + '_ {
self.entries.iter().map(|((db, tid, _storage), tombstone)| {
(*db, *tid, tombstone.bare_name.as_str(), tombstone.purge_lsn)
})
}
pub fn extend(&mut self, other: TombstoneSet) {
for (slot, tombstone) in other.entries {
self.merge(slot, tombstone);
}
}
}
pub struct DatabaseTombstones<'a> {
set: &'a TombstoneSet,
database_id: u64,
}
impl DatabaseTombstones<'_> {
pub fn is_tombstoned(&self, tenant_id: u64, storage_name: &str, lsn: u64) -> bool {
self.set
.is_tombstoned(self.database_id, tenant_id, storage_name, lsn)
}
}
impl std::ops::Deref for DatabaseTombstones<'_> {
type Target = TombstoneSet;
fn deref(&self) -> &Self::Target {
self.set
}
}
pub fn extract_tombstones(records: &[WalRecord]) -> crate::Result<TombstoneSet> {
Ok(extract_replay_filters(records)?.tombstones)
}
#[derive(Debug, Default, Clone)]
pub struct ReplayFilters {
pub tombstones: TombstoneSet,
pub aborted: AbortedWrites,
}
pub fn extract_replay_filters(records: &[WalRecord]) -> crate::Result<ReplayFilters> {
let mut filters = ReplayFilters::default();
for record in records {
let Some(kind) = RecordType::from_raw(record.logical_record_type()) else {
continue;
};
match kind {
RecordType::CollectionTombstoned => {
reject_if_encrypted(record, "collection-tombstone extraction")?;
let payload = CollectionTombstonePayload::from_bytes(&record.payload)?;
filters.tombstones.insert(
CollectionKey::from_bare(
DatabaseId::new(record.header.database_id),
&payload.collection,
),
record.header.tenant_id,
payload.purge_lsn,
);
}
RecordType::WriteAborted => {
reject_if_encrypted(record, "write-abort extraction")?;
let payload = WriteAbortedPayload::from_bytes(&record.payload)?;
filters.aborted.insert(payload.aborted_lsn);
}
_ => continue,
}
}
Ok(filters)
}
fn reject_if_encrypted(record: &WalRecord, site: &'static str) -> crate::Result<()> {
if !record.is_encrypted() {
return Ok(());
}
let err = crate::WalError::EncryptedRecordWithoutKey {
lsn: record.header.lsn,
context: site,
};
crate::diag::encrypted_record_without_key(&err, record.header.lsn, site);
Err(err)
}
pub fn drop_aborted_records(records: Vec<WalRecord>, aborted: &AbortedWrites) -> Vec<WalRecord> {
if aborted.is_empty() {
return records;
}
let mut kept = records;
kept.retain(|record| !aborted.contains(record.header.lsn));
kept
}
#[cfg(test)]
mod tests {
use nodedb_types::QualifiedCollection;
use super::*;
use crate::record::{WalRecord, WalRecordArgs};
fn tombstone_record(
database: u64,
tenant: u64,
name: &str,
purge_lsn: u64,
record_lsn: u64,
) -> WalRecord {
let payload = CollectionTombstonePayload::new(name, purge_lsn)
.to_bytes()
.unwrap();
WalRecord::new(WalRecordArgs {
record_type: RecordType::CollectionTombstoned as u32,
lsn: record_lsn,
tenant_id: tenant,
vshard_id: 0,
database_id: database,
payload,
encryption_key: None,
preamble_bytes: None,
})
.unwrap()
}
fn key(database: u64, name: &str) -> CollectionKey<'_> {
CollectionKey::from_bare(DatabaseId::new(database), name)
}
#[test]
fn is_tombstoned_shadows_older_writes() {
let mut set = TombstoneSet::new();
set.insert(key(7, "users"), 1, 100);
assert!(set.is_tombstoned(7, 1, "7/users", 50));
assert!(!set.is_tombstoned(7, 1, "7/users", 100));
assert!(!set.is_tombstoned(7, 1, "7/users", 200));
assert!(!set.is_tombstoned(7, 1, "7/other", 50));
assert!(!set.is_tombstoned(7, 2, "7/users", 50));
assert!(!set.is_tombstoned(8, 1, "8/users", 50));
}
#[test]
fn bare_tombstone_shadows_qualified_record_in_named_database() {
let db = DatabaseId::new(1024);
let mut set = TombstoneSet::new();
set.insert(CollectionKey::from_bare(db, "orders"), 1, 100);
let storage = QualifiedCollection::new(db, "orders");
assert_eq!(storage.as_str(), "1024/orders");
assert!(set.is_tombstoned(1024, 1, storage.as_str(), 99));
assert!(
set.for_database(1024)
.is_tombstoned(1, storage.as_str(), 99)
);
assert!(!set.is_tombstoned(1024, 1, storage.as_str(), 100));
assert!(!set.is_tombstoned(1024, 1, "orders", 99));
assert!(set.is_key_tombstoned(CollectionKey::from_bare(db, "orders"), 1, 99));
}
#[test]
fn default_database_tombstone_shadows_bare_record() {
let mut set = TombstoneSet::new();
set.insert(
CollectionKey::from_bare(DatabaseId::DEFAULT, "orders"),
1,
100,
);
assert!(set.is_tombstoned(0, 1, "orders", 99));
assert!(set.for_database(0).is_tombstoned(1, "orders", 99));
assert!(!set.is_tombstoned(0, 1, "orders", 100));
assert!(set.is_key_tombstoned(
CollectionKey::from_bare(DatabaseId::DEFAULT, "orders"),
1,
99
));
}
#[test]
fn extracted_tombstone_shadows_qualified_records() {
let set = extract_tombstones(&[tombstone_record(1024, 1, "orders", 100, 101)]).unwrap();
assert!(set.is_tombstoned(1024, 1, "1024/orders", 50));
assert!(!set.is_tombstoned(1024, 1, "orders", 50));
}
#[test]
fn iter_yields_bare_names() {
let mut set = TombstoneSet::new();
set.insert(key(1024, "orders"), 1, 100);
let rows: Vec<_> = set.iter().collect();
assert_eq!(rows, vec![(1024, 1, "orders", 100)]);
}
#[test]
fn insert_keeps_highest_purge_lsn() {
let mut set = TombstoneSet::new();
set.insert(key(7, "users"), 1, 100);
set.insert(key(7, "users"), 1, 50);
assert_eq!(set.purge_lsn(key(7, "users"), 1), Some(100));
set.insert(key(7, "users"), 1, 200);
assert_eq!(set.purge_lsn(key(7, "users"), 1), Some(200));
}
#[test]
fn extract_from_record_stream() {
let records = vec![
tombstone_record(7, 1, "users", 100, 10),
tombstone_record(7, 1, "orders", 150, 11),
tombstone_record(8, 1, "users", 200, 12),
];
let set = extract_tombstones(&records).unwrap();
assert_eq!(set.len(), 3);
assert_eq!(set.purge_lsn(key(7, "users"), 1), Some(100));
assert_eq!(set.purge_lsn(key(7, "orders"), 1), Some(150));
assert_eq!(set.purge_lsn(key(8, "users"), 1), Some(200));
}
#[test]
fn extract_ignores_non_tombstone_records() {
let records = vec![
tombstone_record(0, 1, "users", 100, 10),
WalRecord::new(WalRecordArgs {
record_type: RecordType::Put as u32,
lsn: 11,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: b"junk".to_vec(),
encryption_key: None,
preamble_bytes: None,
})
.unwrap(),
WalRecord::new(WalRecordArgs {
record_type: RecordType::Noop as u32,
lsn: 12,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: vec![],
encryption_key: None,
preamble_bytes: None,
})
.unwrap(),
];
let set = extract_tombstones(&records).unwrap();
assert_eq!(set.len(), 1);
}
#[test]
fn extract_rejects_corrupt_payload() {
let bogus = WalRecord::new(WalRecordArgs {
record_type: RecordType::CollectionTombstoned as u32,
lsn: 5,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: vec![0xFF, 0xFF, 0xFF], encryption_key: None,
preamble_bytes: None,
})
.unwrap();
let records = vec![bogus, tombstone_record(0, 1, "users", 100, 10)];
assert!(
extract_tombstones(&records).is_err(),
"malformed tombstone must fail replay rather than lose a purge boundary"
);
}
fn abort_record(aborted_lsn: u64, record_lsn: u64) -> WalRecord {
WalRecord::new(WalRecordArgs {
record_type: RecordType::WriteAborted as u32,
lsn: record_lsn,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: WriteAbortedPayload::new(aborted_lsn).to_bytes().to_vec(),
encryption_key: None,
preamble_bytes: None,
})
.unwrap()
}
fn put_record(lsn: u64) -> WalRecord {
WalRecord::new(WalRecordArgs {
record_type: RecordType::Put as u32,
lsn,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: b"row".to_vec(),
encryption_key: None,
preamble_bytes: None,
})
.unwrap()
}
#[test]
fn one_pass_extracts_tombstones_and_aborted_lsns() {
let records = vec![
put_record(10),
tombstone_record(0, 1, "users", 100, 11),
put_record(12),
abort_record(10, 13),
abort_record(12, 14),
];
let filters = extract_replay_filters(&records).unwrap();
assert_eq!(filters.tombstones.purge_lsn(key(0, "users"), 1), Some(100));
assert_eq!(filters.aborted.len(), 2);
assert!(filters.aborted.contains(10));
assert!(filters.aborted.contains(12));
assert!(!filters.aborted.contains(11));
}
#[test]
fn extract_rejects_corrupt_abort_payload() {
let bogus = WalRecord::new(WalRecordArgs {
record_type: RecordType::WriteAborted as u32,
lsn: 5,
tenant_id: 1,
vshard_id: 0,
database_id: 0,
payload: vec![0x01, 0x02],
encryption_key: None,
preamble_bytes: None,
})
.unwrap();
assert!(extract_replay_filters(&[bogus]).is_err());
}
#[test]
fn gate_drops_exactly_the_aborted_records() {
let records = vec![
put_record(10),
put_record(11),
put_record(12),
abort_record(11, 13),
];
let filters = extract_replay_filters(&records).unwrap();
let kept = drop_aborted_records(records, &filters.aborted);
let lsns: Vec<u64> = kept.iter().map(|r| r.header.lsn).collect();
assert_eq!(lsns, vec![10, 12, 13]);
}
#[test]
fn gate_is_a_no_op_without_aborts() {
let records = vec![put_record(1), put_record(2)];
let kept = drop_aborted_records(records, &AbortedWrites::new());
assert_eq!(kept.len(), 2);
}
#[test]
fn extend_merges_sets() {
let mut a = TombstoneSet::new();
a.insert(key(7, "users"), 1, 100);
let mut b = TombstoneSet::new();
b.insert(key(7, "users"), 1, 150);
b.insert(key(8, "orders"), 1, 200);
a.extend(b);
assert_eq!(a.purge_lsn(key(7, "users"), 1), Some(150));
assert_eq!(a.purge_lsn(key(8, "orders"), 1), Some(200));
}
}