wip: limit use of self ref struct in accounts index

accounts-db/src/accounts.rs

@@ -131,13 +131,15 @@ impl Accounts {
         address_table_lookup: &MessageAddressTableLookup,
         slot_hashes: &SlotHashes,
     ) -> std::result::Result<LoadedAddresses, AddressLookupError> {
+        self.accounts_db.accounts_index.valid.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+
         let table_account = self
             .accounts_db
             .load_with_fixed_root(ancestors, &address_table_lookup.account_key)
             .map(|(account, _rent)| account)
             .ok_or(AddressLookupError::LookupTableAccountNotFound)?;
 
-        if table_account.owner() == &address_lookup_table::program::id() {
+        let r = if table_account.owner() == &address_lookup_table::program::id() {
             let current_slot = ancestors.max_slot();
             let lookup_table = AddressLookupTable::deserialize(table_account.data())
                 .map_err(|_ix_err| AddressLookupError::InvalidAccountData)?;
@@ -156,7 +158,9 @@ impl Accounts {
             })
         } else {
             Err(AddressLookupError::InvalidAccountOwner)
-        }
+        };
+        self.accounts_db.accounts_index.valid.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        r
     }
 
     /// Slow because lock is held for 1 operation instead of many

accounts-db/src/accounts_db.rs

@@ -5059,36 +5059,36 @@ impl AccountsDb {
         max_root: Option<Slot>,
         clone_in_lock: bool,
     ) -> Option<(Slot, StorageLocation, Option<LoadedAccountAccessor<'a>>)> {
-        let (lock, index) = match self.accounts_index.get(pubkey, Some(ancestors), max_root) {
-            AccountIndexGetResult::Found(lock, index) => (lock, index),
-            // we bail out pretty early for missing.
-            AccountIndexGetResult::NotFound => {
-                return None;
-            }
-        };
-
-        let slot_list = lock.slot_list();
-        let (slot, info) = slot_list[index];
-        let storage_location = info.storage_location();
-        let some_from_slow_path = if clone_in_lock {
-            // the fast path must have failed.... so take the slower approach
-            // of copying potentially large Account::data inside the lock.
-
-            // calling check_and_get_loaded_account is safe as long as we're guaranteed to hold
-            // the lock during the time and there should be no purge thanks to alive ancestors
-            // held by our caller.
-            Some(self.get_account_accessor(slot, pubkey, &storage_location))
-        } else {
-            None
-        };
-
-        Some((slot, storage_location, some_from_slow_path))
-        // `lock` is dropped here rather pretty quickly with clone_in_lock = false,
-        // so the entry could be raced for mutation by other subsystems,
-        // before we actually provision an account data for caller's use from now on.
-        // This is traded for less contention and resultant performance, introducing fair amount of
-        // delicate handling in retry_to_get_account_accessor() below ;)
-        // you're warned!
+        if !clone_in_lock {
+            let (slot, info) = self.accounts_index.get_internal(pubkey, Some(ancestors), max_root)?;
+
+            let storage_location = info.storage_location();
+            let some_from_slow_path = None;
+
+            Some((slot, storage_location, some_from_slow_path))
+        }
+        else {
+            self.accounts_index.get_and_lock(pubkey, Some(ancestors), max_root,|(slot, info)| {
+            let storage_location = info.storage_location();
+            let some_from_slow_path = 
+                // the fast path must have failed.... so take the slower approach
+                // of copying potentially large Account::data inside the lock.
+    
+                // calling check_and_get_loaded_account is safe as long as we're guaranteed to hold
+                // the lock during the time and there should be no purge thanks to alive ancestors
+                // held by our caller.
+                Some(self.get_account_accessor(slot, pubkey, &storage_location));
+
+            (slot, storage_location, some_from_slow_path)
+            // `lock` is dropped here rather pretty quickly with clone_in_lock = false,
+            // so the entry could be raced for mutation by other subsystems,
+            // before we actually provision an account data for caller's use from now on.
+            // This is traded for less contention and resultant performance, introducing fair amount of
+            // delicate handling in retry_to_get_account_accessor() below ;)
+            // you're warned!
+            })
+        }
+
     }
 
     fn retry_to_get_account_accessor<'a>(

accounts-db/src/accounts_index.rs

@@ -1,3 +1,5 @@
+use crate::account_info::StorageLocation;
+
 use {
     crate::{
         accounts_index_storage::{AccountsIndexStorage, Startup},
@@ -677,6 +679,9 @@ pub enum AccountsIndexScanResult {
 /// T: account info type to interact in in-memory items
 /// U: account info type to be persisted to disk
 pub struct AccountsIndex<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> {
+    pub valid: AtomicUsize,
+    pub count: AtomicUsize,
+
     pub account_maps: LockMapType<T, U>,
     pub bin_calculator: PubkeyBinCalculator24,
     program_id_index: SecondaryIndex<DashMapSecondaryIndexEntry>,
@@ -725,6 +730,8 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
             .and_then(|config| config.scan_results_limit_bytes);
         let (account_maps, bin_calculator, storage) = Self::allocate_accounts_index(config, exit);
         Self {
+            valid: AtomicUsize::default(),
+            count: AtomicUsize::default(),
             account_maps,
             bin_calculator,
             program_id_index: SecondaryIndex::<DashMapSecondaryIndexEntry>::new(
@@ -1134,6 +1141,12 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         pubkey: &Pubkey,
         lock: &AccountMaps<'_, T, U>,
     ) -> Option<ReadAccountMapEntry<T>> {
+        if self.valid.load(Ordering::Relaxed) == 0 {
+            // this is only called at startup
+            panic!("{}, {}", line!(), self.count.load(Ordering::Relaxed));
+        } else {
+            self.count.fetch_add(1, Ordering::Relaxed);
+        }
         lock.get(pubkey)
             .map(ReadAccountMapEntry::from_account_map_entry)
     }
@@ -1437,6 +1450,11 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         max_root: Option<Slot>,
     ) -> AccountIndexGetResult<T> {
         let read_lock = self.get_bin(pubkey);
+        if self.valid.load(Ordering::Relaxed) == 0 {
+            panic!("{}, {}", line!(), self.count.load(Ordering::Relaxed));
+        } else {
+            self.count.fetch_add(1, Ordering::Relaxed);
+        }
         let account = read_lock
             .get(pubkey)
             .map(ReadAccountMapEntry::from_account_map_entry);
@@ -1454,6 +1472,57 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         }
     }
 
+    /// Get an account
+    /// The latest account that appears in `ancestors` or `roots` is returned.
+    pub fn get_and_lock<R>(
+        &self,
+        pubkey: &Pubkey,
+        ancestors: Option<&Ancestors>,
+        max_root: Option<Slot>,
+        callback: impl Fn((Slot, T)) -> R,
+    ) -> Option<R> {
+        let read_lock = self.get_bin(pubkey);
+        if self.valid.load(Ordering::Relaxed) == 0 {
+            panic!("{}, {}", line!(), self.count.load(Ordering::Relaxed));
+        } else {
+            self.count.fetch_add(1, Ordering::Relaxed);
+        }
+        read_lock.get_internal(pubkey, |account| {
+            account
+                .map(|account| {
+                    let read = account.slot_list.read().unwrap();
+                    let slot_list = &read;
+                    self.latest_slot(ancestors, slot_list, max_root)
+                        .map(|found_index| (false, Some(callback(slot_list[found_index]))))
+                })
+                .flatten()
+                .unwrap_or((false, None))
+        })
+    }
+
+    /// Get an account
+    /// The latest account that appears in `ancestors` or `roots` is returned.
+    pub fn get_internal(
+        &self,
+        pubkey: &Pubkey,
+        ancestors: Option<&Ancestors>,
+        max_root: Option<Slot>,
+    ) -> Option<(Slot, T)> {
+        let read_lock = self.get_bin(pubkey);
+        // note this still gets the refcount on the entry. We don't need that.
+        read_lock.get_internal(pubkey, |account| {
+            account
+                .map(|account| {
+                    let slot_list = account.slot_list.read().unwrap();
+                    let found_index = self.latest_slot(ancestors, &slot_list, max_root);
+                    found_index.map(|found_index| slot_list[found_index])
+                })
+                .flatten()
+                .map(|r| (false, Some(r)))
+                .unwrap_or((false, None))
+        })
+    }
+
     // Get the maximum root <= `max_allowed_root` from the given `slice`
     fn get_newest_root_in_slot_list(
         alive_roots: &RollingBitField,

runtime/src/bank.rs

@@ -1213,6 +1213,8 @@ impl Bank {
         reward_calc_tracer: Option<impl RewardCalcTracer>,
         new_bank_options: NewBankOptions,
     ) -> Self {
+        parent.rc.accounts.accounts_db.accounts_index.valid.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+
         let mut time = Measure::start("bank::new_from_parent");
         let NewBankOptions { vote_only_bank } = new_bank_options;
 
@@ -1473,6 +1475,10 @@ impl Bank {
             .submit(parent.slot());
 
         new.loaded_programs_cache.write().unwrap().stats.reset();
+        parent.rc.accounts.accounts_db.accounts_index.valid.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+        error!("jwash: {}, {}", line!(), parent.rc.accounts.accounts_db.accounts_index.count.swap(0, Ordering::Relaxed));
+
+
         new
     }
 
@@ -1758,6 +1764,7 @@ impl Bank {
         debug_do_not_add_builtins: bool,
         accounts_data_size_initial: u64,
     ) -> Self {
+        bank_rc.accounts.accounts_db.accounts_index.valid.fetch_add(1, Ordering::Relaxed);
         let now = Instant::now();
         let ancestors = Ancestors::from(&fields.ancestors);
         // For backward compatibility, we can only serialize and deserialize
@@ -1895,6 +1902,7 @@ impl Bank {
                 i64
             ),
         );
+        bank.rc.accounts.accounts_db.accounts_index.valid.fetch_sub(1, Ordering::Relaxed);
         bank
     }
 
@@ -2169,6 +2177,8 @@ impl Bank {
     }
 
     fn update_slot_hashes(&self) {
+        self.rc.accounts.accounts_db.accounts_index.valid.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+
         self.update_sysvar_account(&sysvar::slot_hashes::id(), |account| {
             let mut slot_hashes = account
                 .as_ref()
@@ -2180,6 +2190,7 @@ impl Bank {
                 self.inherit_specially_retained_account_fields(account),
             )
         });
+        self.rc.accounts.accounts_db.accounts_index.valid.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
     }
 
     pub fn get_slot_history(&self) -> SlotHistory {

runtime/src/snapshot_bank_utils.rs

@@ -843,8 +843,13 @@ fn verify_slot_deltas(
     slot_deltas: &[BankSlotDelta],
     bank: &Bank,
 ) -> std::result::Result<(), VerifySlotDeltasError> {
+    bank.rc.accounts.accounts_db.accounts_index.valid.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+
     let info = verify_slot_deltas_structural(slot_deltas, bank.slot())?;
-    verify_slot_deltas_with_history(&info.slots, &bank.get_slot_history(), bank.slot())
+    let r = verify_slot_deltas_with_history(&info.slots, &bank.get_slot_history(), bank.slot());
+    bank.rc.accounts.accounts_db.accounts_index.valid.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
+
+    r
 }
 
 /// Verify that the snapshot's slot deltas are not corrupt/invalid