Create benchmark for the system::set_code instrisic

frame/system/Cargo.toml

@@ -13,6 +13,7 @@ readme = "README.md"
 targets = ["x86_64-unknown-linux-gnu"]
 
 [dependencies]
+cfg-if = "1.0"
 codec = { package = "parity-scale-codec", version = "3.2.2", default-features = false, features = ["derive"] }
 log = { version = "0.4.17", default-features = false }
 scale-info = { version = "2.5.0", default-features = false, features = ["derive"] }

frame/system/benchmarking/Cargo.toml

@@ -24,6 +24,8 @@ sp-std = { version = "5.0.0", default-features = false, path = "../../../primiti
 
 [dev-dependencies]
 sp-io = { version = "7.0.0", path = "../../../primitives/io" }
+sp-externalities = { version = "0.13.0", path = "../../../primitives/externalities" }
+sp-version = { version = "5.0.0", path = "../../../primitives/version" }
 
 [features]
 default = ["std"]

frame/system/benchmarking/res/README.md

@@ -0,0 +1,5 @@
+These runtimes are used for benchmarking the `set_code` intrinsic.
+
+**Don't use them in production environments!**
+
+To update the just copy the new runtime from `target/release/wbuild/kitchensink-runtime/kitchensink_runtime.compact.compressed.wasm` to here.

frame/system/benchmarking/src/lib.rs

@@ -48,9 +48,12 @@ benchmarks! {
 	set_heap_pages {
 	}: _(RawOrigin::Root, Default::default())
 
-	// `set_code` was not benchmarked because it is pretty hard to come up with a real
-	// Wasm runtime to test the upgrade with. But this is okay because we will make
-	// `set_code` take a full block anyway.
+	set_code {
+		let runtime_blob = include_bytes!("../res/kitchensink_runtime.compact.compressed.wasm");
+	}: _(RawOrigin::Root, runtime_blob.to_vec())
+	verify {
+		System::<T>::assert_last_event(frame_system::Event::<T>::CodeUpdated.into());
+	}
 
 	#[extra]
 	set_code_without_checks {

frame/system/benchmarking/src/mock.rs

@@ -19,6 +19,7 @@
 
 #![cfg(test)]
 
+use codec::Encode;
 use sp_runtime::traits::IdentityLookup;
 
 type AccountId = u64;
@@ -67,7 +68,29 @@ impl frame_system::Config for Test {
 
 impl crate::Config for Test {}
 
+struct MockedReadRuntimeVersion(Vec<u8>);
+
+impl sp_core::traits::ReadRuntimeVersion for MockedReadRuntimeVersion {
+	fn read_runtime_version(
+		&self,
+		_wasm_code: &[u8],
+		_ext: &mut dyn sp_externalities::Externalities,
+	) -> Result<Vec<u8>, String> {
+		Ok(self.0.clone())
+	}
+}
+
 pub fn new_test_ext() -> sp_io::TestExternalities {
 	let t = frame_system::GenesisConfig::default().build_storage::<Test>().unwrap();
-	sp_io::TestExternalities::new(t)
+
+	let version = sp_version::RuntimeVersion {
+		spec_name: "spec_name".into(),
+		spec_version: 123,
+		impl_version: 456,
+		..Default::default()
+	};
+	let read_runtime_version = MockedReadRuntimeVersion(version.encode());
+	let mut ext = sp_io::TestExternalities::new(t);
+	ext.register_extension(sp_core::traits::ReadRuntimeVersionExt::new(read_runtime_version));
+	ext
 }

frame/system/src/lib.rs

@@ -373,7 +373,6 @@ pub mod pallet {
 	impl<T: Config> Pallet<T> {
 		/// Make some on-chain remark.
 		///
-		/// ## Complexity
 		/// - `O(1)`
 		#[pallet::call_index(0)]
 		#[pallet::weight(T::SystemWeightInfo::remark(_remark.len() as u32))]
@@ -394,30 +393,34 @@ pub mod pallet {
 
 		/// Set the new runtime code.
 		///
-		/// ## Complexity
 		/// - `O(C + S)` where `C` length of `code` and `S` complexity of `can_set_code`
+		/// - 1 call to `can_set_code`: `O(S)` (calls `sp_io::misc::runtime_version` which is
+		///   expensive).
+		/// The weight of this function is dependent on the runtime, but generally this is very
+		/// expensive.
 		#[pallet::call_index(2)]
-		#[pallet::weight((T::BlockWeights::get().max_block, DispatchClass::Operational))]
+		#[pallet::weight((T::SystemWeightInfo::set_code(), DispatchClass::Operational))]
 		pub fn set_code(origin: OriginFor<T>, code: Vec<u8>) -> DispatchResultWithPostInfo {
 			ensure_root(origin)?;
 			Self::can_set_code(&code)?;
 			T::OnSetCode::set_code(code)?;
-			Ok(().into())
+			// consume the rest of the block to prevent further transactions
+			Ok(Some(T::BlockWeights::get().max_block).into())
 		}
 
 		/// Set the new runtime code without doing any checks of the given `code`.
 		///
-		/// ## Complexity
 		/// - `O(C)` where `C` length of `code`
+		/// The weight of this function is dependent on the runtime.
 		#[pallet::call_index(3)]
-		#[pallet::weight((T::BlockWeights::get().max_block, DispatchClass::Operational))]
+		#[pallet::weight((T::SystemWeightInfo::set_code(), DispatchClass::Operational))]
 		pub fn set_code_without_checks(
 			origin: OriginFor<T>,
 			code: Vec<u8>,
 		) -> DispatchResultWithPostInfo {
 			ensure_root(origin)?;
 			T::OnSetCode::set_code(code)?;
-			Ok(().into())
+			Ok(Some(T::BlockWeights::get().max_block).into())
 		}
 
 		/// Set some items of storage.
@@ -1405,7 +1408,6 @@ impl<T: Config> Pallet<T> {
 
 	/// Deposits a log and ensures it matches the block's log data.
 	///
-	/// ## Complexity
 	/// - `O(1)`
 	pub fn deposit_log(item: generic::DigestItem) {
 		<Digest<T>>::append(item);
@@ -1604,15 +1606,23 @@ impl<T: Config> Pallet<T> {
 			.and_then(|v| RuntimeVersion::decode(&mut &v[..]).ok())
 			.ok_or(Error::<T>::FailedToExtractRuntimeVersion)?;
 
-		if new_version.spec_name != current_version.spec_name {
-			return Err(Error::<T>::InvalidSpecName.into())
-		}
+		cfg_if::cfg_if! {
+			 if #[cfg(all(feature = "runtime-benchmarks", not(test)))] {
+					// Let's ensure the compiler doesn't optimize our fetching of the runtime version away.
+					core::hint::black_box((new_version, current_version));
+					Ok(())
+			  } else {
+				  if new_version.spec_name != current_version.spec_name {
+					  return Err(Error::<T>::InvalidSpecName.into())
+				   }
 
-		if new_version.spec_version <= current_version.spec_version {
-			return Err(Error::<T>::SpecVersionNeedsToIncrease.into())
-		}
+				   if new_version.spec_version <= current_version.spec_version {
+						return Err(Error::<T>::SpecVersionNeedsToIncrease.into())
+				   }
 
-		Ok(())
+				   Ok(())
+			  }
+		}
 	}
 }
 

frame/system/src/tests.rs

@@ -550,7 +550,12 @@ fn set_code_checks_works() {
 		("test", 1, 2, Err(Error::<Test>::SpecVersionNeedsToIncrease)),
 		("test", 1, 1, Err(Error::<Test>::SpecVersionNeedsToIncrease)),
 		("test2", 1, 1, Err(Error::<Test>::InvalidSpecName)),
-		("test", 2, 1, Ok(PostDispatchInfo::default())),
+		(
+			"test",
+			2,
+			1,
+			Ok(Some(<mock::Test as pallet::Config>::BlockWeights::get().max_block).into()),
+		),
 		("test", 0, 1, Err(Error::<Test>::SpecVersionNeedsToIncrease)),
 		("test", 1, 0, Err(Error::<Test>::SpecVersionNeedsToIncrease)),
 	];