mirror of
https://github.com/0glabs/0g-storage-node.git
synced 2025-04-04 15:35:18 +00:00

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
* 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.
220 lines
6.6 KiB
Rust
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
|
|
}
|
|
}
|