Find max root and purge roots below it. (#4645)

* Test for forking accounts

* Find max root and purge roots below it.

runtime/src/accounts_db.rs

@@ -1219,14 +1219,7 @@ mod tests {
             .clone();
         //fork 0 is behind root, but it is not root, therefore it is purged
         accounts.add_root(1);
-        {
+        assert!(accounts.accounts_index.read().unwrap().is_purged(0));
-            let accounts_index = accounts.accounts_index.read().unwrap();
-            assert!(AccountsIndex::<AccountInfo>::is_purged(
-                &accounts_index.roots,
-                accounts_index.last_root,
-                0
-            ));
-        }
 
         //fork is still there, since gc is lazy
         assert!(accounts.storage.read().unwrap().0[&0]

runtime/src/accounts_index.rs

@@ -39,6 +39,16 @@ impl<T: Clone> AccountsIndex<T> {
         rv
     }
 
+    pub fn get_max_root(roots: &HashSet<Fork>, fork_vec: &[(Fork, T)]) -> Fork {
+        let mut max_root = 0;
+        for (f, _) in fork_vec.iter() {
+            if *f > max_root && roots.contains(f) {
+                max_root = *f;
+            }
+        }
+        max_root
+    }
+
     pub fn insert(
         &mut self,
         fork: Fork,
@@ -46,7 +56,6 @@ impl<T: Clone> AccountsIndex<T> {
         account_info: T,
         reclaims: &mut Vec<(Fork, T)>,
     ) {
-        let last_root = self.last_root;
         let roots = &self.roots;
         let fork_vec = self
             .account_maps
@@ -60,13 +69,15 @@ impl<T: Clone> AccountsIndex<T> {
         // add the new entry
         fork_vec.push((fork, account_info));
 
+        let max_root = Self::get_max_root(roots, fork_vec);
+
         reclaims.extend(
             fork_vec
                 .iter()
-                .filter(|(fork, _)| Self::is_purged(roots, last_root, *fork))
+                .filter(|(fork, _)| Self::can_purge(max_root, *fork))
                 .cloned(),
         );
-        fork_vec.retain(|(fork, _)| !Self::is_purged(roots, last_root, *fork));
+        fork_vec.retain(|(fork, _)| !Self::can_purge(max_root, *fork));
     }
 
     pub fn add_index(&mut self, fork: Fork, pubkey: &Pubkey, account_info: T) {
@@ -74,8 +85,12 @@ impl<T: Clone> AccountsIndex<T> {
         entry.push((fork, account_info));
     }
 
-    pub fn is_purged(roots: &HashSet<Fork>, last_root: Fork, fork: Fork) -> bool {
+    pub fn is_purged(&self, fork: Fork) -> bool {
-        !roots.contains(&fork) && fork < last_root
+        fork < self.last_root
+    }
+
+    pub fn can_purge(max_root: Fork, fork: Fork) -> bool {
+        fork < max_root
     }
 
     pub fn is_root(&self, fork: Fork) -> bool {
@@ -170,17 +185,9 @@ mod tests {
     #[test]
     fn test_is_purged() {
         let mut index = AccountsIndex::<bool>::default();
-        assert!(!AccountsIndex::<bool>::is_purged(
+        assert!(!index.is_purged(0));
-            &index.roots,
-            index.last_root,
-            0
-        ));
         index.add_root(1);
-        assert!(AccountsIndex::<bool>::is_purged(
+        assert!(index.is_purged(0));
-            &index.roots,
-            index.last_root,
-            0
-        ));
     }
 
     #[test]
@@ -257,10 +264,15 @@ mod tests {
         let mut gc = Vec::new();
         index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
-        index.add_root(1);
         index.insert(1, &key.pubkey(), false, &mut gc);
-        assert_eq!(gc, vec![(0, true)]);
+        index.insert(2, &key.pubkey(), true, &mut gc);
+        index.insert(3, &key.pubkey(), true, &mut gc);
+        index.add_root(0);
+        index.add_root(1);
+        index.add_root(3);
+        index.insert(4, &key.pubkey(), true, &mut gc);
+        assert_eq!(gc, vec![(0, true), (1, false), (2, true)]);
         let ancestors = vec![].into_iter().collect();
-        assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&false, 1)));
+        assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&true, 3)));
     }
 }

runtime/src/bank.rs

@@ -2046,6 +2046,49 @@ mod tests {
         assert_eq!(parent.get_balance(&key1.pubkey()), 1);
     }
 
+    #[test]
+    fn test_bank_get_account_in_parent_after_squash2() {
+        solana_logger::setup();
+        let (genesis_block, mint_keypair) = create_genesis_block(500);
+        let bank0 = Arc::new(Bank::new(&genesis_block));
+
+        let key1 = Keypair::new();
+
+        bank0.transfer(1, &mint_keypair, &key1.pubkey()).unwrap();
+        assert_eq!(bank0.get_balance(&key1.pubkey()), 1);
+
+        let bank1 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 1));
+        bank1.transfer(3, &mint_keypair, &key1.pubkey()).unwrap();
+        let bank2 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 2));
+        bank2.transfer(2, &mint_keypair, &key1.pubkey()).unwrap();
+        let bank3 = Arc::new(Bank::new_from_parent(&bank1, &Pubkey::default(), 3));
+        bank1.squash();
+
+        assert_eq!(bank0.get_balance(&key1.pubkey()), 1);
+        assert_eq!(bank3.get_balance(&key1.pubkey()), 4);
+        assert_eq!(bank2.get_balance(&key1.pubkey()), 3);
+        bank3.squash();
+        assert_eq!(bank1.get_balance(&key1.pubkey()), 4);
+
+        let bank4 = Arc::new(Bank::new_from_parent(&bank3, &Pubkey::default(), 4));
+        bank4.transfer(4, &mint_keypair, &key1.pubkey()).unwrap();
+        assert_eq!(bank4.get_balance(&key1.pubkey()), 8);
+        assert_eq!(bank3.get_balance(&key1.pubkey()), 4);
+        bank4.squash();
+        let bank5 = Arc::new(Bank::new_from_parent(&bank4, &Pubkey::default(), 5));
+        bank5.squash();
+        let bank6 = Arc::new(Bank::new_from_parent(&bank5, &Pubkey::default(), 6));
+        bank6.squash();
+
+        // This picks up the values from 4 which is the highest root:
+        // TODO: if we need to access rooted banks older than this,
+        // need to fix the lookup.
+        assert_eq!(bank3.get_balance(&key1.pubkey()), 8);
+        assert_eq!(bank2.get_balance(&key1.pubkey()), 8);
+
+        assert_eq!(bank4.get_balance(&key1.pubkey()), 8);
+    }
+
     #[test]
     fn test_bank_epoch_vote_accounts() {
         let leader_pubkey = Pubkey::new_rand();