Fix getProgramAccounts RPC (#5024)

* Use scan_accounts to load accounts by program_id

* Add bank test

* Use get_program_accounts in RPC

core/src/rpc.rs

@@ -73,7 +73,7 @@ impl JsonRpcRequestProcessor {
     pub fn get_program_accounts(&self, program_id: &Pubkey) -> Result<Vec<(String, Account)>> {
         Ok(self
             .bank()
-            .get_program_accounts_modified_since_parent(&program_id)
+            .get_program_accounts(&program_id)
             .into_iter()
             .map(|(pubkey, account)| (pubkey.to_string(), account))
             .collect())

runtime/src/accounts.rs

@@ -338,7 +338,7 @@ impl Accounts {
             .collect()
     }
 
-    pub fn load_by_program(&self, fork: Fork, program_id: &Pubkey) -> Vec<(Pubkey, Account)> {
+    pub fn load_by_program_fork(&self, fork: Fork, program_id: &Pubkey) -> Vec<(Pubkey, Account)> {
         self.scan_fork(fork, |stored_account| {
             if stored_account.balance.owner == *program_id {
                 Some((stored_account.meta.pubkey, stored_account.clone_account()))
@@ -348,6 +348,24 @@ impl Accounts {
         })
     }
 
+    pub fn load_by_program(
+        &self,
+        ancestors: &HashMap<Fork, usize>,
+        program_id: &Pubkey,
+    ) -> Vec<(Pubkey, Account)> {
+        self.accounts_db.scan_accounts(
+            ancestors,
+            |collector: &mut Vec<(Pubkey, Account)>, option| {
+                if let Some(data) = option
+                    .filter(|(_, account, _)| account.owner == *program_id && account.lamports != 0)
+                    .map(|(pubkey, account, _fork)| (*pubkey, account))
+                {
+                    collector.push(data)
+                }
+            },
+        )
+    }
+
     /// Slow because lock is held for 1 operation instead of many
     pub fn store_slow(&self, fork: Fork, pubkey: &Pubkey, account: &Account) {
         let mut accounts = HashMap::new();
@@ -1050,7 +1068,7 @@ mod tests {
     }
 
     #[test]
-    fn test_load_by_program() {
+    fn test_load_by_program_fork() {
         let accounts = Accounts::new(None);
 
         // Load accounts owned by various programs into AccountsDB
@@ -1064,11 +1082,11 @@ mod tests {
         let account2 = Account::new(1, 0, &Pubkey::new(&[3; 32]));
         accounts.store_slow(0, &pubkey2, &account2);
 
-        let loaded = accounts.load_by_program(0, &Pubkey::new(&[2; 32]));
+        let loaded = accounts.load_by_program_fork(0, &Pubkey::new(&[2; 32]));
         assert_eq!(loaded.len(), 2);
-        let loaded = accounts.load_by_program(0, &Pubkey::new(&[3; 32]));
+        let loaded = accounts.load_by_program_fork(0, &Pubkey::new(&[3; 32]));
         assert_eq!(loaded, vec![(pubkey2, account2)]);
-        let loaded = accounts.load_by_program(0, &Pubkey::new(&[4; 32]));
+        let loaded = accounts.load_by_program_fork(0, &Pubkey::new(&[4; 32]));
         assert_eq!(loaded, vec![]);
     }
 

runtime/src/bank.rs

@@ -1206,11 +1206,19 @@ impl Bank {
             .map(|(account, _)| account)
     }
 
+    pub fn get_program_accounts(&self, program_id: &Pubkey) -> Vec<(Pubkey, Account)> {
+        self.rc
+            .accounts
+            .load_by_program(&self.ancestors, program_id)
+    }
+
     pub fn get_program_accounts_modified_since_parent(
         &self,
         program_id: &Pubkey,
     ) -> Vec<(Pubkey, Account)> {
-        self.rc.accounts.load_by_program(self.slot(), program_id)
+        self.rc
+            .accounts
+            .load_by_program_fork(self.slot(), program_id)
     }
 
     pub fn get_account_modified_since_parent(&self, pubkey: &Pubkey) -> Option<(Account, Fork)> {
@@ -2731,4 +2739,51 @@ mod tests {
             (1.0, 1.0)
         );
     }
+
+    #[test]
+    fn test_bank_get_program_accounts() {
+        let (genesis_block, _mint_keypair) = create_genesis_block(500);
+        let parent = Arc::new(Bank::new(&genesis_block));
+
+        let bank0 = Arc::new(new_from_parent(&parent));
+
+        let pubkey0 = Pubkey::new_rand();
+        let program_id = Pubkey::new(&[2; 32]);
+        let account0 = Account::new(1, 0, &program_id);
+        bank0.store_account(&pubkey0, &account0);
+
+        assert_eq!(
+            bank0.get_program_accounts_modified_since_parent(&program_id),
+            vec![(pubkey0, account0.clone())]
+        );
+
+        let bank1 = Arc::new(new_from_parent(&bank0));
+        bank1.squash();
+        assert_eq!(
+            bank0.get_program_accounts(&program_id),
+            vec![(pubkey0, account0.clone())]
+        );
+        assert_eq!(
+            bank1.get_program_accounts(&program_id),
+            vec![(pubkey0, account0.clone())]
+        );
+        assert_eq!(
+            bank1.get_program_accounts_modified_since_parent(&program_id),
+            vec![]
+        );
+
+        let bank2 = Arc::new(new_from_parent(&bank1));
+        let pubkey1 = Pubkey::new_rand();
+        let account1 = Account::new(3, 0, &program_id);
+        bank2.store_account(&pubkey1, &account1);
+        // Accounts with 0 lamports should be filtered out by Accounts::load_by_program()
+        let pubkey2 = Pubkey::new_rand();
+        let account2 = Account::new(0, 0, &program_id);
+        bank2.store_account(&pubkey2, &account2);
+
+        let bank3 = Arc::new(new_from_parent(&bank2));
+        bank3.squash();
+        assert_eq!(bank1.get_program_accounts(&program_id).len(), 2);
+        assert_eq!(bank3.get_program_accounts(&program_id).len(), 2);
+    }
 }