hold lock to speed up insert (#17194)

* hold lock to speed up insert

* add tests

runtime/src/accounts_db.rs

@@ -5034,22 +5034,33 @@ impl AccountsDb {
                         let dirty_keys =
                             accounts_map.iter().map(|(pubkey, _info)| *pubkey).collect();
                         self.uncleaned_pubkeys.insert(*slot, dirty_keys);
-                        for (pubkey, account_infos) in accounts_map.into_iter() {
-                            for (_, (store_id, stored_account)) in account_infos.into_iter() {
+                        let mut lock = self.accounts_index.get_account_maps_write_lock();
+                        for (pubkey, account_infos) in accounts_map.iter() {
+                            for (_, (store_id, stored_account)) in account_infos.iter() {
                                 let account_info = AccountInfo {
-                                    store_id,
+                                    store_id: *store_id,
                                     offset: stored_account.offset,
                                     stored_size: stored_account.stored_size,
                                     lamports: stored_account.account_meta.lamports,
                                 };
-                                self.accounts_index.insert_new_if_missing(
-                                    *slot,
+                                self.accounts_index
+                                    .insert_new_if_missing_into_primary_index(
+                                        *slot,
+                                        &pubkey,
+                                        account_info,
+                                        &mut _reclaims,
+                                        &mut lock,
+                                    );
+                            }
+                        }
+                        drop(lock);
+                        for (pubkey, account_infos) in accounts_map.into_iter() {
+                            for (_, (_store_id, stored_account)) in account_infos.iter() {
+                                self.accounts_index.update_secondary_indexes(
                                     &pubkey,
                                     &stored_account.account_meta.owner,
                                     &stored_account.data,
                                     &self.account_indexes,
-                                    account_info,
-                                    &mut _reclaims,
                                 );
                             }
                         }

runtime/src/accounts_index.rs

@@ -849,6 +849,15 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
     fn insert_new_entry_if_missing(&self, pubkey: &Pubkey) -> (WriteAccountMapEntry<T>, bool) {
         let new_entry = Self::new_entry();
         let mut w_account_maps = self.get_account_maps_write_lock();
+        self.insert_new_entry_if_missing_with_lock(pubkey, &mut w_account_maps, new_entry)
+    }
+
+    fn insert_new_entry_if_missing_with_lock(
+        &self,
+        pubkey: &Pubkey,
+        w_account_maps: &mut ReadWriteLockMapType<T>,
+        new_entry: AccountMapEntry<T>,
+    ) -> (WriteAccountMapEntry<T>, bool) {
         let mut is_newly_inserted = false;
         let account_entry = w_account_maps.entry(*pubkey).or_insert_with(|| {
             is_newly_inserted = true;
@@ -1093,7 +1102,7 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         max_root
     }
 
-    fn update_secondary_indexes(
+    pub(crate) fn update_secondary_indexes(
         &self,
         pubkey: &Pubkey,
         account_owner: &Pubkey,
@@ -1147,31 +1156,29 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         }
     }
 
-    fn get_account_maps_write_lock(&self) -> ReadWriteLockMapType<T> {
+    pub(crate) fn get_account_maps_write_lock(&self) -> ReadWriteLockMapType<T> {
         self.account_maps.write().unwrap()
     }
 
     // Same functionally to upsert, but doesn't take the read lock
     // initially on the accounts_map
-    // Can save time when inserting lots of new keys
-    pub fn insert_new_if_missing(
+    // Can save time when inserting lots of new keys.
+    // But, does NOT update secondary index
+    pub(crate) fn insert_new_if_missing_into_primary_index(
         &self,
         slot: Slot,
         pubkey: &Pubkey,
-        account_owner: &Pubkey,
-        account_data: &[u8],
-        account_indexes: &AccountSecondaryIndexes,
         account_info: T,
         reclaims: &mut SlotList<T>,
+        w_account_maps: &mut ReadWriteLockMapType<T>,
     ) {
-        {
-            let (mut w_account_entry, _is_new) = self.insert_new_entry_if_missing(pubkey);
-            if account_info.is_zero_lamport() {
-                self.zero_lamport_pubkeys.insert(*pubkey);
-            }
-            w_account_entry.update(slot, account_info, reclaims);
+        let new_entry = Self::new_entry();
+        let (mut w_account_entry, _is_new) =
+            self.insert_new_entry_if_missing_with_lock(pubkey, w_account_maps, new_entry);
+        if account_info.is_zero_lamport() {
+            self.zero_lamport_pubkeys.insert(*pubkey);
         }
-        self.update_secondary_indexes(pubkey, account_owner, account_data, account_indexes);
+        w_account_entry.update(slot, account_info, reclaims);
     }
 
     // Updates the given pubkey at the given slot with the new account information.
@@ -2180,6 +2187,121 @@ pub mod tests {
         assert_eq!(num, 0);
     }
 
+    type AccountInfoTest = f64;
+
+    impl IsCached for AccountInfoTest {
+        fn is_cached(&self) -> bool {
+            true
+        }
+    }
+
+    impl ZeroLamport for AccountInfoTest {
+        fn is_zero_lamport(&self) -> bool {
+            true
+        }
+    }
+    #[test]
+    fn test_insert_new_with_lock_no_ancestors() {
+        let key = Keypair::new();
+        let mut gc = Vec::new();
+        let slot = 0;
+
+        let index = AccountsIndex::<bool>::default();
+        let mut w_account_maps = index.get_account_maps_write_lock();
+        let account_info = true;
+        index.insert_new_if_missing_into_primary_index(
+            slot,
+            &key.pubkey(),
+            account_info,
+            &mut gc,
+            &mut w_account_maps,
+        );
+        drop(w_account_maps);
+        assert!(gc.is_empty());
+
+        assert!(index.zero_lamport_pubkeys().is_empty());
+
+        let mut ancestors = Ancestors::default();
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_none());
+        assert!(index.get(&key.pubkey(), None, None).is_none());
+
+        let mut num = 0;
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 0);
+        ancestors.insert(slot, 0);
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_some());
+        assert_eq!(index.ref_count_from_storage(&key.pubkey()), 1);
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 1);
+
+        // not zero lamports
+        let mut gc = Vec::new();
+        let index = AccountsIndex::<AccountInfoTest>::default();
+        let mut w_account_maps = index.get_account_maps_write_lock();
+        let account_info: AccountInfoTest = 0 as AccountInfoTest;
+        index.insert_new_if_missing_into_primary_index(
+            slot,
+            &key.pubkey(),
+            account_info,
+            &mut gc,
+            &mut w_account_maps,
+        );
+        drop(w_account_maps);
+        assert!(gc.is_empty());
+
+        assert!(!index.zero_lamport_pubkeys().is_empty());
+
+        let mut ancestors = Ancestors::default();
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_none());
+        assert!(index.get(&key.pubkey(), None, None).is_none());
+
+        let mut num = 0;
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 0);
+        ancestors.insert(slot, 0);
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_some());
+        assert_eq!(index.ref_count_from_storage(&key.pubkey()), 0); // cached, so 0
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 1);
+    }
+
+    #[test]
+    fn test_insert_with_lock_no_ancestors() {
+        let key = Keypair::new();
+        let index = AccountsIndex::<bool>::default();
+        let mut gc = Vec::new();
+        let slot = 0;
+
+        let new_entry = AccountsIndex::new_entry();
+        assert_eq!(new_entry.ref_count.load(Ordering::Relaxed), 0);
+        assert!(new_entry.slot_list.read().unwrap().is_empty());
+        assert_eq!(new_entry.slot_list.read().unwrap().capacity(), 1);
+        let mut w_account_maps = index.get_account_maps_write_lock();
+        let (mut write, insert) = index.insert_new_entry_if_missing_with_lock(
+            &key.pubkey(),
+            &mut w_account_maps,
+            new_entry,
+        );
+        assert!(insert);
+        drop(w_account_maps);
+        let account_info = true;
+        write.update(slot, account_info, &mut gc);
+        assert!(gc.is_empty());
+        drop(write);
+
+        let mut ancestors = Ancestors::default();
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_none());
+        assert!(index.get(&key.pubkey(), None, None).is_none());
+
+        let mut num = 0;
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 0);
+        ancestors.insert(slot, 0);
+        assert!(index.get(&key.pubkey(), Some(&ancestors), None).is_some());
+        index.unchecked_scan_accounts("", &ancestors, |_pubkey, _index| num += 1);
+        assert_eq!(num, 1);
+    }
+
     #[test]
     fn test_insert_wrong_ancestors() {
         let key = Keypair::new();