ProofRecorder: Implement transactional support

client/block-builder/src/lib.rs

@@ -312,7 +312,9 @@ mod tests {
 	use sp_blockchain::HeaderBackend;
 	use sp_core::Blake2Hasher;
 	use sp_state_machine::Backend;
-	use substrate_test_runtime_client::{DefaultTestClientBuilderExt, TestClientBuilderExt};
+	use substrate_test_runtime_client::{
+		runtime::Extrinsic, DefaultTestClientBuilderExt, TestClientBuilderExt,
+	};
 
 	#[test]
 	fn block_building_storage_proof_does_not_include_runtime_by_default() {
@@ -345,4 +347,62 @@ mod tests {
 			.unwrap_err()
 			.contains("Database missing expected key"),);
 	}
+
+	#[test]
+	fn failing_extrinsic_rolls_back_changes_in_storage_proof() {
+		let builder = substrate_test_runtime_client::TestClientBuilder::new();
+		let backend = builder.backend();
+		let client = builder.build();
+
+		let mut block_builder = BlockBuilder::new(
+			&client,
+			client.info().best_hash,
+			client.info().best_number,
+			RecordProof::Yes,
+			Default::default(),
+			&*backend,
+		)
+		.unwrap();
+
+		block_builder.push(Extrinsic::ReadAndPanic(8)).unwrap_err();
+
+		let block = block_builder.build().unwrap();
+
+		let proof_with_panic = block.proof.expect("Proof is build on request").encoded_size();
+
+		let mut block_builder = BlockBuilder::new(
+			&client,
+			client.info().best_hash,
+			client.info().best_number,
+			RecordProof::Yes,
+			Default::default(),
+			&*backend,
+		)
+		.unwrap();
+
+		block_builder.push(Extrinsic::Read(8)).unwrap();
+
+		let block = block_builder.build().unwrap();
+
+		let proof_without_panic = block.proof.expect("Proof is build on request").encoded_size();
+
+		let block = BlockBuilder::new(
+			&client,
+			client.info().best_hash,
+			client.info().best_number,
+			RecordProof::Yes,
+			Default::default(),
+			&*backend,
+		)
+		.unwrap()
+		.build()
+		.unwrap();
+
+		let proof_empty_block = block.proof.expect("Proof is build on request").encoded_size();
+
+		// Ensure that we rolled back the changes of the panicked transaction.
+		assert!(proof_without_panic > proof_with_panic);
+		assert!(proof_without_panic > proof_empty_block);
+		assert_eq!(proof_empty_block, proof_with_panic);
+	}
 }

primitives/api/proc-macro/src/impl_runtime_apis.rs

@@ -236,7 +236,8 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 				&self,
 				call: F,
 			) -> R where Self: Sized {
-				#crate_::OverlayedChanges::start_transaction(&mut std::cell::RefCell::borrow_mut(&self.changes));
+				self.start_transaction();
+
 				*std::cell::RefCell::borrow_mut(&self.commit_on_success) = false;
 				let res = call(self);
 				*std::cell::RefCell::borrow_mut(&self.commit_on_success) = true;
@@ -340,18 +341,51 @@ fn generate_runtime_api_base_structures() -> Result<TokenStream> {
 					transactions; qed";
 				if *std::cell::RefCell::borrow(&self.commit_on_success) {
 					let res = if commit {
-						#crate_::OverlayedChanges::commit_transaction(
+						let res = if let Some(recorder) = &self.recorder {
+							#crate_::ProofRecorder::<Block>::commit_transaction(&recorder)
+						} else {
+							Ok(())
+						};
+
+						let res2 = #crate_::OverlayedChanges::commit_transaction(
 							&mut std::cell::RefCell::borrow_mut(&self.changes)
-						)
+						);
+
+						// Will panic on an `Err` below, however we should call commit
+						// on the recorder and the changes together.
+						std::result::Result::and(res, std::result::Result::map_err(res2, drop))
 					} else {
-						#crate_::OverlayedChanges::rollback_transaction(
+						let res = if let Some(recorder) = &self.recorder {
+							#crate_::ProofRecorder::<Block>::rollback_transaction(&recorder)
+						} else {
+							Ok(())
+						};
+
+						let res2 = #crate_::OverlayedChanges::rollback_transaction(
 							&mut std::cell::RefCell::borrow_mut(&self.changes)
-						)
+						);
+
+						// Will panic on an `Err` below, however we should call commit
+						// on the recorder and the changes together.
+						std::result::Result::and(res, std::result::Result::map_err(res2, drop))
 					};
 
 					std::result::Result::expect(res, proof);
 				}
 			}
+
+			fn start_transaction(&self) {
+				if !*std::cell::RefCell::borrow(&self.commit_on_success) {
+					return
+				}
+
+				#crate_::OverlayedChanges::start_transaction(
+					&mut std::cell::RefCell::borrow_mut(&self.changes)
+				);
+				if let Some(recorder) = &self.recorder {
+					#crate_::ProofRecorder::<Block>::start_transaction(&recorder);
+				}
+			}
 		}
 	))
 }
@@ -443,11 +477,7 @@ impl<'a> ApiRuntimeImplToApiRuntimeApiImpl<'a> {
 				params: std::vec::Vec<u8>,
 				fn_name: &dyn Fn(#crate_::RuntimeVersion) -> &'static str,
 			) -> std::result::Result<std::vec::Vec<u8>, #crate_::ApiError> {
-				if *std::cell::RefCell::borrow(&self.commit_on_success) {
-					#crate_::OverlayedChanges::start_transaction(
-						&mut std::cell::RefCell::borrow_mut(&self.changes)
-					);
-				}
+				self.start_transaction();
 
 				let res = (|| {
 					let version = #crate_::CallApiAt::<__SrApiBlock__>::runtime_version_at(