batch insert account_index items in generate_index (#17290)

runtime/src/accounts_db.rs

@@ -5065,14 +5065,13 @@ impl AccountsDb {
 
                     if !accounts_map.is_empty() {
                         let mut _reclaims: Vec<(u64, AccountInfo)> = vec![];
-                        let dirty_keys =
+                        let len = accounts_map.len();
-                            accounts_map.iter().map(|(pubkey, _info)| *pubkey).collect();
-                        self.uncleaned_pubkeys.insert(*slot, dirty_keys);
 
-                        let infos: Vec<_> = accounts_map
+                        let mut items = Vec::with_capacity(len);
+                        let dirty_keys = accounts_map
                             .iter()
                             .map(|(pubkey, (_, store_id, stored_account))| {
-                                (
+                                items.push((
                                     pubkey,
                                     AccountInfo {
                                         store_id: *store_id,
@@ -5080,21 +5079,14 @@ impl AccountsDb {
                                         stored_size: stored_account.stored_size,
                                         lamports: stored_account.account_meta.lamports,
                                     },
-                                )
+                                ));
+                                *pubkey
                             })
-                            .collect::<Vec<_>>(); // we want this collection to occur before the lock below
+                            .collect::<Vec<_>>();
-                        let mut lock = self.accounts_index.get_account_maps_write_lock();
+                        self.uncleaned_pubkeys.insert(*slot, dirty_keys);
-                        infos.into_iter().for_each(|(pubkey, account_info)| {
+
-                            self.accounts_index
+                        self.accounts_index
-                                .insert_new_if_missing_into_primary_index(
+                            .insert_new_if_missing_into_primary_index(*slot, items);
-                                    *slot,
-                                    &pubkey,
-                                    account_info,
-                                    &mut _reclaims,
-                                    &mut lock,
-                                );
-                        });
-                        drop(lock);
                         if !self.account_indexes.is_empty() {
                             for (pubkey, (_, _store_id, stored_account)) in accounts_map.iter() {
                                 self.accounts_index.update_secondary_indexes(

runtime/src/accounts_index.rs

@@ -1179,30 +1179,49 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         }
     }
 
-    pub(crate) fn get_account_maps_write_lock(&self) -> AccountMapsWriteLock<T> {
+    fn get_account_maps_write_lock(&self) -> AccountMapsWriteLock<T> {
         self.account_maps.write().unwrap()
     }
 
-    // Same functionally to upsert, but doesn't take the read lock
+    // Same functionally to upsert, but:
-    // initially on the accounts_map
+    // 1. operates on a batch of items
+    // 2. holds the write lock for the duration of adding the items
     // Can save time when inserting lots of new keys.
     // But, does NOT update secondary index
+    // This is designed to be called at startup time.
+    #[allow(clippy::needless_collect)]
     pub(crate) fn insert_new_if_missing_into_primary_index(
         &self,
         slot: Slot,
-        pubkey: &Pubkey,
+        items: Vec<(&Pubkey, T)>,
-        account_info: T,
-        reclaims: &mut SlotList<T>,
-        w_account_maps: &mut AccountMapsWriteLock<T>,
     ) {
-        let account_entry =
+        let potentially_new_items = items
-            self.insert_new_entry_if_missing(pubkey, slot, &account_info, Some(w_account_maps));
+            .iter()
-        if account_info.is_zero_lamport() {
+            .map(|(_pubkey, account_info)| {
-            self.zero_lamport_pubkeys.insert(*pubkey);
+                // this value is equivalent to what update() below would have created if we inserted a new item
-        }
+                WriteAccountMapEntry::new_entry_after_update(slot, account_info)
-        if let Some(mut w_account_entry) = account_entry {
+            })
-            w_account_entry.update(slot, account_info, reclaims);
+            .collect::<Vec<_>>(); // collect here so we have created all data prior to obtaining lock
-        }
+
+        let mut _reclaims = SlotList::new();
+
+        let mut w_account_maps = self.get_account_maps_write_lock();
+        items
+            .into_iter()
+            .zip(potentially_new_items.into_iter())
+            .for_each(|((pubkey, account_info), new_item)| {
+                let account_entry = self.insert_new_entry_if_missing_with_lock(
+                    pubkey,
+                    &mut w_account_maps,
+                    new_item,
+                );
+                if account_info.is_zero_lamport() {
+                    self.zero_lamport_pubkeys.insert(*pubkey);
+                }
+                if let Some(mut w_account_entry) = account_entry {
+                    w_account_entry.update(slot, account_info, &mut _reclaims);
+                }
+            });
     }
 
     // Updates the given pubkey at the given slot with the new account information.
@@ -2232,21 +2251,12 @@ pub mod tests {
     fn test_insert_new_with_lock_no_ancestors() {
         let key = Keypair::new();
         let pubkey = &key.pubkey();
-        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(
+        let items = vec![(pubkey, account_info)];
-            slot,
+        index.insert_new_if_missing_into_primary_index(slot, items);
-            pubkey,
-            account_info,
-            &mut gc,
-            &mut w_account_maps,
-        );
-        drop(w_account_maps);
-        assert!(gc.is_empty());
 
         assert!(index.zero_lamport_pubkeys().is_empty());
 
@@ -2264,19 +2274,10 @@ pub mod tests {
         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(
+        let items = vec![(pubkey, account_info)];
-            slot,
+        index.insert_new_if_missing_into_primary_index(slot, items);
-            pubkey,
-            account_info,
-            &mut gc,
-            &mut w_account_maps,
-        );
-        drop(w_account_maps);
-        assert!(gc.is_empty());
 
         assert!(!index.zero_lamport_pubkeys().is_empty());
 
@@ -2320,6 +2321,27 @@ pub mod tests {
         );
     }
 
+    #[test]
+    fn test_batch_insert() {
+        let slot0 = 0;
+        let key0 = Keypair::new().pubkey();
+        let key1 = Keypair::new().pubkey();
+
+        let index = AccountsIndex::<bool>::default();
+        let account_infos = [true, false];
+
+        index.insert_new_if_missing_into_primary_index(
+            slot0,
+            vec![(&key0, account_infos[0]), (&key1, account_infos[1])],
+        );
+
+        for (i, key) in [key0, key1].iter().enumerate() {
+            let entry = index.get_account_read_entry(key).unwrap();
+            assert_eq!(entry.ref_count().load(Ordering::Relaxed), 1);
+            assert_eq!(entry.slot_list().to_vec(), vec![(slot0, account_infos[i]),]);
+        }
+    }
+
     fn test_new_entry_code_paths_helper<
         T: 'static + Clone + IsCached + ZeroLamport + std::cmp::PartialEq + std::fmt::Debug,
     >(
@@ -2346,13 +2368,9 @@ pub mod tests {
                 &mut gc,
             );
         } else {
-            let mut lock = index.get_account_maps_write_lock();
             index.insert_new_if_missing_into_primary_index(
                 slot0,
-                &key,
+                vec![(&key, account_infos[0].clone())],
-                account_infos[0].clone(),
-                &mut gc,
-                &mut lock,
             );
         }
         assert!(gc.is_empty());
@@ -2385,13 +2403,9 @@ pub mod tests {
                 &mut gc,
             );
         } else {
-            let mut lock = index.get_account_maps_write_lock();
             index.insert_new_if_missing_into_primary_index(
                 slot1,
-                &key,
+                vec![(&key, account_infos[1].clone())],
-                account_infos[1].clone(),
-                &mut gc,
-                &mut lock,
             );
         }
         assert!(gc.is_empty());