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

substrate/Cargo.lock

@@ -4347,6 +4347,8 @@ dependencies = [
  "sr-arithmetic 2.0.0",
  "sr-io 2.0.0",
  "sr-std 2.0.0",
+ "srml-support 2.0.0",
+ "srml-system 2.0.0",
  "substrate-application-crypto 2.0.0",
  "substrate-inherents 2.0.0",
  "substrate-offchain 2.0.0",

substrate/core/sr-primitives/Cargo.toml

@@ -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 = []

substrate/core/sr-primitives/src/testing.rs

@@ -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)?;

substrate/core/sr-primitives/src/weights.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)
+	}
+}

substrate/node/runtime/src/lib.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);
+	}
 }

substrate/srml/example/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.

substrate/srml/executive/Cargo.toml

@@ -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]

substrate/srml/executive/src/lib.rs

@@ -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);
+		})
+	}
 }

substrate/srml/support/src/dispatch.rs

@@ -26,7 +26,7 @@ pub use srml_metadata::{
 pub use sr_primitives::{
 	weights::{
 		SimpleDispatchInfo, GetDispatchInfo, DispatchInfo, WeighData, ClassifyDispatch,
-		TransactionPriority
+		TransactionPriority, Weight, WeighBlock,
 	},
 	traits::{Dispatchable, DispatchResult, ModuleDispatchError},
 	DispatchError,
@@ -320,6 +320,7 @@ macro_rules! decl_module {
 			"`deposit_event` function is reserved and must follow the syntax: `$vis:vis fn deposit_event() = default;`"
 		);
 	};
+	// Add on_finalize, without a given weight.
 	(@normalize
 		$(#[$attr:meta])*
 		pub struct $mod_type:ident<$trait_instance:ident: $trait_name:ident$(<I>, I: $instantiable:path $(= $module_default_instance:path)?)?>
@@ -343,7 +344,10 @@ macro_rules! decl_module {
 			{ $( $other_where_bounds )* }
 			{ $( $deposit_event )* }
 			{ $( $on_initialize )* }
-			{ fn on_finalize( $( $param_name : $param ),* ) { $( $impl )* } }
+			{
+				#[weight = $crate::dispatch::SimpleDispatchInfo::zero()]
+				fn on_finalize( $( $param_name : $param ),* ) { $( $impl )* }
+			}
 			{ $( $offchain )* }
 			{ $( $constants )* }
 			{ $( $error_type )* }
@@ -351,12 +355,51 @@ macro_rules! decl_module {
 			$($rest)*
 		);
 	};
+	// Add on_finalize, given weight.
 	(@normalize
 		$(#[$attr:meta])*
 		pub struct $mod_type:ident<$trait_instance:ident: $trait_name:ident$(<I>, I: $instantiable:path $(= $module_default_instance:path)?)?>
 		for enum $call_type:ident where origin: $origin_type:ty, system = $system:ident
 		{ $( $other_where_bounds:tt )* }
 		{ $( $deposit_event:tt )* }
+		{ $( $on_initialize:tt )* }
+		{}
+		{ $( $offchain:tt )* }
+		{ $( $constants:tt )* }
+		{ $( $error_type:tt )* }
+		[ $( $dispatchables:tt )* ]
+		$(#[doc = $doc_attr:tt])*
+		#[weight = $weight:expr]
+		fn on_finalize( $( $param_name:ident : $param:ty ),* $(,)? ) { $( $impl:tt )* }
+		$($rest:tt)*
+	) => {
+		$crate::decl_module!(@normalize
+			$(#[$attr])*
+			pub struct $mod_type<$trait_instance: $trait_name$(<I>, I: $instantiable $(= $module_default_instance)?)?>
+			for enum $call_type where origin: $origin_type, system = $system
+			{ $( $other_where_bounds )* }
+			{ $( $deposit_event )* }
+			{ $( $on_initialize )* }
+			{
+				#[weight = $weight]
+				fn on_finalize( $( $param_name : $param ),* ) { $( $impl )* }
+			}
+			{ $( $offchain )* }
+			{ $( $constants )* }
+			{ $( $error_type )* }
+			[ $( $dispatchables )* ]
+			$($rest)*
+		);
+	};
+	// Add on_initialize, without a given weight.
+	(@normalize
+		$(#[$attr:meta])*
+		pub struct $mod_type:ident<
+			$trait_instance:ident: $trait_name:ident$(<I>, I: $instantiable:path $(= $module_default_instance:path)?)?
+		>
+		for enum $call_type:ident where origin: $origin_type:ty, system = $system:ident
+		{ $( $other_where_bounds:tt )* }
+		{ $( $deposit_event:tt )* }
 		{}
 		{ $( $on_finalize:tt )* }
 		{ $( $offchain:tt )* }
@@ -373,7 +416,48 @@ macro_rules! decl_module {
 			for enum $call_type where origin: $origin_type, system = $system
 			{ $( $other_where_bounds )* }
 			{ $( $deposit_event )* }
-			{ fn on_initialize( $( $param_name : $param ),* ) { $( $impl )* } }
+			{
+				#[weight = $crate::dispatch::SimpleDispatchInfo::zero()]
+				fn on_initialize( $( $param_name : $param ),* ) { $( $impl )* }
+			}
+			{ $( $on_finalize )* }
+			{ $( $offchain )* }
+			{ $( $constants )* }
+			{ $( $error_type )* }
+			[ $( $dispatchables )* ]
+			$($rest)*
+		);
+	};
+	// Add on_initialize, given weight.
+	(@normalize
+		$(#[$attr:meta])*
+		pub struct $mod_type:ident<
+			$trait_instance:ident: $trait_name:ident$(<I>, I: $instantiable:path $(= $module_default_instance:path)?)?
+		>
+		for enum $call_type:ident where origin: $origin_type:ty, system = $system:ident
+		{ $( $other_where_bounds:tt )* }
+		{ $( $deposit_event:tt )* }
+		{}
+		{ $( $on_finalize:tt )* }
+		{ $( $offchain:tt )* }
+		{ $( $constants:tt )* }
+		{ $( $error_type:tt )* }
+		[ $( $dispatchables:tt )* ]
+		$(#[doc = $doc_attr:tt])*
+		#[weight = $weight:expr]
+		fn on_initialize( $( $param_name:ident : $param:ty ),* $(,)? ) { $( $impl:tt )* }
+		$($rest:tt)*
+	) => {
+		$crate::decl_module!(@normalize
+			$(#[$attr])*
+			pub struct $mod_type<$trait_instance: $trait_name$(<I>, I: $instantiable $(= $module_default_instance)?)?>
+			for enum $call_type where origin: $origin_type, system = $system
+			{ $( $other_where_bounds )* }
+			{ $( $deposit_event )* }
+			{
+				#[weight = $weight]
+				fn on_initialize( $( $param_name : $param ),* ) { $( $impl )* }
+			}
 			{ $( $on_finalize )* }
 			{ $( $offchain )* }
 			{ $( $constants )* }
@@ -770,6 +854,7 @@ macro_rules! decl_module {
 	(@impl_on_initialize
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
 		{ $( $other_where_bounds:tt )* }
+		#[weight = $weight:expr]
 		fn on_initialize() { $( $impl:tt )* }
 	) => {
 		impl<$trait_instance: $trait_name$(<I>, $instance: $instantiable)?>
@@ -783,6 +868,7 @@ macro_rules! decl_module {
 	(@impl_on_initialize
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
 		{ $( $other_where_bounds:tt )* }
+		#[weight = $weight:expr]
 		fn on_initialize($param:ident : $param_ty:ty) { $( $impl:tt )* }
 	) => {
 		impl<$trait_instance: $trait_name$(<I>, $instance: $instantiable)?>
@@ -806,6 +892,7 @@ macro_rules! decl_module {
 	(@impl_on_finalize
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
 		{ $( $other_where_bounds:tt )* }
+		#[weight = $weight:expr]
 		fn on_finalize() { $( $impl:tt )* }
 	) => {
 		impl<$trait_instance: $trait_name$(<I>, $instance: $instantiable)?>
@@ -819,6 +906,7 @@ macro_rules! decl_module {
 	(@impl_on_finalize
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
 		{ $( $other_where_bounds:tt )* }
+		#[weight = $weight:expr]
 		fn on_finalize($param:ident : $param_ty:ty) { $( $impl:tt )* }
 	) => {
 		impl<$trait_instance: $trait_name$(<I>, $instance: $instantiable)?>
@@ -840,6 +928,35 @@ macro_rules! decl_module {
 		}
 	};
 
+	(@impl_block_hooks_weight
+		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
+		{ $( $other_where_bounds:tt )* }
+		@init $(
+			#[weight = $weight_initialize:expr]
+			fn on_initialize($( $param_initialize:ident : $param_ty_initialize:ty )*) { $( $impl_initialize:tt )* }
+		)?
+		@fin $(
+			#[weight = $weight_finalize:expr]
+			fn on_finalize($( $param_finalize:ident : $param_ty_finalize:ty )*) { $( $impl_finalize:tt )* }
+		)?
+	) => {
+		impl<$trait_instance: $trait_name$(<I>, $instance: $instantiable)?>
+		$crate::dispatch::WeighBlock<$trait_instance::BlockNumber> for $module<$trait_instance$(, $instance)?> where
+			$( $other_where_bounds )*
+		{
+			$(
+				fn on_initialize(n: $trait_instance::BlockNumber) -> $crate::dispatch::Weight {
+					<dyn $crate::dispatch::WeighData<$trait_instance::BlockNumber>>::weigh_data(&$weight_initialize, n)
+				}
+			)?
+			$(
+				fn on_finalize(n: $trait_instance::BlockNumber) -> $crate::dispatch::Weight {
+					<dyn $crate::dispatch::WeighData<$trait_instance::BlockNumber>>::weigh_data(&$weight_finalize, n)
+				}
+			)?
+		}
+	};
+
 	(@impl_offchain
 		$module:ident<$trait_instance:ident: $trait_name:ident$(<I>, $instance:ident: $instantiable:path)?>;
 		{ $( $other_where_bounds:tt )* }
@@ -1083,6 +1200,14 @@ macro_rules! decl_module {
 			$( $on_finalize )*
 		}
 
+		$crate::decl_module! {
+			@impl_block_hooks_weight
+			$mod_type<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?>;
+			{ $( $other_where_bounds )* }
+			@init $( $on_initialize )*
+			@fin $( $on_finalize )*
+		}
+
 		$crate::decl_module! {
 			@impl_offchain
 			$mod_type<$trait_instance: $trait_name $(<I>, $instance: $instantiable)?>;
@@ -1727,6 +1852,14 @@ mod tests {
 		}
 	}
 
+	struct BLockWeight;
+	impl<BlockNumber: Into<u32>> WeighData<BlockNumber> for BLockWeight {
+		fn weigh_data(&self, target: BlockNumber) -> Weight {
+			let target: u32 = target.into();
+			if target % 2 == 0 { 10 } else { 0 }
+		}
+	}
+
 	decl_module! {
 		pub struct Module<T: Trait> for enum Call where origin: T::Origin, T::AccountId: From<u32> {
 			/// Hi, this is a comment.
@@ -1738,7 +1871,9 @@ mod tests {
 			fn aux_4(_origin, _data: i32) -> Result { unreachable!() }
 			fn aux_5(_origin, _data: i32, #[compact] _data2: u32,) -> Result { unreachable!() }
 
+			#[weight = SimpleDispatchInfo::FixedNormal(7)]
 			fn on_initialize(n: T::BlockNumber,) { if n.into() == 42 { panic!("on_initialize") } }
+			#[weight = BLockWeight]
 			fn on_finalize(n: T::BlockNumber) { if n.into() == 42 { panic!("on_finalize") } }
 			fn offchain_worker() {}
 
@@ -1899,4 +2034,16 @@ mod tests {
 			DispatchInfo { weight: 3, class: DispatchClass::Normal },
 		);
 	}
+
+	#[test]
+	fn weight_for_block_hooks() {
+		// independent of block number
+		assert_eq!(<Test as WeighBlock<u32>>::on_initialize(0), 7);
+		assert_eq!(<Test as WeighBlock<u32>>::on_initialize(10), 7);
+		assert_eq!(<Test as WeighBlock<u32>>::on_initialize(100), 7);
+
+		// dependent
+		assert_eq!(<Test as WeighBlock<u32>>::on_finalize(2), 10);
+		assert_eq!(<Test as WeighBlock<u32>>::on_finalize(3), 0);
+	}
 }

substrate/srml/system/src/lib.rs

@@ -606,6 +606,29 @@ impl<T: Trait> Module<T> {
 		AllExtrinsicsLen::get().unwrap_or_default()
 	}
 
+	/// Inform the system module of some additional weight that should be accounted for, in the
+	/// current block.
+	///
+	/// NOTE: use with extra care; this function is made public only be used for certain modules
+	/// that need it. A runtime that does not have dynamic calls should never need this and should
+	/// stick to static weights. A typical use case for this is inner calls or smart contract calls.
+	/// Furthermore, it only makes sense to use this when it is presumably  _cheap_ to provide the
+	/// argument `weight`; In other words, if this function is to be used to account for some
+	/// unknown, user provided call's weight, it would only make sense to use it if you are sure you
+	/// can rapidly compute the weight of the inner call.
+	///
+	/// Even more dangerous is to note that this function does NOT take any action, if the new sum
+	/// of block weight is more than the block weight limit. This is what the _unchecked_.
+	///
+	/// Another potential use-case could be for the `on_initialise` and `on_finalize` hooks.
+	///
+	/// If no previous weight exists, the function initializes the weight to zero.
+	pub fn register_extra_weight_unchecked(weight: Weight) {
+		let current_weight = AllExtrinsicsWeight::get().unwrap_or_default();
+		let next_weight = current_weight.saturating_add(weight).min(T::MaximumBlockWeight::get());
+		AllExtrinsicsWeight::put(next_weight);
+	}
+
 	/// Start the execution of a particular block.
 	pub fn initialize(
 		number: &T::BlockNumber,