Skip to content
This repository has been archived by the owner on Jan 22, 2025. It is now read-only.

Commit

Permalink
Adds get_with_and_then() to AccountsIndex (#35307)
Browse files Browse the repository at this point in the history
  • Loading branch information
brooksprumo authored Feb 23, 2024
1 parent 1e47aac commit 54706a8
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 68 deletions.
65 changes: 34 additions & 31 deletions accounts-db/src/accounts_db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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![];
Expand Down Expand Up @@ -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());
Expand All @@ -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]
Expand Down Expand Up @@ -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);
Expand Down
156 changes: 119 additions & 37 deletions accounts-db/src/accounts_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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(
"",
Expand All @@ -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(
Expand All @@ -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(
"",
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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(
Expand All @@ -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]
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 54706a8

Please sign in to comment.