Fix on_runtime_upgrade weight recording (#7480)

* Fix on_runtime_upgrade weight recording

* fix naming

* Update lib.rs

* fix line width

* fix line width again

Author: shawntabrizi

src/lib.rs

@@ -234,12 +234,12 @@ where
 		extrinsics_root: &System::Hash,
 		digest: &Digest<System::Hash>,
 	) {
+		let mut weight = 0;
 		if Self::runtime_upgraded() {
 			// System is not part of `AllModules`, so we need to call this manually.
-			let mut weight = <frame_system::Module::<System> as OnRuntimeUpgrade>::on_runtime_upgrade();
+			weight = weight.saturating_add(<frame_system::Module::<System> as OnRuntimeUpgrade>::on_runtime_upgrade());
 			weight = weight.saturating_add(COnRuntimeUpgrade::on_runtime_upgrade());
 			weight = weight.saturating_add(<AllModules as OnRuntimeUpgrade>::on_runtime_upgrade());
-			<frame_system::Module<System>>::register_extra_weight_unchecked(weight, DispatchClass::Mandatory);
 		}
 		<frame_system::Module<System>>::initialize(
 			block_number,
@@ -248,8 +248,10 @@ where
 			digest,
 			frame_system::InitKind::Full,
 		);
-		<frame_system::Module<System> as OnInitialize<System::BlockNumber>>::on_initialize(*block_number);
-		let weight = <AllModules as OnInitialize<System::BlockNumber>>::on_initialize(*block_number)
+		weight = weight.saturating_add(
+			<frame_system::Module<System> as OnInitialize<System::BlockNumber>>::on_initialize(*block_number)
+		);
+		weight = weight.saturating_add(<AllModules as OnInitialize<System::BlockNumber>>::on_initialize(*block_number))
 			.saturating_add(<System::BlockExecutionWeight as frame_support::traits::Get<_>>::get());
 		<frame_system::Module::<System>>::register_extra_weight_unchecked(weight, DispatchClass::Mandatory);
 
@@ -543,7 +545,7 @@ mod tests {
 
 				fn on_runtime_upgrade() -> Weight {
 					sp_io::storage::set(super::TEST_KEY, "module".as_bytes());
-					0
+					200
 				}
 			}
 		}
@@ -675,7 +677,7 @@ mod tests {
 		fn on_runtime_upgrade() -> Weight {
 			sp_io::storage::set(TEST_KEY, "custom_upgrade".as_bytes());
 			sp_io::storage::set(CUSTOM_ON_RUNTIME_KEY, &true.encode());
-			0
+			100
 		}
 	}
 
@@ -810,7 +812,7 @@ mod tests {
 		let xt = TestXt::new(Call::Balances(BalancesCall::transfer(33, 0)), sign_extra(1, 0, 0));
 		let encoded = xt.encode();
 		let encoded_len = encoded.len() as Weight;
-		// Block execution weight + on_initialize weight
+		// on_initialize weight + block execution weight
 		let base_block_weight = 175 + <Runtime as frame_system::Trait>::BlockExecutionWeight::get();
 		let limit = AvailableBlockRatio::get() * MaximumBlockWeight::get() - base_block_weight;
 		let num_to_exhaust_block = limit / (encoded_len + 5);
@@ -1073,4 +1075,44 @@ mod tests {
 			assert_eq!(sp_io::storage::get(CUSTOM_ON_RUNTIME_KEY).unwrap(), true.encode());
 		});
 	}
+
+	#[test]
+	fn all_weights_are_recorded_correctly() {
+		new_test_ext(1).execute_with(|| {
+			// Make sure `on_runtime_upgrade` is called for maximum complexity
+			RUNTIME_VERSION.with(|v| *v.borrow_mut() = sp_version::RuntimeVersion {
+				spec_version: 1,
+				..Default::default()
+			});
+
+			let block_number = 1;
+
+			Executive::initialize_block(&Header::new(
+				block_number,
+				H256::default(),
+				H256::default(),
+				[69u8; 32].into(),
+				Digest::default(),
+			));
+
+			// All weights that show up in the `initialize_block_impl`
+			let frame_system_upgrade_weight = frame_system::Module::<Runtime>::on_runtime_upgrade();
+			let custom_runtime_upgrade_weight = CustomOnRuntimeUpgrade::on_runtime_upgrade();
+			let runtime_upgrade_weight = <AllModules as OnRuntimeUpgrade>::on_runtime_upgrade();
+			let frame_system_on_initialize_weight = frame_system::Module::<Runtime>::on_initialize(block_number);
+			let on_initialize_weight = <AllModules as OnInitialize<u64>>::on_initialize(block_number);
+			let base_block_weight = <Runtime as frame_system::Trait>::BlockExecutionWeight::get();
+
+			// Weights are recorded correctly
+			assert_eq!(
+				frame_system::Module::<Runtime>::block_weight().total(),
+				frame_system_upgrade_weight +
+				custom_runtime_upgrade_weight +
+				runtime_upgrade_weight +
+				frame_system_on_initialize_weight +
+				on_initialize_weight +
+				base_block_weight,
+			);
+		});
+	}
 }