Keep block height with state roots to be pruned

Keep the block height with the state root to be pruned: This prevents
roots from being pruned if there is a lot of fork swapping, and also
allows the abandoned chains' roots to be pruned at the prune depth as
well.

This required the state_root_pruning_queue to be made a binary heap for
efficient access.

Additionally:

- renamed `decimate` to `execute`, which now takes an `at_depth`
parameter

Signed-off-by: Peter Schwarz <pschwarz@bitwise.io>

Author: peterschwarz

validator/src/journal/chain.rs

@@ -283,7 +283,7 @@ impl<BC: BlockCache + 'static, BV: BlockValidator + 'static, CW: ChainWriter + '
                 &result
                     .current_chain
                     .iter()
-                    .map(|block| block.state_root_hash())
+                    .map(|block| (block.block_num(), block.state_root_hash()))
                     .collect::<Vec<_>>(),
             );
 
@@ -359,16 +359,17 @@ impl<BC: BlockCache + 'static, BV: BlockValidator + 'static, CW: ChainWriter + '
 
             let chain_head_block_num = state.chain_head.as_ref().unwrap().block_num();
             if chain_head_block_num + 1 > self.state_pruning_block_depth as u64 {
-                let prune_at =  chain_head_block_num - (self.state_pruning_block_depth as u64);
+                let prune_at = chain_head_block_num - (self.state_pruning_block_depth as u64);
                 match state.chain_reader.get_block_by_block_num(prune_at) {
-                    Ok(Some(block)) => state.state_pruning_manager.add_to_queue(block.state_root_hash()),
+                    Ok(Some(block)) => state.state_pruning_manager.add_to_queue(block.block_num(),block.state_root_hash()),
                     Ok(None) => warn!("No block at block height {}; ignoring...", prune_at),
                     Err(err) => error!("Unable to fetch block at height {}: {:?}", prune_at, err)
                 }
+
+                // Execute pruning:
+                state.state_pruning_manager.execute(prune_at)
             }
 
-            // Execute pruning:
-            state.state_pruning_manager.decimate()
         } else {
             info!("Rejected new chain head: {}", new_block);
         }

validator/src/state/state_pruning_manager.rs

@@ -14,7 +14,8 @@
  * limitations under the License.
  * ------------------------------------------------------------------------------
  */
-use std::collections::VecDeque;
+use std::cmp::Ordering;
+use std::collections::BinaryHeap;
 
 use database::lmdb::LmdbDatabase;
 use metrics;
@@ -33,76 +34,105 @@ lazy_static! {
 /// abandoned and then re-chosen as the primary chain.
 pub struct StatePruningManager {
     // Contains the state root hashes slated for pruning
-    state_root_prune_queue: VecDeque<String>,
+    state_root_prune_queue: BinaryHeap<PruneCandidate>,
     state_database: LmdbDatabase,
 }
 
+#[derive(Eq, PartialEq, Debug, Ord)]
+struct PruneCandidate(u64, String);
+
+impl PartialOrd for PruneCandidate {
+    fn partial_cmp(&self, other: &PruneCandidate) -> Option<Ordering> {
+        Some(Ordering::reverse(self.0.cmp(&other.0)))
+    }
+}
+
 impl StatePruningManager {
     pub fn new(state_database: LmdbDatabase) -> Self {
         StatePruningManager {
-            state_root_prune_queue: VecDeque::new(),
+            state_root_prune_queue: BinaryHeap::new(),
             state_database,
         }
     }
 
     /// Updates the pruning queue.  Abandoned roots will be added to the queue.
     /// Added roots will be removed from the queue.  This ensures that the state
     /// roots won't be removed, regardless of the chain state.
-    pub fn update_queue(&mut self, added_roots: &[&str], _abandoned_roots: &[&str]) {
+    pub fn update_queue(&mut self, added_roots: &[&str], abandoned_roots: &[(u64, &str)]) {
+        // add the roots that have been abandoned.
+        for (height, state_root_hash) in abandoned_roots {
+            self.add_to_queue(*height, state_root_hash);
+        }
         // Remove any state root hashes from the pruning queue that we may have switched
         // back too from an alternate chain
-        added_roots.iter().for_each(|state_root_hash| {
-            self.state_root_prune_queue.retain(|hash| {
-                if hash != state_root_hash {
+        let mut new_queue = BinaryHeap::with_capacity(0);
+        ::std::mem::swap(&mut self.state_root_prune_queue, &mut new_queue);
+        self.state_root_prune_queue = new_queue
+            .into_iter()
+            .filter(|candidate| {
+                if !added_roots.contains(&candidate.1.as_str()) {
                     true
                 } else {
-                    debug!("Removing {} from pruning queue", state_root_hash);
+                    debug!("Removing {} from pruning queue", candidate.1);
                     false
                 }
             })
-        });
+            .collect();
     }
 
     /// Add a single state root to the pruning queue.
-    pub fn add_to_queue(&mut self, state_root_hash: &str) {
+    pub fn add_to_queue(&mut self, height: u64, state_root_hash: &str) {
         let state_root_hash = state_root_hash.into();
-        if !self.state_root_prune_queue.contains(&state_root_hash) {
+        if !self.state_root_prune_queue
+            .iter()
+            .any(|candidate| candidate.1 == state_root_hash)
+        {
             debug!("Adding {} to pruning queue", state_root_hash);
-            self.state_root_prune_queue.push_back(state_root_hash);
+            self.state_root_prune_queue
+                .push(PruneCandidate(height, state_root_hash));
         }
     }
 
-    /// Remove the first ten percent (floored) of the prune queue.
-    pub fn decimate(&mut self) {
-        let prune_count = (self.state_root_prune_queue.len() as f64 / 10.0).floor() as usize;
+    /// Executes prune on any state root hash at or below the given depth.
+    pub fn execute(&mut self, at_depth: u64) {
+        let mut prune_candidates = vec![];
 
-        let mut total_pruned_entries: usize = 0;
-        for _ in 0..prune_count {
-            if let Some(state_root) = self.state_root_prune_queue.pop_front() {
-                match MerkleDatabase::prune(&self.state_database, &state_root) {
-                    Ok(removed_keys) => {
-                        total_pruned_entries += removed_keys.len();
-
-                        let mut state_roots_pruned_count =
-                            COLLECTOR.counter("StatePruneManager.state_roots_pruned", None, None);
-                        state_roots_pruned_count.inc();
-
-                        // the state root was not pruned (it is likely the root of a
-                        // fork), so push it back into the queue.
-                        if removed_keys.is_empty() {
-                            self.state_root_prune_queue.push_back(state_root);
-                        }
-                    }
-                    Err(err) => {
-                        error!("Unable to prune state root {}: {:?}", state_root, err);
-                        self.state_root_prune_queue.push_back(state_root);
-                    }
+        loop {
+            if let Some(candidate) = self.state_root_prune_queue.pop() {
+                if candidate.0 <= at_depth {
+                    prune_candidates.push(candidate);
+                } else {
+                    self.state_root_prune_queue.push(candidate);
+                    break;
                 }
             } else {
                 break;
             }
         }
 
+        let mut total_pruned_entries: usize = 0;
+        for candidate in prune_candidates {
+            match MerkleDatabase::prune(&self.state_database, &candidate.1) {
+                Ok(removed_keys) => {
+                    total_pruned_entries += removed_keys.len();
+
+                    let mut state_roots_pruned_count =
+                        COLLECTOR.counter("StatePruneManager.state_roots_pruned", None, None);
+                    state_roots_pruned_count.inc();
+
+                    // the state root was not pruned (it is likely the root of a
+                    // fork), so push it back into the queue.
+                    if removed_keys.is_empty() {
+                        self.state_root_prune_queue.push(candidate);
+                    }
+                }
+                Err(err) => {
+                    error!("Unable to prune state root {}: {:?}", candidate.1, err);
+                    self.state_root_prune_queue.push(candidate);
+                }
+            }
+        }
+
         if total_pruned_entries > 0 {
             info!(
                 "Pruned {} keys from the Global state Database",
@@ -111,3 +141,24 @@ impl StatePruningManager {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn ordering_candidates() {
+        let mut heap = ::std::collections::BinaryHeap::new();
+
+        heap.push(PruneCandidate(2, "two".into()));
+        heap.push(PruneCandidate(3, "three".into()));
+        heap.push(PruneCandidate(3, "another_three".into()));
+        heap.push(PruneCandidate(4, "four".into()));
+
+        assert_eq!(heap.pop(), Some(PruneCandidate(2, "two".into())));
+        assert_eq!(heap.pop(), Some(PruneCandidate(3, "another_three".into())));
+        assert_eq!(heap.pop(), Some(PruneCandidate(3, "three".into())));
+        assert_eq!(heap.pop(), Some(PruneCandidate(4, "four".into())));
+        assert_eq!(heap.pop(), None);
+    }
+}