Pin states in memory so that they are not pruned away while still referenced (#2761)

* State pinning in client

* Canonicalization queue

* Fixed prioritization queue

* possible fix of "hash mismatch"

* Check for pinned discarded states

* Release state before finalization

* Style

* Style

substrate/core/client/db/src/lib.rs

@@ -72,6 +72,97 @@ const MIN_BLOCKS_TO_KEEP_CHANGES_TRIES_FOR: u32 = 32768;
 /// DB-backed patricia trie state, transaction type is an overlay of changes to commit.
 pub type DbState = state_machine::TrieBackend<Arc<state_machine::Storage<Blake2Hasher>>, Blake2Hasher>;
 
+pub struct RefTrackingState<Block: BlockT> {
+	state: DbState,
+	storage: Arc<StorageDb<Block>>,
+	parent_hash: Option<Block::Hash>,
+}
+
+impl<B: BlockT> RefTrackingState<B> {
+	fn new(state: DbState, storage: Arc<StorageDb<B>>, parent_hash: Option<B::Hash>) -> RefTrackingState<B> {
+		if let Some(hash) = &parent_hash {
+			storage.state_db.pin(hash);
+		}
+		RefTrackingState {
+			state,
+			parent_hash,
+			storage,
+		}
+	}
+}
+
+impl<B: BlockT> Drop for RefTrackingState<B> {
+	fn drop(&mut self) {
+		if let Some(hash) = &self.parent_hash {
+			self.storage.state_db.unpin(hash);
+		}
+	}
+}
+
+impl<B: BlockT> StateBackend<Blake2Hasher> for RefTrackingState<B> {
+	type Error =  <DbState as StateBackend<Blake2Hasher>>::Error;
+	type Transaction = <DbState as StateBackend<Blake2Hasher>>::Transaction;
+	type TrieBackendStorage = <DbState as StateBackend<Blake2Hasher>>::TrieBackendStorage;
+
+	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.storage(key)
+	}
+
+	fn storage_hash(&self, key: &[u8]) -> Result<Option<H256>, Self::Error> {
+		self.state.storage_hash(key)
+	}
+
+	fn child_storage(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.child_storage(storage_key, key)
+	}
+
+	fn exists_storage(&self, key: &[u8]) -> Result<bool, Self::Error> {
+		self.state.exists_storage(key)
+	}
+
+	fn exists_child_storage(&self, storage_key: &[u8], key: &[u8]) -> Result<bool, Self::Error> {
+		self.state.exists_child_storage(storage_key, key)
+	}
+
+	fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
+		self.state.for_keys_with_prefix(prefix, f)
+	}
+
+	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {
+		self.state.for_keys_in_child_storage(storage_key, f)
+	}
+
+	fn storage_root<I>(&self, delta: I) -> (H256, Self::Transaction)
+		where
+			I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>
+	{
+		self.state.storage_root(delta)
+	}
+
+	fn child_storage_root<I>(&self, storage_key: &[u8], delta: I) -> (Vec<u8>, bool, Self::Transaction)
+		where
+			I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>,
+	{
+		self.state.child_storage_root(storage_key, delta)
+	}
+
+	fn pairs(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
+		self.state.pairs()
+	}
+
+	fn keys(&self, prefix: &[u8]) -> Vec<Vec<u8>> {
+		self.state.keys(prefix)
+	}
+
+	fn child_keys(&self, child_key: &[u8], prefix: &[u8]) -> Vec<Vec<u8>> {
+		self.state.child_keys(child_key, prefix)
+	}
+
+	fn as_trie_backend(&mut self) -> Option<&state_machine::TrieBackend<Self::TrieBackendStorage, Blake2Hasher>> {
+		self.state.as_trie_backend()
+	}
+}
+
 /// Database settings.
 pub struct DatabaseSettings {
 	/// Cache size in bytes. If `None` default is used.
@@ -270,7 +361,7 @@ impl<Block: BlockT> client::blockchain::ProvideCache<Block> for BlockchainDb<Blo
 
 /// Database transaction
 pub struct BlockImportOperation<Block: BlockT, H: Hasher> {
-	old_state: CachingState<Blake2Hasher, DbState, Block>,
+	old_state: CachingState<Blake2Hasher, RefTrackingState<Block>, Block>,
 	db_updates: PrefixedMemoryDB<H>,
 	storage_updates: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	changes_trie_updates: MemoryDB<H>,
@@ -295,7 +386,7 @@ impl<Block> client::backend::BlockImportOperation<Block, Blake2Hasher>
 for BlockImportOperation<Block, Blake2Hasher>
 where Block: BlockT<Hash=H256>,
 {
-	type State = CachingState<Blake2Hasher, DbState, Block>;
+	type State = CachingState<Blake2Hasher, RefTrackingState<Block>, Block>;
 
 	fn state(&self) -> Result<Option<&Self::State>, client::error::Error> {
 		Ok(Some(&self.old_state))
@@ -922,6 +1013,8 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			let changes_trie_updates = operation.changes_trie_updates;
 
 			self.changes_tries_storage.commit(&mut transaction, changes_trie_updates);
+			let cache = operation.old_state.release(); // release state reference so that it can be finalized
+
 
 			if finalized {
 				// TODO: ensure best chain contains this block.
@@ -953,7 +1046,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 
 			meta_updates.push((hash, number, pending_block.leaf_state.is_best(), finalized));
 
-			Some((number, hash, enacted, retracted, displaced_leaf, is_best))
+			Some((number, hash, enacted, retracted, displaced_leaf, is_best, cache))
 		} else {
 			None
 		};
@@ -975,7 +1068,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 
 		let write_result = self.storage.db.write(transaction).map_err(db_err);
 
-		if let Some((number, hash, enacted, retracted, displaced_leaf, is_best)) = imported {
+		if let Some((number, hash, enacted, retracted, displaced_leaf, is_best, mut cache)) = imported {
 			if let Err(e) = write_result {
 				let mut leaves = self.blockchain.leaves.write();
 				let mut undo = leaves.undo();
@@ -990,7 +1083,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 				return Err(e)
 			}
 
-			operation.old_state.sync_cache(
+			cache.sync_cache(
 				&enacted,
 				&retracted,
 				operation.storage_updates,
@@ -1090,7 +1183,7 @@ impl<Block> client::backend::AuxStore for Backend<Block> where Block: BlockT<Has
 impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> where Block: BlockT<Hash=H256> {
 	type BlockImportOperation = BlockImportOperation<Block, Blake2Hasher>;
 	type Blockchain = BlockchainDb<Block>;
-	type State = CachingState<Blake2Hasher, DbState, Block>;
+	type State = CachingState<Blake2Hasher, RefTrackingState<Block>, Block>;
 	type ChangesTrieStorage = DbChangesTrieStorage<Block>;
 
 	fn begin_operation(&self) -> Result<Self::BlockImportOperation, client::error::Error> {
@@ -1217,7 +1310,8 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 			BlockId::Hash(h) if h == Default::default() => {
 				let genesis_storage = DbGenesisStorage::new();
 				let root = genesis_storage.0.clone();
-				let state = DbState::new(Arc::new(genesis_storage), root);
+				let db_state = DbState::new(Arc::new(genesis_storage), root);
+				let state = RefTrackingState::new(db_state, self.storage.clone(), None);
 				return Ok(CachingState::new(state, self.shared_cache.clone(), None));
 			},
 			_ => {}
@@ -1228,7 +1322,8 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 				let hash = hdr.hash();
 				if !self.storage.state_db.is_pruned(&hash, (*hdr.number()).saturated_into::<u64>()) {
 					let root = H256::from_slice(hdr.state_root().as_ref());
-					let state = DbState::new(self.storage.clone(), root);
+					let db_state = DbState::new(self.storage.clone(), root);
+					let state = RefTrackingState::new(db_state, self.storage.clone(), Some(hash.clone()));
 					Ok(CachingState::new(state, self.shared_cache.clone(), Some(hash)))
 				} else {
 					Err(client::error::Error::UnknownBlock(format!("State already discarded for {:?}", block)))
@@ -1243,10 +1338,10 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 		!self.storage.state_db.is_pruned(hash, number.saturated_into::<u64>())
 	}
 
-	fn destroy_state(&self, mut state: Self::State) -> Result<(), client::error::Error> {
-		if let Some(hash) = state.parent_hash.clone() {
+	fn destroy_state(&self, state: Self::State) -> Result<(), client::error::Error> {
+		if let Some(hash) = state.cache.parent_hash.clone() {
 			let is_best = || self.blockchain.meta.read().best_hash == hash;
-			state.sync_cache(&[], &[], vec![], None, None, is_best);
+			state.release().sync_cache(&[], &[], vec![], None, None, is_best);
 		}
 		Ok(())
 	}