From 6a1d67649bd86f50d21fbf96813d8ceab8627ae7 Mon Sep 17 00:00:00 2001 From: Ashwin Sekar Date: Tue, 10 Jan 2023 16:04:45 -0800 Subject: [PATCH] Track pruned subtrees in repair weight --- core/src/heaviest_subtree_fork_choice.rs | 38 +- core/src/repair_weight.rs | 716 ++++++++++++++++------- 2 files changed, 539 insertions(+), 215 deletions(-) diff --git a/core/src/heaviest_subtree_fork_choice.rs b/core/src/heaviest_subtree_fork_choice.rs index 4d325e3ecb6a83..e6eb71b1d78e35 100644 --- a/core/src/heaviest_subtree_fork_choice.rs +++ b/core/src/heaviest_subtree_fork_choice.rs @@ -16,8 +16,10 @@ use { }, std::{ borrow::Borrow, + cmp::Reverse, collections::{ - btree_set::Iter, hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet, VecDeque, + btree_set::Iter, hash_map::Entry, BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, + VecDeque, }, sync::{Arc, RwLock}, time::Instant, @@ -309,6 +311,40 @@ impl HeaviestSubtreeForkChoice { self.last_root_time = Instant::now(); } + /// Prunes all slots not descending from `new_root`. + /// Returns slots that were removed (`< new_root`) and the newly pruned subtrees with only the slots > `new_root` + pub fn prune(&mut self, new_root: SlotHashKey) -> (Vec, Vec) { + let mut pruned = vec![]; + let mut removed = vec![]; + // Find the pruned trees + let mut to_visit = BinaryHeap::from([Reverse(self.tree_root)]); + while !to_visit.is_empty() { + let cur_slot = to_visit.pop().unwrap().0; + if cur_slot == new_root { + continue; + } + if cur_slot < new_root { + for child in (&*self) + .children(&cur_slot) + .expect("slot was discovered earlier, must exist") + { + to_visit.push(Reverse(*child)) + } + removed.push(cur_slot); + } else { + // The start of a pruned subtree. Split it out and stop traversing this subtree. + pruned.push(self.split_off(&cur_slot)); + } + } + + for slot in removed.iter() { + self.fork_infos + .remove(slot) + .expect("Slots reachable from old root must exist in tree"); + } + (removed, pruned) + } + pub fn add_root_parent(&mut self, root_parent: SlotHashKey) { assert!(root_parent.0 < self.tree_root.0); assert!(self.fork_infos.get(&root_parent).is_none()); diff --git a/core/src/repair_weight.rs b/core/src/repair_weight.rs index 9ea3b82f112ef1..4786663a5f8f99 100644 --- a/core/src/repair_weight.rs +++ b/core/src/repair_weight.rs @@ -5,7 +5,6 @@ use { repair_service::{BestRepairsStats, RepairTiming}, repair_weighted_traversal, serve_repair::ShredRepairType, - tree_diff::TreeDiff, }, solana_ledger::{ ancestor_iterator::AncestorIterator, blockstore::Blockstore, blockstore_meta::SlotMeta, @@ -18,36 +17,51 @@ use { hash::Hash, pubkey::Pubkey, }, - std::collections::{BTreeSet, HashMap, HashSet, VecDeque}, + std::collections::{HashMap, HashSet, VecDeque}, }; +#[derive(PartialEq, Eq, Copy, Clone, Hash, Debug)] +enum TreeRoot { + Root(Slot), + PrunedRoot(Slot), +} + +impl TreeRoot { + pub fn is_pruned(&self) -> bool { + matches!(self, Self::PrunedRoot(_)) + } +} + +impl From for Slot { + fn from(val: TreeRoot) -> Self { + match val { + TreeRoot::Root(slot) => slot, + TreeRoot::PrunedRoot(slot) => slot, + } + } +} pub struct RepairWeight { // Map from root -> a subtree rooted at that `root` trees: HashMap, + // Map from root -> pruned subtree + // We manage the size by removing slots < root + pruned_trees: HashMap, + // Maps each slot to the root of the tree that contains it - slot_to_tree: HashMap, - // Prevents spam attacks that send votes for long chains of - // unrooted slots, which would incur high cost of blockstore - // lookup/iteration to resolve ancestry. - - // Must maintain invariant that any slot in `unrooted_slots` - // does not exist in `slot_to_tree` and any descendants of - // slots in the set `unrooted_slots` must also be in the set - unrooted_slots: BTreeSet, + slot_to_tree: HashMap, root: Slot, } impl RepairWeight { pub fn new(root: Slot) -> Self { let root_tree = HeaviestSubtreeForkChoice::new((root, Hash::default())); - let slot_to_tree: HashMap = vec![(root, root)].into_iter().collect(); - let trees: HashMap = - vec![(root, root_tree)].into_iter().collect(); + let slot_to_tree = HashMap::from([(root, TreeRoot::Root(root))]); + let trees = HashMap::from([(root, root_tree)]); Self { trees, slot_to_tree, root, - unrooted_slots: BTreeSet::new(), + pruned_trees: HashMap::new(), } } @@ -60,15 +74,15 @@ impl RepairWeight { ) where I: Iterator)>, { - let mut all_subtree_updates: HashMap> = HashMap::new(); + let mut all_subtree_updates: HashMap> = HashMap::new(); for (slot, pubkey_votes) in votes { - if slot < self.root || self.unrooted_slots.contains(&slot) { + if slot < self.root { continue; } - let mut tree_root = self.slot_to_tree.get(&slot).cloned(); + let mut tree_root = self.slot_to_tree.get(&slot).copied(); let mut new_ancestors = VecDeque::new(); // If we don't know know how this slot chains to any existing trees - // in `self.trees`, then use `blockstore` to see if this chains + // in `self.trees` or `self.pruned_trees`, then use `blockstore` to see if this chains // any existing trees in `self.trees` if tree_root.is_none() { let (discovered_ancestors, existing_subtree_root) = @@ -77,33 +91,42 @@ impl RepairWeight { tree_root = existing_subtree_root; } - let should_create_new_subtree = tree_root.is_none(); - - let tree_root = tree_root.unwrap_or( - // If `tree_root` is still None, then there is no known - // subtree that contains `slot`. Thus, create a new - // subtree rooted at the earliest known ancestor of `slot` - *new_ancestors.front().unwrap_or(&slot), - ); - - // If the discovered root of this tree is unrooted, mark all - // ancestors in `new_ancestors` and this slot as unrooted - if self.is_unrooted_slot(tree_root) { - for ancestor in new_ancestors.into_iter().chain(std::iter::once(slot)) { - self.unrooted_slots.insert(ancestor); + let (tree_root, tree) = { + match (tree_root, *new_ancestors.front().unwrap_or(&slot)) { + (Some(tree_root), _) if !tree_root.is_pruned() => ( + tree_root, + self.trees + .get_mut(&tree_root.into()) + .expect("If tree root was found, it must exist in `self.trees`"), + ), + (Some(tree_root), _) => ( + tree_root, + self.pruned_trees.get_mut(&tree_root.into()).expect( + "If a pruned tree root was found, it must exist in `self.pruned_trees`", + ), + ), + (None, earliest_ancestor) => { + // There is no known subtree that contains `slot`. Thus, create a new + // subtree rooted at the earliest known ancestor of `slot`. + // If this earliest known ancestor is not part of the rooted path, create a new + // pruned tree from the ancestor that is `> self.root` instead. + if earliest_ancestor < self.root { + let next_earliest_ancestor = *new_ancestors.get(1).unwrap_or(&slot); + self.insert_new_pruned_tree(next_earliest_ancestor); + ( + TreeRoot::PrunedRoot(next_earliest_ancestor), + self.pruned_trees.get_mut(&next_earliest_ancestor).unwrap(), + ) + } else { + self.insert_new_tree(earliest_ancestor); + ( + TreeRoot::Root(earliest_ancestor), + self.trees.get_mut(&earliest_ancestor).unwrap(), + ) + } + } } - - continue; - } - - if should_create_new_subtree { - self.insert_new_tree(tree_root); - } - - let tree = self - .trees - .get_mut(&tree_root) - .expect("If tree root was found, it must exist in `self.trees`"); + }; // First element in `ancestors` must be either: // 1) Leaf of some existing subtree @@ -132,9 +155,8 @@ impl RepairWeight { for (tree_root, updates) in all_subtree_updates { let tree = self - .trees - .get_mut(&tree_root) - .expect("`slot_to_tree` and `self.trees` must be in sync"); + .get_tree(&tree_root) + .expect("Tree for `tree_root` must exist here"); let updates: Vec<_> = updates.into_iter().collect(); tree.add_votes( updates @@ -249,14 +271,16 @@ impl RepairWeight { } /// Split `slot` and descendants into an orphan tree in repair weighting - /// These orphaned slots should be removed from `unrooted_slots` as on proper repair these slots might - /// now be part of the rooted path + /// Even if these orphaned slots were previously pruned, they should be added back properly as + /// on repair they might now be part of the rooted path. + /// Assumes that `slot` is greater than `self.root` pub fn split_off(&mut self, slot: Slot) { + assert!(slot >= self.root); if slot == self.root { error!("Trying to orphan root of repair tree {}", slot); return; } - if let Some(subtree_root) = self.slot_to_tree.get(&slot) { + if let Some(TreeRoot::Root(subtree_root)) = self.slot_to_tree.get(&slot) { if *subtree_root == slot { info!("{} is already orphan, skipping", slot); return; @@ -264,13 +288,37 @@ impl RepairWeight { let subtree = self .trees .get_mut(subtree_root) - .expect("subtree must exist"); + .expect("`self.slot_to_tree` and `self.trees` must be in sync"); let orphaned_tree = subtree.split_off(&(slot, Hash::default())); - for ((orphaned_slot, _), _) in orphaned_tree.all_slots_stake_voted_subtree() { - self.unrooted_slots.remove(orphaned_slot); - self.slot_to_tree.insert(*orphaned_slot, slot); - } + self.rename_tree_root(&orphaned_tree, TreeRoot::Root(slot)); self.trees.insert(slot, orphaned_tree); + } else if let Some(TreeRoot::PrunedRoot(subtree_root)) = + self.slot_to_tree.get(&slot).copied() + { + info!( + "Dumping pruned slot {} of tree {} in repair", + slot, subtree_root + ); + let mut subtree = self + .pruned_trees + .remove(&subtree_root) + .expect("`self.slot_to_tree` and `self.pruned_trees` must be in sync"); + + if subtree_root == slot { + // In this case we simply unprune the entire subtree + self.rename_tree_root(&subtree, TreeRoot::Root(subtree_root)); + self.trees.insert(subtree_root, subtree); + } else { + let orphaned_tree = subtree.split_off(&(slot, Hash::default())); + self.pruned_trees.insert(subtree_root, subtree); + self.rename_tree_root(&orphaned_tree, TreeRoot::Root(slot)); + self.trees.insert(slot, orphaned_tree); + } + } else { + warn!( + "Trying to split off slot {} which doesn't currently exist in repair", + slot + ); } } @@ -283,9 +331,21 @@ impl RepairWeight { } // Root slot of the tree that contains `new_root`, if one exists - let new_root_tree_root = self.slot_to_tree.get(&new_root).cloned(); + let new_root_tree_root = self + .slot_to_tree + .get(&new_root) + .copied() + .map(|root| match root { + TreeRoot::Root(r) => r, + TreeRoot::PrunedRoot(r) => panic!( + "New root {} chains to a pruned tree with root {}", + new_root, r + ), + }); // Purge outdated trees from `self.trees` + // These are all subtrees that have a `tree_root` < `new_root` and are not part of the + // rooted subtree let subtrees_to_purge: Vec<_> = self .trees .keys() @@ -295,42 +355,42 @@ impl RepairWeight { .map(|new_root_tree_root| **subtree_root != new_root_tree_root) .unwrap_or(true) }) - .cloned() + .copied() .collect(); for subtree_root in subtrees_to_purge { let subtree = self .trees .remove(&subtree_root) .expect("Must exist, was found in `self.trees` above"); - self.remove_tree_slots( - subtree - .all_slots_stake_voted_subtree() - .map(|((slot, _), _)| slot), - new_root, - ); + + // Track these trees as part of the pruned set + self.rename_tree_root(&subtree, TreeRoot::PrunedRoot(subtree_root)); + self.pruned_trees.insert(subtree_root, subtree); } if let Some(new_root_tree_root) = new_root_tree_root { let mut new_root_tree = self .trees .remove(&new_root_tree_root) - .expect("Found slot root earlier in self.slot_to_trees, treee must exist"); + .expect("Found slot root earlier in self.slot_to_trees, tree must exist"); + // Find all descendants of `self.root` that are not reachable from `new_root`. - // These are exactly the unrooted slots, which can be purged and added to - // `self.unrooted_slots`. - let unrooted_slots = (&new_root_tree).subtree_diff( - (new_root_tree_root, Hash::default()), - (new_root, Hash::default()), - ); - self.remove_tree_slots( - unrooted_slots.iter().map(|slot_hash| &slot_hash.0), - new_root, - ); + // Prune these out and add to `self.pruned_trees` + trace!("pruning tree {} with {}", new_root_tree_root, new_root); + let (removed, pruned) = new_root_tree.prune((new_root, Hash::default())); + for pruned_tree in pruned { + let pruned_tree_root = pruned_tree.tree_root().0; + self.rename_tree_root(&pruned_tree, TreeRoot::PrunedRoot(pruned_tree_root)); + self.pruned_trees.insert(pruned_tree_root, pruned_tree); + } - new_root_tree.set_tree_root((new_root, Hash::default())); + for (slot, _) in removed { + self.slot_to_tree.remove(&slot); + } // Update `self.slot_to_tree` to reflect new root - self.rename_tree_root(&new_root_tree, new_root); + new_root_tree.set_tree_root((new_root, Hash::default())); + self.rename_tree_root(&new_root_tree, TreeRoot::Root(new_root)); // Insert the tree for the new root self.trees.insert(new_root, new_root_tree); @@ -338,10 +398,35 @@ impl RepairWeight { self.insert_new_tree(new_root); } - // Purge `self.unrooted_slots` of slots less than `new_root` as we know any - // unrooted votes for slots < `new_root` will now be rejected, so we won't - // need to check `self.unrooted_slots` to see if those slots are unrooted. - self.unrooted_slots = self.unrooted_slots.split_off(&new_root); + // Clean up the pruned set by splitting out slots that are less than `new_root` + self.pruned_trees = self + .pruned_trees + .drain() + .flat_map(|(tree_root, mut pruned_tree)| { + if tree_root < new_root { + trace!("pruning tree {} with {}", tree_root, new_root); + let (removed, pruned) = pruned_tree.prune((new_root, Hash::default())); + for (slot, _) in removed { + self.slot_to_tree.remove(&slot); + } + pruned + .into_iter() + .map(|new_pruned_subtree| { + let new_pruned_tree_root = new_pruned_subtree.tree_root().0; + // Resync `self.slot_to_tree` + for ((slot, _), _) in new_pruned_subtree.all_slots_stake_voted_subtree() + { + *self.slot_to_tree.get_mut(slot).unwrap() = + TreeRoot::PrunedRoot(new_pruned_tree_root); + } + (new_pruned_tree_root, new_pruned_subtree) + }) + .collect() + } else { + vec![(tree_root, pruned_tree)] + } + }) + .collect::>(); self.root = new_root; } @@ -519,15 +604,18 @@ impl RepairWeight { for ancestor in new_ancestors.iter().skip(num_skip).rev() { // We temporarily use orphan_tree_root as the tree root and later // rename tree root to either the parent_tree_root or the earliest_ancestor - self.slot_to_tree.insert(*ancestor, orphan_tree_root); - heaviest_tree.add_root_parent((*ancestor, Hash::default())); + if *ancestor >= self.root { + self.slot_to_tree + .insert(*ancestor, TreeRoot::Root(orphan_tree_root)); + heaviest_tree.add_root_parent((*ancestor, Hash::default())); + } } } if let Some(parent_tree_root) = parent_tree_root { // If another subtree is discovered to be the parent // of this subtree, merge the two. self.merge_trees( - orphan_tree_root, + TreeRoot::Root(orphan_tree_root), parent_tree_root, *new_ancestors .front() @@ -535,42 +623,60 @@ impl RepairWeight { epoch_stakes, epoch_schedule, ); - orphan_tree_root = parent_tree_root; + if let TreeRoot::Root(parent_tree_root) = parent_tree_root { + orphan_tree_root = parent_tree_root; + } else { + // This orphan is now part of a pruned tree, no need to chain further + return None; + } } else { // If there's no other subtree to merge with, then return // the current known earliest ancestor of this branch if let Some(earliest_ancestor) = new_ancestors.front() { + assert!(*earliest_ancestor != self.root); // In this case we would merge above let orphan_tree = self .trees .remove(&orphan_tree_root) .expect("orphan tree must exist"); - self.rename_tree_root(&orphan_tree, *earliest_ancestor); - assert!(self.trees.insert(*earliest_ancestor, orphan_tree).is_none()); - orphan_tree_root = *earliest_ancestor; + if *earliest_ancestor > self.root { + self.rename_tree_root(&orphan_tree, TreeRoot::Root(*earliest_ancestor)); + assert!(self.trees.insert(*earliest_ancestor, orphan_tree).is_none()); + orphan_tree_root = *earliest_ancestor; + } else { + // In this case we should create a new pruned subtree + let next_earliest_ancestor = + new_ancestors.get(1).unwrap_or(&orphan_tree_root); + self.rename_tree_root( + &orphan_tree, + TreeRoot::PrunedRoot(*next_earliest_ancestor), + ); + assert!(self + .pruned_trees + .insert(*next_earliest_ancestor, orphan_tree) + .is_none()); + return None; + } } break; } } - if self.is_unrooted_slot(orphan_tree_root) { - // If this orphan branch chained down to an unrooted - // slot, then purge the entire subtree as we know the - // entire subtree is unrooted - let orphan_tree = self - .trees - .remove(&orphan_tree_root) - .expect("Must exist, was found in `self.trees` above"); - self.remove_tree_slots( - orphan_tree - .all_slots_stake_voted_subtree() - .map(|((slot, _), _)| slot), - self.root, - ); - None - } else { - // Return the (potentially new in the case of some merges) - // root of this orphan subtree - Some(orphan_tree_root) + // Return the (potentially new in the case of some merges) + // root of this orphan subtree + Some(orphan_tree_root) + } + + fn get_tree(&mut self, tree_root: &TreeRoot) -> Option<&mut HeaviestSubtreeForkChoice> { + match tree_root { + TreeRoot::Root(r) => self.trees.get_mut(r), + TreeRoot::PrunedRoot(r) => self.pruned_trees.get_mut(r), + } + } + + fn remove_tree(&mut self, tree_root: &TreeRoot) -> Option { + match tree_root { + TreeRoot::Root(r) => self.trees.remove(r), + TreeRoot::PrunedRoot(r) => self.pruned_trees.remove(r), } } @@ -578,64 +684,73 @@ impl RepairWeight { assert!(!self.trees.contains_key(&new_tree_root)); // Update `self.slot_to_tree` - self.slot_to_tree.insert(new_tree_root, new_tree_root); + self.slot_to_tree + .insert(new_tree_root, TreeRoot::Root(new_tree_root)); self.trees.insert( new_tree_root, HeaviestSubtreeForkChoice::new((new_tree_root, Hash::default())), ); } + fn insert_new_pruned_tree(&mut self, new_pruned_tree_root: Slot) { + assert!(!self.pruned_trees.contains_key(&new_pruned_tree_root)); + + // Update `self.slot_to_tree` + self.slot_to_tree.insert( + new_pruned_tree_root, + TreeRoot::PrunedRoot(new_pruned_tree_root), + ); + self.pruned_trees.insert( + new_pruned_tree_root, + HeaviestSubtreeForkChoice::new((new_pruned_tree_root, Hash::default())), + ); + } + + /// Finds any ancestors avaiable from `blockstore` for `slot`, stopping if it chains to any + /// tree in `self.trees` or `self.pruned_trees`. + /// Returns ancestors, and possibly the root of a tree that `slot` belongs to fn find_ancestor_subtree_of_slot( &self, blockstore: &Blockstore, slot: Slot, - ) -> (VecDeque, Option) { + ) -> (VecDeque, Option) { let ancestors = AncestorIterator::new(slot, blockstore); let mut ancestors_to_add = VecDeque::new(); - let mut tree_to_merge = None; + let mut tree_root = None; // This means `heaviest_tree_root` has not been // chained back to slots currently being replayed // in BankForks. Try to see if blockstore has sufficient // information to link this slot back for a in ancestors { ancestors_to_add.push_front(a); - // If this tree chains to some unrooted fork, purge it - if self.is_unrooted_slot(a) { - break; - } - // If an ancestor chains back to another subtree, then return - let other_tree_root = self.slot_to_tree.get(&a).cloned(); - tree_to_merge = other_tree_root; - if tree_to_merge.is_some() { + // If an ancestor chains back to another subtree or pruned, then return + tree_root = self.slot_to_tree.get(&a).copied(); + if tree_root.is_some() || a < self.root { break; } } - (ancestors_to_add, tree_to_merge) - } - - fn is_unrooted_slot(&self, slot: Slot) -> bool { - slot < self.root || self.unrooted_slots.contains(&slot) + (ancestors_to_add, tree_root) } - // Attaches `tree1` rooted at `root1` to `tree2` rooted at `root2` - // at the leaf in `tree2` given by `merge_leaf` + /// Attaches `tree1` rooted at `root1` to `tree2` rooted at `root2` + /// at the leaf in `tree2` given by `merge_leaf` + /// Expects `root1` and `root2` to exist in `self.trees` or `self.pruned_trees` fn merge_trees( &mut self, - root1: Slot, - root2: Slot, + root1: TreeRoot, + root2: TreeRoot, merge_leaf: Slot, epoch_stakes: &HashMap, epoch_schedule: &EpochSchedule, ) { // Update self.slot_to_tree to reflect the merge - let tree1 = self.trees.remove(&root1).expect("tree to merge must exist"); + let tree1 = self.remove_tree(&root1).expect("tree to merge must exist"); self.rename_tree_root(&tree1, root2); // Merge trees let tree2 = self - .trees - .get_mut(&root2) + .get_tree(&root2) .expect("tree to be merged into must exist"); tree2.merge( @@ -647,7 +762,7 @@ impl RepairWeight { } // Update all slots in the `tree1` to point to `root2`, - fn rename_tree_root(&mut self, tree1: &HeaviestSubtreeForkChoice, root2: Slot) { + fn rename_tree_root(&mut self, tree1: &HeaviestSubtreeForkChoice, root2: TreeRoot) { let all_slots = tree1.all_slots_stake_voted_subtree(); for ((slot, _), _) in all_slots { *self @@ -657,23 +772,6 @@ impl RepairWeight { } } - // For all slots `s` in `tree1` - // 1) If `s` < min, purge `s` from `self.slot_to_tree` - // 2) Else add `s` to `self.unrooted_slots` - fn remove_tree_slots<'a, I>(&'a mut self, slots_to_remove: I, min: Slot) - where - I: Iterator, - { - for slot in slots_to_remove { - self.slot_to_tree - .remove(slot) - .expect("Item in tree must exist in `slot_to_tree`"); - if *slot >= min { - self.unrooted_slots.insert(*slot); - } - } - } - // Heavier, smaller slots come first fn sort_by_stake_weight_slot(slot_stake_voted: &mut [(Slot, u64)]) { slot_stake_voted.sort_by(|(slot, stake_voted), (slot_, stake_voted_)| { @@ -716,9 +814,13 @@ mod test { // Try to add a vote for slot < root and a slot that is unrooted for old_slot in &[2, 4] { if *old_slot > root { - assert!(repair_weight.unrooted_slots.contains(old_slot)); + assert!(repair_weight + .slot_to_tree + .get(old_slot) + .unwrap() + .is_pruned()); } else { - assert!(!repair_weight.unrooted_slots.contains(old_slot)); + assert!(!repair_weight.slot_to_tree.contains(old_slot)); } let votes = vec![(*old_slot, vec![Pubkey::default()])]; repair_weight.add_votes( @@ -727,8 +829,17 @@ mod test { bank.epoch_stakes_map(), bank.epoch_schedule(), ); - assert!(!repair_weight.trees.contains_key(old_slot)); - assert!(!repair_weight.slot_to_tree.contains_key(old_slot)); + if *old_slot > root { + assert!(repair_weight.pruned_trees.contains_key(old_slot)); + assert!(repair_weight + .slot_to_tree + .get(old_slot) + .unwrap() + .is_pruned()); + } else { + assert!(!repair_weight.trees.contains_key(old_slot)); + assert!(!repair_weight.slot_to_tree.contains_key(old_slot)); + } } } @@ -758,7 +869,10 @@ mod test { vec![(0, Hash::default())] ); for i in &[0, 1] { - assert_eq!(*repair_weight.slot_to_tree.get(i).unwrap(), 0); + assert_eq!( + *repair_weight.slot_to_tree.get(i).unwrap(), + TreeRoot::Root(0) + ); } // Add slots 6 and 4 with the same set of pubkeys, @@ -795,7 +909,10 @@ mod test { vec![5, 3, 1, 0] ); for slot in 0..=6 { - assert_eq!(*repair_weight.slot_to_tree.get(&slot).unwrap(), 0); + assert_eq!( + *repair_weight.slot_to_tree.get(&slot).unwrap(), + TreeRoot::Root(0) + ); let stake_voted_at = repair_weight .trees .get(&0) @@ -920,10 +1037,16 @@ mod test { ); for slot in &[8, 10, 11] { - assert_eq!(*repair_weight.slot_to_tree.get(slot).unwrap(), 8); + assert_eq!( + *repair_weight.slot_to_tree.get(slot).unwrap(), + TreeRoot::Root(8) + ); } for slot in 0..=1 { - assert_eq!(*repair_weight.slot_to_tree.get(&slot).unwrap(), 0); + assert_eq!( + *repair_weight.slot_to_tree.get(&slot).unwrap(), + TreeRoot::Root(0) + ); } // Call `update_orphan_ancestors` to resolve orphan @@ -935,7 +1058,10 @@ mod test { ); for slot in &[8, 10, 11] { - assert_eq!(*repair_weight.slot_to_tree.get(slot).unwrap(), 0); + assert_eq!( + *repair_weight.slot_to_tree.get(slot).unwrap(), + TreeRoot::Root(0) + ); } assert_eq!(repair_weight.trees.len(), 1); assert!(repair_weight.trees.contains_key(&0)); @@ -1178,12 +1304,18 @@ mod test { // Set root at 1 repair_weight.set_root(1); check_old_root_purged_verify_new_root(0, 1, &repair_weight); - assert!(repair_weight.unrooted_slots.is_empty()); + assert!(repair_weight.pruned_trees.is_empty()); // Other later slots in the fork should be updated to map to the // the new root - assert_eq!(*repair_weight.slot_to_tree.get(&1).unwrap(), 1); - assert_eq!(*repair_weight.slot_to_tree.get(&2).unwrap(), 1); + assert_eq!( + *repair_weight.slot_to_tree.get(&1).unwrap(), + TreeRoot::Root(1) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&2).unwrap(), + TreeRoot::Root(1) + ); // Trees tracked should be updated assert_eq!(repair_weight.trees.get(&1).unwrap().tree_root().0, 1); @@ -1194,7 +1326,10 @@ mod test { repair_weight.trees.get(orphan).unwrap().tree_root().0, *orphan ); - assert_eq!(repair_weight.slot_to_tree.get(orphan).unwrap(), orphan); + assert_eq!( + *repair_weight.slot_to_tree.get(orphan).unwrap(), + TreeRoot::Root(*orphan) + ); } } @@ -1221,7 +1356,10 @@ mod test { repair_weight.trees.get(orphan).unwrap().tree_root().0, *orphan ); - assert_eq!(repair_weight.slot_to_tree.get(orphan).unwrap(), orphan); + assert_eq!( + *repair_weight.slot_to_tree.get(orphan).unwrap(), + TreeRoot::Root(*orphan) + ); } } @@ -1243,11 +1381,14 @@ mod test { // Other higher orphan branch rooted at slot `20` remains unchanged assert_eq!(repair_weight.trees.get(&20).unwrap().tree_root().0, 20); - assert_eq!(*repair_weight.slot_to_tree.get(&20).unwrap(), 20); + assert_eq!( + *repair_weight.slot_to_tree.get(&20).unwrap(), + TreeRoot::Root(20) + ); } #[test] - fn test_set_root_check_unrooted_slots() { + fn test_set_root_check_pruned_slots() { let (blockstore, bank, mut repair_weight) = setup_orphan_repair_weight(); // Chain orphan 8 back to the main fork, but don't @@ -1266,59 +1407,84 @@ mod test { repair_weight.set_root(3); check_old_root_purged_verify_new_root(0, 3, &repair_weight); - // Setting root at 3 should purge all descendants of slot `2`, and - // add any purged slots > 3 to the the `unrooted` set - let purged_slots = vec![0, 1, 2, 4, 8, 10]; - let mut expected_unrooted_len = 0; - for purged_slot in &purged_slots { - assert!(!repair_weight.slot_to_tree.contains_key(purged_slot)); - assert!(!repair_weight.trees.contains_key(purged_slot)); - if *purged_slot > 3 { - assert!(repair_weight.unrooted_slots.contains(purged_slot)); - expected_unrooted_len += 1; - } + // Setting root at 3 should purge all slots < 3 and prune slots > 3 that are not part of + // the rooted path. Additionally slot 4 should now be a standalone pruned tree + for purged_slot in 0..3 { + assert!(!repair_weight.slot_to_tree.contains_key(&purged_slot)); + assert!(!repair_weight.trees.contains_key(&purged_slot)); } - - assert_eq!(repair_weight.unrooted_slots.len(), expected_unrooted_len); + for pruned_slot in &[4, 8, 10] { + assert!(repair_weight + .slot_to_tree + .get(pruned_slot) + .unwrap() + .is_pruned()); + } + assert_eq!( + repair_weight.pruned_trees.keys().copied().collect_vec(), + vec![4] + ); // Orphan 20 should still exist assert_eq!(repair_weight.trees.len(), 2); assert_eq!(repair_weight.trees.get(&20).unwrap().tree_root().0, 20); - assert_eq!(*repair_weight.slot_to_tree.get(&20).unwrap(), 20); + assert_eq!( + *repair_weight.slot_to_tree.get(&20).unwrap(), + TreeRoot::Root(20) + ); // Now set root at a slot 30 that doesnt exist in `repair_weight`, but is // higher than the remaining orphan assert!(!repair_weight.slot_to_tree.contains_key(&30)); repair_weight.set_root(30); check_old_root_purged_verify_new_root(3, 30, &repair_weight); - assert_eq!(repair_weight.trees.len(), 1); - // Orphan 20 should be purged. All other old slots in `unrooted_slots` - // should have been purged - assert!(repair_weight.unrooted_slots.is_empty()); + // Previously pruned tree should now also be purged + assert_eq!(repair_weight.pruned_trees.len(), 0); // Trees tracked should be updated assert_eq!(repair_weight.trees.len(), 1); + assert_eq!(repair_weight.slot_to_tree.len(), 1); + assert_eq!(repair_weight.root, 30); + assert_eq!(repair_weight.trees.get(&30).unwrap().tree_root().0, 30); } #[test] fn test_add_votes_update_orphans_unrooted() { let root = 3; // Test chaining back to slots that were purged when the root 3 was set. - // Two cases: - // 1) Slot 2 < root should not be in the `unrooted_slots` set - // 2) Slot 4 > root should be in the `unrooted_slots` set - // Both cases should purge any branches that chain to them + // 1) Slot 2 < root should be purged and subsequent children purged + // 2) Slot 4 > root should be pruned and subsequent children pruned for old_parent in &[2, 4] { let (blockstore, bank, mut repair_weight) = setup_orphan_repair_weight(); // Set a root at 3 repair_weight.set_root(root); // Check + assert_eq!(repair_weight.pruned_trees.len(), 1); if *old_parent > root { - assert!(repair_weight.unrooted_slots.contains(old_parent)); + assert!(repair_weight + .slot_to_tree + .get(old_parent) + .unwrap() + .is_pruned()); + assert_eq!( + repair_weight + .pruned_trees + .get(old_parent) + .unwrap() + .tree_root() + .0, + *old_parent + ); + assert_eq!( + *repair_weight.slot_to_tree.get(old_parent).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); } else { - assert!(!repair_weight.unrooted_slots.contains(old_parent)); + assert!(!repair_weight.slot_to_tree.contains_key(old_parent)); + assert!(!repair_weight.trees.contains_key(old_parent)); + assert!(!repair_weight.pruned_trees.contains_key(old_parent)); } // Chain orphan 8 back to slot `old_parent` @@ -1335,20 +1501,74 @@ mod test { bank.epoch_schedule(), ); - // Should mark entire branch as unrooted and purge them - for purged_slot in &[*old_parent, 8, 20] { - assert!(!repair_weight.slot_to_tree.contains_key(purged_slot)); - assert!(!repair_weight.slot_to_tree.contains_key(purged_slot)); - if *purged_slot > root { - assert!(repair_weight.unrooted_slots.contains(purged_slot)); - assert!(repair_weight.unrooted_slots.contains(purged_slot)); + // Should prune entire branch + for slot in &[*old_parent, 8, 10, 20] { + if *slot > root { + assert!(repair_weight.slot_to_tree.get(slot).unwrap().is_pruned()); } else { - assert!(!repair_weight.unrooted_slots.contains(purged_slot)); - assert!(!repair_weight.unrooted_slots.contains(purged_slot)); + assert!(!repair_weight.slot_to_tree.contains_key(slot)); } } + if *old_parent > root { + assert_eq!(repair_weight.pruned_trees.len(), 1); + + assert_eq!( + repair_weight + .pruned_trees + .get(old_parent) + .unwrap() + .tree_root() + .0, + *old_parent + ); + assert_eq!( + *repair_weight.slot_to_tree.get(old_parent).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&8).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&10).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&20).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); + } else { + // Should create a new pruned tree alongside 4 + assert_eq!(repair_weight.pruned_trees.len(), 2); + + assert_eq!(repair_weight.pruned_trees.get(&4).unwrap().tree_root().0, 4); + assert_eq!( + *repair_weight.slot_to_tree.get(&4).unwrap(), + TreeRoot::PrunedRoot(4) + ); + + assert_eq!(repair_weight.pruned_trees.get(&8).unwrap().tree_root().0, 8); + assert_eq!( + *repair_weight.slot_to_tree.get(&8).unwrap(), + TreeRoot::PrunedRoot(8) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&10).unwrap(), + TreeRoot::PrunedRoot(8) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&20).unwrap(), + TreeRoot::PrunedRoot(8) + ); + + // Old parent remains missing + assert!(!repair_weight.slot_to_tree.contains_key(old_parent)); + assert!(!repair_weight.trees.contains_key(old_parent)); + assert!(!repair_weight.pruned_trees.contains_key(old_parent)); + } - // Add a vote that chains back to `old_parent`, should be purged + // Add a vote that chains back to `old_parent`, should be added to appropriate pruned + // tree let new_vote_slot = 100; blockstore.add_tree( tr(*old_parent) / tr(new_vote_slot), @@ -1364,50 +1584,115 @@ mod test { bank.epoch_schedule(), ); - assert!(!repair_weight.slot_to_tree.contains_key(&new_vote_slot)); - // Because new_vote_slot > root (otherwise vote isn't processsed) - assert!(repair_weight.unrooted_slots.contains(&new_vote_slot)); + assert!(repair_weight + .slot_to_tree + .get(&new_vote_slot) + .unwrap() + .is_pruned()); + if *old_parent > root { + // Adds to new tree + assert_eq!(repair_weight.pruned_trees.len(), 1); + assert_eq!( + repair_weight + .pruned_trees + .get(old_parent) + .unwrap() + .tree_root() + .0, + *old_parent + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&new_vote_slot).unwrap(), + TreeRoot::PrunedRoot(*old_parent) + ); + } else { + // Creates new tree + assert_eq!(repair_weight.pruned_trees.len(), 3); + + assert_eq!(repair_weight.pruned_trees.get(&4).unwrap().tree_root().0, 4); + assert_eq!( + *repair_weight.slot_to_tree.get(&4).unwrap(), + TreeRoot::PrunedRoot(4) + ); + + assert_eq!(repair_weight.pruned_trees.get(&8).unwrap().tree_root().0, 8); + assert_eq!( + *repair_weight.slot_to_tree.get(&8).unwrap(), + TreeRoot::PrunedRoot(8) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&10).unwrap(), + TreeRoot::PrunedRoot(8) + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&20).unwrap(), + TreeRoot::PrunedRoot(8) + ); + + assert_eq!( + repair_weight + .pruned_trees + .get(&new_vote_slot) + .unwrap() + .tree_root() + .0, + new_vote_slot + ); + assert_eq!( + *repair_weight.slot_to_tree.get(&new_vote_slot).unwrap(), + TreeRoot::PrunedRoot(new_vote_slot) + ); + + // Old parent remains missing + assert!(!repair_weight.slot_to_tree.contains_key(old_parent)); + assert!(!repair_weight.trees.contains_key(old_parent)); + assert!(!repair_weight.pruned_trees.contains_key(old_parent)); + } + + // TODO: add more purge scenarios from here / add a new test } } #[test] fn test_find_ancestor_subtree_of_slot() { let (blockstore, _, mut repair_weight) = setup_orphan_repair_weight(); + // TODO: add more PurgedRoot scenario // Ancestor of slot 4 is slot 2, with an existing subtree rooted at 0 // because there wass a vote for a descendant assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 4), - (vec![2].into_iter().collect::>(), Some(0)) + (VecDeque::from([2]), Some(TreeRoot::Root(0))) ); // Ancestors of 5 are [1, 3], with an existing subtree rooted at 0 // because there wass a vote for a descendant assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 5), - (vec![1, 3].into_iter().collect::>(), Some(0)) + (VecDeque::from([1, 3]), Some(TreeRoot::Root(0))) ); // Ancestors of slot 23 are [20, 22], with an existing subtree of 20 // because there wass a vote for 20 assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 23), - (vec![20, 22].into_iter().collect::>(), Some(20)) + (VecDeque::from([20, 22]), Some(TreeRoot::Root(20))) ); - // Add 22 to `unrooted_slots`, ancestor search should now - // terminate early and return no valid existing subtree - repair_weight.unrooted_slots.insert(22); + // Add 22 to `pruned_trees`, ancestor search should now + // chain it back to 20 + repair_weight.remove_tree(&TreeRoot::Root(20)).unwrap(); + repair_weight.insert_new_pruned_tree(20); assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 23), - (vec![22].into_iter().collect::>(), None) + (VecDeque::from([20, 22]), Some(TreeRoot::PrunedRoot(20))) ); // Ancestors of slot 31 are [30], with no existing subtree blockstore.add_tree(tr(30) / (tr(31)), true, true, 2, Hash::default()); assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 31), - (vec![30].into_iter().collect::>(), None) + (VecDeque::from([30]), None), ); // Set a root at 5 @@ -1417,15 +1702,16 @@ mod test { // at slot 5 assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 6), - (vec![5].into_iter().collect::>(), Some(5)) + (VecDeque::from([5]), Some(TreeRoot::Root(5))) ); // Chain orphan 8 back to slot 4 on a different fork, ancestor search // should not return ancestors earlier than the root + // TODO: original behavior does nto match this comment should we fix blockstore.add_tree(tr(4) / (tr(8)), true, true, 2, Hash::default()); assert_eq!( repair_weight.find_ancestor_subtree_of_slot(&blockstore, 8), - (vec![4].into_iter().collect::>(), None) + (VecDeque::from([4]), None), ); } @@ -1594,7 +1880,10 @@ mod test { repair_weight.trees.get(orphan).unwrap().tree_root().0, *orphan ); - assert_eq!(repair_weight.slot_to_tree.get(orphan).unwrap(), orphan); + assert_eq!( + *repair_weight.slot_to_tree.get(orphan).unwrap(), + TreeRoot::Root(*orphan) + ); } (blockstore, bank, repair_weight) } @@ -1607,7 +1896,6 @@ mod test { // Check old root is purged assert!(!repair_weight.trees.contains_key(&old_root)); assert!(!repair_weight.slot_to_tree.contains_key(&old_root)); - assert!(!repair_weight.unrooted_slots.contains(&old_root)); // Validate new root assert_eq!( @@ -1616,7 +1904,7 @@ mod test { ); assert_eq!( *repair_weight.slot_to_tree.get(&new_root).unwrap(), - new_root + TreeRoot::Root(new_root) ); assert_eq!(repair_weight.root, new_root); }