fix(client): fix trie_node_touched bug (#3025)

Change some code so that Trie instances are no longer persisted.

Test plan
---------
existing tests, new regression test

chain/chain/src/test_utils.rs

@@ -266,7 +266,7 @@ impl RuntimeAdapter for KeyValueRuntime {
         self.tries.clone()
     }
 
-    fn get_trie_for_shard(&self, shard_id: ShardId) -> Arc<Trie> {
+    fn get_trie_for_shard(&self, shard_id: ShardId) -> Trie {
         self.tries.get_trie_for_shard(shard_id)
     }
 

chain/chain/src/types.rs

@@ -137,7 +137,7 @@ pub trait RuntimeAdapter: Send + Sync {
     fn get_tries(&self) -> ShardTries;
 
     /// Returns trie.
-    fn get_trie_for_shard(&self, shard_id: ShardId) -> Arc<Trie>;
+    fn get_trie_for_shard(&self, shard_id: ShardId) -> Trie;
 
     /// Verify block producer validity
     fn verify_block_signature(&self, header: &BlockHeader) -> Result<(), Error>;

core/store/src/db.rs

@@ -12,6 +12,7 @@ use rocksdb::{
 use strum_macros::EnumIter;
 
 use near_primitives::version::DbVersion;
+use std::marker::PhantomPinned;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct DBError(rocksdb::Error);
@@ -226,6 +227,7 @@ impl DBTransaction {
 pub struct RocksDB {
     db: DB,
     cfs: Vec<*const ColumnFamily>,
+    _pin: PhantomPinned,
 }
 
 // DB was already Send+Sync. cf and read_options are const pointers using only functions in
@@ -404,7 +406,7 @@ impl RocksDB {
         let db = DB::open_cf_for_read_only(&options, path, cf_names.iter(), false)?;
         let cfs =
             cf_names.iter().map(|n| db.cf_handle(n).unwrap() as *const ColumnFamily).collect();
-        Ok(Self { db, cfs })
+        Ok(Self { db, cfs, _pin: PhantomPinned })
     }
 
     pub fn new<P: AsRef<std::path::Path>>(path: P) -> Result<Self, DBError> {
@@ -416,7 +418,7 @@ impl RocksDB {
         let db = DB::open_cf_descriptors(&options, path, cf_descriptors)?;
         let cfs =
             cf_names.iter().map(|n| db.cf_handle(n).unwrap() as *const ColumnFamily).collect();
-        Ok(Self { db, cfs })
+        Ok(Self { db, cfs, _pin: PhantomPinned })
     }
 }
 

core/store/src/lib.rs

@@ -32,18 +32,20 @@ pub use crate::trie::{
     update::TrieUpdateValuePtr, KeyForStateChanges, PartialStorage, ShardTries, Trie, TrieChanges,
     WrappedTrieChanges,
 };
+use std::ops::Deref;
+use std::pin::Pin;
 
 mod db;
 pub mod migrations;
 pub mod test_utils;
 mod trie;
 
 pub struct Store {
-    storage: Arc<dyn Database>,
+    storage: Pin<Arc<dyn Database>>,
 }
 
 impl Store {
-    pub fn new(storage: Arc<dyn Database>) -> Store {
+    pub fn new(storage: Pin<Arc<dyn Database>>) -> Store {
         Store { storage }
     }
 
@@ -134,14 +136,14 @@ impl Store {
 
 /// Keeps track of current changes to the database and can commit all of them to the database.
 pub struct StoreUpdate {
-    storage: Arc<dyn Database>,
+    storage: Pin<Arc<dyn Database>>,
     transaction: DBTransaction,
     /// Optionally has reference to the trie to clear cache on the commit.
     tries: Option<ShardTries>,
 }
 
 impl StoreUpdate {
-    pub fn new(storage: Arc<dyn Database>) -> Self {
+    pub fn new(storage: Pin<Arc<dyn Database>>) -> Self {
         let transaction = storage.transaction();
         StoreUpdate { storage, transaction, tries: None }
     }
@@ -178,8 +180,8 @@ impl StoreUpdate {
                 self.tries = Some(tries);
             } else {
                 debug_assert_eq!(
-                    self.tries.as_ref().unwrap().tries.as_ref() as *const _,
-                    tries.tries.as_ref() as *const _
+                    self.tries.as_ref().unwrap().caches.as_ref() as *const _,
+                    tries.caches.as_ref() as *const _
                 );
             }
         }
@@ -215,8 +217,8 @@ impl StoreUpdate {
         );
         if let Some(tries) = self.tries {
             assert_eq!(
-                tries.get_store().storage.as_ref() as *const _,
-                self.storage.as_ref() as *const _
+                tries.get_store().storage.deref() as *const _,
+                self.storage.deref() as *const _
             );
             tries.update_cache(&self.transaction)?;
         }
@@ -255,7 +257,7 @@ pub fn read_with_cache<'a, T: BorshDeserialize + 'a>(
 }
 
 pub fn create_store(path: &str) -> Arc<Store> {
-    let db = Arc::new(RocksDB::new(path).expect("Failed to open the database"));
+    let db = Arc::pin(RocksDB::new(path).expect("Failed to open the database"));
     Arc::new(Store::new(db))
 }
 

core/store/src/test_utils.rs

@@ -5,12 +5,13 @@ use rand::seq::SliceRandom;
 use rand::Rng;
 
 use crate::db::TestDB;
-use crate::{ShardTries, Store, Trie};
+use crate::{ShardTries, Store};
 use near_primitives::hash::CryptoHash;
+use near_primitives::types::ShardId;
 
 /// Creates an in-memory database.
 pub fn create_test_store() -> Arc<Store> {
-    let db = Arc::new(TestDB::new());
+    let db = Arc::pin(TestDB::new());
     Arc::new(Store::new(db))
 }
 
@@ -21,12 +22,13 @@ pub fn create_tries() -> ShardTries {
 }
 
 pub fn test_populate_trie(
-    trie: Arc<Trie>,
+    tries: &ShardTries,
     root: &CryptoHash,
+    shard_id: ShardId,
     changes: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 ) -> CryptoHash {
+    let trie = tries.get_trie_for_shard(shard_id);
     assert_eq!(trie.storage.as_caching_storage().unwrap().shard_id, 0);
-    let tries = Arc::new(ShardTries { tries: Arc::new(vec![trie.clone()]) });
     let trie_changes = trie.update(root, changes.iter().cloned()).unwrap();
     let (store_update, root) = tries.apply_all(&trie_changes, 0).unwrap();
     store_update.commit().unwrap();