use std::collections::BTreeMap;
use tracing::{debug, instrument};
use crate::merkle::MerkleNode;
use crate::storage::traits::StorageError;
use super::SyncEngine;
#[derive(Debug, Clone)]
pub struct MerkleDiff {
pub divergent_paths: Vec<String>,
pub local_only: Vec<String>,
pub remote_only: Vec<String>,
}
impl SyncEngine {
#[instrument(skip(self))]
pub async fn get_merkle_root(&self) -> Result<Option<[u8; 32]>, StorageError> {
if let Some(ref sql_merkle) = self.sql_merkle {
return sql_merkle.root_hash().await;
}
if let Some(ref merkle_cache) = self.merkle_cache {
return merkle_cache.get_hash("").await;
}
Ok(None)
}
#[instrument(skip(self))]
pub async fn get_merkle_hash(&self, path: &str) -> Result<Option<[u8; 32]>, StorageError> {
if self.is_merkle_synced().await? {
if let Some(ref merkle_cache) = self.merkle_cache {
return merkle_cache.get_hash(path).await;
}
}
if let Some(ref sql_merkle) = self.sql_merkle {
return sql_merkle.get_hash(path).await;
}
Ok(None)
}
#[instrument(skip(self))]
pub async fn get_merkle_node(&self, path: &str) -> Result<Option<MerkleNode>, StorageError> {
if self.is_merkle_synced().await? {
if let Some(ref merkle_cache) = self.merkle_cache {
return merkle_cache.get_node(path).await;
}
}
if let Some(ref sql_merkle) = self.sql_merkle {
return sql_merkle.get_node(path).await;
}
Ok(None)
}
#[instrument(skip(self))]
pub async fn get_merkle_children(&self, path: &str) -> Result<BTreeMap<String, [u8; 32]>, StorageError> {
if self.is_merkle_synced().await? {
if let Some(ref merkle_cache) = self.merkle_cache {
return merkle_cache.get_children(path).await;
}
}
if let Some(ref sql_merkle) = self.sql_merkle {
return sql_merkle.get_children(path).await;
}
Ok(BTreeMap::new())
}
#[instrument(skip(self, remote_nodes), fields(remote_count = remote_nodes.len()))]
pub async fn find_divergent_paths(
&self,
remote_nodes: &[(String, [u8; 32])],
) -> Result<MerkleDiff, StorageError> {
let mut divergent_paths = Vec::new();
let remote_only = Vec::new(); let local_only = Vec::new();
for (path, remote_hash) in remote_nodes {
let local_hash = self.get_merkle_hash(path).await?;
match local_hash {
Some(local) if local == *remote_hash => {
debug!(path = %path, "Merkle path synced");
}
Some(local) => {
debug!(
path = %path,
local = %hex::encode(local),
remote = %hex::encode(remote_hash),
"Merkle path diverged"
);
divergent_paths.push(path.clone());
}
None => {
debug!(path = %path, "Path exists on remote only");
divergent_paths.push(path.clone());
}
}
}
Ok(MerkleDiff {
divergent_paths,
local_only,
remote_only,
})
}
#[instrument(skip(self))]
pub async fn is_merkle_synced(&self) -> Result<bool, StorageError> {
let sql_root = if let Some(ref sql_merkle) = self.sql_merkle {
sql_merkle.root_hash().await?
} else {
return Ok(false);
};
let redis_root = if let Some(ref merkle_cache) = self.merkle_cache {
merkle_cache.get_hash("").await?
} else {
return Ok(false);
};
Ok(sql_root == redis_root)
}
#[instrument(skip(self, remote_children))]
pub async fn drill_down_divergence<F, Fut>(
&self,
start_path: &str,
remote_children: F,
) -> Result<Vec<String>, StorageError>
where
F: Fn(String) -> Fut,
Fut: std::future::Future<Output = Result<BTreeMap<String, [u8; 32]>, StorageError>>,
{
let mut divergent_leaves = Vec::new();
let mut paths_to_check = vec![start_path.to_string()];
while let Some(path) = paths_to_check.pop() {
let local_children = self.get_merkle_children(&path).await?;
let remote = remote_children(path.clone()).await?;
for (segment, remote_hash) in &remote {
let child_path = if path.is_empty() {
segment.clone()
} else {
format!("{}.{}", path, segment)
};
match local_children.get(segment) {
Some(local_hash) if local_hash == remote_hash => {
}
Some(_) => {
let node = self.get_merkle_node(&child_path).await?;
if node.map(|n| n.is_leaf).unwrap_or(true) {
divergent_leaves.push(child_path);
} else {
paths_to_check.push(child_path);
}
}
None => {
divergent_leaves.push(child_path);
}
}
}
for segment in local_children.keys() {
if !remote.contains_key(segment) {
let child_path = if path.is_empty() {
segment.clone()
} else {
format!("{}.{}", path, segment)
};
divergent_leaves.push(child_path);
}
}
}
Ok(divergent_leaves)
}
#[instrument(skip(self))]
pub async fn branch_dirty_count(&self, prefix: &str) -> Result<u64, StorageError> {
if let Some(ref sql_store) = self.sql_store {
return sql_store.branch_dirty_count(prefix).await;
}
Ok(0)
}
#[instrument(skip(self))]
pub async fn get_dirty_prefixes(&self) -> Result<Vec<String>, StorageError> {
if let Some(ref sql_store) = self.sql_store {
return sql_store.get_dirty_prefixes().await;
}
Ok(Vec::new())
}
#[instrument(skip(self))]
pub async fn get_top_level_branches(&self) -> Result<BTreeMap<String, [u8; 32]>, StorageError> {
self.get_merkle_children("").await
}
#[instrument(skip(self))]
pub async fn get_clean_branches(&self) -> Result<BTreeMap<String, [u8; 32]>, StorageError> {
let all_branches = self.get_top_level_branches().await?;
let dirty_prefixes = self.get_dirty_prefixes().await?;
let clean: BTreeMap<String, [u8; 32]> = all_branches
.into_iter()
.filter(|(branch, _)| !dirty_prefixes.contains(branch))
.collect();
debug!(
total = clean.len() + dirty_prefixes.len(),
clean = clean.len(),
dirty = dirty_prefixes.len(),
"Branch hygiene check"
);
Ok(clean)
}
#[instrument(skip(self))]
pub async fn is_fully_clean(&self) -> Result<bool, StorageError> {
if let Some(ref sql_store) = self.sql_store {
return sql_store.has_dirty_merkle().await.map(|has_dirty| !has_dirty);
}
Ok(true)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_merkle_diff_struct() {
let diff = MerkleDiff {
divergent_paths: vec!["uk.nhs".to_string()],
local_only: vec!["uk.private".to_string()],
remote_only: vec!["us.medicare".to_string()],
};
assert_eq!(diff.divergent_paths.len(), 1);
assert_eq!(diff.local_only.len(), 1);
assert_eq!(diff.remote_only.len(), 1);
}
}