diff --git a/accounts-db/src/accounts_index.rs b/accounts-db/src/accounts_index.rs index 1dc1646e90b0f6..029ef8ee8924f9 100644 --- a/accounts-db/src/accounts_index.rs +++ b/accounts-db/src/accounts_index.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "dev-context-only-utils")] +use qualifier_attr::qualifiers; pub(crate) mod in_mem_accounts_index; use { crate::{ @@ -487,32 +489,16 @@ pub struct AccountsIndexRootsStats { type RangeItemVec = Vec<(Pubkey, AccountMapEntry)>; -pub struct AccountsIndexIterator<'a, T: IndexValue, U: DiskIndexValue + From + Into> { +struct AccountsIndexIteratorInner<'a, T: IndexValue, U: DiskIndexValue + From + Into> { account_maps: &'a LockMapTypeSlice, bin_calculator: &'a PubkeyBinCalculator24, start_bound: Bound, end_bound: Bound, - is_finished: bool, - returns_items: AccountsIndexIteratorReturnsItems, - last_bin_range: Option<(usize, RangeItemVec)>, } -impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> AccountsIndexIterator<'a, T, U> { - fn range( - map: &AccountMaps, - range: R, - returns_items: AccountsIndexIteratorReturnsItems, - ) -> RangeItemVec - where - R: RangeBounds + std::fmt::Debug, - { - let mut result = map.items(&range); - if returns_items == AccountsIndexIteratorReturnsItems::Sorted { - result.sort_unstable_by(|a, b| a.0.cmp(&b.0)); - } - result - } - +impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> + AccountsIndexIteratorInner<'a, T, U> +{ fn clone_bound(bound: Bound<&Pubkey>) -> Bound { match bound { Unbounded => Unbounded, @@ -556,11 +542,7 @@ impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> AccountsIndexIter (start_bin, bin_range) } - pub fn new( - index: &'a AccountsIndex, - range: Option<&R>, - returns_items: AccountsIndexIteratorReturnsItems, - ) -> Self + pub fn new(index: &'a AccountsIndex, range: Option<&R>) -> Self where R: RangeBounds, { @@ -574,10 +556,7 @@ impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> AccountsIndexIter .map(|r| Self::clone_bound(r.end_bound())) .unwrap_or(Unbounded), account_maps: &index.account_maps, - is_finished: false, bin_calculator: &index.bin_calculator, - returns_items, - last_bin_range: None, } } @@ -598,17 +577,82 @@ impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> AccountsIndexIter } } +pub struct AccountsIndexIteratorSorted<'a, T: IndexValue, U: DiskIndexValue + From + Into> { + inner: AccountsIndexIteratorInner<'a, T, U>, + is_finished: bool, + last_bin_range: Option<(usize, RangeItemVec)>, +} + +impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> + AccountsIndexIteratorSorted<'a, T, U> +{ + fn range(map: &AccountMaps, range: R) -> RangeItemVec + where + R: RangeBounds + std::fmt::Debug, + { + let mut result = map.items(&range); + result.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + result + } + + pub fn new(index: &'a AccountsIndex, range: Option<&R>) -> Self + where + R: RangeBounds, + { + let inner = AccountsIndexIteratorInner::new(index, range); + Self { + inner, + is_finished: false, + last_bin_range: None, + } + } +} + +pub struct AccountsIndexIteratorUnsorted<'a, T: IndexValue, U: DiskIndexValue + From + Into> { + inner: AccountsIndexIteratorInner<'a, T, U>, +} + +impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> + AccountsIndexIteratorUnsorted<'a, T, U> +{ + fn range(map: &AccountMaps, range: R) -> RangeItemVec + where + R: RangeBounds + std::fmt::Debug, + { + map.items(&range) + } + + pub fn new(index: &'a AccountsIndex, range: Option<&R>) -> Self + where + R: RangeBounds, + { + let inner = AccountsIndexIteratorInner::new(index, range); + Self { inner } + } + + pub fn hold_range_in_memory(&self, range: &R, start_holding: bool, thread_pool: &ThreadPool) + where + R: RangeBounds + Debug + Sync, + { + self.inner + .hold_range_in_memory(range, start_holding, thread_pool) + } +} + impl + Into> Iterator - for AccountsIndexIterator<'_, T, U> + for AccountsIndexIteratorSorted<'_, T, U> { type Item = Vec<(Pubkey, AccountMapEntry)>; fn next(&mut self) -> Option { + let inner = &mut self.inner; + if self.is_finished { return None; } - let (start_bin, bin_range) = self.bin_start_and_range(); + + let (start_bin, bin_range) = inner.bin_start_and_range(); let mut chunk = Vec::with_capacity(ITER_BATCH_SIZE); - 'outer: for (i, map) in self + 'outer: for (i, map) in inner .account_maps .iter() .skip(start_bin) @@ -623,11 +667,11 @@ impl + Into> Iterator } _ => { // else load the new bin - Self::range(&map, (self.start_bound, self.end_bound), self.returns_items) + Self::range(&map, (inner.start_bound, inner.end_bound)) } }; for (count, (pubkey, account_map_entry)) in range.iter().enumerate() { - if chunk.len() >= ITER_BATCH_SIZE && self.returns_items.is_sorted() { + if chunk.len() >= ITER_BATCH_SIZE { range.drain(0..count); self.last_bin_range = Some((bin, range)); break 'outer; @@ -639,13 +683,92 @@ impl + Into> Iterator if chunk.is_empty() { self.is_finished = true; - return None; - } else if self.returns_items == AccountsIndexIteratorReturnsItems::Unsorted { - self.is_finished = true; + None + } else { + inner.start_bound = Excluded(chunk.last().unwrap().0); + Some(chunk) } + } +} + +impl + Into> Iterator + for AccountsIndexIteratorUnsorted<'_, T, U> +{ + type Item = Vec<(Pubkey, AccountMapEntry)>; + fn next(&mut self) -> Option { + let inner = &mut self.inner; + let (start_bin, bin_range) = inner.bin_start_and_range(); + let mut chunk = Vec::with_capacity(ITER_BATCH_SIZE); + for map in inner.account_maps.iter().skip(start_bin).take(bin_range) { + let mut range = Self::range(&map, (inner.start_bound, inner.end_bound)); + chunk.append(&mut range); + } + if chunk.is_empty() { + None + } else { + Some(chunk) + } + } +} + +pub enum AccountsIndexIterator<'a, T: IndexValue, U: DiskIndexValue + From + Into> { + Sorted(AccountsIndexIteratorSorted<'a, T, U>), + Unsorted(AccountsIndexIteratorUnsorted<'a, T, U>), +} - self.start_bound = Excluded(chunk.last().unwrap().0); - Some(chunk) +impl<'a, T: IndexValue, U: DiskIndexValue + From + Into> AccountsIndexIterator<'a, T, U> { + pub fn new( + index: &'a AccountsIndex, + range: Option<&R>, + returns_items: AccountsIndexIteratorReturnsItems, + ) -> Self + where + R: RangeBounds, + { + match returns_items { + AccountsIndexIteratorReturnsItems::Sorted => { + AccountsIndexIterator::Sorted(AccountsIndexIteratorSorted::new(index, range)) + } + AccountsIndexIteratorReturnsItems::Unsorted => { + AccountsIndexIterator::Unsorted(AccountsIndexIteratorUnsorted::new(index, range)) + } + } + } +} + +impl + Into> Iterator + for AccountsIndexIterator<'_, T, U> +{ + type Item = Vec<(Pubkey, AccountMapEntry)>; + fn next(&mut self) -> Option { + match self { + AccountsIndexIterator::Sorted(iter) => iter.next(), + AccountsIndexIterator::Unsorted(iter) => iter.next(), + } + } +} + +#[cfg(feature = "dev-context-only-utils")] +impl + Into> AccountsIndexIterator<'_, T, U> { + pub fn bin_start_and_range(&self) -> (usize, usize) { + match self { + AccountsIndexIterator::Sorted(iter) => iter.inner.bin_start_and_range(), + AccountsIndexIterator::Unsorted(iter) => iter.inner.bin_start_and_range(), + } + } + + pub fn start_bin(&self) -> usize { + match self { + AccountsIndexIterator::Sorted(iter) => iter.inner.start_bin(), + AccountsIndexIterator::Unsorted(iter) => iter.inner.start_bin(), + } + } + + pub fn end_bin_inclusive(&self) -> usize { + match self { + AccountsIndexIterator::Sorted(iter) => iter.inner.end_bin_inclusive(), + AccountsIndexIterator::Unsorted(iter) => iter.inner.end_bin_inclusive(), + } } } @@ -819,6 +942,21 @@ impl + Into> AccountsIndex { AccountsIndexIterator::new(self, range, returns_items) } + #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))] + fn sorted_iter(&self, range: Option<&R>) -> AccountsIndexIteratorSorted + where + R: RangeBounds, + { + AccountsIndexIteratorSorted::new(self, range) + } + + fn unsorted_iter(&self, range: Option<&R>) -> AccountsIndexIteratorUnsorted + where + R: RangeBounds, + { + AccountsIndexIteratorUnsorted::new(self, range) + } + /// is the accounts index using disk as a backing store pub fn is_disk_index_enabled(&self) -> bool { self.storage.storage.is_disk_index_enabled() @@ -1442,7 +1580,7 @@ impl + Into> AccountsIndex { where R: RangeBounds + Debug + Sync, { - let iter = self.iter(Some(range), AccountsIndexIteratorReturnsItems::Unsorted); + let iter = self.unsorted_iter(Some(range)); iter.hold_range_in_memory(range, start_holding, thread_pool); } @@ -3956,10 +4094,7 @@ pub mod tests { } // Create an iterator for the whole pubkey range. - let mut iter = index.iter( - None::<&Range>, - AccountsIndexIteratorReturnsItems::Sorted, - ); + let mut iter = index.sorted_iter(None::<&Range>); // First iter.next() should return the first batch of pubkeys (1000 // pubkeys) out of the 2000 pubkeys in the first bin. And the remaining // 1000 pubkeys from the first bin should be cached in