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

Adds get_with_and_then() to AccountsIndex #35307

Merged
merged 4 commits into from
Feb 23, 2024
Merged

Adds get_with_and_then() to AccountsIndex #35307

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
293a926
Adds get_with_and_then() to AccountsIndex
brooksprumo Feb 17, 2024
97c3313
pr: pub(crate)
brooksprumo Feb 23, 2024
02b209e
pr: mark as dead code
brooksprumo Feb 23, 2024
24acb5f
pr: use #[cfg(test)]
brooksprumo Feb 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 34 additions & 31 deletions accounts-db/src/accounts_db.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11642,13 +11642,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 @@ -11703,13 +11706,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 @@ -11833,10 +11846,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 @@ -12022,10 +12032,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 @@ -12069,10 +12076,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 @@ -12081,10 +12085,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 @@ -12159,13 +12160,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
154 changes: 117 additions & 37 deletions accounts-db/src/accounts_index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1141,6 +1141,26 @@ 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
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)
})
}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's the main new function. In subsequent PRs more calls to get() will be replaced with calls to get_with_and_then().


/// Gets the index's entry for `pubkey` and clones it
///
/// Prefer `get_and_then()` whenever possible.
Expand All @@ -1154,6 +1174,17 @@ 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?
pub fn contains_with(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pub(crate)?

Copy link
Contributor Author

@brooksprumo brooksprumo Feb 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I can do that. For the most part, all of AccountsIndex really should not be public to the world, only the accounts db crate, right?

Another way to handle this is to remove the pub when importing the module in accounts-db/src/lib.rs.

(Mostly trying to avoid (1) unnecessary CI runs and (2) creating inconsistency within the accounts index code.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 97c3313.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My thought process is that contains() and get() et al really are the public APIs for AccountsIndex. So pub makes sense in my head.

Simultaneously we do not want others to call into AccountsIndex.

So only AccountsDb would have visibility to AccountsIndex, and all others would not. Thus only AccountsDb could call AccountsIndex functions.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah bummer... Because this PR doesn't use get_with_and_then() nor contains_with() outside of tests, clippy is throwing errors saying these functions are dead code 🫠

Do you have a preference for resolution? I could mark these functions are dead code in this PR, and remove the annotations in future PRs. Or I could change the vis to pub instead of pub(crate). Or something else?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. Marked as dead code in 02b209e.

I'm ok with marked as dead code. Isn't there a cfg[test] or something? Not a huge deal. Just trying to get it right with each commit.

Copy link
Contributor Author

@brooksprumo brooksprumo Feb 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I could decorate the functions with #[cfg(test)] too. Would you prefer I change to use #[cfg(test)]?

Just trying to get it right with each commit.

Yah, I prefer that way too. IMO I still think pub is the right visibility here for these functions 😸

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#[cfg(test)]

yes, that's what I'd prefer.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still think pub is the right visibility

I'm inclined to say make as little pub as possible. If this can be not pub, don't make it pub. Then we can get rid of it or change it later without any concern.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Decorated with #[cfg(test)] in 24acb5f.

&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 @@ -2158,8 +2189,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 @@ -2287,8 +2318,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 @@ -2370,10 +2401,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 @@ -2384,9 +2413,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 @@ -2408,10 +2435,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 @@ -2422,9 +2447,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 @@ -2686,8 +2709,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 @@ -2698,7 +2721,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 @@ -2726,7 +2749,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 @@ -2851,8 +2874,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 @@ -3076,8 +3109,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 @@ -3146,9 +3183,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 @@ -3162,8 +3208,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 @@ -3195,11 +3251,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 @@ -3265,8 +3341,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
Loading