use std::collections::{hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet};
use anyhow::{bail, Result};
use crate::{
metrics::DIEM_JELLYFISH_STORAGE_READS,
node_type::{Node, NodeKey},
storage::{
NodeBatch, NodeStats, StaleNodeIndex, StaleNodeIndexBatch, TreeReader, TreeUpdateBatch,
},
types::{Version, PRE_GENESIS_VERSION},
RootHash,
};
struct FrozenTreeCache {
node_cache: NodeBatch,
stale_node_index_cache: StaleNodeIndexBatch,
node_stats: Vec<NodeStats>,
root_hashes: Vec<RootHash>,
}
impl FrozenTreeCache {
fn new() -> Self {
Self {
node_cache: BTreeMap::new(),
stale_node_index_cache: BTreeSet::new(),
node_stats: Vec::new(),
root_hashes: Vec::new(),
}
}
}
pub struct TreeCache<'a, R> {
root_node_key: NodeKey,
next_version: Version,
node_cache: HashMap<NodeKey, Node>,
num_new_leaves: usize,
stale_node_index_cache: HashSet<NodeKey>,
num_stale_leaves: usize,
frozen_cache: FrozenTreeCache,
reader: &'a R,
}
impl<'a, R> TreeCache<'a, R>
where
R: 'a + TreeReader,
{
pub async fn new(reader: &'a R, next_version: Version) -> Result<TreeCache<'a, R>> {
let mut node_cache = HashMap::new();
let root_node_key = if next_version == 0 {
let pre_genesis_root_key = NodeKey::new_empty_path(PRE_GENESIS_VERSION);
let pre_genesis_root = reader.get_node_option(&pre_genesis_root_key).await?;
match pre_genesis_root {
Some(_) => {
pre_genesis_root_key
}
None => {
let genesis_root_key = NodeKey::new_empty_path(0);
node_cache.insert(genesis_root_key.clone(), Node::new_null());
genesis_root_key
}
}
} else {
NodeKey::new_empty_path(next_version - 1)
};
Ok(Self {
node_cache,
stale_node_index_cache: HashSet::new(),
frozen_cache: FrozenTreeCache::new(),
root_node_key,
next_version,
reader,
num_stale_leaves: 0,
num_new_leaves: 0,
})
}
pub async fn get_node(&self, node_key: &NodeKey) -> Result<Node> {
Ok(if let Some(node) = self.node_cache.get(node_key) {
node.clone()
} else if let Some(node) = self.frozen_cache.node_cache.get(node_key) {
node.clone()
} else {
DIEM_JELLYFISH_STORAGE_READS.inc();
self.reader.get_node_option(node_key).await?.unwrap()
})
}
pub fn get_root_node_key(&self) -> &NodeKey {
&self.root_node_key
}
pub fn set_root_node_key(&mut self, root_node_key: NodeKey) {
self.root_node_key = root_node_key;
}
pub fn put_node(&mut self, node_key: NodeKey, new_node: Node) -> Result<()> {
match self.node_cache.entry(node_key) {
Entry::Vacant(o) => {
if new_node.is_leaf() {
self.num_new_leaves += 1
}
o.insert(new_node);
}
Entry::Occupied(o) => bail!("Node with key {:?} already exists in NodeBatch", o.key()),
};
Ok(())
}
pub fn delete_node(&mut self, old_node_key: &NodeKey, is_leaf: bool) {
if self.node_cache.remove(old_node_key).is_none() {
let is_new_entry = self.stale_node_index_cache.insert(old_node_key.clone());
assert!(is_new_entry, "Node gets stale twice unexpectedly.");
if is_leaf {
self.num_stale_leaves += 1;
}
} else if is_leaf {
self.num_new_leaves -= 1;
}
}
pub async fn freeze(&mut self) {
let root_node_key = self.get_root_node_key();
let root_hash = self
.get_node(root_node_key)
.await
.unwrap_or_else(|_| unreachable!("Root node with key {:?} must exist", root_node_key))
.hash();
self.frozen_cache.root_hashes.push(RootHash(root_hash));
let node_stats = NodeStats {
new_nodes: self.node_cache.len(),
new_leaves: self.num_new_leaves,
stale_nodes: self.stale_node_index_cache.len(),
stale_leaves: self.num_stale_leaves,
};
self.frozen_cache.node_stats.push(node_stats);
self.frozen_cache.node_cache.extend(self.node_cache.drain());
let stale_since_version = self.next_version;
self.frozen_cache
.stale_node_index_cache
.extend(
self.stale_node_index_cache
.drain()
.map(|node_key| StaleNodeIndex {
stale_since_version,
node_key,
}),
);
self.num_stale_leaves = 0;
self.num_new_leaves = 0;
self.next_version += 1;
}
}
impl<'a, R> From<TreeCache<'a, R>> for (Vec<RootHash>, TreeUpdateBatch)
where
R: 'a + TreeReader,
{
fn from(tree_cache: TreeCache<'a, R>) -> Self {
(
tree_cache.frozen_cache.root_hashes,
TreeUpdateBatch {
node_batch: tree_cache.frozen_cache.node_cache,
stale_node_index_batch: tree_cache.frozen_cache.stale_node_index_cache,
node_stats: tree_cache.frozen_cache.node_stats,
},
)
}
}