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

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

@@ -99,6 +99,17 @@ struct LocalCache<H: Hasher> {
 	hashes: HashMap<StorageKey, Option<H::Out>>,
 }
 
+/// Cache changes.
+pub struct CacheChanges<H: Hasher, B: Block> {
+	/// Shared canonical state cache.
+	shared_cache: SharedCache<B, H>,
+	/// Local cache of values for this state.
+	local_cache: RwLock<LocalCache<H>>,
+	/// Hash of the block on top of which this instance was created or
+	/// `None` if cache is disabled
+	pub parent_hash: Option<B::Hash>,
+}
+
 /// State abstraction.
 /// Manages shared global state cache which reflects the canonical
 /// state as it is on the disk.
@@ -109,56 +120,11 @@ struct LocalCache<H: Hasher> {
 pub struct CachingState<H: Hasher, S: StateBackend<H>, B: Block> {
 	/// Backing state.
 	state: S,
-	/// Shared canonical state cache.
-	shared_cache: SharedCache<B, H>,
-	/// Local cache of values for this state.
-	local_cache: RwLock<LocalCache<H>>,
-	/// Hash of the block on top of which this instance was created or
-	/// `None` if cache is disabled
-	pub parent_hash: Option<B::Hash>,
+	/// Cache data.
+	pub cache: CacheChanges<H, B>
 }
 
-impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
-	/// Create a new instance wrapping generic State and shared cache.
-	pub fn new(state: S, shared_cache: SharedCache<B, H>, parent_hash: Option<B::Hash>) -> CachingState<H, S, B> {
-		CachingState {
-			state,
-			shared_cache,
-			local_cache: RwLock::new(LocalCache {
-				storage: Default::default(),
-				hashes: Default::default(),
-			}),
-			parent_hash: parent_hash,
-		}
-	}
-
-	fn storage_insert(cache: &mut Cache<B, H>, k: StorageValue, v: Option<StorageValue>) {
-		if let Some(v_) = &v {
-			while cache.storage_used_size + v_.len() > cache.shared_cache_size {
-				// pop until space constraint satisfied
-				match cache.storage.remove_lru() {
-					Some((_, Some(popped_v))) =>
-						cache.storage_used_size = cache.storage_used_size - popped_v.len(),
-					Some((_, None)) => continue,
-					None => break,
-				};
-			}
-			cache.storage_used_size = cache.storage_used_size + v_.len();
-		}
-		cache.storage.insert(k, v);
-	}
-
-	fn storage_remove(
-		storage: &mut LruCache<StorageKey, Option<StorageValue>>,
-		k: &StorageKey,
-		storage_used_size: &mut usize,
-	) {
-		let v = storage.remove(k);
-		if let Some(Some(v_)) = v {
-			*storage_used_size = *storage_used_size - v_.len();
-		}
-	}
-
+impl<H: Hasher, B: Block> CacheChanges<H, B> {
 	/// Propagate local cache into the shared cache and synchronize
 	/// the shared cache with the best block state.
 	/// This function updates the shared cache by removing entries
@@ -189,7 +155,7 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 					m.is_canon = true;
 					for a in &m.storage {
 						trace!("Reverting enacted key {:?}", a);
-						CachingState::<H, S, B>::storage_remove(&mut cache.storage, a, &mut cache.storage_used_size);
+						CacheChanges::<H, B>::storage_remove(&mut cache.storage, a, &mut cache.storage_used_size);
 					}
 					false
 				} else {
@@ -205,7 +171,7 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 					m.is_canon = false;
 					for a in &m.storage {
 						trace!("Retracted key {:?}", a);
-						CachingState::<H, S, B>::storage_remove(&mut cache.storage, a, &mut cache.storage_used_size);
+						CacheChanges::<H, B>::storage_remove(&mut cache.storage, a, &mut cache.storage_used_size);
 					}
 					false
 				} else {
@@ -228,7 +194,7 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 			if is_best {
 				trace!("Committing {} local, {} hashes, {} modified entries", local_cache.storage.len(), local_cache.hashes.len(), changes.len());
 				for (k, v) in local_cache.storage.drain() {
-					CachingState::<H, S, B>::storage_insert(cache, k, v);
+					CacheChanges::<H, B>::storage_insert(cache, k, v);
 				}
 				for (k, v) in local_cache.hashes.drain() {
 					cache.hashes.insert(k, v);
@@ -248,7 +214,7 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 				modifications.insert(k.clone());
 				if is_best {
 					cache.hashes.remove(&k);
-					CachingState::<H, S, B>::storage_insert(cache, k, v);
+					CacheChanges::<H, B>::storage_insert(cache, k, v);
 				}
 			}
 			// Save modified storage. These are ordered by the block number.
@@ -272,6 +238,50 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 		}
 	}
 
+	fn storage_insert(cache: &mut Cache<B, H>, k: StorageValue, v: Option<StorageValue>) {
+		if let Some(v_) = &v {
+			while cache.storage_used_size + v_.len() > cache.shared_cache_size {
+				// pop until space constraint satisfied
+				match cache.storage.remove_lru() {
+					Some((_, Some(popped_v))) =>
+						cache.storage_used_size = cache.storage_used_size - popped_v.len(),
+						Some((_, None)) => continue,
+						None => break,
+				};
+			}
+			cache.storage_used_size = cache.storage_used_size + v_.len();
+		}
+		cache.storage.insert(k, v);
+	}
+
+	fn storage_remove(
+		storage: &mut LruCache<StorageKey, Option<StorageValue>>,
+		k: &StorageKey,
+		storage_used_size: &mut usize,
+		) {
+		let v = storage.remove(k);
+		if let Some(Some(v_)) = v {
+			*storage_used_size = *storage_used_size - v_.len();
+		}
+	}
+}
+
+impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
+	/// Create a new instance wrapping generic State and shared cache.
+	pub fn new(state: S, shared_cache: SharedCache<B, H>, parent_hash: Option<B::Hash>) -> CachingState<H, S, B> {
+		CachingState {
+			state,
+			cache: CacheChanges {
+				shared_cache,
+				local_cache: RwLock::new(LocalCache {
+					storage: Default::default(),
+					hashes: Default::default(),
+				}),
+				parent_hash: parent_hash,
+			},
+		}
+	}
+
 	/// Check if the key can be returned from cache by matching current block parent hash against canonical
 	/// state and filtering out entries modified in later blocks.
 	fn is_allowed(
@@ -312,6 +322,11 @@ impl<H: Hasher, S: StateBackend<H>, B: Block> CachingState<H, S, B> {
 		trace!("Cache lookup skipped for {:?}: parent hash is unknown", key);
 		false
 	}
+
+	/// Dispose state and return cache data.
+	pub fn release(self) -> CacheChanges<H, B> {
+		self.cache
+	}
 }
 
 impl<H: Hasher, S: StateBackend<H>, B:Block> StateBackend<H> for CachingState<H, S, B> {
@@ -320,13 +335,13 @@ impl<H: Hasher, S: StateBackend<H>, B:Block> StateBackend<H> for CachingState<H,
 	type TrieBackendStorage = S::TrieBackendStorage;
 
 	fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		let local_cache = self.local_cache.upgradable_read();
+		let local_cache = self.cache.local_cache.upgradable_read();
 		if let Some(entry) = local_cache.storage.get(key).cloned() {
 			trace!("Found in local cache: {:?}", key);
 			return Ok(entry)
 		}
-		let mut cache = self.shared_cache.lock();
-		if Self::is_allowed(key, &self.parent_hash, &cache.modifications) {
+		let mut cache = self.cache.shared_cache.lock();
+		if Self::is_allowed(key, &self.cache.parent_hash, &cache.modifications) {
 			if let Some(entry) = cache.storage.get_mut(key).map(|a| a.clone()) {
 				trace!("Found in shared cache: {:?}", key);
 				return Ok(entry)
@@ -339,13 +354,13 @@ impl<H: Hasher, S: StateBackend<H>, B:Block> StateBackend<H> for CachingState<H,
 	}
 
 	fn storage_hash(&self, key: &[u8]) -> Result<Option<H::Out>, Self::Error> {
-		let local_cache = self.local_cache.upgradable_read();
+		let local_cache = self.cache.local_cache.upgradable_read();
 		if let Some(entry) = local_cache.hashes.get(key).cloned() {
 			trace!("Found hash in local cache: {:?}", key);
 			return Ok(entry)
 		}
-		let mut cache = self.shared_cache.lock();
-		if Self::is_allowed(key, &self.parent_hash, &cache.modifications) {
+		let mut cache = self.cache.shared_cache.lock();
+		if Self::is_allowed(key, &self.cache.parent_hash, &cache.modifications) {
 			if let Some(entry) = cache.hashes.get_mut(key).map(|a| a.clone()) {
 				trace!("Found hash in shared cache: {:?}", key);
 				return Ok(entry)
@@ -405,8 +420,8 @@ impl<H: Hasher, S: StateBackend<H>, B:Block> StateBackend<H> for CachingState<H,
 		self.state.child_keys(child_key, prefix)
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
-		self.state.try_into_trie_backend()
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
+		self.state.as_trie_backend()
 	}
 }
 
@@ -436,22 +451,22 @@ mod tests {
 		// blocks  [ 3a(c) 2a(c) 2b 1b 1a(c) 0 ]
 		// state   [ 5     5     4  3  2     2 ]
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(root_parent.clone()));
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![2]))], Some(h0.clone()), Some(0), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![2]))], Some(h0.clone()), Some(0), || true);
 
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h0.clone()));
-		s.sync_cache(&[], &[], vec![], Some(h1a.clone()), Some(1), || true);
+		s.cache.sync_cache(&[], &[], vec![], Some(h1a.clone()), Some(1), || true);
 
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h0.clone()));
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![3]))], Some(h1b.clone()), Some(1), || false);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![3]))], Some(h1b.clone()), Some(1), || false);
 
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h1b.clone()));
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![4]))], Some(h2b.clone()), Some(2), || false);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![4]))], Some(h2b.clone()), Some(2), || false);
 
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h1a.clone()));
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![5]))], Some(h2a.clone()), Some(2), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![5]))], Some(h2a.clone()), Some(2), || true);
 
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h2a.clone()));
-		s.sync_cache(&[], &[], vec![], Some(h3a.clone()), Some(3), || true);
+		s.cache.sync_cache(&[], &[], vec![], Some(h3a.clone()), Some(3), || true);
 
 		let s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h3a.clone()));
 		assert_eq!(s.storage(&key).unwrap().unwrap(), vec![5]);
@@ -468,7 +483,14 @@ mod tests {
 		// reorg to 3b
 		// blocks  [ 3b(c) 3a 2a 2b(c) 1b 1a 0 ]
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h2b.clone()));
-		s.sync_cache(&[h1b.clone(), h2b.clone(), h3b.clone()], &[h1a.clone(), h2a.clone(), h3a.clone()], vec![], Some(h3b.clone()), Some(3), || true);
+		s.cache.sync_cache(
+			&[h1b.clone(), h2b.clone(), h3b.clone()],
+			&[h1a.clone(), h2a.clone(), h3a.clone()],
+			vec![],
+			Some(h3b.clone()),
+			Some(3),
+			|| true
+		);
 		let s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(h3a.clone()));
 		assert!(s.storage(&key).unwrap().is_none());
 	}
@@ -482,11 +504,11 @@ mod tests {
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(root_parent.clone()));
 
 		let key = H256::random()[..].to_vec();
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2, 3]))], Some(h0.clone()), Some(0), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2, 3]))], Some(h0.clone()), Some(0), || true);
 		assert_eq!(shared.lock().used_storage_cache_size(), 3 /* bytes */);
 
 		let key = H256::random()[..].to_vec();
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2]))], Some(h0.clone()), Some(0), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2]))], Some(h0.clone()), Some(0), || true);
 		assert_eq!(shared.lock().used_storage_cache_size(), 5 /* bytes */);
 	}
 
@@ -499,11 +521,11 @@ mod tests {
 		let mut s = CachingState::new(InMemory::<Blake2Hasher>::default(), shared.clone(), Some(root_parent.clone()));
 
 		let key = H256::random()[..].to_vec();
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2, 3, 4]))], Some(h0.clone()), Some(0), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2, 3, 4]))], Some(h0.clone()), Some(0), || true);
 		assert_eq!(shared.lock().used_storage_cache_size(), 4 /* bytes */);
 
 		let key = H256::random()[..].to_vec();
-		s.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2]))], Some(h0.clone()), Some(0), || true);
+		s.cache.sync_cache(&[], &[], vec![(key.clone(), Some(vec![1, 2]))], Some(h0.clone()), Some(0), || true);
 		assert_eq!(shared.lock().used_storage_cache_size(), 2 /* bytes */);
 	}
 }

substrate/core/client/src/call_executor.rs

@@ -117,17 +117,17 @@ where
 	/// No changes are made.
 	fn prove_at_state<S: state_machine::Backend<H>>(
 		&self,
-		state: S,
+		mut state: S,
 		overlay: &mut OverlayedChanges,
 		method: &str,
 		call_data: &[u8]
 	) -> Result<(Vec<u8>, Vec<Vec<u8>>), error::Error> {
-		let trie_state = state.try_into_trie_backend()
+		let trie_state = state.as_trie_backend()
 			.ok_or_else(||
 				Box::new(state_machine::ExecutionError::UnableToGenerateProof)
 					as Box<state_machine::Error>
 			)?;
-		self.prove_at_trie_state(&trie_state, overlay, method, call_data)
+		self.prove_at_trie_state(trie_state, overlay, method, call_data)
 	}
 
 	/// Execute a call to a contract on top of given trie state, gathering execution proof.
@@ -239,18 +239,18 @@ where
 			_ => {},
 		}
 
-		let state = self.backend.state_at(*at)?;
+		let mut state = self.backend.state_at(*at)?;
 
 		match recorder {
 			Some(recorder) => {
-				let trie_state = state.try_into_trie_backend()
+				let trie_state = state.as_trie_backend()
 					.ok_or_else(||
 						Box::new(state_machine::ExecutionError::UnableToGenerateProof)
 							as Box<state_machine::Error>
 					)?;
 
 				let backend = state_machine::ProvingBackend::new_with_recorder(
-					&trie_state,
+					trie_state,
 					recorder.clone()
 				);
 

substrate/core/client/src/cht.rs

@@ -101,14 +101,14 @@ pub fn build_proof<Header, Hasher, BlocksI, HashesI>(
 		.into_iter()
 		.map(|(k, v)| (None, k, Some(v)))
 		.collect::<Vec<_>>();
-	let storage = InMemoryState::<Hasher>::default().update(transaction);
-	let trie_storage = storage.try_into_trie_backend()
-		.expect("InMemoryState::try_into_trie_backend always returns Some; qed");
+	let mut storage = InMemoryState::<Hasher>::default().update(transaction);
+	let trie_storage = storage.as_trie_backend()
+		.expect("InMemoryState::as_trie_backend always returns Some; qed");
 	let mut total_proof = HashSet::new();
 	for block in blocks.into_iter() {
 		debug_assert_eq!(block_to_cht_number(cht_size, block), Some(cht_num));
 
-		let (value, proof) = prove_read_on_trie_backend(&trie_storage, &encode_cht_key(block))?;
+		let (value, proof) = prove_read_on_trie_backend(trie_storage, &encode_cht_key(block))?;
 		assert!(value.is_some(), "we have just built trie that includes the value for block");
 		total_proof.extend(proof);
 	}

substrate/core/client/src/client.rs

@@ -1277,7 +1277,7 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 
 	/// Prepare in-memory header that is used in execution environment.
 	fn prepare_environment_block(&self, parent: &BlockId<Block>) -> error::Result<Block::Header> {
-		let parent_header = self.backend().blockchain().expect_header(*parent)?;
+		let parent_header = self.backend.blockchain().expect_header(*parent)?;
 		Ok(<<Block as BlockT>::Header as HeaderT>::new(
 			self.backend.blockchain().expect_block_number_from_id(parent)? + One::one(),
 			Default::default(),

substrate/core/client/src/light/backend.rs

@@ -38,13 +38,13 @@ use consensus::well_known_cache_keys;
 const IN_MEMORY_EXPECT_PROOF: &str = "InMemory state backend has Void error type and always suceeds; qed";
 
 /// Light client backend.
-pub struct Backend<S, F, H> {
+pub struct Backend<S, F, H: Hasher> {
 	blockchain: Arc<Blockchain<S, F>>,
 	genesis_state: RwLock<Option<InMemoryState<H>>>,
 }
 
 /// Light block (header and justification) import operation.
-pub struct ImportOperation<Block: BlockT, S, F, H> {
+pub struct ImportOperation<Block: BlockT, S, F, H: Hasher> {
 	header: Option<Block::Header>,
 	cache: HashMap<well_known_cache_keys::Id, Vec<u8>>,
 	leaf_state: NewBlockState,
@@ -64,14 +64,14 @@ pub struct OnDemandState<Block: BlockT, S, F> {
 }
 
 /// On-demand or in-memory genesis state.
-pub enum OnDemandOrGenesisState<Block: BlockT, S, F, H> {
+pub enum OnDemandOrGenesisState<Block: BlockT, S, F, H: Hasher> {
 	/// On-demand state - storage values are fetched from remote nodes.
 	OnDemand(OnDemandState<Block, S, F>),
 	/// Genesis state - storage values are stored in-memory.
 	Genesis(InMemoryState<H>),
 }
 
-impl<S, F, H> Backend<S, F, H> {
+impl<S, F, H: Hasher> Backend<S, F, H> {
 	/// Create new light backend.
 	pub fn new(blockchain: Arc<Blockchain<S, F>>) -> Self {
 		Self {
@@ -86,7 +86,7 @@ impl<S, F, H> Backend<S, F, H> {
 	}
 }
 
-impl<S: AuxStore, F, H> AuxStore for Backend<S, F, H> {
+impl<S: AuxStore, F, H: Hasher> AuxStore for Backend<S, F, H> {
 	fn insert_aux<
 		'a,
 		'b: 'a,
@@ -387,7 +387,7 @@ where
 		Vec::new()
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		None
 	}
 }
@@ -482,10 +482,10 @@ where
 		}
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		match self {
-			OnDemandOrGenesisState::OnDemand(state) => state.try_into_trie_backend(),
-			OnDemandOrGenesisState::Genesis(state) => state.try_into_trie_backend(),
+			OnDemandOrGenesisState::OnDemand(ref mut state) => state.as_trie_backend(),
+			OnDemandOrGenesisState::Genesis(ref mut state) => state.as_trie_backend(),
 		}
 	}
 }

substrate/core/client/src/light/call_executor.rs

@@ -388,7 +388,7 @@ impl<Block, B, Remote, Local> CallExecutor<Block, Blake2Hasher> for
 /// Method is executed using passed header as environment' current block.
 /// Proof includes both environment preparation proof and method execution proof.
 pub fn prove_execution<Block, S, E>(
-	state: S,
+	mut state: S,
 	header: Block::Header,
 	executor: &E,
 	method: &str,
@@ -399,13 +399,13 @@ pub fn prove_execution<Block, S, E>(
 		S: StateBackend<Blake2Hasher>,
 		E: CallExecutor<Block, Blake2Hasher>,
 {
-	let trie_state = state.try_into_trie_backend()
+	let trie_state = state.as_trie_backend()
 		.ok_or_else(|| Box::new(state_machine::ExecutionError::UnableToGenerateProof) as Box<state_machine::Error>)?;
 
 	// prepare execution environment + record preparation proof
 	let mut changes = Default::default();
 	let (_, init_proof) = executor.prove_at_trie_state(
-		&trie_state,
+		trie_state,
 		&mut changes,
 		"Core_initialize_block",
 		&header.encode(),

substrate/core/state-db/src/lib.rs

@@ -37,7 +37,7 @@ use std::fmt;
 use parking_lot::RwLock;
 use parity_codec as codec;
 use codec::Codec;
-use std::collections::HashSet;
+use std::collections::{VecDeque, HashMap, hash_map::Entry};
 use noncanonical::NonCanonicalOverlay;
 use pruning::RefWindow;
 use log::trace;
@@ -78,6 +78,8 @@ pub enum Error<E: fmt::Debug> {
 	InvalidBlockNumber,
 	/// Trying to insert block with unknown parent.
 	InvalidParent,
+	/// Canonicalization would discard pinned state.
+	DiscardingPinned,
 }
 
 impl<E: fmt::Debug> fmt::Debug for Error<E> {
@@ -88,6 +90,7 @@ impl<E: fmt::Debug> fmt::Debug for Error<E> {
 			Error::InvalidBlock => write!(f, "Trying to canonicalize invalid block"),
 			Error::InvalidBlockNumber => write!(f, "Trying to insert block with invalid number"),
 			Error::InvalidParent => write!(f, "Trying to insert block with unknown parent"),
+			Error::DiscardingPinned => write!(f, "Trying to discard pinned state"),
 		}
 	}
 }
@@ -112,7 +115,7 @@ pub struct CommitSet<H: Hash> {
 }
 
 /// Pruning constraints. If none are specified pruning is
-#[derive(Default, Debug, Clone)]
+#[derive(Default, Debug, Clone, Eq, PartialEq)]
 pub struct Constraints {
 	/// Maximum blocks. Defaults to 0 when unspecified, effectively keeping only non-canonical states.
 	pub max_blocks: Option<u32>,
@@ -121,7 +124,7 @@ pub struct Constraints {
 }
 
 /// Pruning mode.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub enum PruningMode {
 	/// Maintain a pruning window.
 	Constrained(Constraints),
@@ -165,13 +168,14 @@ fn to_meta_key<S: Codec>(suffix: &[u8], data: &S) -> Vec<u8> {
 struct StateDbSync<BlockHash: Hash, Key: Hash> {
 	mode: PruningMode,
 	non_canonical: NonCanonicalOverlay<BlockHash, Key>,
+	canonicalization_queue: VecDeque<BlockHash>,
 	pruning: Option<RefWindow<BlockHash, Key>>,
-	pinned: HashSet<BlockHash>,
+	pinned: HashMap<BlockHash, u32>,
 }
 
 impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 	pub fn new<D: MetaDb>(mode: PruningMode, db: &D) -> Result<StateDbSync<BlockHash, Key>, Error<D::Error>> {
-		trace!("StateDb settings: {:?}", mode);
+		trace!(target: "state-db", "StateDb settings: {:?}", mode);
 		let non_canonical: NonCanonicalOverlay<BlockHash, Key> = NonCanonicalOverlay::new(db)?;
 		let pruning: Option<RefWindow<BlockHash, Key>> = match mode {
 			PruningMode::Constrained(Constraints {
@@ -186,6 +190,7 @@ impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 			non_canonical,
 			pruning,
 			pinned: Default::default(),
+			canonicalization_queue: Default::default(),
 		})
 	}
 
@@ -206,21 +211,30 @@ impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 	}
 
 	pub fn canonicalize_block<E: fmt::Debug>(&mut self, hash: &BlockHash) -> Result<CommitSet<Key>, Error<E>> {
-		let mut commit = match self.mode {
-			PruningMode::ArchiveAll => {
-				CommitSet::default()
-			},
-			PruningMode::ArchiveCanonical => {
-				let mut commit = self.non_canonical.canonicalize(hash)?;
-				commit.data.deleted.clear();
-				commit
-			},
-			PruningMode::Constrained(_) => {
-				self.non_canonical.canonicalize(hash)?
-			},
-		};
-		if let Some(ref mut pruning) = self.pruning {
-			pruning.note_canonical(hash, &mut commit);
+		let mut commit = CommitSet::default();
+		if self.mode == PruningMode::ArchiveAll {
+			return Ok(commit)
+		}
+		self.canonicalization_queue.push_back(hash.clone());
+		while let Some(hash) = self.canonicalization_queue.front().cloned() {
+			if self.pinned.contains_key(&hash) {
+				break;
+			}
+			match self.non_canonical.canonicalize(&hash, &self.pinned, &mut commit) {
+				Ok(()) => {
+					self.canonicalization_queue.pop_front();
+					if self.mode == PruningMode::ArchiveCanonical {
+						commit.data.deleted.clear();
+					}
+				}
+				Err(Error::DiscardingPinned) => {
+					break;
+				}
+				Err(e) => return Err(e),
+			};
+			if let Some(ref mut pruning) = self.pruning {
+				pruning.note_canonical(&hash, &mut commit);
+			}
 		}
 		self.prune(&mut commit);
 		Ok(commit)
@@ -255,7 +269,7 @@ impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 				}
 
 				let pinned = &self.pinned;
-				if pruning.next_hash().map_or(false, |h| pinned.contains(&h)) {
+				if pruning.next_hash().map_or(false, |h| pinned.contains_key(&h)) {
 					break;
 				}
 				pruning.prune_one(commit);
@@ -278,11 +292,23 @@ impl<BlockHash: Hash, Key: Hash> StateDbSync<BlockHash, Key> {
 	}
 
 	pub fn pin(&mut self, hash: &BlockHash) {
-		self.pinned.insert(hash.clone());
+		trace!(target: "state-db", "Pinned block: {:?}", hash);
+		*self.pinned.entry(hash.clone()).or_default() += 1;
 	}
 
 	pub fn unpin(&mut self, hash: &BlockHash) {
-		self.pinned.remove(hash);
+		match self.pinned.entry(hash.clone()) {
+			Entry::Occupied(mut entry) => {
+				*entry.get_mut() -= 1;
+				if *entry.get() == 0 {
+					trace!(target: "state-db", "Unpinned block: {:?}", hash);
+					entry.remove();
+				} else {
+					trace!(target: "state-db", "Releasing reference for {:?}", hash);
+				}
+			},
+			Entry::Vacant(_) => {},
+		}
 	}
 
 	pub fn get<D: NodeDb>(&self, key: &Key, db: &D) -> Result<Option<DBValue>, Error<D::Error>>

substrate/core/state-db/src/noncanonical.rs

@@ -230,13 +230,20 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 		Ok(commit)
 	}
 
-	fn discard_journals(&self, level_index: usize, discarded_journals: &mut Vec<Vec<u8>>, hash: &BlockHash) {
+	fn discard_journals(
+		&self,
+		level_index: usize,
+		discarded_journals: &mut Vec<Vec<u8>>,
+		discarded_blocks: &mut Vec<BlockHash>,
+		hash: &BlockHash
+	) {
 		if let Some(level) = self.levels.get(level_index) {
 			level.iter().for_each(|overlay| {
 				let parent = self.parents.get(&overlay.hash).expect("there is a parent entry for each entry in levels; qed").clone();
 				if parent == *hash {
 					discarded_journals.push(overlay.journal_key.clone());
-					self.discard_journals(level_index + 1, discarded_journals, &overlay.hash);
+					discarded_blocks.push(overlay.hash.clone());
+					self.discard_journals(level_index + 1, discarded_journals, discarded_blocks, &overlay.hash);
 				}
 			});
 		}
@@ -268,7 +275,12 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 
 	/// Select a top-level root and canonicalized it. Discards all sibling subtrees and the root.
 	/// Returns a set of changes that need to be added to the DB.
-	pub fn canonicalize<E: fmt::Debug>(&mut self, hash: &BlockHash) -> Result<CommitSet<Key>, Error<E>> {
+	pub fn canonicalize<E: fmt::Debug>(
+		&mut self,
+		hash: &BlockHash,
+		pinned: &HashMap<BlockHash, u32>,
+		commit: &mut CommitSet<Key>,
+	) -> Result<(), Error<E>> {
 		trace!(target: "state-db", "Canonicalizing {:?}", hash);
 		let level = self.levels.get(self.pending_canonicalizations.len()).ok_or_else(|| Error::InvalidBlock)?;
 		let index = level
@@ -276,26 +288,40 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 			.position(|overlay| overlay.hash == *hash)
 			.ok_or_else(|| Error::InvalidBlock)?;
 
-		let mut commit = CommitSet::default();
 		let mut discarded_journals = Vec::new();
-		for (i, overlay) in level.into_iter().enumerate() {
-			if i == index {
-				// that's the one we need to canonicalize
-				commit.data.inserted = overlay.inserted.iter()
-					.map(|k| (k.clone(), self.values.get(k).expect("For each key in verlays there's a value in values").1.clone()))
-					.collect();
-				commit.data.deleted = overlay.deleted.clone();
-			} else {
-				self.discard_journals(self.pending_canonicalizations.len() + 1, &mut discarded_journals, &overlay.hash);
+		let mut discarded_blocks = Vec::new();
+		for (i, overlay) in level.iter().enumerate() {
+			if i != index {
+				self.discard_journals(
+					self.pending_canonicalizations.len() + 1,
+					&mut discarded_journals,
+					&mut discarded_blocks,
+					&overlay.hash
+				);
 			}
 			discarded_journals.push(overlay.journal_key.clone());
+			discarded_blocks.push(overlay.hash.clone());
 		}
+
+		for hash in discarded_blocks.into_iter() {
+			if pinned.contains_key(&hash) {
+				trace!(target: "state-db", "Refusing to discard pinned state {:?}", hash);
+				return Err(Error::DiscardingPinned)
+			}
+		}
+
+		// get the one we need to canonicalize
+		let overlay = &level[index];
+		commit.data.inserted.extend(overlay.inserted.iter()
+			.map(|k| (k.clone(), self.values.get(k).expect("For each key in overlays there's a value in values").1.clone())));
+		commit.data.deleted.extend(overlay.deleted.clone());
+
 		commit.meta.deleted.append(&mut discarded_journals);
 		let canonicalized = (hash.clone(), self.front_block_number() + self.pending_canonicalizations.len() as u64);
 		commit.meta.inserted.push((to_meta_key(LAST_CANONICAL, &()), canonicalized.encode()));
 		trace!(target: "state-db", "Discarding {} records", commit.meta.deleted.len());
 		self.pending_canonicalizations.push(hash.clone());
-		Ok(commit)
+		Ok(())
 	}
 
 	fn apply_canonicalizations(&mut self) {
@@ -385,10 +411,10 @@ impl<BlockHash: Hash, Key: Hash> NonCanonicalOverlay<BlockHash, Key> {
 
 #[cfg(test)]
 mod tests {
-	use std::io;
+	use std::{collections::HashMap, io};
 	use primitives::H256;
 	use super::{NonCanonicalOverlay, to_journal_key};
-	use crate::ChangeSet;
+	use crate::{ChangeSet, CommitSet};
 	use crate::test::{make_db, make_changeset};
 
 	fn contains(overlay: &NonCanonicalOverlay<H256, H256>, key: u64) -> bool {
@@ -409,7 +435,8 @@ mod tests {
 	fn canonicalize_empty_panics() {
 		let db = make_db(&[]);
 		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
-		overlay.canonicalize::<io::Error>(&H256::default()).unwrap();
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&H256::default(), &HashMap::default(), &mut commit).unwrap();
 	}
 
 	#[test]
@@ -453,7 +480,8 @@ mod tests {
 		let db = make_db(&[]);
 		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
 		overlay.insert::<io::Error>(&h1, 1, &H256::default(), ChangeSet::default()).unwrap();
-		overlay.canonicalize::<io::Error>(&h2).unwrap();
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
 	}
 
 	#[test]
@@ -468,7 +496,8 @@ mod tests {
 		assert_eq!(insertion.meta.inserted.len(), 2);
 		assert_eq!(insertion.meta.deleted.len(), 0);
 		db.commit(&insertion);
-		let finalization = overlay.canonicalize::<io::Error>(&h1).unwrap();
+		let mut finalization = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut finalization).unwrap();
 		assert_eq!(finalization.data.inserted.len(), changeset.inserted.len());
 		assert_eq!(finalization.data.deleted.len(), changeset.deleted.len());
 		assert_eq!(finalization.meta.inserted.len(), 1);
@@ -501,7 +530,9 @@ mod tests {
 		let mut overlay = NonCanonicalOverlay::<H256, H256>::new(&db).unwrap();
 		db.commit(&overlay.insert::<io::Error>(&h1, 10, &H256::default(), make_changeset(&[3, 4], &[2])).unwrap());
 		db.commit(&overlay.insert::<io::Error>(&h2, 11, &h1, make_changeset(&[5], &[3])).unwrap());
-		db.commit(&overlay.canonicalize::<io::Error>(&h1).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 1);
 
@@ -526,7 +557,9 @@ mod tests {
 		assert!(contains(&overlay, 5));
 		assert_eq!(overlay.levels.len(), 2);
 		assert_eq!(overlay.parents.len(), 2);
-		db.commit(&overlay.canonicalize::<io::Error>(&h1).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		assert!(contains(&overlay, 5));
 		assert_eq!(overlay.levels.len(), 2);
 		assert_eq!(overlay.parents.len(), 2);
@@ -535,7 +568,9 @@ mod tests {
 		assert_eq!(overlay.parents.len(), 1);
 		assert!(!contains(&overlay, 5));
 		assert!(contains(&overlay, 7));
-		db.commit(&overlay.canonicalize::<io::Error>(&h2).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 0);
 		assert_eq!(overlay.parents.len(), 0);
@@ -552,7 +587,9 @@ mod tests {
 		db.commit(&overlay.insert::<io::Error>(&h_1, 1, &H256::default(), c_1).unwrap());
 		db.commit(&overlay.insert::<io::Error>(&h_2, 1, &H256::default(), c_2).unwrap());
 		assert!(contains(&overlay, 1));
-		db.commit(&overlay.canonicalize::<io::Error>(&h_1).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_1, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		assert!(contains(&overlay, 1));
 		overlay.apply_pending();
 		assert!(!contains(&overlay, 1));
@@ -569,8 +606,10 @@ mod tests {
 		db.commit(&overlay.insert::<io::Error>(&h1, 1, &H256::default(), changeset.clone()).unwrap());
 		db.commit(&overlay.insert::<io::Error>(&h2, 2, &h1, changeset.clone()).unwrap());
 		overlay.apply_pending();
-		db.commit(&overlay.canonicalize::<io::Error>(&h1).unwrap());
-		db.commit(&overlay.canonicalize::<io::Error>(&h2).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h1, &HashMap::default(), &mut commit).unwrap();
+		overlay.canonicalize::<io::Error>(&h2, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		db.commit(&overlay.insert::<io::Error>(&h3, 3, &h2, changeset.clone()).unwrap());
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 1);
@@ -639,7 +678,9 @@ mod tests {
 		assert_eq!(overlay.last_canonicalized, overlay2.last_canonicalized);
 
 		// canonicalize 1. 2 and all its children should be discarded
-		db.commit(&overlay.canonicalize::<io::Error>(&h_1).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_1, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 2);
 		assert_eq!(overlay.parents.len(), 6);
@@ -657,8 +698,15 @@ mod tests {
 		assert!(db.get_meta(&to_journal_key(2, 2)).unwrap().is_none());
 		assert!(db.get_meta(&to_journal_key(2, 3)).unwrap().is_none());
 
+		// check that discarding pinned state produces an error.
+		let mut commit = CommitSet::default();
+		let pinned = vec![(h_1_1_1, 1)].into_iter().collect();
+		assert!(overlay.canonicalize::<io::Error>(&h_1_2, &pinned, &mut commit).is_err());
+
 		// canonicalize 1_2. 1_1 and all its children should be discarded
-		db.commit(&overlay.canonicalize::<io::Error>(&h_1_2).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_1_2, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 1);
 		assert_eq!(overlay.parents.len(), 3);
@@ -673,7 +721,9 @@ mod tests {
 		assert!(!overlay.have_block(&h_1_1_1));
 
 		// canonicalize 1_2_2
-		db.commit(&overlay.canonicalize::<io::Error>(&h_1_2_2).unwrap());
+		let mut commit = CommitSet::default();
+		overlay.canonicalize::<io::Error>(&h_1_2_2, &HashMap::default(), &mut commit).unwrap();
+		db.commit(&commit);
 		overlay.apply_pending();
 		assert_eq!(overlay.levels.len(), 0);
 		assert_eq!(overlay.parents.len(), 0);

substrate/core/state-machine/src/backend.rs

@@ -106,7 +106,7 @@ pub trait Backend<H: Hasher> {
 	}
 
 	/// Try convert into trie backend.
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>>;
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>>;
 
 	/// Calculate the storage root, with given delta over what is already stored
 	/// in the backend, and produce a "transaction" that can be used to commit.
@@ -185,31 +185,33 @@ impl error::Error for Void {
 
 /// In-memory backend. Fully recomputes tries on each commit but useful for
 /// tests.
-#[derive(Eq)]
-pub struct InMemory<H> {
+pub struct InMemory<H: Hasher> {
 	inner: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>,
+	trie: Option<TrieBackend<MemoryDB<H>, H>>,
 	_hasher: PhantomData<H>,
 }
 
-impl<H> Default for InMemory<H> {
+impl<H: Hasher> Default for InMemory<H> {
 	fn default() -> Self {
 		InMemory {
 			inner: Default::default(),
+			trie: None,
 			_hasher: PhantomData,
 		}
 	}
 }
 
-impl<H> Clone for InMemory<H> {
+impl<H: Hasher> Clone for InMemory<H> {
 	fn clone(&self) -> Self {
 		InMemory {
 			inner: self.inner.clone(),
+			trie: None,
 			_hasher: PhantomData,
 		}
 	}
 }
 
-impl<H> PartialEq for InMemory<H> {
+impl<H: Hasher> PartialEq for InMemory<H> {
 	fn eq(&self, other: &Self) -> bool {
 		self.inner.eq(&other.inner)
 	}
@@ -230,27 +232,29 @@ impl<H: Hasher> InMemory<H> {
 	}
 }
 
-impl<H> From<HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>> for InMemory<H> {
+impl<H: Hasher> From<HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>> for InMemory<H> {
 	fn from(inner: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>) -> Self {
 		InMemory {
 			inner: inner,
+			trie: None,
 			_hasher: PhantomData,
 		}
 	}
 }
 
-impl<H> From<HashMap<Vec<u8>, Vec<u8>>> for InMemory<H> {
+impl<H: Hasher> From<HashMap<Vec<u8>, Vec<u8>>> for InMemory<H> {
 	fn from(inner: HashMap<Vec<u8>, Vec<u8>>) -> Self {
 		let mut expanded = HashMap::new();
 		expanded.insert(None, inner);
 		InMemory {
 			inner: expanded,
+			trie: None,
 			_hasher: PhantomData,
 		}
 	}
 }
 
-impl<H> From<Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>> for InMemory<H> {
+impl<H: Hasher> From<Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>> for InMemory<H> {
 	fn from(inner: Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>) -> Self {
 		let mut expanded: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>> = HashMap::new();
 		for (child_key, key, value) in inner {
@@ -365,16 +369,14 @@ impl<H: Hasher> Backend<H> for InMemory<H> {
 			.collect()
 	}
 
-	fn try_into_trie_backend(
-		self
-	)-> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&mut self)-> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		let mut mdb = MemoryDB::default();
 		let mut root = None;
 		let mut new_child_roots = Vec::new();
 		let mut root_map = None;
-		for (storage_key, map) in self.inner {
+		for (storage_key, map) in &self.inner {
 			if let Some(storage_key) = storage_key.as_ref() {
-				let ch = insert_into_memory_db::<H, _>(&mut mdb, map.into_iter())?;
+				let ch = insert_into_memory_db::<H, _>(&mut mdb, map.clone().into_iter())?;
 				new_child_roots.push((storage_key.clone(), ch.as_ref().into()));
 			} else {
 				root_map = Some(map);
@@ -384,14 +386,15 @@ impl<H: Hasher> Backend<H> for InMemory<H> {
 		if let Some(map) = root_map.take() {
 			root = Some(insert_into_memory_db::<H, _>(
 				&mut mdb,
-				map.into_iter().chain(new_child_roots.into_iter())
+				map.clone().into_iter().chain(new_child_roots.into_iter())
 			)?);
 		}
 		let root = match root {
 			Some(root) => root,
 			None => insert_into_memory_db::<H, _>(&mut mdb, ::std::iter::empty())?,
 		};
-		Some(TrieBackend::new(mdb, root))
+		self.trie = Some(TrieBackend::new(mdb, root));
+		self.trie.as_ref()
 	}
 }
 

substrate/core/state-machine/src/lib.rs

@@ -669,7 +669,7 @@ impl<'a, H, N, B, T, O, Exec> StateMachine<'a, H, N, B, T, O, Exec> where
 
 /// Prove execution using the given state backend, overlayed changes, and call executor.
 pub fn prove_execution<B, H, Exec>(
-	backend: B,
+	mut backend: B,
 	overlay: &mut OverlayedChanges,
 	exec: &Exec,
 	method: &str,
@@ -681,9 +681,9 @@ where
 	Exec: CodeExecutor<H>,
 	H::Out: Ord + 'static,
 {
-	let trie_backend = backend.try_into_trie_backend()
+	let trie_backend = backend.as_trie_backend()
 		.ok_or_else(|| Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>)?;
-	prove_execution_on_trie_backend(&trie_backend, overlay, exec, method, call_data)
+	prove_execution_on_trie_backend(trie_backend, overlay, exec, method, call_data)
 }
 
 /// Prove execution using the given trie backend, overlayed changes, and call executor.
@@ -778,7 +778,7 @@ where
 
 /// Generate storage read proof.
 pub fn prove_read<B, H>(
-	backend: B,
+	mut backend: B,
 	key: &[u8]
 ) -> Result<(Option<Vec<u8>>, Vec<Vec<u8>>), Box<Error>>
 where
@@ -786,16 +786,16 @@ where
 	H: Hasher,
 	H::Out: Ord
 {
-	let trie_backend = backend.try_into_trie_backend()
+	let trie_backend = backend.as_trie_backend()
 		.ok_or_else(
 			||Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>
 		)?;
-	prove_read_on_trie_backend(&trie_backend, key)
+	prove_read_on_trie_backend(trie_backend, key)
 }
 
 /// Generate child storage read proof.
 pub fn prove_child_read<B, H>(
-	backend: B,
+	mut backend: B,
 	storage_key: &[u8],
 	key: &[u8],
 ) -> Result<(Option<Vec<u8>>, Vec<Vec<u8>>), Box<Error>>
@@ -804,9 +804,9 @@ where
 	H: Hasher,
 	H::Out: Ord
 {
-	let trie_backend = backend.try_into_trie_backend()
+	let trie_backend = backend.as_trie_backend()
 		.ok_or_else(|| Box::new(ExecutionError::UnableToGenerateProof) as Box<Error>)?;
-	prove_child_read_on_trie_backend(&trie_backend, storage_key, key)
+	prove_child_read_on_trie_backend(trie_backend, storage_key, key)
 }
 
 
@@ -1100,7 +1100,8 @@ mod tests {
 			b"abc".to_vec() => b"2".to_vec(),
 			b"bbb".to_vec() => b"3".to_vec()
 		];
-		let backend = InMemory::<Blake2Hasher>::from(initial).try_into_trie_backend().unwrap();
+		let mut state = InMemory::<Blake2Hasher>::from(initial);
+		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges {
 			committed: map![
 				b"aba".to_vec() => OverlayedValue::from(Some(b"1312".to_vec())),
@@ -1115,7 +1116,7 @@ mod tests {
 
 		{
 			let changes_trie_storage = InMemoryChangesTrieStorage::<Blake2Hasher, u64>::new();
-			let mut ext = Ext::new(&mut overlay, &backend, Some(&changes_trie_storage), NeverOffchainExt::new());
+			let mut ext = Ext::new(&mut overlay, backend, Some(&changes_trie_storage), NeverOffchainExt::new());
 			ext.clear_prefix(b"ab");
 		}
 		overlay.commit_prospective();
@@ -1136,12 +1137,13 @@ mod tests {
 
 	#[test]
 	fn set_child_storage_works() {
-		let backend = InMemory::<Blake2Hasher>::default().try_into_trie_backend().unwrap();
+		let mut state = InMemory::<Blake2Hasher>::default();
+		let backend = state.as_trie_backend().unwrap();
 		let changes_trie_storage = InMemoryChangesTrieStorage::<Blake2Hasher, u64>::new();
 		let mut overlay = OverlayedChanges::default();
 		let mut ext = Ext::new(
 			&mut overlay,
-			&backend,
+			backend,
 			Some(&changes_trie_storage),
 			NeverOffchainExt::new()
 		);

substrate/core/state-machine/src/proving_backend.rs

@@ -184,7 +184,7 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		self.backend.child_storage_root(storage_key, delta)
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		None
 	}
 }
@@ -269,16 +269,16 @@ mod tests {
 	fn proof_recorded_and_checked() {
 		let contents = (0..64).map(|i| (None, vec![i], Some(vec![i]))).collect::<Vec<_>>();
 		let in_memory = InMemory::<Blake2Hasher>::default();
-		let in_memory = in_memory.update(contents);
+		let mut in_memory = in_memory.update(contents);
 		let in_memory_root = in_memory.storage_root(::std::iter::empty()).0;
 		(0..64).for_each(|i| assert_eq!(in_memory.storage(&[i]).unwrap().unwrap(), vec![i]));
 
-		let trie = in_memory.try_into_trie_backend().unwrap();
+		let trie = in_memory.as_trie_backend().unwrap();
 		let trie_root = trie.storage_root(::std::iter::empty()).0;
 		assert_eq!(in_memory_root, trie_root);
 		(0..64).for_each(|i| assert_eq!(trie.storage(&[i]).unwrap().unwrap(), vec![i]));
 
-		let proving = ProvingBackend::new(&trie);
+		let proving = ProvingBackend::new(trie);
 		assert_eq!(proving.storage(&[42]).unwrap().unwrap(), vec![42]);
 
 		let proof = proving.extract_proof();
@@ -302,7 +302,7 @@ mod tests {
 			.chain((10..15).map(|i| (Some(own2.clone()), vec![i], Some(vec![i]))))
 			.collect::<Vec<_>>();
 		let in_memory = InMemory::<Blake2Hasher>::default();
-		let in_memory = in_memory.update(contents);
+		let mut in_memory = in_memory.update(contents);
 		let in_memory_root = in_memory.full_storage_root::<_, Vec<_>, _>(
 			::std::iter::empty(),
       in_memory.child_storage_keys().map(|k|(k.to_vec(), Vec::new()))
@@ -320,7 +320,7 @@ mod tests {
 			vec![i]
 		));
 
-		let trie = in_memory.try_into_trie_backend().unwrap();
+		let trie = in_memory.as_trie_backend().unwrap();
 		let trie_root = trie.storage_root(::std::iter::empty()).0;
 		assert_eq!(in_memory_root, trie_root);
 		(0..64).for_each(|i| assert_eq!(
@@ -328,7 +328,7 @@ mod tests {
 			vec![i]
 		));
 
-		let proving = ProvingBackend::new(&trie);
+		let proving = ProvingBackend::new(trie);
 		assert_eq!(proving.storage(&[42]).unwrap().unwrap(), vec![42]);
 
 		let proof = proving.extract_proof();
@@ -343,7 +343,7 @@ mod tests {
 		assert_eq!(proof_check.storage(&[41]).unwrap().unwrap(), vec![41]);
 		assert_eq!(proof_check.storage(&[64]).unwrap(), None);
 
-		let proving = ProvingBackend::new(&trie);
+		let proving = ProvingBackend::new(trie);
 		assert_eq!(proving.child_storage(&own1[..], &[64]), Ok(Some(vec![64])));
 
 		let proof = proving.extract_proof();

substrate/core/state-machine/src/trie_backend.rs

@@ -179,7 +179,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 		(root, is_default, write_overlay)
 	}
 
-	fn try_into_trie_backend(self) -> Option<TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		Some(self)
 	}
 }