perf: do not read whole snapshot for ancestor check in find_new_block_arrivals

Author: kantai

stackslib/src/chainstate/burn/db/sortdb.rs

@@ -1264,13 +1264,25 @@ impl<'a> SortitionHandleTx<'a> {
         burn_header_hash: &BurnchainHeaderHash,
         chain_tip: &SortitionId,
     ) -> Result<Option<BlockSnapshot>, db_error> {
+        let Some(sortition_id) = self.get_sortition_id_for_bhh(burn_header_hash, chain_tip)? else {
+            return Ok(None);
+        };
+
+        SortitionDB::get_block_snapshot(self.tx(), &sortition_id)
+    }
+
+    fn get_sortition_id_for_bhh(
+        &mut self,
+        burn_header_hash: &BurnchainHeaderHash,
+        chain_tip: &SortitionId,
+    ) -> Result<Option<SortitionId>, db_error> {
         let sortition_identifier_key = db_keys::sortition_id_for_bhh(burn_header_hash);
         let sortition_id = match self.get_indexed(chain_tip, &sortition_identifier_key)? {
             None => return Ok(None),
             Some(x) => SortitionId::from_hex(&x).expect("FATAL: bad Sortition ID stored in DB"),
         };
 
-        SortitionDB::get_block_snapshot(self.tx(), &sortition_id)
+        Ok(Some(sortition_id))
     }
 
     /// Get a leader key at a specific location in the burn chain's fork history, given the
@@ -6528,25 +6540,25 @@ impl SortitionHandleTx<'_> {
             }
 
             // must be an ancestor of this tip, or must be this tip
-            if let Some(sn) =
-                self.get_block_snapshot(&arrival_sn.burn_header_hash, &parent_tip.sortition_id)?
+            if let Some(sortition_id) = self
+                .get_sortition_id_for_bhh(&arrival_sn.burn_header_hash, &parent_tip.sortition_id)?
             {
-                if !sn.pox_valid || sn != arrival_sn {
+                if sortition_id != arrival_sn.sortition_id {
                     continue;
                 }
 
                 debug!(
                     "New Stacks anchored block arrived: block {}/{} ({}) ari={} tip={}",
-                    &sn.consensus_hash,
-                    &sn.winning_stacks_block_hash,
-                    sn.stacks_block_height,
+                    &arrival_sn.consensus_hash,
+                    &arrival_sn.winning_stacks_block_hash,
+                    arrival_sn.stacks_block_height,
                     ari,
                     &parent_tip.burn_header_hash
                 );
                 new_block_arrivals.push((
-                    sn.consensus_hash,
-                    sn.winning_stacks_block_hash,
-                    sn.stacks_block_height,
+                    arrival_sn.consensus_hash,
+                    arrival_sn.winning_stacks_block_hash,
+                    arrival_sn.stacks_block_height,
                 ));
             } else {
                 // this block did not arrive on an ancestor block

stackslib/src/chainstate/stacks/index/cache.rs

@@ -15,6 +15,7 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 use std::char::from_digit;
+use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::hash::{Hash, Hasher};
 use std::io::{BufWriter, Cursor, Read, Seek, SeekFrom, Write};
@@ -137,6 +138,11 @@ impl<T: MarfTrieId> TrieCacheState<T> {
         self.block_hash_cache.get(&block_id).cloned()
     }
 
+    /// Return the HashMap Entry for a block_id
+    pub fn entry_block_hash(&mut self, block_id: u32) -> Entry<'_, u32, T> {
+        self.block_hash_cache.entry(block_id)
+    }
+
     /// Cache a block hash, given its ID
     pub fn store_block_hash(&mut self, block_id: u32, block_hash: T) {
         assert!(!self.block_hash_cache.contains_key(&block_id));
@@ -309,6 +315,11 @@ impl<T: MarfTrieId> TrieCache<T> {
         self.state_mut().load_block_hash(block_id)
     }
 
+    /// Get entry for a block hash, given its ID
+    pub fn entry_block_hash<'a>(&'a mut self, block_id: u32) -> Entry<'a, u32, T> {
+        self.state_mut().entry_block_hash(block_id)
+    }
+
     /// Store a block's ID and hash to teh cache.
     pub fn store_block_hash(&mut self, block_id: u32, block_hash: T) {
         self.state_mut().store_block_hash(block_id, block_hash)

stackslib/src/chainstate/stacks/index/storage.rs

@@ -15,6 +15,7 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 use std::char::from_digit;
+use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::hash::{Hash, Hasher};
 use std::io::{BufWriter, Cursor, Read, Seek, SeekFrom, Write};
@@ -108,12 +109,15 @@ impl<T: MarfTrieId> BlockMap for TrieStorageConnection<'_, T> {
         trie_sql::get_block_hash(&self.db, id)
     }
 
-    fn get_block_hash_caching(&mut self, id: u32) -> Result<&T, Error> {
-        if !self.is_block_hash_cached(id) {
-            let block_hash = self.get_block_hash(id)?;
-            self.cache.store_block_hash(id, block_hash);
+    fn get_block_hash_caching<'a>(&'a mut self, id: u32) -> Result<&'a T, Error> {
+        match self.cache.entry_block_hash(id) {
+            Entry::Occupied(occupied_entry) => Ok(occupied_entry.into_mut()),
+            Entry::Vacant(vacant_entry) => {
+                let block_hash = trie_sql::get_block_hash(&self.db, id)?;
+                let block_hash_ref = vacant_entry.insert(block_hash);
+                Ok(block_hash_ref)
+            }
         }
-        self.cache.ref_block_hash(id).ok_or(Error::NotFoundError)
     }
 
     fn is_block_hash_cached(&self, id: u32) -> bool {