Add a storage accounts cache to Bank (#4578)

runtime/Cargo.toml

@@ -29,6 +29,7 @@ solana-noop-program = { path = "../programs/noop_program", version = "0.16.0" }
 solana-sdk = { path = "../sdk", version = "0.16.0" }
 solana-stake-api = { path = "../programs/stake_api", version = "0.16.0" }
 solana-stake-program = { path = "../programs/stake_program", version = "0.16.0" }
+solana-storage-api = { path = "../programs/storage_api", version = "0.16.0" }
 solana-vote-api = { path = "../programs/vote_api", version = "0.16.0" }
 solana-vote-program = { path = "../programs/vote_program", version = "0.16.0" }
 sys-info = "0.5.7"

runtime/src/bank.rs

@@ -14,6 +14,8 @@ use crate::serde_utils::{
 };
 use crate::stakes::Stakes;
 use crate::status_cache::StatusCache;
+use crate::storage_utils;
+use crate::storage_utils::StorageAccounts;
 use bincode::{deserialize_from, serialize, serialize_into, serialized_size};
 use log::*;
 use serde::{Deserialize, Serialize};
@@ -232,6 +234,9 @@ pub struct Bank {
     /// cache of vote_account and stake_account state for this fork
     stakes: RwLock<Stakes>,
 
+    /// cache of validator and replicator storage accounts for this fork
+    storage_accounts: RwLock<StorageAccounts>,
+
     /// staked nodes on epoch boundaries, saved off when a bank.slot() is at
     ///   a leader schedule calculation boundary
     epoch_stakes: HashMap<u64, Stakes>,
@@ -287,6 +292,7 @@ impl Bank {
         bank.transaction_count
             .store(parent.transaction_count() as usize, Ordering::Relaxed);
         bank.stakes = RwLock::new(parent.stakes.read().unwrap().clone());
+        bank.storage_accounts = RwLock::new(parent.storage_accounts.read().unwrap().clone());
 
         bank.tick_height
             .store(parent.tick_height.load(Ordering::SeqCst), Ordering::SeqCst);
@@ -936,7 +942,7 @@ impl Bank {
             .accounts
             .store_accounts(self.slot(), txs, executed, loaded_accounts);
 
-        self.store_stakes(txs, executed, loaded_accounts);
+        self.update_cached_accounts(txs, executed, loaded_accounts);
 
         // once committed there is no way to unroll
         let write_elapsed = now.elapsed();
@@ -1005,6 +1011,11 @@ impl Bank {
 
         if Stakes::is_stake(account) {
             self.stakes.write().unwrap().store(pubkey, account);
+        } else if storage_utils::is_storage(account) {
+            self.storage_accounts
+                .write()
+                .unwrap()
+                .store(pubkey, account);
         }
     }
 
@@ -1129,7 +1140,7 @@ impl Bank {
     }
 
     /// a bank-level cache of vote accounts
-    fn store_stakes(
+    fn update_cached_accounts(
         &self,
         txs: &[Transaction],
         res: &[Result<()>],
@@ -1143,17 +1154,31 @@ impl Bank {
             let message = &txs[i].message();
             let acc = raccs.as_ref().unwrap();
 
-            for (pubkey, account) in message
-                .account_keys
-                .iter()
-                .zip(acc.0.iter())
-                .filter(|(_, account)| Stakes::is_stake(account))
+            for (pubkey, account) in
+                message
+                    .account_keys
+                    .iter()
+                    .zip(acc.0.iter())
+                    .filter(|(_, account)| {
+                        (Stakes::is_stake(account)) || storage_utils::is_storage(account)
+                    })
             {
-                self.stakes.write().unwrap().store(pubkey, account);
+                if Stakes::is_stake(account) {
+                    self.stakes.write().unwrap().store(pubkey, account);
+                } else if storage_utils::is_storage(account) {
+                    self.storage_accounts
+                        .write()
+                        .unwrap()
+                        .store(pubkey, account);
+                }
             }
         }
     }
 
+    pub fn storage_accounts(&self) -> StorageAccounts {
+        self.storage_accounts.read().unwrap().clone()
+    }
+
     /// current vote accounts for this bank along with the stake
     ///   attributed to each account
     pub fn vote_accounts(&self) -> HashMap<Pubkey, (u64, Account)> {

runtime/src/lib.rs

@@ -15,6 +15,7 @@ mod native_loader;
 mod serde_utils;
 pub mod stakes;
 pub mod status_cache;
+pub mod storage_utils;
 mod system_instruction_processor;
 
 #[macro_use]

runtime/src/storage_utils.rs

@@ -0,0 +1,116 @@
+use crate::bank::Bank;
+use solana_sdk::account::Account;
+use solana_sdk::account_utils::State;
+use solana_sdk::pubkey::Pubkey;
+use solana_storage_api::storage_contract::StorageContract;
+use std::collections::{HashMap, HashSet};
+
+#[derive(Default, Clone, PartialEq, Debug, Deserialize, Serialize)]
+pub struct StorageAccounts {
+    /// validator storage accounts
+    validator_accounts: HashSet<Pubkey>,
+
+    /// replicator storage accounts
+    replicator_accounts: HashSet<Pubkey>,
+}
+
+pub fn is_storage(account: &Account) -> bool {
+    solana_storage_api::check_id(&account.owner)
+}
+
+impl StorageAccounts {
+    pub fn store(&mut self, pubkey: &Pubkey, account: &Account) {
+        if let Ok(storage_state) = account.state() {
+            if let StorageContract::ReplicatorStorage { .. } = storage_state {
+                if account.lamports == 0 {
+                    self.replicator_accounts.remove(pubkey);
+                } else {
+                    self.replicator_accounts.insert(*pubkey);
+                }
+            } else if let StorageContract::ValidatorStorage { .. } = storage_state {
+                if account.lamports == 0 {
+                    self.validator_accounts.remove(pubkey);
+                } else {
+                    self.validator_accounts.insert(*pubkey);
+                }
+            }
+        };
+    }
+}
+
+pub fn validator_accounts(bank: &Bank) -> HashMap<Pubkey, Account> {
+    bank.storage_accounts()
+        .validator_accounts
+        .iter()
+        .filter_map(|account_id| {
+            bank.get_account(account_id)
+                .and_then(|account| Some((*account_id, account)))
+        })
+        .collect()
+}
+
+pub fn replicator_accounts(bank: &Bank) -> HashMap<Pubkey, Account> {
+    bank.storage_accounts()
+        .replicator_accounts
+        .iter()
+        .filter_map(|account_id| {
+            bank.get_account(account_id)
+                .and_then(|account| Some((*account_id, account)))
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::bank_client::BankClient;
+    use solana_sdk::client::SyncClient;
+    use solana_sdk::genesis_block::create_genesis_block;
+    use solana_sdk::message::Message;
+    use solana_sdk::signature::{Keypair, KeypairUtil};
+    use solana_storage_api::{storage_instruction, storage_processor};
+    use std::sync::Arc;
+
+    #[test]
+    fn test_store_and_recover() {
+        let (genesis_block, mint_keypair) = create_genesis_block(1000);
+        let mint_pubkey = mint_keypair.pubkey();
+        let replicator_keypair = Keypair::new();
+        let replicator_pubkey = replicator_keypair.pubkey();
+        let validator_keypair = Keypair::new();
+        let validator_pubkey = validator_keypair.pubkey();
+        let mut bank = Bank::new(&genesis_block);
+        bank.add_instruction_processor(
+            solana_storage_api::id(),
+            storage_processor::process_instruction,
+        );
+
+        let bank = Arc::new(bank);
+        let bank_client = BankClient::new_shared(&bank);
+
+        bank_client
+            .transfer(10, &mint_keypair, &replicator_pubkey)
+            .unwrap();
+        let message = Message::new(storage_instruction::create_replicator_storage_account(
+            &mint_pubkey,
+            &Pubkey::default(),
+            &replicator_pubkey,
+            1,
+        ));
+        bank_client.send_message(&[&mint_keypair], message).unwrap();
+
+        bank_client
+            .transfer(10, &mint_keypair, &validator_pubkey)
+            .unwrap();
+        let message = Message::new(storage_instruction::create_validator_storage_account(
+            &mint_pubkey,
+            &Pubkey::default(),
+            &validator_pubkey,
+            1,
+        ));
+        bank_client.send_message(&[&mint_keypair], message).unwrap();
+
+        assert_eq!(validator_accounts(bank.as_ref()).len(), 1);
+        assert_eq!(replicator_accounts(bank.as_ref()).len(), 1);
+    }
+}