Skip to content

Commit

Permalink
Update or insert nodes (#545)
Browse files Browse the repository at this point in the history 
Signed-off-by: Richard Pringle <[email protected]>
Co-authored-by: Ron Kuris <[email protected]>
  • Loading branch information
@richardpringle @rkuris
richardpringle and rkuris authored Feb 21, 2024
1 parent f7d7414 commit 56e2d2d
Show file tree
Hide file tree
Showing 4 changed files with 236 additions and 44 deletions.
258 changes: 221 additions & 37 deletions firewood/src/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::db::{MutStore, SharedStore};
// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
// See the file LICENSE.md for licensing terms.
use crate::nibbles::Nibbles;
use crate::shale::{self, disk_address::DiskAddress, ObjWriteError, ShaleError, ShaleStore};
use crate::shale::{self, disk_address::DiskAddress, ObjWriteSizeError, ShaleError, ShaleStore};
use crate::v2::api;
use futures::{StreamExt, TryStreamExt};
use sha3::Digest;
Expand Down Expand Up @@ -49,7 +49,7 @@ pub enum MerkleError {
#[error("removing internal node references failed")]
UnsetInternal,
#[error("error updating nodes: {0}")]
WriteError(#[from] ObjWriteError),
WriteError(#[from] ObjWriteSizeError),
#[error("merkle serde error: {0}")]
BinarySerdeError(String),
}
Expand Down Expand Up @@ -523,10 +523,11 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {
match (overlap.unique_a.len(), overlap.unique_b.len()) {
// same node, overwrite the data
(0, 0) => {
node.write(|node| {
node.inner.set_data(Data(val));
node.rehash();
})?;
self.update_data_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
Data(val),
)?;
}

// new node is a child of the old node
Expand Down Expand Up @@ -571,12 +572,13 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {

let new_branch_path = overlap.shared.to_vec();

node.write(move |old_leaf| {
*old_leaf.inner.path_mut() = PartialPath(old_leaf_path.to_vec());
old_leaf.rehash();
})?;

let old_leaf = node.as_ptr();
let old_leaf = self
.update_path_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
PartialPath(old_leaf_path.to_vec()),
)?
.as_ptr();

let mut new_branch = BranchNode {
path: PartialPath(new_branch_path),
Expand Down Expand Up @@ -607,18 +609,19 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {

let new_branch_path = overlap.shared.to_vec();

node.write(move |old_leaf| {
*old_leaf.inner.path_mut() = PartialPath(old_leaf_path.to_vec());
old_leaf.rehash();
})?;
let old_leaf = self
.update_path_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
PartialPath(old_leaf_path.to_vec()),
)?
.as_ptr();

let new_leaf = Node::from_leaf(LeafNode::new(
PartialPath(new_leaf_path),
Data(val),
));

let old_leaf = node.as_ptr();

let new_leaf = self.put_node(new_leaf)?.as_ptr();

let mut new_branch = BranchNode {
Expand Down Expand Up @@ -680,11 +683,11 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {
match (overlap.unique_a.len(), overlap.unique_b.len()) {
// same node, overwrite the data
(0, 0) => {
node.write(|node| {
node.inner.set_data(Data(val));
node.rehash();
})?;

self.update_data_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
Data(val),
)?;
break None;
}

Expand Down Expand Up @@ -732,13 +735,13 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {

let new_branch_path = overlap.shared.to_vec();

node.write(move |old_branch| {
*old_branch.inner.path_mut() =
PartialPath(old_branch_path.to_vec());
old_branch.rehash();
})?;

let old_branch = node.as_ptr();
let old_branch = self
.update_path_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
PartialPath(old_branch_path.to_vec()),
)?
.as_ptr();

let mut new_branch = BranchNode {
path: PartialPath(new_branch_path),
Expand Down Expand Up @@ -771,19 +774,19 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {

let new_branch_path = overlap.shared.to_vec();

node.write(move |old_branch| {
*old_branch.inner.path_mut() =
PartialPath(old_branch_path.to_vec());
old_branch.rehash();
})?;
let old_branch = self
.update_path_and_move_node_if_larger(
(&mut parents, &mut deleted),
node,
PartialPath(old_branch_path.to_vec()),
)?
.as_ptr();

let new_leaf = Node::from_leaf(LeafNode::new(
PartialPath(new_leaf_path),
Data(val),
));

let old_branch = node.as_ptr();

let new_leaf = self.put_node(new_leaf)?.as_ptr();

let mut new_branch = BranchNode {
Expand Down Expand Up @@ -1474,6 +1477,56 @@ impl<S: ShaleStore<Node> + Send + Sync, T> Merkle<S, T> {
last_key_proof,
}))
}

/// Try to update the [NodeObjRef]'s path in-place. If the update fails because the node can no longer fit at its old address,
/// then the old address is marked for deletion and the [Node] (with its update) is inserted at a new address.
fn update_path_and_move_node_if_larger<'a>(
&'a self,
(parents, to_delete): (&mut [(NodeObjRef, u8)], &mut Vec<DiskAddress>),
mut node: NodeObjRef<'a>,
path: PartialPath,
) -> Result<NodeObjRef<'a>, MerkleError> {
let write_result = node.write(|node| {
node.inner_mut().set_path(path);
node.rehash();
});

self.move_node_if_write_failed((parents, to_delete), node, write_result)
}

/// Try to update the [NodeObjRef]'s data/value in-place. If the update fails because the node can no longer fit at its old address,
/// then the old address is marked for deletion and the [Node] (with its update) is inserted at a new address.
fn update_data_and_move_node_if_larger<'a>(
&'a self,
(parents, to_delete): (&mut [(NodeObjRef, u8)], &mut Vec<DiskAddress>),
mut node: NodeObjRef<'a>,
data: Data,
) -> Result<NodeObjRef, MerkleError> {
let write_result = node.write(|node| {
node.inner_mut().set_data(data);
node.rehash();
});

self.move_node_if_write_failed((parents, to_delete), node, write_result)
}

/// Checks if the `write_result` is an [ObjWriteSizeError]. If it is, then the `node` is moved to a new address and the old address is marked for deletion.
fn move_node_if_write_failed<'a>(
&'a self,
(parents, deleted): (&mut [(NodeObjRef, u8)], &mut Vec<DiskAddress>),
mut node: NodeObjRef<'a>,
write_result: Result<(), ObjWriteSizeError>,
) -> Result<NodeObjRef<'a>, MerkleError> {
if let Err(ObjWriteSizeError) = write_result {
let old_node_address = node.as_ptr();
node = self.put_node(node.clone())?;
deleted.push(old_node_address);

set_parent(node.as_ptr(), parents);
}

Ok(node)
}
}

fn set_parent(new_chd: DiskAddress, parents: &mut [(NodeObjRef, u8)]) {
Expand Down Expand Up @@ -2270,6 +2323,7 @@ mod tests {
}

#[test]
#[ignore]
fn single_key_proof_with_one_node() {
let mut merkle = create_test_merkle();
let root = merkle.init_root().unwrap();
Expand Down Expand Up @@ -2306,4 +2360,134 @@ mod tests {

assert_eq!(verified.as_deref(), Some(key1.as_slice()));
}

#[test]
#[ignore]
fn update_leaf_with_larger_path() -> Result<(), MerkleError> {
let path = vec![0x00];
let data = vec![0x00];

let double_path = path
.clone()
.into_iter()
.chain(path.clone())
.collect::<Vec<_>>();

let node = Node::from_leaf(LeafNode {
path: PartialPath::from(path),
data: Data(data.clone()),
});

check_node_update(node, double_path, data)
}

#[test]
#[ignore]
fn update_leaf_with_larger_data() -> Result<(), MerkleError> {
let path = vec![0x00];
let data = vec![0x00];

let double_data = data
.clone()
.into_iter()
.chain(data.clone())
.collect::<Vec<_>>();

let node = Node::from_leaf(LeafNode {
path: PartialPath::from(path.clone()),
data: Data(data),
});

check_node_update(node, path, double_data)
}

#[test]
#[ignore]
fn update_branch_with_larger_path() -> Result<(), MerkleError> {
let path = vec![0x00];
let data = vec![0x00];

let double_path = path
.clone()
.into_iter()
.chain(path.clone())
.collect::<Vec<_>>();

let node = Node::from_branch(BranchNode {
path: PartialPath::from(path.clone()),
children: Default::default(),
value: Some(Data(data.clone())),
children_encoded: Default::default(),
});

check_node_update(node, double_path, data)
}

#[test]
#[ignore]
fn update_branch_with_larger_data() -> Result<(), MerkleError> {
let path = vec![0x00];
let data = vec![0x00];

let double_data = data
.clone()
.into_iter()
.chain(data.clone())
.collect::<Vec<_>>();

let node = Node::from_branch(BranchNode {
path: PartialPath::from(path.clone()),
children: Default::default(),
value: Some(Data(data)),
children_encoded: Default::default(),
});

check_node_update(node, path, double_data)
}

fn check_node_update(
node: Node,
new_path: Vec<u8>,
new_data: Vec<u8>,
) -> Result<(), MerkleError> {
let merkle = create_test_merkle();
let root = merkle.init_root()?;
let root = merkle.get_node(root)?;

let mut node_ref = merkle.put_node(node)?;
let addr = node_ref.as_ptr();

// make sure that doubling the path length will fail on a normal write
let write_result = node_ref.write(|node| {
node.inner_mut().set_path(PartialPath(new_path.clone()));
node.inner_mut().set_data(Data(new_data.clone()));
node.rehash();
});

assert!(matches!(write_result, Err(ObjWriteSizeError)));

let mut to_delete = vec![];
// could be any branch node, convenient to use the root.
let mut parents = vec![(root, 0)];

let node = merkle.update_path_and_move_node_if_larger(
(&mut parents, &mut to_delete),
node_ref,
PartialPath(new_path.clone()),
)?;

assert_ne!(node.as_ptr(), addr);
assert_eq!(&to_delete[0], &addr);

let (path, data) = match node.inner() {
NodeType::Leaf(leaf) => (&leaf.path, Some(&leaf.data)),
NodeType::Branch(branch) => (&branch.path, branch.value.as_ref()),
_ => unreachable!(),
};

assert_eq!(path, &PartialPath(new_path));
assert_eq!(data, Some(&Data(new_data)));

Ok(())
}
}
8 changes: 8 additions & 0 deletions firewood/src/merkle/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,14 @@ impl NodeType {
}
}

pub fn set_path(&mut self, path: PartialPath) {
match self {
NodeType::Branch(u) => u.path = path,
NodeType::Leaf(node) => node.path = path,
NodeType::Extension(node) => node.path = path,
}
}

pub fn set_data(&mut self, data: Data) {
match self {
NodeType::Branch(u) => u.value = Some(data),
Expand Down
4 changes: 2 additions & 2 deletions firewood/src/merkle/proof.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
use std::cmp::Ordering;
use std::collections::HashMap;

use crate::shale::ObjWriteError;
use crate::shale::ObjWriteSizeError;
use crate::shale::{disk_address::DiskAddress, ShaleError, ShaleStore};
use crate::v2::api::HashKey;
use nix::errno::Errno;
Expand Down Expand Up @@ -58,7 +58,7 @@ pub enum ProofError {
#[error("invalid root hash")]
InvalidRootHash,
#[error("{0}")]
WriteError(#[from] ObjWriteError),
WriteError(#[from] ObjWriteSizeError),
}

impl From<DataStoreError> for ProofError {
Expand Down
Loading

0 comments on commit 56e2d2d

Please sign in to comment.