Fetch runtime code from storage cache when using proofing backend (paritytech#9611)

Before we fetched the runtime code from the `TrieBackend` and this lead
to not using the storage cache. Thus, we recalculated the storage hash
for the runtime code on every call into the runtime and this killed the
performance on parachains block authoring. The solution is to fetch the
runtime code from the storage cache, to make sure we use the cached
storage cache.

(cherry picked from commit b79e2b9)

Author: bkchr

client/api/src/cht.rs

@@ -116,7 +116,7 @@ where
 		.into_iter()
 		.map(|(k, v)| (k, Some(v)))
 		.collect::<Vec<_>>();
-	let mut storage = InMemoryBackend::<Hasher>::default().update(vec![(None, transaction)]);
+	let storage = InMemoryBackend::<Hasher>::default().update(vec![(None, transaction)]);
 	let trie_storage = storage
 		.as_trie_backend()
 		.expect("InMemoryState::as_trie_backend always returns Some; qed");

client/db/src/bench.rs

@@ -454,12 +454,6 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 			.map_or(Default::default(), |s| s.child_keys(child_info, prefix))
 	}
 
-	fn as_trie_backend(
-		&mut self,
-	) -> Option<&sp_state_machine::TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
-		None
-	}
-
 	fn commit(
 		&self,
 		storage_root: <HashFor<B> as Hasher>::Out,

client/db/src/lib.rs

@@ -274,7 +274,7 @@ impl<B: BlockT> StateBackend<HashFor<B>> for RefTrackingState<B> {
 	}
 
 	fn as_trie_backend(
-		&mut self,
+		&self,
 	) -> Option<&sp_state_machine::TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
 		self.state.as_trie_backend()
 	}

client/db/src/storage_cache.rs

@@ -703,7 +703,7 @@ impl<S: StateBackend<HashFor<B>>, B: BlockT> StateBackend<HashFor<B>> for Cachin
 		self.state.child_keys(child_info, prefix)
 	}
 
-	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
+	fn as_trie_backend(&self) -> Option<&TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
 		self.state.as_trie_backend()
 	}
 
@@ -901,9 +901,9 @@ impl<S: StateBackend<HashFor<B>>, B: BlockT> StateBackend<HashFor<B>>
 		self.caching_state().child_keys(child_info, prefix)
 	}
 
-	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
+	fn as_trie_backend(&self) -> Option<&TrieBackend<Self::TrieBackendStorage, HashFor<B>>> {
 		self.caching_state
-			.as_mut()
+			.as_ref()
 			.expect("`caching_state` is valid for the lifetime of the object; qed")
 			.as_trie_backend()
 	}

client/light/src/backend.rs

@@ -569,9 +569,9 @@ where
 		sp_state_machine::UsageInfo::empty()
 	}
 
-	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		match self {
-			GenesisOrUnavailableState::Genesis(ref mut state) => state.as_trie_backend(),
+			GenesisOrUnavailableState::Genesis(ref state) => state.as_trie_backend(),
 			GenesisOrUnavailableState::Unavailable => None,
 		}
 	}

client/service/src/client/call_executor.rs

@@ -212,29 +212,30 @@ where
 			backend::changes_tries_state_at_block(at, self.backend.changes_trie_storage())?;
 		let mut storage_transaction_cache = storage_transaction_cache.map(|c| c.borrow_mut());
 
-		let mut state = self.backend.state_at(*at)?;
+		let state = self.backend.state_at(*at)?;
 
 		let changes = &mut *changes.borrow_mut();
 
 		let at_hash = self.backend.blockchain().block_hash_from_id(at)?.ok_or_else(|| {
 			sp_blockchain::Error::UnknownBlock(format!("Could not find block hash for {:?}", at))
 		})?;
 
+		// It is important to extract the runtime code here before we create the proof
+		// recorder to not record it. We also need to fetch the runtime code from `state` to
+		// make sure we use the caching layers.
+		let state_runtime_code = sp_state_machine::backend::BackendRuntimeCode::new(&state);
+
+		let runtime_code =
+			state_runtime_code.runtime_code().map_err(sp_blockchain::Error::RuntimeCode)?;
+		let runtime_code = self.check_override(runtime_code, at)?;
+
 		match recorder {
 			Some(recorder) => {
 				let trie_state = state.as_trie_backend().ok_or_else(|| {
 					Box::new(sp_state_machine::ExecutionError::UnableToGenerateProof)
 						as Box<dyn sp_state_machine::Error>
 				})?;
 
-				let state_runtime_code =
-					sp_state_machine::backend::BackendRuntimeCode::new(trie_state);
-				// It is important to extract the runtime code here before we create the proof
-				// recorder.
-				let runtime_code =
-					state_runtime_code.runtime_code().map_err(sp_blockchain::Error::RuntimeCode)?;
-				let runtime_code = self.check_override(runtime_code, at)?;
-
 				let backend = sp_state_machine::ProvingBackend::new_with_recorder(
 					trie_state,
 					recorder.clone(),
@@ -260,11 +261,6 @@ where
 				)
 			},
 			None => {
-				let state_runtime_code = sp_state_machine::backend::BackendRuntimeCode::new(&state);
-				let runtime_code =
-					state_runtime_code.runtime_code().map_err(sp_blockchain::Error::RuntimeCode)?;
-				let runtime_code = self.check_override(runtime_code, at)?;
-
 				let mut state_machine = StateMachine::new(
 					&state,
 					changes_trie_state,
@@ -310,7 +306,7 @@ where
 		method: &str,
 		call_data: &[u8],
 	) -> sp_blockchain::Result<(Vec<u8>, StorageProof)> {
-		let mut state = self.backend.state_at(*at)?;
+		let state = self.backend.state_at(*at)?;
 
 		let trie_backend = state.as_trie_backend().ok_or_else(|| {
 			Box::new(sp_state_machine::ExecutionError::UnableToGenerateProof)

primitives/state-machine/src/backend.rs

@@ -173,7 +173,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	}
 
 	/// Try convert into trie backend.
-	fn as_trie_backend(&mut self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
+	fn as_trie_backend(&self) -> Option<&TrieBackend<Self::TrieBackendStorage, H>> {
 		None
 	}
 

primitives/state-machine/src/in_memory_backend.rs

@@ -175,7 +175,7 @@ mod tests {
 		let storage = new_in_mem::<BlakeTwo256>();
 		let child_info = ChildInfo::new_default(b"1");
 		let child_info = &child_info;
-		let mut storage = storage
+		let storage = storage
 			.update(vec![(Some(child_info.clone()), vec![(b"2".to_vec(), Some(b"3".to_vec()))])]);
 		let trie_backend = storage.as_trie_backend().unwrap();
 		assert_eq!(trie_backend.child_storage(child_info, b"2").unwrap(), Some(b"3".to_vec()));

primitives/state-machine/src/lib.rs

@@ -714,7 +714,7 @@ mod execution {
 	}
 
 	/// Generate storage read proof.
-	pub fn prove_read<B, H, I>(mut backend: B, keys: I) -> Result<StorageProof, Box<dyn Error>>
+	pub fn prove_read<B, H, I>(backend: B, keys: I) -> Result<StorageProof, Box<dyn Error>>
 	where
 		B: Backend<H>,
 		H: Hasher,
@@ -730,7 +730,7 @@ mod execution {
 
 	/// Generate range storage read proof.
 	pub fn prove_range_read_with_size<B, H>(
-		mut backend: B,
+		backend: B,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
 		size_limit: usize,
@@ -789,7 +789,7 @@ mod execution {
 
 	/// Generate child storage read proof.
 	pub fn prove_child_read<B, H, I>(
-		mut backend: B,
+		backend: B,
 		child_info: &ChildInfo,
 		keys: I,
 	) -> Result<StorageProof, Box<dyn Error>>
@@ -1192,7 +1192,7 @@ mod tests {
 			b"abc".to_vec() => b"2".to_vec(),
 			b"bbb".to_vec() => b"3".to_vec()
 		];
-		let mut state = InMemoryBackend::<BlakeTwo256>::from(initial);
+		let state = InMemoryBackend::<BlakeTwo256>::from(initial);
 		let backend = state.as_trie_backend().unwrap();
 
 		let mut overlay = OverlayedChanges::default();
@@ -1345,7 +1345,7 @@ mod tests {
 	fn set_child_storage_works() {
 		let child_info = ChildInfo::new_default(b"sub1");
 		let child_info = &child_info;
-		let mut state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 		let mut cache = StorageTransactionCache::default();
@@ -1367,7 +1367,7 @@ mod tests {
 	fn append_storage_works() {
 		let reference_data = vec![b"data1".to_vec(), b"2".to_vec(), b"D3".to_vec(), b"d4".to_vec()];
 		let key = b"key".to_vec();
-		let mut state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 		let mut cache = StorageTransactionCache::default();
@@ -1422,7 +1422,7 @@ mod tests {
 
 		let key = b"events".to_vec();
 		let mut cache = StorageTransactionCache::default();
-		let mut state = new_in_mem::<BlakeTwo256>();
+		let state = new_in_mem::<BlakeTwo256>();
 		let backend = state.as_trie_backend().unwrap();
 		let mut overlay = OverlayedChanges::default();
 
@@ -1793,7 +1793,7 @@ mod tests {
 			b"aaa".to_vec() => b"0".to_vec(),
 			b"bbb".to_vec() => b"".to_vec()
 		];
-		let mut state = InMemoryBackend::<BlakeTwo256>::from(initial);
+		let state = InMemoryBackend::<BlakeTwo256>::from(initial);
 		let backend = state.as_trie_backend().unwrap();
 
 		let mut overlay = OverlayedChanges::default();

primitives/state-machine/src/proving_backend.rs

@@ -433,7 +433,7 @@ mod tests {
 	fn proof_recorded_and_checked() {
 		let contents = (0..64).map(|i| (vec![i], Some(vec![i]))).collect::<Vec<_>>();
 		let in_memory = InMemoryBackend::<BlakeTwo256>::default();
-		let mut in_memory = in_memory.update(vec![(None, contents)]);
+		let in_memory = in_memory.update(vec![(None, 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]));
 
@@ -464,7 +464,7 @@ mod tests {
 			(Some(child_info_2.clone()), (10..15).map(|i| (vec![i], Some(vec![i]))).collect()),
 		];
 		let in_memory = InMemoryBackend::<BlakeTwo256>::default();
-		let mut in_memory = in_memory.update(contents);
+		let in_memory = in_memory.update(contents);
 		let child_storage_keys = vec![child_info_1.to_owned(), child_info_2.to_owned()];
 		let in_memory_root = in_memory
 			.full_storage_root(