Accounts index opt (#4621)

* Add accounts_index bench

* Don't take the accounts index lock unless needed

* Accounts_index remove insert return vec and add capacity stats

* Use hashbrown hashmap for accounts_index

runtime/benches/accounts_index.rs

@@ -0,0 +1,46 @@
+#![feature(test)]
+
+extern crate test;
+
+use rand::{thread_rng, Rng};
+use solana_runtime::accounts_db::AccountInfo;
+use solana_runtime::accounts_index::AccountsIndex;
+use solana_sdk::pubkey::Pubkey;
+use test::Bencher;
+
+#[bench]
+fn bench_accounts_index(bencher: &mut Bencher) {
+    const NUM_PUBKEYS: usize = 10_000;
+    let pubkeys: Vec<_> = (0..NUM_PUBKEYS)
+        .into_iter()
+        .map(|_| Pubkey::new_rand())
+        .collect();
+
+    const NUM_FORKS: u64 = 16;
+
+    let mut reclaims = vec![];
+    let mut index = AccountsIndex::<AccountInfo>::default();
+    for f in 0..NUM_FORKS {
+        for _p in 0..NUM_PUBKEYS {
+            index.insert(f, &pubkeys[_p], AccountInfo::default(), &mut reclaims);
+        }
+    }
+
+    let mut fork = NUM_FORKS;
+    let mut root = 0;
+    bencher.iter(|| {
+        for _p in 0..NUM_PUBKEYS {
+            let pubkey = thread_rng().gen_range(0, NUM_PUBKEYS);
+            index.insert(
+                fork,
+                &pubkeys[pubkey],
+                AccountInfo::default(),
+                &mut reclaims,
+            );
+            reclaims.clear();
+        }
+        index.add_root(root);
+        root += 1;
+        fork += 1;
+    });
+}

runtime/src/accounts_db.rs

@@ -474,14 +474,14 @@ impl AccountsDB {
         fork_id: Fork,
         infos: Vec<AccountInfo>,
         accounts: &[(&Pubkey, &Account)],
-    ) -> Vec<(Fork, AccountInfo)> {
+    ) -> (Vec<(Fork, AccountInfo)>, u64) {
+        let mut reclaims = Vec::with_capacity(infos.len() * 2);
         let mut index = self.accounts_index.write().unwrap();
-        let mut reclaims = vec![];
         for (i, info) in infos.into_iter().enumerate() {
             let key = &accounts[i].0;
-            reclaims.extend(index.insert(fork_id, key, info).into_iter())
+            index.insert(fork_id, key, info, &mut reclaims);
         }
-        reclaims
+        (reclaims, index.last_root)
     }
 
     fn remove_dead_accounts(&self, reclaims: Vec<(Fork, AccountInfo)>) -> HashSet<Fork> {
@@ -516,24 +516,30 @@ impl AccountsDB {
         dead_forks
     }
 
-    fn cleanup_dead_forks(&self, dead_forks: &mut HashSet<Fork>) {
+    fn cleanup_dead_forks(&self, dead_forks: &mut HashSet<Fork>, last_root: u64) {
-        let mut index = self.accounts_index.write().unwrap();
         // a fork is not totally dead until it is older than the root
-        dead_forks.retain(|fork| *fork < index.last_root);
+        dead_forks.retain(|fork| *fork < last_root);
-        for fork in dead_forks.iter() {
+        if !dead_forks.is_empty() {
-            index.cleanup_dead_fork(*fork);
+            let mut index = self.accounts_index.write().unwrap();
+            for fork in dead_forks.iter() {
+                index.cleanup_dead_fork(*fork);
+            }
         }
     }
 
     /// Store the account update.
     pub fn store(&self, fork_id: Fork, accounts: &[(&Pubkey, &Account)]) {
         let infos = self.store_accounts(fork_id, accounts);
-        let reclaims = self.update_index(fork_id, infos, accounts);
+
+        let (reclaims, last_root) = self.update_index(fork_id, infos, accounts);
         trace!("reclaim: {}", reclaims.len());
+
         let mut dead_forks = self.remove_dead_accounts(reclaims);
         trace!("dead_forks: {}", dead_forks.len());
-        self.cleanup_dead_forks(&mut dead_forks);
+
+        self.cleanup_dead_forks(&mut dead_forks, last_root);
         trace!("purge_forks: {}", dead_forks.len());
+
         for fork in dead_forks {
             self.purge_fork(fork);
         }
@@ -1182,7 +1188,14 @@ 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

@@ -1,7 +1,9 @@
+use hashbrown::HashMap;
 use log::*;
 use serde::{Deserialize, Serialize};
 use solana_sdk::pubkey::Pubkey;
-use std::collections::{HashMap, HashSet};
+use std::collections;
+use std::collections::HashSet;
 
 pub type Fork = u64;
 
@@ -19,7 +21,11 @@ pub struct AccountsIndex<T> {
 impl<T: Clone> AccountsIndex<T> {
     /// Get an account
     /// The latest account that appears in `ancestors` or `roots` is returned.
-    pub fn get(&self, pubkey: &Pubkey, ancestors: &HashMap<Fork, usize>) -> Option<(&T, Fork)> {
+    pub fn get(
+        &self,
+        pubkey: &Pubkey,
+        ancestors: &collections::HashMap<Fork, usize>,
+    ) -> Option<(&T, Fork)> {
         let list = self.account_maps.get(pubkey)?;
         let mut max = 0;
         let mut rv = None;
@@ -33,35 +39,34 @@ impl<T: Clone> AccountsIndex<T> {
         rv
     }
 
-    /// Insert a new fork.
+    pub fn insert(
-    /// @retval - The return value contains any squashed accounts that can freed from storage.
+        &mut self,
-    pub fn insert(&mut self, fork: Fork, pubkey: &Pubkey, account_info: T) -> Vec<(Fork, T)> {
+        fork: Fork,
-        let mut rv = vec![];
+        pubkey: &Pubkey,
-        let mut fork_vec: Vec<(Fork, T)> = vec![];
+        account_info: T,
-        {
+        reclaims: &mut Vec<(Fork, T)>,
-            let entry = self.account_maps.entry(*pubkey).or_insert_with(|| vec![]);
+    ) {
-            std::mem::swap(entry, &mut fork_vec);
+        let last_root = self.last_root;
-        };
+        let roots = &self.roots;
+        let fork_vec = self
+            .account_maps
+            .entry(*pubkey)
+            .or_insert_with(|| (Vec::with_capacity(32)));
 
         // filter out old entries
-        rv.extend(fork_vec.iter().filter(|(f, _)| *f == fork).cloned());
+        reclaims.extend(fork_vec.iter().filter(|(f, _)| *f == fork).cloned());
         fork_vec.retain(|(f, _)| *f != fork);
 
         // add the new entry
         fork_vec.push((fork, account_info));
 
-        rv.extend(
+        reclaims.extend(
             fork_vec
                 .iter()
-                .filter(|(fork, _)| self.is_purged(*fork))
+                .filter(|(fork, _)| Self::is_purged(roots, last_root, *fork))
                 .cloned(),
         );
-        fork_vec.retain(|(fork, _)| !self.is_purged(*fork));
+        fork_vec.retain(|(fork, _)| !Self::is_purged(roots, last_root, *fork));
-        {
-            let entry = self.account_maps.entry(*pubkey).or_insert_with(|| vec![]);
-            std::mem::swap(entry, &mut fork_vec);
-        };
-        rv
     }
 
     pub fn add_index(&mut self, fork: Fork, pubkey: &Pubkey, account_info: T) {
@@ -69,12 +74,14 @@ impl<T: Clone> AccountsIndex<T> {
         entry.push((fork, account_info));
     }
 
-    pub fn is_purged(&self, fork: Fork) -> bool {
+    pub fn is_purged(roots: &HashSet<Fork>, last_root: Fork, fork: Fork) -> bool {
-        !self.is_root(fork) && fork < self.last_root
+        !roots.contains(&fork) && fork < last_root
     }
+
     pub fn is_root(&self, fork: Fork) -> bool {
         self.roots.contains(&fork)
     }
+
     pub fn add_root(&mut self, fork: Fork) {
         assert!(
             (self.last_root == 0 && fork == 0) || (fork > self.last_root),
@@ -99,7 +106,7 @@ mod tests {
     fn test_get_empty() {
         let key = Keypair::new();
         let index = AccountsIndex::<bool>::default();
-        let ancestors = HashMap::new();
+        let ancestors = collections::HashMap::new();
         assert_eq!(index.get(&key.pubkey(), &ancestors), None);
     }
 
@@ -107,10 +114,11 @@ mod tests {
     fn test_insert_no_ancestors() {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
 
-        let ancestors = HashMap::new();
+        let ancestors = collections::HashMap::new();
         assert_eq!(index.get(&key.pubkey(), &ancestors), None);
     }
 
@@ -118,7 +126,8 @@ mod tests {
     fn test_insert_wrong_ancestors() {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
 
         let ancestors = vec![(1, 1)].into_iter().collect();
@@ -129,7 +138,8 @@ mod tests {
     fn test_insert_with_ancestors() {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
 
         let ancestors = vec![(0, 0)].into_iter().collect();
@@ -148,7 +158,8 @@ mod tests {
     fn test_insert_with_root() {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
 
         let ancestors = vec![].into_iter().collect();
@@ -159,9 +170,17 @@ mod tests {
     #[test]
     fn test_is_purged() {
         let mut index = AccountsIndex::<bool>::default();
-        assert!(!index.is_purged(0));
+        assert!(!AccountsIndex::<bool>::is_purged(
+            &index.roots,
+            index.last_root,
+            0
+        ));
         index.add_root(1);
-        assert!(index.is_purged(0));
+        assert!(AccountsIndex::<bool>::is_purged(
+            &index.roots,
+            index.last_root,
+            0
+        ));
     }
 
     #[test]
@@ -205,11 +224,13 @@ mod tests {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
         let ancestors = vec![(0, 0)].into_iter().collect();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
         assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&true, 0)));
 
-        let gc = index.insert(0, &key.pubkey(), false);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), false, &mut gc);
         assert_eq!(gc, vec![(0, true)]);
         assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&false, 0)));
     }
@@ -219,9 +240,10 @@ mod tests {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
         let ancestors = vec![(0, 0)].into_iter().collect();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
-        let gc = index.insert(1, &key.pubkey(), false);
+        index.insert(1, &key.pubkey(), false, &mut gc);
         assert!(gc.is_empty());
         assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&true, 0)));
         let ancestors = vec![(1, 0)].into_iter().collect();
@@ -232,10 +254,11 @@ mod tests {
     fn test_update_gc_purged_fork() {
         let key = Keypair::new();
         let mut index = AccountsIndex::<bool>::default();
-        let gc = index.insert(0, &key.pubkey(), true);
+        let mut gc = Vec::new();
+        index.insert(0, &key.pubkey(), true, &mut gc);
         assert!(gc.is_empty());
         index.add_root(1);
-        let gc = index.insert(1, &key.pubkey(), false);
+        index.insert(1, &key.pubkey(), false, &mut gc);
         assert_eq!(gc, vec![(0, true)]);
         let ancestors = vec![].into_iter().collect();
         assert_eq!(index.get(&key.pubkey(), &ancestors), Some((&false, 1)));

runtime/src/lib.rs

@@ -1,6 +1,6 @@
 mod accounts;
 pub mod accounts_db;
-mod accounts_index;
+pub mod accounts_index;
 pub mod append_vec;
 pub mod bank;
 pub mod bank_client;