[unstable2507] Backport #9756

prdoc/pr_9756.prdoc

@@ -0,0 +1,19 @@
+title: 'FRAME: Register `on_initialize` after each pallet'
+doc:
+- audience: Runtime Dev
+  description: |-
+    Before this pull request, FRAME was executing all pallets `on_initialize` and then register the weight, including the weight of `on_runtime_upgrade`. Thus, other pallets were not aware on how much weight was already used when they were executing their `on_initialize` code. As some pallets are doing some work in `on_initialize`, they need to be aware of how much weight is still left.
+    To register the weight after each `on_initialize` call, a new trait is added. This new trait is implemented for tuples of types that implement `OnInitialize` and then it registers the weight after each call to `on_initialize`.
+
+    `pallet-scheduler` is changed to take the remaining weight into account and to not just assume that its configured weight is always available.
+crates:
+- name: frame-support
+  bump: patch
+- name: frame-executive
+  bump: patch
+- name: pallet-scheduler
+  bump: patch
+- name: frame-system
+  bump: patch
+- name: sp-weights
+  bump: patch

substrate/frame/executive/src/lib.rs

@@ -161,8 +161,9 @@ use frame_support::{
 	migrations::MultiStepMigrator,
 	pallet_prelude::InvalidTransaction,
 	traits::{
-		BeforeAllRuntimeMigrations, ExecuteBlock, IsInherent, OffchainWorker, OnFinalize, OnIdle,
-		OnInitialize, OnPoll, OnRuntimeUpgrade, PostInherents, PostTransactions, PreInherents,
+		BeforeAllRuntimeMigrations, ExecuteBlock, Get, IsInherent, OffchainWorker, OnFinalize,
+		OnIdle, OnInitialize, OnPoll, OnRuntimeUpgrade, PostInherents, PostTransactions,
+		PreInherents,
 	},
 	weights::{Weight, WeightMeter},
 	MAX_EXTRINSIC_DEPTH,
@@ -260,7 +261,7 @@ impl<
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
 			+ BeforeAllRuntimeMigrations
-			+ OnInitialize<BlockNumberFor<System>>
+			+ OnInitializeWithWeightRegistration<System>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
 			+ OffchainWorker<BlockNumberFor<System>>
@@ -299,7 +300,7 @@ impl<
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
 			+ BeforeAllRuntimeMigrations
-			+ OnInitialize<BlockNumberFor<System>>
+			+ OnInitializeWithWeightRegistration<System>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
 			+ OffchainWorker<BlockNumberFor<System>>
@@ -493,6 +494,38 @@ where
 	}
 }
 
+/// Extension trait for [`OnInitialize`].
+///
+/// It takes care to register the weight of each pallet directly after executing its
+/// `on_initialize`.
+///
+/// The trait is sealed.
+pub trait OnInitializeWithWeightRegistration<T: frame_system::Config> {
+	/// The actual logic that calls `on_initialize` and registers the weight.
+	fn on_initialize_with_weight_registration(_n: BlockNumberFor<T>) -> Weight;
+}
+
+frame_support::impl_for_tuples_attr! {
+	#[tuple_types_custom_trait_bound(OnInitialize<frame_system::pallet_prelude::BlockNumberFor<T>>)]
+	impl<T: frame_system::Config> OnInitializeWithWeightRegistration<T> for Tuple {
+		fn on_initialize_with_weight_registration(n: BlockNumberFor<T>) -> Weight {
+			let mut weight = Weight::zero();
+			for_tuples!( #(
+				let individual_weight = Tuple::on_initialize(n);
+
+				<frame_system::Pallet<T>>::register_extra_weight_unchecked(
+					individual_weight,
+					DispatchClass::Mandatory,
+				);
+
+				weight = weight.saturating_add(individual_weight);
+			)* );
+
+			weight
+		}
+	}
+}
+
 impl<
 		System: frame_system::Config + IsInherent<Block::Extrinsic>,
 		Block: traits::Block<
@@ -503,7 +536,7 @@ impl<
 		UnsignedValidator,
 		AllPalletsWithSystem: OnRuntimeUpgrade
 			+ BeforeAllRuntimeMigrations
-			+ OnInitialize<BlockNumberFor<System>>
+			+ OnInitializeWithWeightRegistration<System>
 			+ OnIdle<BlockNumberFor<System>>
 			+ OnFinalize<BlockNumberFor<System>>
 			+ OffchainWorker<BlockNumberFor<System>>
@@ -571,7 +604,7 @@ where
 	) {
 		// Reset events before apply runtime upgrade hook.
 		// This is required to preserve events from runtime upgrade hook.
-		// This means the format of all the event related storages must always be compatible.
+		// This means the format of all the event related storage must always be compatible.
 		<frame_system::Pallet<System>>::reset_events();
 
 		let mut weight = Weight::zero();
@@ -585,17 +618,24 @@ where
 			);
 		}
 		<frame_system::Pallet<System>>::initialize(block_number, parent_hash, digest);
-		weight = weight.saturating_add(<AllPalletsWithSystem as OnInitialize<
-			BlockNumberFor<System>,
-		>>::on_initialize(*block_number));
-		weight = weight.saturating_add(
-			<System::BlockWeights as frame_support::traits::Get<_>>::get().base_block,
-		);
+
+		weight = System::BlockWeights::get().base_block.saturating_add(weight);
+		// Register the base block weight and optional `on_runtime_upgrade` weight.
 		<frame_system::Pallet<System>>::register_extra_weight_unchecked(
 			weight,
 			DispatchClass::Mandatory,
 		);
 
+		weight = weight
+			.saturating_add(<AllPalletsWithSystem as OnInitializeWithWeightRegistration<
+			System,
+		>>::on_initialize_with_weight_registration(*block_number));
+
+		log::debug!(
+			target: LOG_TARGET,
+			"[{block_number:?}]: Block initialization weight consumption: {weight:?}",
+		);
+
 		frame_system::Pallet::<System>::note_finished_initialize();
 		<System as frame_system::Config>::PreInherents::pre_inherents();
 	}

substrate/frame/scheduler/src/lib.rs

@@ -432,7 +432,8 @@ pub mod pallet {
 		/// Execute the scheduled calls
 		fn on_initialize(_now: SystemBlockNumberFor<T>) -> Weight {
 			let now = T::BlockNumberProvider::current_block_number();
-			let mut weight_counter = WeightMeter::with_limit(T::MaximumWeight::get());
+			let mut weight_counter = frame_system::Pallet::<T>::remaining_block_weight()
+				.limit_to(T::MaximumWeight::get());
 			Self::service_agendas(&mut weight_counter, now, u32::MAX);
 			weight_counter.consumed()
 		}

substrate/frame/scheduler/src/mock.rs

@@ -26,6 +26,7 @@ use frame_support::{
 };
 use frame_system::{EnsureRoot, EnsureSignedBy};
 use sp_runtime::{BuildStorage, Perbill};
+use sp_weights::constants::WEIGHT_REF_TIME_PER_SECOND;
 
 // Logger module to track execution.
 #[frame_support::pallet]
@@ -134,14 +135,15 @@ impl Contains<RuntimeCall> for BaseFilter {
 parameter_types! {
 	pub BlockWeights: frame_system::limits::BlockWeights =
 		frame_system::limits::BlockWeights::simple_max(
-			Weight::from_parts(2_000_000_000_000, u64::MAX),
+			Weight::from_parts(WEIGHT_REF_TIME_PER_SECOND * 2, u64::MAX),
 		);
 }
 
 #[derive_impl(frame_system::config_preludes::TestDefaultConfig)]
 impl system::Config for Test {
 	type BaseCallFilter = BaseFilter;
 	type Block = Block;
+	type BlockWeights = BlockWeights;
 }
 impl logger::Config for Test {
 	type RuntimeEvent = RuntimeEvent;

substrate/frame/scheduler/src/tests.rs

@@ -1639,7 +1639,7 @@ fn on_initialize_weight_is_correct() {
 		let now = 1;
 		<Test as Config>::BlockNumberProvider::set_block_number(now);
 		assert_eq!(
-			Scheduler::on_initialize(42), // BN unused
+			Scheduler::on_initialize(42), // block number unused
 			TestWeightInfo::service_agendas_base() +
 				TestWeightInfo::service_agenda_base(1) +
 				<TestWeightInfo as MarginalWeightInfo>::service_task(None, true, true) +
@@ -1653,7 +1653,7 @@ fn on_initialize_weight_is_correct() {
 		let now = 2;
 		<Test as Config>::BlockNumberProvider::set_block_number(now);
 		assert_eq!(
-			Scheduler::on_initialize(123), // BN unused
+			Scheduler::on_initialize(123), // block number unused
 			TestWeightInfo::service_agendas_base() +
 				TestWeightInfo::service_agenda_base(2) +
 				<TestWeightInfo as MarginalWeightInfo>::service_task(None, false, true) +
@@ -1670,7 +1670,7 @@ fn on_initialize_weight_is_correct() {
 		let now = 3;
 		<Test as Config>::BlockNumberProvider::set_block_number(now);
 		assert_eq!(
-			Scheduler::on_initialize(555), // BN unused
+			Scheduler::on_initialize(555), // block number unused
 			TestWeightInfo::service_agendas_base() +
 				TestWeightInfo::service_agenda_base(1) +
 				<TestWeightInfo as MarginalWeightInfo>::service_task(None, true, false) +
@@ -1686,12 +1686,20 @@ fn on_initialize_weight_is_correct() {
 		// Will contain none
 		let now = 4;
 		<Test as Config>::BlockNumberProvider::set_block_number(now);
-		let actual_weight = Scheduler::on_initialize(444); // BN unused
+		let actual_weight = Scheduler::on_initialize(444); // block number unused
 		assert_eq!(
 			actual_weight,
 			TestWeightInfo::service_agendas_base() + TestWeightInfo::service_agenda_base(0)
 		);
 		assert_eq!(IncompleteSince::<Test>::get(), Some(now + 1));
+
+		frame_system::Pallet::<Test>::register_extra_weight_unchecked(
+			BlockWeights::get().max_block,
+			frame_support::dispatch::DispatchClass::Mandatory,
+		);
+
+		let actual_weight = Scheduler::on_initialize(444); // block number unused
+		assert!(actual_weight.is_zero());
 	});
 }
 

substrate/frame/support/src/lib.rs

@@ -54,6 +54,7 @@ pub mod __private {
 	};
 	pub use codec;
 	pub use frame_metadata as metadata;
+	pub use impl_trait_for_tuples;
 	pub use log;
 	pub use paste;
 	pub use scale_info;
@@ -878,6 +879,108 @@ macro_rules! hypothetically_ok {
 	};
 }
 
+/// Puts the [`impl_for_tuples`](impl_trait_for_tuples::impl_for_tuples) attribute above the given
+/// code.
+///
+/// The main purpose of this macro is to handle the `tuples-*` feature which informs the attribute
+/// about the maximum size of the tuple to generate. Besides that, there is no difference to use the
+/// attribute directly.
+///
+/// # Example
+///
+/// ```rust
+/// trait ILoveTuples {
+///     fn really_hard();
+/// }
+///
+/// frame_support::impl_for_tuples_attr! {
+///     impl ILoveTuples for Tuple {
+///         fn really_hard() {
+///             for_tuples! { #(
+///                 // Print it for each tuple
+///                 println!("I LOVE TUPLES");
+///             )* }
+///         }
+///     }
+/// }
+/// ```
+#[cfg(all(not(feature = "tuples-96"), not(feature = "tuples-128")))]
+#[macro_export]
+macro_rules! impl_for_tuples_attr {
+	( $( $input:tt )* ) => {
+		#[$crate::__private::impl_trait_for_tuples::impl_for_tuples(64)]
+		$( $input )*
+	}
+}
+
+/// Puts the [`impl_for_tuples`](impl_trait_for_tuples::impl_for_tuples) attribute above the given
+/// code.
+///
+/// The main purpose of this macro is to handle the `tuples-*` feature which informs the attribute
+/// about the maximum size of the tuple to generate. Besides that, there is no difference to use the
+/// attribute directly.
+///
+/// # Example
+///
+/// ```rust
+/// trait ILoveTuples {
+///     fn really_hard();
+/// }
+///
+/// frame_support::impl_for_tuples_attr! {
+///     impl ILoveTuples for Tuple {
+///         fn really_hard() {
+///             for_tuples! { #(
+///                 // Print it for each tuple
+///                 println!("I LOVE TUPLES");
+///             )* }
+///         }
+///     }
+/// }
+/// ```
+#[cfg(all(feature = "tuples-96", not(feature = "tuples-128")))]
+#[macro_export]
+macro_rules! impl_for_tuples_attr {
+	( $( $input:tt )* ) => {
+		#[$crate::__private::impl_trait_for_tuples::impl_for_tuples(96)]
+		$( $input )*
+	}
+}
+
+/// Puts the [`impl_for_tuples`](impl_trait_for_tuples::impl_for_tuples) attribute above the given
+/// code.
+///
+/// The main purpose of this macro is to handle the `tuples-*` feature which informs the attribute
+/// about the maximum size of the tuple to generate. Besides that, there is no difference to use the
+/// attribute directly.
+///
+/// # Example
+///
+/// ```rust
+/// trait ILoveTuples {
+///     fn really_hard();
+/// }
+///
+/// frame_support::impl_for_tuples_attr! {
+///     impl ILoveTuples for Tuple {
+///         fn really_hard() {
+///             for_tuples! { #(
+///                 // Print it for each tuple
+///                 println!("I LOVE TUPLES");
+///             )* }
+///         }
+///     }
+/// }
+/// ```
+#[cfg(feature = "tuples-128")]
+#[macro_export]
+macro_rules! impl_for_tuples_attr {
+	( $( $input:tt )* ) => {
+		#[$crate::__private::impl_trait_for_tuples::impl_for_tuples(128)]
+		$( $input )*
+	}
+}
+
 #[doc(hidden)]
 pub use serde::{Deserialize, Serialize};