feat: improved shard cache

chain/chain/src/chain.rs

@@ -4533,7 +4533,7 @@ impl<'a> ChainUpdate<'a> {
                         &result.shard_uid,
                         new_chunk_extra,
                     );
-                    self.chain_store_update.save_trie_changes(result.trie_changes);
+                    self.chain_store_update.save_trie_changes(result.trie_changes)?;
                 }
                 assert_eq!(sum_gas_used, total_gas_used);
                 assert_eq!(sum_balance_burnt, total_balance_burnt);
@@ -4580,7 +4580,7 @@ impl<'a> ChainUpdate<'a> {
                         apply_result.total_balance_burnt,
                     ),
                 );
-                self.chain_store_update.save_trie_changes(apply_result.trie_changes);
+                self.chain_store_update.save_trie_changes(apply_result.trie_changes)?;
                 self.chain_store_update.save_outgoing_receipt(
                     &block_hash,
                     shard_id,
@@ -4614,7 +4614,7 @@ impl<'a> ChainUpdate<'a> {
                 *new_extra.state_root_mut() = apply_result.new_root;
 
                 self.chain_store_update.save_chunk_extra(&block_hash, &shard_uid, new_extra);
-                self.chain_store_update.save_trie_changes(apply_result.trie_changes);
+                self.chain_store_update.save_trie_changes(apply_result.trie_changes)?;
 
                 if let Some(apply_results_or_state_changes) = apply_split_result_or_state_changes {
                     self.process_split_state(
@@ -4966,7 +4966,7 @@ impl<'a> ChainUpdate<'a> {
 
         self.chain_store_update.save_chunk(chunk);
 
-        self.chain_store_update.save_trie_changes(apply_result.trie_changes);
+        self.chain_store_update.save_trie_changes(apply_result.trie_changes)?;
         let chunk_extra = ChunkExtra::new(
             &apply_result.new_root,
             outcome_root,
@@ -5045,7 +5045,7 @@ impl<'a> ChainUpdate<'a> {
             Default::default(),
         )?;
 
-        self.chain_store_update.save_trie_changes(apply_result.trie_changes);
+        self.chain_store_update.save_trie_changes(apply_result.trie_changes)?;
 
         let mut new_chunk_extra = ChunkExtra::clone(&chunk_extra);
         *new_chunk_extra.state_root_mut() = apply_result.new_root;

chain/chain/src/store.rs

@@ -1818,8 +1818,16 @@ impl<'a> ChainStoreUpdate<'a> {
         self.chain_store_cache_update.outcome_ids.insert((*block_hash, shard_id), outcome_ids);
     }
 
-    pub fn save_trie_changes(&mut self, trie_changes: WrappedTrieChanges) {
+    // fix me: catch error everywhere
+    pub fn save_trie_changes(&mut self, trie_changes: WrappedTrieChanges) -> Result<(), Error> {
+        // Pop deleted trie nodes from the shard cache.
+        // TODO: get rid of creating `StoreUpdate`
+        let mut store_update = self.store().store_update();
+        trie_changes.deletions_into(&mut store_update);
+        store_update.update_cache()?;
+
         self.trie_changes.push(trie_changes);
+        Ok(())
     }
 
     pub fn add_state_changes_for_split_states(

chain/chain/src/tests/gc.rs

@@ -137,7 +137,7 @@ fn do_fork(
                 Default::default(),
                 *block.hash(),
             );
-            store_update.save_trie_changes(wrapped_trie_changes);
+            store_update.save_trie_changes(wrapped_trie_changes).unwrap();
 
             prev_state_roots[shard_id as usize] = new_root;
             trie_changes_shards.push(trie_changes_data);

core/store/src/lib.rs

@@ -356,6 +356,17 @@ impl StoreUpdate {
         self.transaction.merge(other.transaction)
     }
 
+    pub fn update_cache(&self) -> io::Result<()> {
+        if let Some(tries) = &self.shard_tries {
+            // Note: avoid comparing wide pointers here to work-around
+            // https://github.com/rust-lang/rust/issues/69757
+            let addr = |arc| Arc::as_ptr(arc) as *const u8;
+            assert_eq!(addr(&tries.get_store().storage), addr(&self.storage),);
+            tries.update_cache(&self.transaction)?;
+        }
+        Ok(())
+    }
+
     pub fn commit(self) -> io::Result<()> {
         debug_assert!(
             {
@@ -376,13 +387,7 @@ impl StoreUpdate {
             "Transaction overwrites itself: {:?}",
             self
         );
-        if let Some(tries) = self.shard_tries {
-            // Note: avoid comparing wide pointers here to work-around
-            // https://github.com/rust-lang/rust/issues/69757
-            let addr = |arc| Arc::as_ptr(arc) as *const u8;
-            assert_eq!(addr(&tries.get_store().storage), addr(&self.storage),);
-            tries.update_cache(&self.transaction)?;
-        }
+        self.update_cache()?;
         let _span = tracing::trace_span!(target: "store", "commit").entered();
         for op in &self.transaction.ops {
             match op {

core/store/src/trie/shard_tries.rs

@@ -44,7 +44,7 @@ impl TrieCacheFactory {
     /// Create new cache for the given shard uid.
     pub fn create_cache(&self, shard_uid: &ShardUId) -> TrieCache {
         match self.capacities.get(shard_uid) {
-            Some(capacity) => TrieCache::with_capacity(*capacity),
+            Some(capacity) => TrieCache::with_capacities(*capacity),
             None => TrieCache::new(),
         }
     }
@@ -325,6 +325,10 @@ impl WrappedTrieChanges {
         self.tries.apply_insertions(&self.trie_changes, self.shard_uid, store_update)
     }
 
+    pub fn deletions_into(&self, store_update: &mut StoreUpdate) {
+        self.tries.apply_deletions(&self.trie_changes, self.shard_uid, store_update)
+    }
+
     /// Save state changes into Store.
     ///
     /// NOTE: the changes are drained from `self`.