Always wipe DB state after benchmark (#12025)

* Add test

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* Always wipe the DB state

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* Silence Clippy

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

substrate/Cargo.lock

@@ -2102,6 +2102,7 @@ dependencies = [
  "serde",
  "sp-api",
  "sp-application-crypto",
+ "sp-core",
  "sp-io",
  "sp-keystore",
  "sp-runtime",

substrate/frame/benchmarking/Cargo.toml

@@ -23,6 +23,7 @@ frame-support = { version = "4.0.0-dev", default-features = false, path = "../su
 frame-system = { version = "4.0.0-dev", default-features = false, path = "../system" }
 sp-api = { version = "4.0.0-dev", default-features = false, path = "../../primitives/api" }
 sp-application-crypto = { version = "6.0.0", default-features = false, path = "../../primitives/application-crypto" }
+sp-core = { version = "6.0.0", default-features = false, path = "../../primitives/core" }
 sp-io = { version = "6.0.0", default-features = false, path = "../../primitives/io" }
 sp-runtime = { version = "6.0.0", default-features = false, path = "../../primitives/runtime" }
 sp-runtime-interface = { version = "6.0.0", default-features = false, path = "../../primitives/runtime-interface" }
@@ -45,6 +46,7 @@ std = [
 	"serde",
 	"sp-api/std",
 	"sp-application-crypto/std",
+	"sp-core/std",
 	"sp-io/std",
 	"sp-runtime-interface/std",
 	"sp-runtime/std",

substrate/frame/benchmarking/src/lib.rs

@@ -38,6 +38,8 @@ pub use log;
 #[doc(hidden)]
 pub use paste;
 #[doc(hidden)]
+pub use sp_core::defer;
+#[doc(hidden)]
 pub use sp_io::storage::root as storage_root;
 #[doc(hidden)]
 pub use sp_runtime::traits::Zero;
@@ -1033,6 +1035,9 @@ macro_rules! impl_benchmark {
 
 				// Always do at least one internal repeat...
 				for _ in 0 .. internal_repeats.max(1) {
+					// Always reset the state after the benchmark.
+					$crate::defer!($crate::benchmarking::wipe_db());
+
 					// Set up the externalities environment for the setup we want to
 					// benchmark.
 					let closure_to_benchmark = <
@@ -1110,9 +1115,6 @@ macro_rules! impl_benchmark {
 						proof_size: diff_pov,
 						keys: read_and_written_keys,
 					});
-
-					// Wipe the DB back to the genesis state.
-					$crate::benchmarking::wipe_db();
 				}
 
 				return Ok(results);
@@ -1175,6 +1177,9 @@ macro_rules! impl_benchmark_test {
 					let execute_benchmark = |
 						c: $crate::Vec<($crate::BenchmarkParameter, u32)>
 					| -> Result<(), $crate::BenchmarkError> {
+						// Always reset the state after the benchmark.
+						$crate::defer!($crate::benchmarking::wipe_db());
+
 						// Set up the benchmark, return execution + verification function.
 						let closure_to_verify = <
 							SelectedBenchmark as $crate::BenchmarkingSetup<T, _>
@@ -1186,12 +1191,7 @@ macro_rules! impl_benchmark_test {
 						}
 
 						// Run execution + verification
-						closure_to_verify()?;
-
-						// Reset the state
-						$crate::benchmarking::wipe_db();
-
-						Ok(())
+						closure_to_verify()
 					};
 
 					if components.is_empty() {

substrate/frame/benchmarking/src/tests.rs

@@ -125,6 +125,8 @@ fn new_test_ext() -> sp_io::TestExternalities {
 	GenesisConfig::default().build_storage().unwrap().into()
 }
 
+// NOTE: This attribute is only needed for the `modify_in_` functions.
+#[allow(unreachable_code)]
 mod benchmarks {
 	use super::{new_test_ext, pallet_test::Value, Test};
 	use crate::{account, BenchmarkError, BenchmarkParameter, BenchmarkResult, BenchmarkingSetup};
@@ -227,6 +229,24 @@ mod benchmarks {
 			// This should never be reached.
 			assert!(value > 100);
 		}
+
+		modify_in_setup_then_error {
+			Value::<T>::set(Some(123));
+			return Err(BenchmarkError::Stop("Should error"));
+		}: { }
+
+		modify_in_call_then_error {
+		}: {
+			Value::<T>::set(Some(123));
+			return Err(BenchmarkError::Stop("Should error"));
+		}
+
+		modify_in_verify_then_error {
+		}: {
+		} verify {
+			Value::<T>::set(Some(123));
+			return Err(BenchmarkError::Stop("Should error"));
+		}
 	}
 
 	#[test]
@@ -350,4 +370,28 @@ mod benchmarks {
 			assert_eq!(Pallet::<Test>::test_benchmark_skip_benchmark(), Err(BenchmarkError::Skip),);
 		});
 	}
+
+	/// An error return of a benchmark test function still causes the db to be wiped.
+	#[test]
+	fn benchmark_error_wipes_storage() {
+		new_test_ext().execute_with(|| {
+			// It resets when the error happens in the setup:
+			assert_err!(
+				Pallet::<Test>::test_benchmark_modify_in_setup_then_error(),
+				"Should error"
+			);
+			assert_eq!(Value::<Test>::get(), None);
+
+			// It resets when the error happens in the call:
+			assert_err!(Pallet::<Test>::test_benchmark_modify_in_call_then_error(), "Should error");
+			assert_eq!(Value::<Test>::get(), None);
+
+			// It resets when the error happens in the verify:
+			assert_err!(
+				Pallet::<Test>::test_benchmark_modify_in_verify_then_error(),
+				"Should error"
+			);
+			assert_eq!(Value::<Test>::get(), None);
+		});
+	}
 }