Refactor Accounts Index and AccountsIndexIterator

accounts-db/src/accounts_db.rs

@@ -8524,15 +8524,17 @@ impl AccountsDb {
         #[allow(clippy::stable_sort_primitive)]
         alive_roots.sort();
         info!("{}: accounts_index alive_roots: {:?}", label, alive_roots,);
-        let full_pubkey_range = Pubkey::from([0; 32])..=Pubkey::from([0xff; 32]);
-
         self.accounts_index.account_maps.iter().for_each(|map| {
-            for (pubkey, account_entry) in map.items(&full_pubkey_range) {
-                info!("  key: {} ref_count: {}", pubkey, account_entry.ref_count(),);
-                info!(
-                    "      slots: {:?}",
-                    *account_entry.slot_list.read().unwrap()
-                );
+            for pubkey in map.keys() {
+                self.accounts_index.get_and_then(&pubkey, |account_entry| {
+                    if let Some(account_entry) = account_entry {
+                        let list_r = &account_entry.slot_list.read().unwrap();
+                        info!(" key: {} ref_count: {}", pubkey, account_entry.ref_count(),);
+                        info!("      slots: {:?}", list_r);
+                    }
+                    let add_to_in_mem_cache = false;
+                    (add_to_in_mem_cache, ())
+                });
             }
         });
     }

accounts-db/src/accounts_index.rs

@@ -727,24 +727,30 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         let mut iterator_elapsed = 0;
         let mut iterator_timer = Measure::start("iterator_elapsed");
 
-        for pubkey_list in self.iter(range.as_ref(), returns_items) {
+        for pubkeys in self.iter(range.as_ref(), returns_items) {
             iterator_timer.stop();
             iterator_elapsed += iterator_timer.as_us();
-            for (pubkey, list) in pubkey_list {
+            for pubkey in pubkeys {
                 num_keys_iterated += 1;
-                let mut read_lock_timer = Measure::start("read_lock");
-                let list_r = &list.slot_list.read().unwrap();
-                read_lock_timer.stop();
-                read_lock_elapsed += read_lock_timer.as_us();
-                let mut latest_slot_timer = Measure::start("latest_slot");
-                if let Some(index) = self.latest_slot(Some(ancestors), list_r, max_root) {
-                    latest_slot_timer.stop();
-                    latest_slot_elapsed += latest_slot_timer.as_us();
-                    let mut load_account_timer = Measure::start("load_account");
-                    func(&pubkey, (&list_r[index].1, list_r[index].0));
-                    load_account_timer.stop();
-                    load_account_elapsed += load_account_timer.as_us();
-                }
+                self.get_and_then(&pubkey, |entry| {
+                    if let Some(list) = entry {
+                        let mut read_lock_timer = Measure::start("read_lock");
+                        let list_r = &list.slot_list.read().unwrap();
+                        read_lock_timer.stop();
+                        read_lock_elapsed += read_lock_timer.as_us();
+                        let mut latest_slot_timer = Measure::start("latest_slot");
+                        if let Some(index) = self.latest_slot(Some(ancestors), list_r, max_root) {
+                            latest_slot_timer.stop();
+                            latest_slot_elapsed += latest_slot_timer.as_us();
+                            let mut load_account_timer = Measure::start("load_account");
+                            func(&pubkey, (&list_r[index].1, list_r[index].0));
+                            load_account_timer.stop();
+                            load_account_elapsed += load_account_timer.as_us();
+                        }
+                    }
+                    let add_to_in_mem_cache = false;
+                    (add_to_in_mem_cache, ())
+                });
                 if config.is_aborted() {
                     return;
                 }

accounts-db/src/accounts_index/in_mem_accounts_index.rs

@@ -250,53 +250,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
         }
     }
 
-    pub fn items<R>(&self, range: &R) -> Vec<(Pubkey, Arc<AccountMapEntry<T>>)>
-    where
-        R: RangeBounds<Pubkey> + std::fmt::Debug,
-    {
-        let start = match range.start_bound() {
-            Bound::Included(bound) | Bound::Excluded(bound) => bound,
-            Bound::Unbounded => &Pubkey::from([0; 32]),
-        };
-
-        let end = match range.end_bound() {
-            Bound::Included(bound) | Bound::Excluded(bound) => bound,
-            Bound::Unbounded => &Pubkey::from([0xff; 32]),
-        };
-
-        if start > &self.highest_pubkey || end < &self.lowest_pubkey {
-            // range does not contain any of the keys in this bin. No need to
-            // load and scan the map.
-            // Example:
-            //    |-------------------|  |-------------------|  |-------------------|
-            //       start          end  low               high  start              end
-            return vec![];
-        }
-
-        let m = Measure::start("items");
-
-        // For simplicity, we check the range for every pubkey in the map. This
-        // can be further optimized for case, such as the range contains lowest
-        // and highest pubkey for this bin. In such case, we can return all
-        // items in the map without range check on item's pubkey. Since the
-        // check is cheap when compared with the cost of reading from disk, we
-        // are not optimizing it for now.
-        self.hold_range_in_memory(range, true);
-        let result = self
-            .map_internal
-            .read()
-            .unwrap()
-            .iter()
-            .filter(|&(k, _v)| range.contains(k))
-            .map(|(k, v)| (*k, Arc::clone(v)))
-            .collect();
-        self.hold_range_in_memory(range, false);
-        Self::update_stat(&self.stats().items, 1);
-        Self::update_time_stat(&self.stats().items_us, m);
-        result
-    }
-
-    // only called in debug code paths
+    /// return all keys in this bin
     pub fn keys(&self) -> Vec<Pubkey> {
         Self::update_stat(&self.stats().keys, 1);
         // easiest implementation is to load everything from disk into cache and return the keys

accounts-db/src/accounts_index/iter.rs

@@ -1,8 +1,5 @@
 use {
-    super::{
-        account_map_entry::AccountMapEntry, in_mem_accounts_index::InMemAccountsIndex,
-        AccountsIndex, DiskIndexValue, IndexValue,
-    },
+    super::{in_mem_accounts_index::InMemAccountsIndex, AccountsIndex, DiskIndexValue, IndexValue},
     solana_pubkey::Pubkey,
     std::{
         ops::{Bound, RangeBounds},
@@ -18,7 +15,7 @@ pub struct AccountsIndexIterator<'a, T: IndexValue, U: DiskIndexValue + From<T>
     end_bound: Bound<&'a Pubkey>,
     start_bin: usize,
     end_bin_inclusive: usize,
-    items: Vec<(Pubkey, Arc<AccountMapEntry<T>>)>,
+    items: Vec<Pubkey>,
     returns_items: AccountsIndexIteratorReturnsItems,
 }
 
@@ -61,18 +58,23 @@ impl<'a, T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndexIter
 impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> Iterator
     for AccountsIndexIterator<'_, T, U>
 {
-    type Item = Vec<(Pubkey, Arc<AccountMapEntry<T>>)>;
+    type Item = Vec<Pubkey>;
     fn next(&mut self) -> Option<Self::Item> {
+        let range = (self.start_bound, self.end_bound);
         while self.items.len() < ITER_BATCH_SIZE {
             if self.start_bin > self.end_bin_inclusive {
                 break;
             }
 
             let bin = self.start_bin;
             let map = &self.account_maps[bin];
-            let mut items = map.items(&(self.start_bound, self.end_bound));
+            let mut items = map
+                .keys()
+                .into_iter()
+                .filter(|k| range.contains(&k))
+                .collect::<Vec<_>>();
             if self.returns_items == AccountsIndexIteratorReturnsItems::Sorted {
-                items.sort_unstable_by(|a, b| a.0.cmp(&b.0));
+                items.sort_unstable();
             }
             self.items.append(&mut items);
             self.start_bin += 1;
@@ -141,7 +143,7 @@ mod tests {
             let x = iter.next().unwrap();
             assert_eq!(x.len(), 2 * ITER_BATCH_SIZE);
             assert_eq!(
-                x.is_sorted_by(|a, b| a.0 < b.0),
+                x.is_sorted(),
                 returns_items == AccountsIndexIteratorReturnsItems::Sorted
             );
             assert_eq!(iter.items.len(), 0); // should be empty.