use std::collections::{HashMap, HashSet};
use crate::error::{DbError, DbResult};
use crate::txn::participant::ClaimKind;
use crate::unique::resolve::{ClaimDecision, decide};
use super::expected::{self, SkippedSource};
use super::registry::{BuiltRegistry, Decl, unique_owner_pair};
use super::repair::RepairOptions;
use super::validate::hex;
fn retain_owner_source_dirty(
reg: &BuiltRegistry,
repaired_sources: &mut HashSet<&'static str>,
owner: &crate::unique::UniqueOwner,
) {
if let Some(entry) = reg
.reg
.entries
.iter()
.find(|entry| crate::unique::collection_tag(entry.name) == owner.collection_id())
{
repaired_sources.remove(entry.name);
}
}
#[cfg(any(test, feature = "test-faults"))]
std::thread_local! {
static FAIL_NEXT_NAMESPACE_REPAIR: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
#[cfg(any(test, feature = "test-faults"))]
pub(crate) fn fail_next_namespace_repair() {
FAIL_NEXT_NAMESPACE_REPAIR.set(true);
}
#[cfg(any(test, feature = "test-faults"))]
fn take_fail_next_namespace_repair() -> bool {
FAIL_NEXT_NAMESPACE_REPAIR.replace(false)
}
#[derive(Debug, Clone, Default)]
pub struct NamespaceRepairStats {
pub namespace: &'static str,
pub written: u64,
pub deleted: u64,
pub declined_orphans: u64,
pub conflicts: Vec<NamespaceConflict>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamespaceConflict {
pub key: String,
pub owners: Vec<(String, String)>,
}
pub(super) fn repair_namespaces(
reg: &BuiltRegistry,
opts: &RepairOptions,
) -> DbResult<(
Vec<NamespaceRepairStats>,
Vec<SkippedSource>,
Vec<&'static str>,
)> {
let (unique_expected, skipped) = expected::build_unique(reg)?;
let mut protected: HashSet<&'static str> = skipped
.iter()
.flat_map(|s| s.targets.iter().copied())
.collect();
for d in ®.reg.decls {
if let Decl::Unique {
source, targets, ..
} = d
&& reg.reg.entries.iter().all(|e| e.name != *source)
{
for t in targets.all() {
protected.insert(t.collection_name());
}
}
}
let handles = expected::declared_unique_handles(reg);
let mut repaired_sources: HashSet<&'static str> =
reg.reg.entries.iter().map(|entry| entry.name).collect();
for d in ®.reg.decls {
let Decl::Unique {
source, targets, ..
} = d
else {
continue;
};
let complete = targets.all().all(|target| {
let name = target.collection_name();
handles.contains_key(name)
&& reg.reg.external.iter().any(|(registered, kind)| {
*registered == name && *kind == armour_core::CollectionKind::UniqueIndex
})
});
if !complete {
repaired_sources.remove(source);
}
}
let mut stored_by_ns: HashMap<&'static str, HashMap<Vec<u8>, crate::UniqueRow>> =
HashMap::new();
let mut corrupt_by_ns: HashMap<&'static str, HashSet<Vec<u8>>> = HashMap::new();
for (ns_name, handle) in &handles {
let stored = stored_by_ns.entry(ns_name).or_default();
let corrupt = corrupt_by_ns.entry(ns_name).or_default();
handle.scan_rows(&mut |key_bytes, row| {
if row.is_well_formed() {
stored.insert(key_bytes.to_vec(), row);
} else {
corrupt.insert(key_bytes.to_vec());
}
Ok(())
})?;
}
#[cfg(any(test, feature = "test-faults"))]
if take_fail_next_namespace_repair() {
return Err(DbError::Internal("injected namespace repair failure"));
}
let mut all_pairs: HashSet<(&'static str, Vec<u8>)> = unique_expected.keys().cloned().collect();
for (ns, rows) in &stored_by_ns {
for key in rows.keys() {
all_pairs.insert((*ns, key.clone()));
}
}
for (ns, keys) in &corrupt_by_ns {
for key in keys {
if let Some(claims) = unique_expected.get(&(*ns, key.clone())) {
for claim in claims {
repaired_sources.remove(claim.source);
}
}
all_pairs.remove(&(*ns, key.clone()));
}
}
let mut stats_by_ns: HashMap<&'static str, NamespaceRepairStats> = HashMap::new();
for (ns, key) in all_pairs {
let Some(handle) = handles.get(ns) else {
continue;
};
let claims = unique_expected.get(&(ns, key.clone()));
let row = stored_by_ns.get(ns).and_then(|m| m.get(&key)).copied();
match claims {
None => {
let Some(row) = row else { continue };
if row.deadline().is_some() {
continue;
}
let stats = stats_by_ns
.entry(ns)
.or_insert_with(|| NamespaceRepairStats {
namespace: ns,
..Default::default()
});
let complete = reg.complete_indices.contains(ns) && !protected.contains(ns);
if complete && opts.delete_orphans {
handle.delete_row(&key)?;
stats.deleted += 1;
} else {
retain_owner_source_dirty(reg, &mut repaired_sources, &row.owner());
stats.declined_orphans += 1;
}
}
Some(claims) if claims.len() == 1 => {
let claim = &claims[0];
let stats = stats_by_ns
.entry(ns)
.or_insert_with(|| NamespaceRepairStats {
namespace: ns,
..Default::default()
});
match decide(row, &claim.owner, ClaimKind::Claim) {
ClaimDecision::NoOp => {}
ClaimDecision::Write(new_row) => {
handle.write_row(&key, new_row)?;
stats.written += 1;
}
ClaimDecision::Conflict(cur) => {
repaired_sources.remove(claim.source);
retain_owner_source_dirty(reg, &mut repaired_sources, &cur);
let (cur_source, cur_key) = unique_owner_pair(reg, &cur);
stats.conflicts.push(NamespaceConflict {
key: hex(&key),
owners: vec![
(claim.source.to_string(), claim.source_key.clone()),
(cur_source, cur_key),
],
});
}
ClaimDecision::Delete => {
return Err(DbError::Internal(
"unique repair: decide(Claim) returned Delete — unreachable",
));
}
}
}
Some(claims) => {
let stats = stats_by_ns
.entry(ns)
.or_insert_with(|| NamespaceRepairStats {
namespace: ns,
..Default::default()
});
let mut owners: Vec<(String, String)> = claims
.iter()
.map(|c| (c.source.to_string(), c.source_key.clone()))
.collect();
for claim in claims {
repaired_sources.remove(claim.source);
}
if let Some(row) = row {
retain_owner_source_dirty(reg, &mut repaired_sources, &row.owner());
}
owners.sort();
stats.conflicts.push(NamespaceConflict {
key: hex(&key),
owners,
});
}
}
}
let mut per_namespace: Vec<NamespaceRepairStats> = stats_by_ns.into_values().collect();
per_namespace.sort_by_key(|s| s.namespace);
for s in &mut per_namespace {
s.conflicts.sort_by(|a, b| a.key.cmp(&b.key));
}
for source in &skipped {
repaired_sources.remove(source.name);
}
let mut repaired_sources: Vec<_> = repaired_sources.into_iter().collect();
repaired_sources.sort_unstable();
Ok((per_namespace, skipped, repaired_sources))
}
#[cfg(test)]
mod tests {
use super::{fail_next_namespace_repair, take_fail_next_namespace_repair};
#[test]
fn namespace_repair_failpoint_is_test_thread_isolated() {
fail_next_namespace_repair();
let other_thread_consumed = std::thread::spawn(take_fail_next_namespace_repair)
.join()
.expect("join failpoint observer");
assert!(
!other_thread_consumed,
"another test thread consumed the failpoint"
);
assert!(
take_fail_next_namespace_repair(),
"the arming thread lost its failpoint"
);
}
}