Merge pull request #6044 from kantai/feat/faster-readonly-marf

Perf: faster readonly MARF

Author: wileyj

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

@@ -14,13 +14,16 @@
 // You should have received a copy of the GNU General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+use std::cell::RefCell;
 use std::cmp::Ordering;
 use std::collections::{HashMap, HashSet};
 use std::io::{ErrorKind, Write};
 use std::ops::{Deref, DerefMut};
 use std::str::FromStr;
+use std::sync::{Arc, LazyLock, Mutex};
 use std::{cmp, fmt, fs};
 
+use clarity::util::lru_cache::LruCache;
 use clarity::vm::ast::ASTRules;
 use clarity::vm::costs::ExecutionCost;
 use clarity::vm::representations::{ClarityName, ContractName};
@@ -95,6 +98,9 @@ pub const REWARD_WINDOW_END: u64 = 144 * 90 + REWARD_WINDOW_START;
 
 pub type BlockHeaderCache = HashMap<ConsensusHash, (Option<BlockHeaderHash>, ConsensusHash)>;
 
+static DESCENDANCY_CACHE: LazyLock<Arc<Mutex<LruCache<(SortitionId, BlockHeaderHash), bool>>>> =
+    LazyLock::new(|| Arc::new(Mutex::new(LruCache::new(2000))));
+
 pub enum FindIter<R> {
     Found(R),
     Continue,
@@ -1112,12 +1118,46 @@ pub trait SortitionHandle {
                 test_debug!("No snapshot at height {}", block_at_burn_height);
                 db_error::NotFoundError
             })?;
+        let top_sortition_id = sn.sortition_id;
+
+        let mut cache = DESCENDANCY_CACHE
+            .lock()
+            .expect("FATAL: lock poisoned in SortitionDB");
 
         while sn.block_height >= earliest_block_height {
+            match cache.get(&(sn.sortition_id, potential_ancestor.clone())) {
+                Ok(Some(result)) => {
+                    if sn.sortition_id != top_sortition_id {
+                        if let Err(_) = cache
+                            .insert_clean((top_sortition_id, potential_ancestor.clone()), result)
+                        {
+                            *cache = LruCache::new(2000);
+                        }
+                    }
+                    return Ok(result);
+                }
+                // not cached, don't need to do anything.
+                Ok(None) => {}
+                // cache is broken, create a new one
+                Err(_) => {
+                    *cache = LruCache::new(2000);
+                }
+            }
+
             if !sn.sortition {
+                if let Err(_) =
+                    cache.insert_clean((top_sortition_id, potential_ancestor.clone()), false)
+                {
+                    *cache = LruCache::new(2000);
+                }
                 return Ok(false);
             }
             if &sn.winning_stacks_block_hash == potential_ancestor {
+                if let Err(_) =
+                    cache.insert_clean((top_sortition_id, potential_ancestor.clone()), true)
+                {
+                    *cache = LruCache::new(2000);
+                }
                 return Ok(true);
             }
 
@@ -1153,6 +1193,9 @@ pub trait SortitionHandle {
                 }
             }
         }
+        if let Err(_) = cache.insert_clean((top_sortition_id, potential_ancestor.clone()), false) {
+            *cache = LruCache::new(2000);
+        }
         return Ok(false);
     }
 }
@@ -2028,15 +2071,15 @@ impl<'a> SortitionHandleConn<'a> {
         connection: &'a SortitionDBConn<'a>,
         chain_tip: &SortitionId,
     ) -> Result<SortitionHandleConn<'a>, db_error> {
-        Ok(SortitionHandleConn {
-            context: SortitionHandleContext {
+        Ok(SortitionHandleConn::new(
+            &connection.index,
+            SortitionHandleContext {
                 chain_tip: chain_tip.clone(),
                 first_block_height: connection.context.first_block_height,
                 pox_constants: connection.context.pox_constants.clone(),
                 dryrun: connection.context.dryrun,
             },
-            index: connection.index,
-        })
+        ))
     }
 
     fn get_tip_indexed(&self, key: &str) -> Result<Option<String>, db_error> {
@@ -3723,15 +3766,15 @@ impl SortitionDBTx<'_> {
 
 impl SortitionDBConn<'_> {
     pub fn as_handle<'b>(&'b self, chain_tip: &SortitionId) -> SortitionHandleConn<'b> {
-        SortitionHandleConn {
-            index: self.index,
-            context: SortitionHandleContext {
+        SortitionHandleConn::new(
+            &self.index,
+            SortitionHandleContext {
                 first_block_height: self.context.first_block_height.clone(),
                 chain_tip: chain_tip.clone(),
                 pox_constants: self.context.pox_constants.clone(),
                 dryrun: self.context.dryrun,
             },
-        }
+        )
     }
 
     /// Given a burnchain consensus hash,

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

@@ -25,6 +25,7 @@ use stacks_common::types::chainstate::{BlockHeaderHash, TrieHash, TRIEHASH_ENCOD
 use stacks_common::util::hash::Sha512Trunc256Sum;
 use stacks_common::util::log;
 
+use super::storage::ReopenedTrieStorageConnection;
 use crate::chainstate::stacks::index::bits::{get_leaf_hash, get_node_hash, read_root_hash};
 use crate::chainstate::stacks::index::node::{
     clear_backptr, is_backptr, set_backptr, CursorError, TrieCursor, TrieNode, TrieNode16,
@@ -251,6 +252,20 @@ impl<T: MarfTrieId> MarfConnection<T> for MarfTransaction<'_, T> {
     }
 }
 
+impl<T: MarfTrieId> MarfConnection<T> for ReopenedTrieStorageConnection<'_, T> {
+    fn with_conn<F, R>(&mut self, exec: F) -> R
+    where
+        F: FnOnce(&mut TrieStorageConnection<T>) -> R,
+    {
+        let mut conn = self.connection();
+        exec(&mut conn)
+    }
+
+    fn sqlite_conn(&self) -> &Connection {
+        self.db_conn()
+    }
+}
+
 impl<T: MarfTrieId> MarfConnection<T> for MARF<T> {
     fn with_conn<F, R>(&mut self, exec: F) -> R
     where
@@ -1620,6 +1635,21 @@ impl<T: MarfTrieId> MARF<T> {
         })
     }
 
+    /// Build a read-only storage connection which can be used for reads without modifying the
+    ///  calling MARF struct (i.e., the tip pointer is only changed in the connection)
+    ///  but reusing self's existing SQLite Connection (avoiding the overhead of
+    ///   `reopen_readonly`).
+    pub fn reopen_connection(&self) -> Result<ReopenedTrieStorageConnection<'_, T>, Error> {
+        if self.open_chain_tip.is_some() {
+            error!(
+                "MARF at {} is already in the process of writing",
+                &self.storage.db_path
+            );
+            return Err(Error::InProgressError);
+        }
+        self.storage.reopen_connection()
+    }
+
     /// Get the root trie hash at a particular block
     pub fn get_root_hash_at(&mut self, block_hash: &T) -> Result<TrieHash, Error> {
         self.storage.connection().get_root_hash_at(block_hash)

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

@@ -1310,6 +1310,50 @@ impl<T: MarfTrieId> TrieStorageTransientData<T> {
     }
 }
 
+pub struct ReopenedTrieStorageConnection<'a, T: MarfTrieId> {
+    pub db_path: &'a str,
+    db: &'a Connection,
+    blobs: Option<TrieFile>,
+    data: TrieStorageTransientData<T>,
+    cache: TrieCache<T>,
+    bench: TrieBenchmark,
+    pub hash_calculation_mode: TrieHashCalculationMode,
+
+    /// row ID of a trie that represents unconfirmed state (i.e. trie state that will never become
+    /// part of the MARF, but nevertheless represents a persistent scratch space).  If this field
+    /// is Some(..), then the storage connection here was used to (re-)open an unconfirmed trie
+    /// (via `open_unconfirmed()` or `open_block()` when `self.unconfirmed()` is `true`), or used
+    /// to create an unconfirmed trie (via `extend_to_unconfirmed_block()`).
+    unconfirmed_block_id: Option<u32>,
+
+    // used in testing in order to short-circuit block-height lookups
+    //   when the trie struct is tested outside of marf.rs usage
+    #[cfg(test)]
+    pub test_genesis_block: Option<T>,
+}
+
+impl<'a, T: MarfTrieId> ReopenedTrieStorageConnection<'a, T> {
+    pub fn db_conn(&self) -> &Connection {
+        self.db
+    }
+
+    pub fn connection(&mut self) -> TrieStorageConnection<'_, T> {
+        TrieStorageConnection {
+            db: SqliteConnection::ConnRef(&self.db),
+            db_path: self.db_path,
+            data: &mut self.data,
+            blobs: self.blobs.as_mut(),
+            cache: &mut self.cache,
+            bench: &mut self.bench,
+            hash_calculation_mode: self.hash_calculation_mode,
+            unconfirmed_block_id: None,
+
+            #[cfg(test)]
+            test_genesis_block: &mut self.test_genesis_block,
+        }
+    }
+}
+
 impl<T: MarfTrieId> TrieFileStorage<T> {
     pub fn connection(&mut self) -> TrieStorageConnection<'_, T> {
         TrieStorageConnection {
@@ -1327,6 +1371,54 @@ impl<T: MarfTrieId> TrieFileStorage<T> {
         }
     }
 
+    /// Build a read-only storage connection which can be used for reads without modifying the
+    ///  calling TrieFileStorage struct (i.e., the tip pointer is only changed in the connection)
+    ///  but reusing the TrieFileStorage's existing SQLite Connection (avoiding the overhead of
+    ///   `reopen_readonly`).
+    pub fn reopen_connection(&self) -> Result<ReopenedTrieStorageConnection<'_, T>, Error> {
+        let data = TrieStorageTransientData {
+            uncommitted_writes: self.data.uncommitted_writes.clone(),
+            cur_block: self.data.cur_block.clone(),
+            cur_block_id: self.data.cur_block_id.clone(),
+
+            read_count: 0,
+            read_backptr_count: 0,
+            read_node_count: 0,
+            read_leaf_count: 0,
+
+            write_count: 0,
+            write_node_count: 0,
+            write_leaf_count: 0,
+
+            trie_ancestor_hash_bytes_cache: None,
+
+            readonly: true,
+            unconfirmed: self.unconfirmed(),
+        };
+        // perf note: should we attempt to clone the cache
+        let cache = TrieCache::default();
+        let blobs = if self.blobs.is_some() {
+            Some(TrieFile::from_db_path(&self.db_path, true)?)
+        } else {
+            None
+        };
+        let bench = TrieBenchmark::new();
+        let hash_calculation_mode = self.hash_calculation_mode;
+        let unconfirmed_block_id = None;
+        Ok(ReopenedTrieStorageConnection {
+            db_path: &self.db_path,
+            db: &self.db,
+            blobs,
+            data,
+            cache,
+            bench,
+            hash_calculation_mode,
+            unconfirmed_block_id,
+            #[cfg(test)]
+            test_genesis_block: self.test_genesis_block.clone(),
+        })
+    }
+
     pub fn transaction(&mut self) -> Result<TrieStorageTransaction<'_, T>, Error> {
         if self.readonly() {
             return Err(Error::ReadOnlyError);

stackslib/src/clarity_vm/database/marf.rs

@@ -263,10 +263,7 @@ impl MarfedKV {
     }
 
     pub fn index_conn<C>(&self, context: C) -> IndexDBConn<'_, C, StacksBlockId> {
-        IndexDBConn {
-            index: &self.marf,
-            context,
-        }
+        IndexDBConn::new(&self.marf, context)
     }
 }
 

stackslib/src/util_lib/db.rs

@@ -613,13 +613,13 @@ impl<'a, C, T: MarfTrieId> IndexDBConn<'a, C, T> {
         ancestor_block_hash: &T,
         tip_block_hash: &T,
     ) -> Result<Option<u64>, Error> {
-        get_ancestor_block_height(self.index, ancestor_block_hash, tip_block_hash)
+        get_ancestor_block_height(&self.index, ancestor_block_hash, tip_block_hash)
     }
 
     /// Get a value from the fork index
     pub fn get_indexed(&self, header_hash: &T, key: &str) -> Result<Option<String>, Error> {
-        let mut ro_index = self.index.reopen_readonly()?;
-        get_indexed(&mut ro_index, header_hash, key)
+        let mut connection = self.index.reopen_connection()?;
+        get_indexed(&mut connection, header_hash, key)
     }
 
     pub fn conn(&self) -> &DBConn {
@@ -727,7 +727,7 @@ pub fn get_ancestor_block_hash<T: MarfTrieId>(
     tip_block_hash: &T,
 ) -> Result<Option<T>, Error> {
     assert!(block_height <= u32::MAX as u64);
-    let mut read_only = index.reopen_readonly()?;
+    let mut read_only = index.reopen_connection()?;
     let bh = read_only.get_block_at_height(block_height as u32, tip_block_hash)?;
     Ok(bh)
 }
@@ -738,7 +738,7 @@ pub fn get_ancestor_block_height<T: MarfTrieId>(
     ancestor_block_hash: &T,
     tip_block_hash: &T,
 ) -> Result<Option<u64>, Error> {
-    let mut read_only = index.reopen_readonly()?;
+    let mut read_only = index.reopen_connection()?;
     let height_opt = read_only
         .get_block_height(ancestor_block_hash, tip_block_hash)?
         .map(|height| height as u64);