diff --git a/accounts-db/src/account_storage.rs b/accounts-db/src/account_storage.rs index da25092a59be09..b1ad6e1b23f7fb 100644 --- a/accounts-db/src/account_storage.rs +++ b/accounts-db/src/account_storage.rs @@ -200,6 +200,29 @@ impl AccountStorage { pub(crate) fn len(&self) -> usize { self.map.len() } + + /// Returns the (slot, storage) tuples where `predicate` returns `true` + /// + /// This function is useful when not all storages are desired, + /// as storages are only Arc::cloned if they pass the predicate. + /// + /// # Panics + /// + /// Panics if `shrink` is in progress. + pub fn get_if( + &self, + predicate: impl Fn(&Slot, &AccountStorageEntry) -> bool, + ) -> Box<[(Slot, Arc)]> { + assert!(self.no_shrink_in_progress()); + self.map + .iter() + .filter_map(|entry| { + let slot = entry.key(); + let storage = &entry.value().storage; + predicate(slot, storage).then(|| (*slot, Arc::clone(storage))) + }) + .collect() + } } /// iterate contents of AccountStorage without exposing internals @@ -291,7 +314,11 @@ impl Default for AccountStorageStatus { #[cfg(test)] pub(crate) mod tests { - use {super::*, crate::accounts_file::AccountsFileProvider, std::path::Path}; + use { + super::*, + crate::accounts_file::AccountsFileProvider, + std::{iter, path::Path}, + }; #[test] fn test_shrink_in_progress() { @@ -568,4 +595,52 @@ pub(crate) mod tests { .is_none()); assert!(storage.get_account_storage_entry(slot, id).is_some()); } + + #[test] + fn test_get_if() { + let storage = AccountStorage::default(); + assert!(storage.get_if(|_, _| true).is_empty()); + + // add some entries + let ids = [123, 456, 789]; + for id in ids { + let slot = id as Slot; + let entry = AccountStorageEntry::new( + Path::new(""), + slot, + id, + 5000, + AccountsFileProvider::AppendVec, + ); + storage.map.insert( + slot, + AccountStorageReference { + id, + storage: entry.into(), + }, + ); + } + + // look 'em up + for id in ids { + let found = storage.get_if(|slot, _| *slot == id as Slot); + assert!(found + .iter() + .map(|(slot, _)| *slot) + .eq(iter::once(id as Slot))); + } + + assert!(storage.get_if(|_, _| false).is_empty()); + assert_eq!(storage.get_if(|_, _| true).len(), ids.len()); + } + + #[test] + #[should_panic(expected = "self.no_shrink_in_progress()")] + fn test_get_if_fail() { + let storage = AccountStorage::default(); + storage + .shrink_in_progress_map + .insert(0, storage.get_test_storage()); + storage.get_if(|_, _| true); + } } diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs index 8eba4a89cf1046..9d904d4bb04bba 100644 --- a/accounts-db/src/accounts_db.rs +++ b/accounts-db/src/accounts_db.rs @@ -8278,44 +8278,27 @@ impl AccountsDb { &self, requested_slots: impl RangeBounds + Sync, ) -> (Vec>, Vec) { - let mut m = Measure::start("get slots"); - let mut slots_and_storages = self + let start = Instant::now(); + let max_alive_root_exclusive = self + .accounts_index + .roots_tracker + .read() + .unwrap() + .alive_roots + .max_exclusive(); + let (slots, storages) = self .storage - .iter() - .filter_map(|(slot, store)| { - requested_slots - .contains(&slot) - .then_some((slot, Some(store))) + .get_if(|slot, storage| { + (*slot < max_alive_root_exclusive) + && requested_slots.contains(slot) + && storage.has_accounts() }) - .collect::>(); - m.stop(); - let mut m2 = Measure::start("filter"); - let chunk_size = 5_000; - let (result, slots): (Vec<_>, Vec<_>) = self.thread_pool_clean.install(|| { - slots_and_storages - .par_chunks_mut(chunk_size) - .map(|slots_and_storages| { - slots_and_storages - .iter_mut() - .filter(|(slot, _)| self.accounts_index.is_alive_root(*slot)) - .filter_map(|(slot, store)| { - let store = std::mem::take(store).unwrap(); - store.has_accounts().then_some((store, *slot)) - }) - .collect::, Slot)>>() - }) - .flatten() - .unzip() - }); - - m2.stop(); - - debug!( - "hash_total: get slots: {}, filter: {}", - m.as_us(), - m2.as_us(), - ); - (result, slots) + .into_vec() + .into_iter() + .unzip(); + let duration = start.elapsed(); + debug!("get_snapshot_storages: {duration:?}"); + (storages, slots) } /// Returns the latest full snapshot slot