Adds StorableAccounts::account_for_geyser()

accounts-db/src/accounts.rs

@@ -594,13 +594,12 @@ impl Accounts {
             for index in 0..accounts.len() {
                 let transaction = transactions
                     .map(|txs| *txs.get(index).expect("txs must be present if provided"));
-                accounts.account(index, |account| {
-                    let account_shared_data = account.take_account();
+                accounts.account_for_geyser(index, |pubkey, account_shared_data| {
                     accounts_db.notify_account_at_accounts_update(
                         slot,
-                        &account_shared_data,
+                        account_shared_data,
                         &transaction,
-                        account.pubkey(),
+                        pubkey,
                         current_write_version,
                     );
                 });

accounts-db/src/accounts_db/tests.rs

@@ -59,6 +59,13 @@ where
     ) -> Ret {
         callback(self.1[index].1.into())
     }
+    fn account_for_geyser<Ret>(
+        &self,
+        _index: usize,
+        _callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        unimplemented!();
+    }
     fn pubkey(&self, index: usize) -> &Pubkey {
         self.1[index].0
     }

accounts-db/src/stake_rewards.rs

@@ -46,7 +46,15 @@ impl<'a> StorableAccounts<'a> for (Slot, &'a [StakeReward]) {
         mut callback: impl for<'local> FnMut(AccountForStorage<'local>) -> Ret,
     ) -> Ret {
         let entry = &self.1[index];
-        callback((&self.1[index].stake_pubkey, &entry.stake_account).into())
+        callback((&entry.stake_pubkey, &entry.stake_account).into())
+    }
+    fn account_for_geyser<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        let entry = &self.1[index];
+        callback(&entry.stake_pubkey, &entry.stake_account)
     }
     fn is_zero_lamport(&self, index: usize) -> bool {
         self.1[index].is_zero_lamport()

accounts-db/src/storable_accounts.rs

@@ -115,6 +115,17 @@ pub trait StorableAccounts<'a>: Sync {
         index: usize,
         callback: impl for<'local> FnMut(AccountForStorage<'local>) -> Ret,
     ) -> Ret;
+    /// Geyser account update notifications need a `&AccountSharedData`. When storing into the
+    /// accounts write cache, we should always have an AccountSharedData, so allow access to it
+    /// when available, which is an optimization to avoid creating a new AccountSharedData only to
+    /// immediately take a reference.
+    /// Note: Only implement this fn if underlying type actually holds an AccountSharedData.
+    /// Otherwise mark it unimplemented!().
+    fn account_for_geyser<Ret>(
+        &self,
+        index: usize,
+        callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret;
     /// whether account at 'index' has zero lamports
     fn is_zero_lamport(&self, index: usize) -> bool;
     /// data length of account at 'index'
@@ -163,6 +174,15 @@ impl<'a: 'b, 'b> StorableAccounts<'a> for (Slot, &'b [(&'a Pubkey, &'a AccountSh
     ) -> Ret {
         callback((self.1[index].0, self.1[index].1).into())
     }
+    fn account_for_geyser<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        let pubkey = self.pubkey(index);
+        let account = self.1[index].1;
+        callback(pubkey, account)
+    }
     fn is_zero_lamport(&self, index: usize) -> bool {
         self.1[index].1.is_zero_lamport()
     }
@@ -192,6 +212,15 @@ impl<'a: 'b, 'b> StorableAccounts<'a> for (Slot, &'b [(Pubkey, AccountSharedData
     ) -> Ret {
         callback((&self.1[index].0, &self.1[index].1).into())
     }
+    fn account_for_geyser<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        let pubkey = self.pubkey(index);
+        let account = &self.1[index].1;
+        callback(pubkey, account)
+    }
     fn is_zero_lamport(&self, index: usize) -> bool {
         self.1[index].1.is_zero_lamport()
     }
@@ -331,6 +360,18 @@ impl<'a> StorableAccounts<'a> for StorableAccountsBySlot<'a> {
         writer.storage = Some(storage);
         ret
     }
+    fn account_for_geyser<Ret>(
+        &self,
+        _index: usize,
+        _callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        // StorableAccountsBySlot does not have an AccountSharedData under the hood, so do not
+        // implement this method.
+        // This is fine because StorableAccountsBySlot is never used to store into the accounts
+        // write cache.  It is only used to store into account storage files.  Thus it'll never
+        // be used for geyser account update notifications.
+        unimplemented!();
+    }
     fn is_zero_lamport(&self, index: usize) -> bool {
         let indexes = self.find_internal_index(index);
         self.slots_and_accounts[indexes.0].1[indexes.1].is_zero_lamport()
@@ -414,6 +455,13 @@ mod tests {
             let account_for_storage = AccountForStorage::StoredAccountInfo(stored_account_info);
             callback(account_for_storage)
         }
+        fn account_for_geyser<Ret>(
+            &self,
+            _index: usize,
+            _callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+        ) -> Ret {
+            unimplemented!();
+        }
         fn is_zero_lamport(&self, index: usize) -> bool {
             self.1[index].is_zero_lamport()
         }
@@ -447,6 +495,13 @@ mod tests {
         ) -> Ret {
             callback((&self.1[index].0, &self.1[index].1).into())
         }
+        fn account_for_geyser<Ret>(
+            &self,
+            _index: usize,
+            _callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+        ) -> Ret {
+            unimplemented!();
+        }
         fn is_zero_lamport(&self, index: usize) -> bool {
             self.1[index].1.lamports() == 0
         }
@@ -481,6 +536,13 @@ mod tests {
             let account_for_storage = AccountForStorage::StoredAccountInfo(stored_account_info);
             callback(account_for_storage)
         }
+        fn account_for_geyser<Ret>(
+            &self,
+            _index: usize,
+            _callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+        ) -> Ret {
+            unimplemented!();
+        }
         fn is_zero_lamport(&self, index: usize) -> bool {
             self.1[index].is_zero_lamport()
         }

runtime/src/bank/partitioned_epoch_rewards/mod.rs

@@ -174,6 +174,15 @@ impl<'a> StorableAccounts<'a> for RewardCommissionAccountsStorable<'a> {
         callback((pubkey, account).into())
     }
 
+    fn account_for_geyser<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl for<'local> FnMut(&'local Pubkey, &'local AccountSharedData) -> Ret,
+    ) -> Ret {
+        let (pubkey, _, account) = &self.reward_commission_accounts.accounts_with_rewards[index];
+        callback(pubkey, account)
+    }
+
     fn is_zero_lamport(&self, index: usize) -> bool {
         self.reward_commission_accounts.accounts_with_rewards[index]
             .2