diff --git a/firewood/src/merkle.rs b/firewood/src/merkle.rs index 9b40aa2ec..168b44500 100644 --- a/firewood/src/merkle.rs +++ b/firewood/src/merkle.rs @@ -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; @@ -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), } @@ -523,10 +523,11 @@ impl + Send + Sync, T> Merkle { 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 @@ -571,12 +572,13 @@ impl + Send + Sync, T> Merkle { 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), @@ -607,18 +609,19 @@ impl + Send + Sync, T> Merkle { 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 { @@ -680,11 +683,11 @@ impl + Send + Sync, T> Merkle { 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; } @@ -732,13 +735,13 @@ impl + Send + Sync, T> Merkle { 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), @@ -771,19 +774,19 @@ impl + Send + Sync, T> Merkle { 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 { @@ -1474,6 +1477,56 @@ impl + Send + Sync, T> Merkle { 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), + mut node: NodeObjRef<'a>, + path: PartialPath, + ) -> Result, 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), + mut node: NodeObjRef<'a>, + data: Data, + ) -> Result { + 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), + mut node: NodeObjRef<'a>, + write_result: Result<(), ObjWriteSizeError>, + ) -> Result, 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)]) { @@ -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(); @@ -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::>(); + + 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::>(); + + 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::>(); + + 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::>(); + + 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, + new_data: Vec, + ) -> 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(()) + } } diff --git a/firewood/src/merkle/node.rs b/firewood/src/merkle/node.rs index 3e3ac0f7a..5df8321e8 100644 --- a/firewood/src/merkle/node.rs +++ b/firewood/src/merkle/node.rs @@ -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), diff --git a/firewood/src/merkle/proof.rs b/firewood/src/merkle/proof.rs index 4eaf0890b..83a29b4a4 100644 --- a/firewood/src/merkle/proof.rs +++ b/firewood/src/merkle/proof.rs @@ -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; @@ -58,7 +58,7 @@ pub enum ProofError { #[error("invalid root hash")] InvalidRootHash, #[error("{0}")] - WriteError(#[from] ObjWriteError), + WriteError(#[from] ObjWriteSizeError), } impl From for ProofError { diff --git a/firewood/src/shale/mod.rs b/firewood/src/shale/mod.rs index 81fbd3b1e..1d3cb79af 100644 --- a/firewood/src/shale/mod.rs +++ b/firewood/src/shale/mod.rs @@ -42,8 +42,8 @@ pub enum ShaleError { // this could probably included with ShaleError, // but keeping it separate for now as Obj/ObjRef might change in the near future #[derive(Debug, Error)] -#[error("write error")] -pub struct ObjWriteError; +#[error("object cannot be written in the space provided")] +pub struct ObjWriteSizeError; pub type SpaceId = u8; pub const INVALID_SPACE_ID: SpaceId = 0xff; @@ -115,13 +115,13 @@ impl Obj { /// Write to the underlying object. Returns `Ok(())` on success. #[inline] - pub fn write(&mut self, modify: impl FnOnce(&mut T)) -> Result<(), ObjWriteError> { + pub fn write(&mut self, modify: impl FnOnce(&mut T)) -> Result<(), ObjWriteSizeError> { modify(self.value.write()); // if `estimate_mem_image` gives overflow, the object will not be written self.dirty = match self.value.estimate_mem_image() { Some(len) => Some(len), - None => return Err(ObjWriteError), + None => return Err(ObjWriteSizeError), }; Ok(()) @@ -181,7 +181,7 @@ impl<'a, T: Storable + Debug> ObjRef<'a, T> { } #[inline] - pub fn write(&mut self, modify: impl FnOnce(&mut T)) -> Result<(), ObjWriteError> { + pub fn write(&mut self, modify: impl FnOnce(&mut T)) -> Result<(), ObjWriteSizeError> { #[allow(clippy::unwrap_used)] let inner = self.inner.as_mut().unwrap(); inner.write(modify)?;