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();