feat: handle fallible LruCache in NonceCache

Author: obycode

stacks-common/src/util/lru_cache.rs

@@ -15,7 +15,7 @@
 
 use std::fmt::Display;
 
-use hashbrown::HashMap;
+use hashbrown::{HashMap, HashSet};
 
 /// Node in the doubly linked list
 struct Node<K, V> {
@@ -75,7 +75,12 @@ impl<K: Display, V: Display> Display for LruCache<K, V> {
 
 impl<K: Eq + std::hash::Hash + Clone, V: Copy> LruCache<K, V> {
     /// Create a new LRU cache with the given capacity
-    pub fn new(capacity: usize) -> Self {
+    pub fn new(mut capacity: usize) -> Self {
+        if capacity == 0 {
+            error!("Capacity must be greater than 0. Defaulting to 1024.");
+            capacity = 1024;
+        }
+
         LruCache {
             capacity,
             cache: HashMap::new(),
@@ -218,18 +223,37 @@ impl<K: Eq + std::hash::Hash + Clone, V: Copy> LruCache<K, V> {
 
     /// Flush all dirty values in the cache, calling the given function, `f`,
     /// for each dirty value.
-    pub fn flush<E>(&mut self, mut f: impl FnMut(&K, V) -> Result<(), E>) -> Result<(), E> {
-        let mut index = self.head;
-        while index != self.capacity {
-            let next = self.order[index].next;
-            if self.order[index].dirty {
-                let value = self.order[index].value;
-                f(&self.order[index].key, value)?;
-                self.order[index].dirty = false;
+    /// Outer result is an error iff the cache is corrupted and should be discarded.
+    /// Inner result is an error iff the function, `f`, returns an error.
+    pub fn flush<E>(
+        &mut self,
+        mut f: impl FnMut(&K, V) -> Result<(), E>,
+    ) -> Result<Result<(), E>, ()> {
+        let mut current = self.head;
+
+        // Keep track of visited nodes to detect cycles
+        let mut visited = HashSet::new();
+
+        while current != self.capacity {
+            // Detect cycles
+            if !visited.insert(current) {
+                return Err(());
             }
-            index = next;
+
+            let node = self.order.get_mut(current).ok_or(())?;
+            let next = node.next;
+            if node.dirty {
+                let value = node.value;
+
+                // Call the flush function
+                match f(&node.key, value) {
+                    Ok(()) => node.dirty = false,
+                    Err(e) => return Ok(Err(e)),
+                }
+            }
+            current = next;
         }
-        Ok(())
+        Ok(Ok(()))
     }
 }
 

stackslib/src/config/mod.rs

@@ -2664,6 +2664,14 @@ impl MinerConfigFile {
             } else {
                 miner_default_config.tenure_cost_limit_per_block_percentage
             };
+
+        let nonce_cache_size = self
+            .nonce_cache_size
+            .unwrap_or(miner_default_config.nonce_cache_size);
+        if nonce_cache_size == 0 {
+            return Err("miner.nonce_cache_size must be greater than 0".to_string());
+        }
+
         Ok(MinerConfig {
             first_attempt_time_ms: self
                 .first_attempt_time_ms

stackslib/src/core/nonce_cache.rs

@@ -39,12 +39,23 @@ use crate::util_lib::db::{query_row, u64_to_sql, DBConn, Error as db_error};
 pub struct NonceCache {
     /// In-memory LRU cache of nonces.
     cache: LruCache<StacksAddress, u64>,
+    max_size: usize,
 }
 
 impl NonceCache {
     pub fn new(max_size: usize) -> Self {
         Self {
             cache: LruCache::new(max_size),
+            max_size,
+        }
+    }
+
+    /// Reset the cache to an empty state and clear the nonce DB.
+    /// This should only be called when the cache is corrupted.
+    fn reset_cache(&mut self, conn: &mut DBConn) {
+        self.cache = LruCache::new(self.max_size);
+        if let Err(e) = conn.execute("DELETE FROM nonces", []) {
+            warn!("error clearing nonces table: {e}");
         }
     }
 
@@ -67,9 +78,14 @@ impl NonceCache {
         C: ClarityConnection,
     {
         // Check in-memory cache
-        if let Some(cached_nonce) = self.cache.get(address) {
-            return cached_nonce;
-        };
+        match self.cache.get(address) {
+            Ok(Some(nonce)) => return nonce,
+            Ok(None) => {}
+            Err(_) => {
+                // The cache is corrupt, reset it
+                self.reset_cache(mempool_db);
+            }
+        }
 
         // Check sqlite cache
         let db_nonce_opt = db_get_nonce(mempool_db, address).unwrap_or_else(|e| {
@@ -79,7 +95,14 @@ impl NonceCache {
         if let Some(db_nonce) = db_nonce_opt {
             // Insert into in-memory cache, but it is not dirty,
             // since we just got it from the database.
-            let evicted = self.cache.insert_clean(address.clone(), db_nonce);
+            let evicted = match self.cache.insert_clean(address.clone(), db_nonce) {
+                Ok(evicted) => evicted,
+                Err(_) => {
+                    // The cache is corrupt, reset it
+                    self.reset_cache(mempool_db);
+                    None
+                }
+            };
             if evicted.is_some() {
                 // If we evicted something, we need to flush the cache.
                 self.flush_with_evicted(mempool_db, evicted);
@@ -97,7 +120,14 @@ impl NonceCache {
     /// Set the nonce for `address` to `value` in the in-memory cache.
     /// If this causes an eviction, flush the in-memory cache to the DB.
     pub fn set(&mut self, address: StacksAddress, value: u64, conn: &mut DBConn) {
-        let evicted = self.cache.insert(address.clone(), value);
+        let evicted = match self.cache.insert(address.clone(), value) {
+            Ok(evicted) => evicted,
+            Err(_) => {
+                // The cache is corrupt, reset it
+                self.reset_cache(conn);
+                Some((address, value))
+            }
+        };
         if evicted.is_some() {
             // If we evicted something, we need to flush the cache.
             self.flush_with_evicted(conn, evicted);
@@ -147,10 +177,18 @@ impl NonceCache {
             tx.execute(sql, params![addr, nonce])?;
         }
 
-        self.cache.flush(|addr, nonce| {
+        match self.cache.flush(|addr, nonce| {
             tx.execute(sql, params![addr, nonce])?;
             Ok::<(), db_error>(())
-        })?;
+        }) {
+            Ok(inner) => inner?,
+            Err(_) => {
+                drop(tx);
+                // The cache is corrupt, reset it and return
+                self.reset_cache(conn);
+                return Ok(());
+            }
+        };
 
         tx.commit()?;