Adds get_with_and_then() to AccountsIndex (#35307)

This commit is contained in:
Brooks 2024-02-23 18:48:36 -05:00 committed by GitHub
parent 1e47aacd0d
commit 54706a885b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 153 additions and 68 deletions

View File

@ -11640,13 +11640,16 @@ pub mod tests {
accounts.add_root_and_flush_write_cache(0);
let ancestors = vec![(0, 0)].into_iter().collect();
let id = {
let (lock, idx) = accounts
.accounts_index
.get_for_tests(&pubkey, Some(&ancestors), None)
.unwrap();
lock.slot_list()[idx].1.store_id()
};
let id = accounts
.accounts_index
.get_with_and_then(
&pubkey,
Some(&ancestors),
None,
false,
|(_slot, account_info)| account_info.store_id(),
)
.unwrap();
accounts.calculate_accounts_delta_hash(0);
//slot is still there, since gc is lazy
@ -11701,13 +11704,23 @@ pub mod tests {
let ancestors = vec![(0, 1)].into_iter().collect();
let (slot1, account_info1) = accounts
.accounts_index
.get_for_tests(&pubkey1, Some(&ancestors), None)
.map(|(account_list1, index1)| account_list1.slot_list()[index1])
.get_with_and_then(
&pubkey1,
Some(&ancestors),
None,
false,
|(slot, account_info)| (slot, account_info),
)
.unwrap();
let (slot2, account_info2) = accounts
.accounts_index
.get_for_tests(&pubkey2, Some(&ancestors), None)
.map(|(account_list2, index2)| account_list2.slot_list()[index2])
.get_with_and_then(
&pubkey2,
Some(&ancestors),
None,
false,
|(slot, account_info)| (slot, account_info),
)
.unwrap();
assert_eq!(slot1, 0);
assert_eq!(slot1, slot2);
@ -11831,10 +11844,7 @@ pub mod tests {
// zero lamport account, should no longer exist in accounts index
// because it has been removed
assert!(accounts
.accounts_index
.get_for_tests(&pubkey, None, None)
.is_none());
assert!(!accounts.accounts_index.contains_with(&pubkey, None, None));
}
#[test]
@ -12020,10 +12030,7 @@ pub mod tests {
// `pubkey1`, a zero lamport account, should no longer exist in accounts index
// because it has been removed by the clean
assert!(accounts
.accounts_index
.get_for_tests(&pubkey1, None, None)
.is_none());
assert!(!accounts.accounts_index.contains_with(&pubkey1, None, None));
// Secondary index should have purged `pubkey1` as well
let mut found_accounts = vec![];
@ -12067,10 +12074,7 @@ pub mod tests {
accounts.clean_accounts(Some(0), false, None, &EpochSchedule::default());
assert_eq!(accounts.alive_account_count_in_slot(0), 1);
assert_eq!(accounts.alive_account_count_in_slot(1), 1);
assert!(accounts
.accounts_index
.get_for_tests(&pubkey, None, None)
.is_some());
assert!(accounts.accounts_index.contains_with(&pubkey, None, None));
// Now the account can be cleaned up
accounts.clean_accounts(Some(1), false, None, &EpochSchedule::default());
@ -12079,10 +12083,7 @@ pub mod tests {
// The zero lamport account, should no longer exist in accounts index
// because it has been removed
assert!(accounts
.accounts_index
.get_for_tests(&pubkey, None, None)
.is_none());
assert!(!accounts.accounts_index.contains_with(&pubkey, None, None));
}
#[test]
@ -12157,13 +12158,15 @@ pub mod tests {
accounts.add_root_and_flush_write_cache(current_slot);
let (slot1, account_info1) = accounts
.accounts_index
.get_for_tests(&pubkey, None, None)
.map(|(account_list1, index1)| account_list1.slot_list()[index1])
.get_with_and_then(&pubkey, None, None, false, |(slot, account_info)| {
(slot, account_info)
})
.unwrap();
let (slot2, account_info2) = accounts
.accounts_index
.get_for_tests(&pubkey2, None, None)
.map(|(account_list2, index2)| account_list2.slot_list()[index2])
.get_with_and_then(&pubkey2, None, None, false, |(slot, account_info)| {
(slot, account_info)
})
.unwrap();
assert_eq!(slot1, current_slot);
assert_eq!(slot1, slot2);

View File

@ -1135,6 +1135,27 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
self.get_bin(pubkey).get_internal(pubkey, callback)
}
/// Gets the index's entry for `pubkey`, with `ancestors` and `max_root`,
/// and applies `callback` to it
#[cfg(test)]
pub(crate) fn get_with_and_then<R>(
&self,
pubkey: &Pubkey,
ancestors: Option<&Ancestors>,
max_root: Option<Slot>,
should_add_to_in_mem_cache: bool,
mut callback: impl FnMut((Slot, T)) -> R,
) -> Option<R> {
self.get_and_then(pubkey, |entry| {
let callback_result = entry.and_then(|entry| {
let slot_list = entry.slot_list.read().unwrap();
self.latest_slot(ancestors, &slot_list, max_root)
.map(|found_index| callback(slot_list[found_index]))
});
(should_add_to_in_mem_cache, callback_result)
})
}
/// Gets the index's entry for `pubkey` and clones it
///
/// Prefer `get_and_then()` whenever possible.
@ -1148,6 +1169,18 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
self.get_and_then(pubkey, |entry| (false, entry.is_some()))
}
/// Is `pubkey`, with `ancestors` and `max_root`, in the index?
#[cfg(test)]
pub(crate) fn contains_with(
&self,
pubkey: &Pubkey,
ancestors: Option<&Ancestors>,
max_root: Option<Slot>,
) -> bool {
self.get_with_and_then(pubkey, ancestors, max_root, false, |_| ())
.is_some()
}
fn slot_list_mut<RT>(
&self,
pubkey: &Pubkey,
@ -2157,8 +2190,8 @@ pub mod tests {
let index = AccountsIndex::<bool, bool>::default_for_tests();
let ancestors = Ancestors::default();
let key = &key;
assert!(index.get_for_tests(key, Some(&ancestors), None).is_none());
assert!(index.get_for_tests(key, None, None).is_none());
assert!(!index.contains_with(key, Some(&ancestors), None));
assert!(!index.contains_with(key, None, None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2286,8 +2319,8 @@ pub mod tests {
assert!(gc.is_empty());
let ancestors = Ancestors::default();
assert!(index.get_for_tests(&key, Some(&ancestors), None).is_none());
assert!(index.get_for_tests(&key, None, None).is_none());
assert!(!index.contains_with(&key, Some(&ancestors), None));
assert!(!index.contains_with(&key, None, None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2356,10 +2389,8 @@ pub mod tests {
index.set_startup(Startup::Normal);
let mut ancestors = Ancestors::default();
assert!(index
.get_for_tests(pubkey, Some(&ancestors), None)
.is_none());
assert!(index.get_for_tests(pubkey, None, None).is_none());
assert!(!index.contains_with(pubkey, Some(&ancestors), None));
assert!(!index.contains_with(pubkey, None, None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2370,9 +2401,7 @@ pub mod tests {
);
assert_eq!(num, 0);
ancestors.insert(slot, 0);
assert!(index
.get_for_tests(pubkey, Some(&ancestors), None)
.is_some());
assert!(index.contains_with(pubkey, Some(&ancestors), None));
assert_eq!(index.ref_count_from_storage(pubkey), 1);
index.unchecked_scan_accounts(
"",
@ -2394,10 +2423,8 @@ pub mod tests {
index.set_startup(Startup::Normal);
let mut ancestors = Ancestors::default();
assert!(index
.get_for_tests(pubkey, Some(&ancestors), None)
.is_none());
assert!(index.get_for_tests(pubkey, None, None).is_none());
assert!(!index.contains_with(pubkey, Some(&ancestors), None));
assert!(!index.contains_with(pubkey, None, None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2408,9 +2435,7 @@ pub mod tests {
);
assert_eq!(num, 0);
ancestors.insert(slot, 0);
assert!(index
.get_for_tests(pubkey, Some(&ancestors), None)
.is_some());
assert!(index.contains_with(pubkey, Some(&ancestors), None));
assert_eq!(index.ref_count_from_storage(pubkey), 1);
index.unchecked_scan_accounts(
"",
@ -2672,8 +2697,8 @@ pub mod tests {
assert_eq!(1, account_maps_stats_len(&index));
let mut ancestors = Ancestors::default();
assert!(index.get_for_tests(&key, Some(&ancestors), None).is_none());
assert!(index.get_for_tests(&key, None, None).is_none());
assert!(!index.contains_with(&key, Some(&ancestors), None));
assert!(!index.contains_with(&key, None, None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2684,7 +2709,7 @@ pub mod tests {
);
assert_eq!(num, 0);
ancestors.insert(slot, 0);
assert!(index.get_for_tests(&key, Some(&ancestors), None).is_some());
assert!(index.contains_with(&key, Some(&ancestors), None));
index.unchecked_scan_accounts(
"",
&ancestors,
@ -2712,7 +2737,7 @@ pub mod tests {
assert!(gc.is_empty());
let ancestors = vec![(1, 1)].into_iter().collect();
assert!(index.get_for_tests(&key, Some(&ancestors), None).is_none());
assert!(!index.contains_with(&key, Some(&ancestors), None));
let mut num = 0;
index.unchecked_scan_accounts(
@ -2837,8 +2862,18 @@ pub mod tests {
assert!(gc.is_empty());
let ancestors = vec![(0, 0)].into_iter().collect();
let (list, idx) = index.get_for_tests(&key, Some(&ancestors), None).unwrap();
assert_eq!(list.slot_list()[idx], (0, true));
index
.get_with_and_then(
&key,
Some(&ancestors),
None,
false,
|(slot, account_info)| {
assert_eq!(slot, 0);
assert!(account_info);
},
)
.unwrap();
let mut num = 0;
let mut found_key = false;
@ -3062,8 +3097,12 @@ pub mod tests {
assert!(gc.is_empty());
index.add_root(0);
let (list, idx) = index.get_for_tests(&key, None, None).unwrap();
assert_eq!(list.slot_list()[idx], (0, true));
index
.get_with_and_then(&key, None, None, false, |(slot, account_info)| {
assert_eq!(slot, 0);
assert!(account_info);
})
.unwrap();
}
#[test]
@ -3132,9 +3171,18 @@ pub mod tests {
UPSERT_POPULATE_RECLAIMS,
);
assert!(gc.is_empty());
let (list, idx) = index.get_for_tests(&key, Some(&ancestors), None).unwrap();
assert_eq!(list.slot_list()[idx], (0, true));
drop(list);
index
.get_with_and_then(
&key,
Some(&ancestors),
None,
false,
|(slot, account_info)| {
assert_eq!(slot, 0);
assert!(account_info);
},
)
.unwrap();
let mut gc = Vec::new();
index.upsert(
@ -3148,8 +3196,18 @@ pub mod tests {
UPSERT_POPULATE_RECLAIMS,
);
assert_eq!(gc, vec![(0, true)]);
let (list, idx) = index.get_for_tests(&key, Some(&ancestors), None).unwrap();
assert_eq!(list.slot_list()[idx], (0, false));
index
.get_with_and_then(
&key,
Some(&ancestors),
None,
false,
|(slot, account_info)| {
assert_eq!(slot, 0);
assert!(!account_info);
},
)
.unwrap();
}
#[test]
@ -3181,11 +3239,31 @@ pub mod tests {
UPSERT_POPULATE_RECLAIMS,
);
assert!(gc.is_empty());
let (list, idx) = index.get_for_tests(&key, Some(&ancestors), None).unwrap();
assert_eq!(list.slot_list()[idx], (0, true));
index
.get_with_and_then(
&key,
Some(&ancestors),
None,
false,
|(slot, account_info)| {
assert_eq!(slot, 0);
assert!(account_info);
},
)
.unwrap();
let ancestors = vec![(1, 0)].into_iter().collect();
let (list, idx) = index.get_for_tests(&key, Some(&ancestors), None).unwrap();
assert_eq!(list.slot_list()[idx], (1, false));
index
.get_with_and_then(
&key,
Some(&ancestors),
None,
false,
|(slot, account_info)| {
assert_eq!(slot, 1);
assert!(!account_info);
},
)
.unwrap();
}
#[test]
@ -3251,8 +3329,12 @@ pub mod tests {
// Updating index should not purge older roots, only purges
// previous updates within the same slot
assert_eq!(gc, vec![]);
let (list, idx) = index.get_for_tests(&key, None, None).unwrap();
assert_eq!(list.slot_list()[idx], (3, true));
index
.get_with_and_then(&key, None, None, false, |(slot, account_info)| {
assert_eq!(slot, 3);
assert!(account_info);
})
.unwrap();
let mut num = 0;
let mut found_key = false;