Weight annotation for block hooks. (#4058)

* Initial version that works with proper tests.

* get rid of todos and grumbles and such.

* Cleanup and fix line-width

* fix test runtime test

Author: kianenigma

Cargo.lock

@@ -22,6 +22,8 @@ inherents = { package = "substrate-inherents", path = "../inherents", default-fe
 serde_json = "1.0.41"
 rand = "0.7.2"
 substrate-offchain = { path = "../offchain" }
+support = { package = "srml-support", path = "../../srml/support" }
+system = { package = "srml-system", path = "../../srml/system" }
 
 [features]
 bench = []

core/sr-primitives/Cargo.toml

@@ -145,7 +145,7 @@ pub type DigestItem = generic::DigestItem<H256>;
 pub type Digest = generic::Digest<H256>;
 
 /// Block Header
-#[derive(PartialEq, Eq, Clone, Serialize, Debug, Encode, Decode)]
+#[derive(PartialEq, Eq, Clone, Serialize, Debug, Encode, Decode, Default)]
 #[serde(rename_all = "camelCase")]
 #[serde(deny_unknown_fields)]
 pub struct Header {
@@ -198,6 +198,16 @@ impl traits::Header for Header {
 	}
 }
 
+impl Header {
+	/// A new header with the given number and default hash for all other fields.
+	pub fn new_from_number(number: <Self as traits::Header>::Number) -> Self {
+		Self {
+			number,
+			..Default::default()
+		}
+	}
+}
+
 impl<'a> Deserialize<'a> for Header {
 	fn deserialize<D: Deserializer<'a>>(de: D) -> Result<Self, D::Error> {
 		let r = <Vec<u8>>::deserialize(de)?;

core/sr-primitives/src/testing.rs

@@ -37,8 +37,9 @@
 
 #[cfg(feature = "std")]
 use serde::{Serialize, Deserialize};
+use impl_trait_for_tuples::impl_for_tuples;
 use codec::{Encode, Decode};
-use arithmetic::traits::Bounded;
+use arithmetic::traits::{Bounded, Zero};
 use crate::RuntimeDebug;
 
 /// Re-export priority as type
@@ -62,6 +63,35 @@ pub trait ClassifyDispatch<T> {
 	fn classify_dispatch(&self, target: T) -> DispatchClass;
 }
 
+/// Means of determining the weight of a block's lifecycle hooks: on_initialize, on_finalize and
+/// such.
+pub trait WeighBlock<BlockNumber> {
+	/// Return the weight of the block's on_initialize hook.
+	fn on_initialize(_: BlockNumber) -> Weight { Zero::zero() }
+	/// Return the weight of the block's on_finalize hook.
+	fn on_finalize(_: BlockNumber) -> Weight { Zero::zero() }
+}
+
+/// Maybe I can do something to remove the duplicate code here.
+#[impl_for_tuples(30)]
+impl<BlockNumber: Copy> WeighBlock<BlockNumber> for SingleModule {
+	fn on_initialize(n: BlockNumber) -> Weight {
+		let mut accumulated_weight: Weight = Zero::zero();
+		for_tuples!(
+			#( accumulated_weight = accumulated_weight.saturating_add(SingleModule::on_initialize(n)); )*
+		);
+		accumulated_weight
+	}
+
+	fn on_finalize(n: BlockNumber) -> Weight {
+		let mut accumulated_weight: Weight = Zero::zero();
+		for_tuples!(
+			#( accumulated_weight = accumulated_weight.saturating_add(SingleModule::on_finalize(n)); )*
+		);
+		accumulated_weight
+	}
+}
+
 /// A generalized group of dispatch types. This is only distinguishing normal, user-triggered transactions
 /// (`Normal`) and anything beyond which serves a higher purpose to the system (`Operational`).
 #[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
@@ -181,3 +211,10 @@ impl Default for SimpleDispatchInfo {
 		SimpleDispatchInfo::FixedNormal(10_000)
 	}
 }
+
+impl SimpleDispatchInfo {
+	/// An _additive zero_ variant of SimpleDispatchInfo.
+	pub fn zero() -> Self {
+		Self::FixedNormal(0)
+	}
+}

core/sr-primitives/src/weights.rs

@@ -697,6 +697,7 @@ impl_runtime_apis! {
 		}
 	}
 }
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -717,4 +718,36 @@ mod tests {
 		let x = SubmitTransaction::default();
 		is_submit_signed_transaction(x);
 	}
+
+	#[test]
+	fn block_hooks_weight_should_not_exceed_limits() {
+		use sr_primitives::weights::WeighBlock;
+		let check_for_block = |b| {
+			let block_hooks_weight =
+				<AllModules as WeighBlock<BlockNumber>>::on_initialize(b) +
+				<AllModules as WeighBlock<BlockNumber>>::on_finalize(b);
+
+			assert_eq!(
+				block_hooks_weight,
+				0,
+				"This test might fail simply because the value being compared to has increased to a \
+				module declaring a new weight for a hook or call. In this case update the test and \
+				happily move on.",
+			);
+
+			// Invariant. Always must be like this to have a sane chain.
+			assert!(block_hooks_weight < MaximumBlockWeight::get());
+
+			// Warning.
+			if block_hooks_weight > MaximumBlockWeight::get() / 2 {
+				println!(
+					"block hooks weight is consuming more than a block's capacity. You probably want \
+					to re-think this. This test will fail now."
+				);
+				assert!(false);
+			}
+		};
+
+		let _ = (0..100_000).for_each(check_for_block);
+	}
 }

node/runtime/src/lib.rs

@@ -498,13 +498,18 @@ decl_module! {
 			<Dummy<T>>::put(new_value);
 		}
 
-		// The signature could also look like: `fn on_initialize()`
+		// The signature could also look like: `fn on_initialize()`.
+		// This function could also very well have a weight annotation, similar to any other. The
+		// only difference being that if it is not annotated, the default is
+		// `SimpleDispatchInfo::zero()`, which resolves into no weight.
+		#[weight = SimpleDispatchInfo::FixedNormal(1000)]
 		fn on_initialize(_n: T::BlockNumber) {
 			// Anything that needs to be done at the start of the block.
 			// We don't do anything here.
 		}
 
 		// The signature could also look like: `fn on_finalize()`
+		#[weight = SimpleDispatchInfo::FixedNormal(2000)]
 		fn on_finalize(_n: T::BlockNumber) {
 			// Anything that needs to be done at the end of the block.
 			// We just kill our dummy storage item.

srml/example/src/lib.rs

@@ -16,8 +16,8 @@ system = { package = "srml-system", path = "../system", default-features = false
 [dev-dependencies]
 hex-literal = "0.2.1"
 primitives = { package = "substrate-primitives",  path = "../../core/primitives" }
-srml-indices = { path = "../indices" }
 balances = { package = "srml-balances", path = "../balances" }
+indices = { package = "srml-indices", path = "../indices" }
 transaction-payment = { package = "srml-transaction-payment", path = "../transaction-payment" }
 
 [features]

srml/executive/Cargo.toml

@@ -78,7 +78,8 @@
 
 use rstd::{prelude::*, marker::PhantomData};
 use sr_primitives::{
-	generic::Digest, ApplyResult, weights::GetDispatchInfo,
+	generic::Digest, ApplyResult,
+	weights::{GetDispatchInfo, WeighBlock},
 	traits::{
 		self, Header, Zero, One, Checkable, Applyable, CheckEqual, OnFinalize, OnInitialize,
 		NumberFor, Block as BlockT, OffchainWorker, Dispatchable,
@@ -110,7 +111,11 @@ impl<
 	Block: traits::Block<Header=System::Header, Hash=System::Hash>,
 	Context: Default,
 	UnsignedValidator,
-	AllModules: OnInitialize<System::BlockNumber> + OnFinalize<System::BlockNumber> + OffchainWorker<System::BlockNumber>,
+	AllModules:
+		OnInitialize<System::BlockNumber> +
+		OnFinalize<System::BlockNumber> +
+		OffchainWorker<System::BlockNumber> +
+		WeighBlock<System::BlockNumber>,
 > ExecuteBlock<Block> for Executive<System, Block, Context, UnsignedValidator, AllModules>
 where
 	Block::Extrinsic: Checkable<Context> + Codec,
@@ -130,7 +135,11 @@ impl<
 	Block: traits::Block<Header=System::Header, Hash=System::Hash>,
 	Context: Default,
 	UnsignedValidator,
-	AllModules: OnInitialize<System::BlockNumber> + OnFinalize<System::BlockNumber> + OffchainWorker<System::BlockNumber>,
+	AllModules:
+		OnInitialize<System::BlockNumber> +
+		OnFinalize<System::BlockNumber> +
+		OffchainWorker<System::BlockNumber> +
+		WeighBlock<System::BlockNumber>,
 > Executive<System, Block, Context, UnsignedValidator, AllModules>
 where
 	Block::Extrinsic: Checkable<Context> + Codec,
@@ -154,6 +163,12 @@ where
 	) {
 		<system::Module<System>>::initialize(block_number, parent_hash, extrinsics_root, digest);
 		<AllModules as OnInitialize<System::BlockNumber>>::on_initialize(*block_number);
+		<system::Module<System>>::register_extra_weight_unchecked(
+			<AllModules as WeighBlock<System::BlockNumber>>::on_initialize(*block_number)
+		);
+		<system::Module<System>>::register_extra_weight_unchecked(
+			<AllModules as WeighBlock<System::BlockNumber>>::on_finalize(*block_number)
+		);
 	}
 
 	fn initial_checks(block: &Block) {
@@ -309,12 +324,48 @@ mod tests {
 		impl_outer_event, impl_outer_origin, parameter_types, impl_outer_dispatch,
 		traits::{Currency, LockIdentifier, LockableCurrency, WithdrawReasons, WithdrawReason},
 	};
-	use system::Call as SystemCall;
+	use system::{Call as SystemCall, ChainContext};
 	use balances::Call as BalancesCall;
 	use hex_literal::hex;
 
+	mod custom {
+		use sr_primitives::weights::SimpleDispatchInfo;
+
+		pub trait Trait: system::Trait {}
+
+		support::decl_module! {
+			pub struct Module<T: Trait> for enum Call where origin: T::Origin {
+				#[weight = SimpleDispatchInfo::FixedNormal(100)]
+				fn some_function(origin) {
+					// NOTE: does not make any different.
+					let _ = system::ensure_signed(origin);
+				}
+				#[weight = SimpleDispatchInfo::FixedOperational(200)]
+				fn some_root_operation(origin) {
+					let _ = system::ensure_root(origin);
+				}
+				#[weight = SimpleDispatchInfo::FreeNormal]
+				fn some_unsigned_message(origin) {
+					let _ = system::ensure_none(origin);
+				}
+
+				// module hooks.
+				// one with block number arg and one without
+				#[weight = SimpleDispatchInfo::FixedNormal(25)]
+				fn on_initialize(n: T::BlockNumber) {
+					println!("on_initialize({})", n);
+				}
+				#[weight = SimpleDispatchInfo::FixedNormal(150)]
+				fn on_finalize() {
+					println!("on_finalize(?)");
+				}
+			}
+		}
+	}
+
 	type System = system::Module<Runtime>;
 	type Balances = balances::Module<Runtime>;
+	type Custom = custom::Module<Runtime>;
 
 	impl_outer_origin! {
 		pub enum Origin for Runtime { }
@@ -386,6 +437,7 @@ mod tests {
 		type WeightToFee = ConvertInto;
 		type FeeMultiplierUpdate = ();
 	}
+	impl custom::Trait for Runtime {}
 
 	#[allow(deprecated)]
 	impl ValidateUnsigned for Runtime {
@@ -409,8 +461,9 @@ mod tests {
 		system::CheckWeight<Runtime>,
 		transaction_payment::ChargeTransactionPayment<Runtime>
 	);
+	type AllModules = (System, Balances, Custom);
 	type TestXt = sr_primitives::testing::TestXt<Call, SignedExtra>;
-	type Executive = super::Executive<Runtime, Block<TestXt>, system::ChainContext<Runtime>, Runtime, ()>;
+	type Executive = super::Executive<Runtime, Block<TestXt>, ChainContext<Runtime>, Runtime, AllModules>;
 
 	fn extra(nonce: u64, fee: u64) -> SignedExtra {
 		(
@@ -534,7 +587,7 @@ mod tests {
 		let xt = sr_primitives::testing::TestXt(sign_extra(1, 0, 0), Call::Balances(BalancesCall::transfer(33, 0)));
 		let encoded = xt.encode();
 		let encoded_len = encoded.len() as Weight;
-		let limit = AvailableBlockRatio::get() * MaximumBlockWeight::get();
+		let limit = AvailableBlockRatio::get() * MaximumBlockWeight::get() - 175;
 		let num_to_exhaust_block = limit / encoded_len;
 		t.execute_with(|| {
 			Executive::initialize_block(&Header::new(
@@ -544,7 +597,8 @@ mod tests {
 				[69u8; 32].into(),
 				Digest::default(),
 			));
-			assert_eq!(<system::Module<Runtime>>::all_extrinsics_weight(), 0);
+			// Initial block weight form the custom module.
+			assert_eq!(<system::Module<Runtime>>::all_extrinsics_weight(), 175);
 
 			for nonce in 0..=num_to_exhaust_block {
 				let xt = sr_primitives::testing::TestXt(
@@ -555,7 +609,7 @@ mod tests {
 					assert!(res.is_ok());
 					assert_eq!(
 						<system::Module<Runtime>>::all_extrinsics_weight(),
-						encoded_len * (nonce + 1),
+						encoded_len * (nonce + 1) + 175,
 					);
 					assert_eq!(<system::Module<Runtime>>::extrinsic_index(), Some(nonce as u32 + 1));
 				} else {
@@ -652,4 +706,15 @@ mod tests {
 		execute_with_lock(WithdrawReasons::all());
 		execute_with_lock(WithdrawReasons::except(WithdrawReason::TransactionPayment));
 	}
+
+	#[test]
+	fn block_hooks_weight_is_stored() {
+		new_test_ext(0).execute_with(|| {
+
+			Executive::initialize_block(&Header::new_from_number(1));
+			// NOTE: might need updates over time if system and balance introduce new weights. For
+			// now only accounts for the custom module.
+			assert_eq!(<system::Module<Runtime>>::all_extrinsics_weight(), 150 + 25);
+		})
+	}
 }