0g-storage-node/common/append_merkle/src/node_manager.rs
peilun-conflux 506d234562
Some checks are pending
abi-consistent-check / build-and-compare (push) Waiting to run
code-coverage / unittest-cov (push) Waiting to run
rust / check (push) Waiting to run
rust / test (push) Waiting to run
rust / lints (push) Waiting to run
functional-test / test (push) Waiting to run
Use LRU to cache MPT nodes. (#227)
* Add trait.

* Update merkle tree trait.

* Use NodeManager.

* fix.

* Use LRU for cache.

* fix clippy.

* Save layer size.

* Initialize LogManager with NodeManager.

* Fix.

* Fix test.

* fix.
2024-10-27 12:52:06 +08:00

220 lines
6.6 KiB
Rust

use crate::HashElement;
use anyhow::Result;
use lru::LruCache;
use std::any::Any;
use std::num::NonZeroUsize;
use std::sync::Arc;
use tracing::error;
pub struct NodeManager<E: HashElement> {
cache: LruCache<(usize, usize), E>,
layer_size: Vec<usize>,
db: Arc<dyn NodeDatabase<E>>,
db_tx: Option<Box<dyn NodeTransaction<E>>>,
}
impl<E: HashElement> NodeManager<E> {
pub fn new(db: Arc<dyn NodeDatabase<E>>, capacity: usize) -> Result<Self> {
let mut layer = 0;
let mut layer_size = Vec::new();
while let Some(size) = db.get_layer_size(layer)? {
layer_size.push(size);
layer += 1;
}
Ok(Self {
cache: LruCache::new(NonZeroUsize::new(capacity).expect("capacity should be non-zero")),
layer_size,
db,
db_tx: None,
})
}
pub fn new_dummy() -> Self {
Self {
cache: LruCache::unbounded(),
layer_size: vec![],
db: Arc::new(EmptyNodeDatabase {}),
db_tx: None,
}
}
pub fn push_node(&mut self, layer: usize, node: E) {
self.add_node(layer, self.layer_size[layer], node);
self.set_layer_size(layer, self.layer_size[layer] + 1);
}
pub fn append_nodes(&mut self, layer: usize, nodes: &[E]) {
let mut pos = self.layer_size[layer];
let mut saved_nodes = Vec::with_capacity(nodes.len());
for node in nodes {
self.cache.put((layer, pos), node.clone());
saved_nodes.push((layer, pos, node));
pos += 1;
}
self.set_layer_size(layer, pos);
self.db_tx().save_node_list(&saved_nodes);
}
pub fn get_node(&self, layer: usize, pos: usize) -> Option<E> {
match self.cache.peek(&(layer, pos)) {
Some(node) => Some(node.clone()),
None => self.db.get_node(layer, pos).unwrap_or_else(|e| {
error!("Failed to get node: {}", e);
None
}),
}
}
pub fn get_nodes(&self, layer: usize, start_pos: usize, end_pos: usize) -> NodeIterator<E> {
NodeIterator {
node_manager: self,
layer,
start_pos,
end_pos,
}
}
pub fn add_node(&mut self, layer: usize, pos: usize, node: E) {
// No need to insert if the value is unchanged.
if self.cache.get(&(layer, pos)) != Some(&node) {
self.db_tx().save_node(layer, pos, &node);
self.cache.put((layer, pos), node);
}
}
pub fn add_layer(&mut self) {
self.layer_size.push(0);
let layer = self.layer_size.len() - 1;
self.db_tx().save_layer_size(layer, 0);
}
pub fn layer_size(&self, layer: usize) -> usize {
self.layer_size[layer]
}
pub fn num_layers(&self) -> usize {
self.layer_size.len()
}
pub fn truncate_nodes(&mut self, layer: usize, pos_end: usize) {
let mut removed_nodes = Vec::new();
for pos in pos_end..self.layer_size[layer] {
self.cache.pop(&(layer, pos));
removed_nodes.push((layer, pos));
}
self.db_tx().remove_node_list(&removed_nodes);
self.set_layer_size(layer, pos_end);
}
pub fn truncate_layer(&mut self, layer: usize) {
self.truncate_nodes(layer, 0);
if layer == self.num_layers() - 1 {
self.layer_size.pop();
self.db_tx().remove_layer_size(layer);
}
}
pub fn start_transaction(&mut self) {
if self.db_tx.is_some() {
error!("start new tx before commit");
panic!("start new tx before commit");
}
self.db_tx = Some(self.db.start_transaction());
}
pub fn commit(&mut self) {
let tx = match self.db_tx.take() {
Some(tx) => tx,
None => {
error!("db_tx is None");
return;
}
};
if let Err(e) = self.db.commit(tx) {
error!("Failed to commit db transaction: {}", e);
}
}
fn db_tx(&mut self) -> &mut dyn NodeTransaction<E> {
(*self.db_tx.as_mut().expect("tx checked")).as_mut()
}
fn set_layer_size(&mut self, layer: usize, size: usize) {
self.layer_size[layer] = size;
self.db_tx().save_layer_size(layer, size);
}
}
pub struct NodeIterator<'a, E: HashElement> {
node_manager: &'a NodeManager<E>,
layer: usize,
start_pos: usize,
end_pos: usize,
}
impl<'a, E: HashElement> Iterator for NodeIterator<'a, E> {
type Item = E;
fn next(&mut self) -> Option<Self::Item> {
if self.start_pos < self.end_pos {
let r = self.node_manager.get_node(self.layer, self.start_pos);
self.start_pos += 1;
r
} else {
None
}
}
}
pub trait NodeDatabase<E: HashElement>: Send + Sync {
fn get_node(&self, layer: usize, pos: usize) -> Result<Option<E>>;
fn get_layer_size(&self, layer: usize) -> Result<Option<usize>>;
fn start_transaction(&self) -> Box<dyn NodeTransaction<E>>;
fn commit(&self, tx: Box<dyn NodeTransaction<E>>) -> Result<()>;
}
pub trait NodeTransaction<E: HashElement>: Send + Sync {
fn save_node(&mut self, layer: usize, pos: usize, node: &E);
/// `nodes` are a list of tuples `(layer, pos, node)`.
fn save_node_list(&mut self, nodes: &[(usize, usize, &E)]);
fn remove_node_list(&mut self, nodes: &[(usize, usize)]);
fn save_layer_size(&mut self, layer: usize, size: usize);
fn remove_layer_size(&mut self, layer: usize);
fn into_any(self: Box<Self>) -> Box<dyn Any>;
}
/// A dummy database structure for in-memory merkle tree that will not read/write db.
pub struct EmptyNodeDatabase {}
pub struct EmptyNodeTransaction {}
impl<E: HashElement> NodeDatabase<E> for EmptyNodeDatabase {
fn get_node(&self, _layer: usize, _pos: usize) -> Result<Option<E>> {
Ok(None)
}
fn get_layer_size(&self, _layer: usize) -> Result<Option<usize>> {
Ok(None)
}
fn start_transaction(&self) -> Box<dyn NodeTransaction<E>> {
Box::new(EmptyNodeTransaction {})
}
fn commit(&self, _tx: Box<dyn NodeTransaction<E>>) -> Result<()> {
Ok(())
}
}
impl<E: HashElement> NodeTransaction<E> for EmptyNodeTransaction {
fn save_node(&mut self, _layer: usize, _pos: usize, _node: &E) {}
fn save_node_list(&mut self, _nodes: &[(usize, usize, &E)]) {}
fn remove_node_list(&mut self, _nodes: &[(usize, usize)]) {}
fn save_layer_size(&mut self, _layer: usize, _size: usize) {}
fn remove_layer_size(&mut self, _layer: usize) {}
fn into_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}