AcctIdx: refactor get to use more refs (#20268)

runtime/src/accounts_index.rs

@@ -1176,12 +1176,7 @@ impl<T: IndexValue> AccountsIndex<T> {
         let read_lock = self.account_maps[self.bin_calculator.bin_from_pubkey(pubkey)]
             .read()
             .unwrap();
-        let get = read_lock.get(pubkey);
+        read_lock.slot_list_mut(pubkey, user)
-        get.map(|entry| {
-            let result = user(&mut entry.slot_list.write().unwrap());
-            entry.set_dirty(true);
-            result
-        })
     }
 
     pub fn handle_dead_keys(

runtime/src/in_mem_accounts_index.rs

@@ -9,7 +9,7 @@ use solana_sdk::{clock::Slot, pubkey::Pubkey};
 use std::collections::{hash_map::Entry, HashMap};
 use std::ops::{Bound, RangeBounds, RangeInclusive};
 use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering};
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, RwLock, RwLockWriteGuard};
 
 use std::fmt::Debug;
 type K = Pubkey;
@@ -133,9 +133,15 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
     }
 
     /// lookup 'pubkey' by only looking in memory. Does not look on disk.
-    fn get_only_in_mem(&self, pubkey: &K) -> Option<AccountMapEntry<T>> {
+    /// callback is called whether pubkey is found or not
+    fn get_only_in_mem<RT>(
+        &self,
+        pubkey: &K,
+        callback: impl for<'a> FnOnce(Option<&'a Arc<AccountMapEntryInner<T>>>) -> RT,
+    ) -> RT {
         let m = Measure::start("get");
-        let result = self.map().read().unwrap().get(pubkey).map(Arc::clone);
+        let map = self.map().read().unwrap();
+        let result = map.get(pubkey);
         let stats = self.stats();
         let (count, time) = if result.is_some() {
             (&stats.gets_from_mem, &stats.get_mem_us)
@@ -145,32 +151,50 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         Self::update_time_stat(time, m);
         Self::update_stat(count, 1);
 
-        if let Some(entry) = result.as_ref() {
+        callback(if let Some(entry) = result {
             entry.set_age(self.storage.future_age_to_flush());
-        }
+            Some(entry)
-        result
+        } else {
+            drop(map);
+            None
+        })
     }
 
     /// lookup 'pubkey' in index (in mem or on disk)
     pub fn get(&self, pubkey: &K) -> Option<AccountMapEntry<T>> {
-        let result = self.get_only_in_mem(pubkey);
+        self.get_internal(pubkey, |entry| entry.map(Arc::clone))
-        if result.is_some() {
-            return result;
     }
 
+    /// lookup 'pubkey' in index.
+    /// call 'callback' whether found or not
+    fn get_internal<RT>(
+        &self,
+        pubkey: &K,
+        callback: impl for<'a> FnOnce(Option<&Arc<AccountMapEntryInner<T>>>) -> RT,
+    ) -> RT {
+        self.get_only_in_mem(pubkey, |entry| {
+            if let Some(entry) = entry {
+                entry.set_age(self.storage.future_age_to_flush());
+                callback(Some(entry))
+            } else {
                 // not in cache, look on disk
                 let stats = &self.stats();
-        let new_entry = self.load_account_entry_from_disk(pubkey)?;
+                let disk_entry = self.load_account_entry_from_disk(pubkey);
+                if disk_entry.is_none() {
+                    return callback(None);
+                }
+                let disk_entry = disk_entry.unwrap();
                 let mut map = self.map().write().unwrap();
                 let entry = map.entry(*pubkey);
-        let result = match entry {
+                match entry {
-            Entry::Occupied(occupied) => Arc::clone(occupied.get()),
+                    Entry::Occupied(occupied) => callback(Some(occupied.get())),
                     Entry::Vacant(vacant) => {
                         stats.insert_or_delete_mem(true, self.bin);
-                Arc::clone(vacant.insert(new_entry))
+                        callback(Some(vacant.insert(disk_entry)))
                     }
-        };
+                }
-        Some(result)
+            }
+        })
     }
 
     fn remove_if_slot_list_empty_value(&self, slot_list: SlotSlice<T>) -> bool {
@@ -245,6 +269,20 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         self.remove_if_slot_list_empty_entry(entry)
     }
 
+    pub fn slot_list_mut<RT>(
+        &self,
+        pubkey: &Pubkey,
+        user: impl for<'a> FnOnce(&mut RwLockWriteGuard<'a, SlotList<T>>) -> RT,
+    ) -> Option<RT> {
+        self.get_internal(pubkey, |entry| {
+            entry.map(|entry| {
+                let result = user(&mut entry.slot_list.write().unwrap());
+                entry.set_dirty(true);
+                result
+            })
+        })
+    }
+
     pub fn upsert(
         &self,
         pubkey: &Pubkey,
@@ -253,17 +291,16 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         previous_slot_entry_was_cached: bool,
     ) {
         // try to get it just from memory first using only a read lock
-        if let Some(get) = self.get_only_in_mem(pubkey) {
+        self.get_only_in_mem(pubkey, |entry| {
+            if let Some(entry) = entry {
                 Self::lock_and_update_slot_list(
-                &get,
+                    entry,
                     new_value.into(),
                     reclaims,
                     previous_slot_entry_was_cached,
                 );
                 Self::update_stat(&self.stats().updates_in_mem, 1);
-            return;
+            } else {
-        }
-
                 let m = Measure::start("entry");
                 let mut map = self.map().write().unwrap();
                 let entry = map.entry(*pubkey);
@@ -309,6 +346,8 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                         self.stats().insert_or_delete(true, self.bin);
                     }
                 }
+            };
+        })
     }
 
     // Try to update an item in the slot list the given `slot` If an item for the slot