Staking async on Basti blocks

cumulus/pallets/parachain-system/src/block_weight/pre_inherents_hook.rs

@@ -38,6 +38,14 @@ where
 	TargetBlockRate: Get<u32>,
 {
 	fn pre_inherents() {
+		log::info!(
+			target: LOG_TARGET,
+			"pre-inherent at block {:?}, target_block_weight: {:?}, consumed: {:?}, is_first_block_in_core: {:?}",
+			frame_system::Pallet::<Config>::block_number(),
+			crate::block_weight::MaxParachainBlockWeight::<Config, TargetBlockRate>::target_block_weight(),
+			frame_system::Pallet::<Config>::remaining_block_weight().consumed(),
+			is_first_block_in_core::<Config>(),
+		);
 		if !block_weight_over_target_block_weight::<Config, TargetBlockRate>() {
 			let new_mode = if Config::MultiBlockMigrator::ongoing() {
 				log::debug!(

substrate/frame/election-provider-multi-block/src/lib.rs

@@ -712,54 +712,88 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_poll(_now: BlockNumberFor<T>, weight_meter: &mut WeightMeter) {
-			// first check we can at least read one storage.
-			if !weight_meter.can_consume(T::DbWeight::get().reads(1)) {
+		fn on_initialize(_now: BlockNumberFor<T>) -> Weight {
+			// Calculate the weight that on_poll will need
+			let current_phase = Self::current_phase();
+			let (self_weight, _) = Self::per_block_exec(current_phase);
+			let (verifier_weight, _) = T::Verifier::per_block_exec();
+
+			// Combined weight for phase transition
+			let combined_weight = self_weight
+				.saturating_add(verifier_weight)
+				.saturating_add(T::DbWeight::get().reads(1)); // Reading current phase
+
+			// Check if there's enough weight left in the block
+			let remaining_weight_meter = frame_system::Pallet::<T>::remaining_block_weight();
+			let remaining_weight = remaining_weight_meter.remaining();
+			let can_execute = remaining_weight.all_gte(combined_weight);
+
+			// TODO: the above doesn't work because the `remaining_block_weight` is already the
+			// mini-basti-block weight that is left, not the potential full-core weight that is
+			// left. We need a way to query the `full_core_weight`, subtract the
+			// `remaining_weight_meter.consumed()` from it, and then check if the result is greater
+			// than or equal to the `combined_weight`. For now, we yolo.
+
+			if true {
+				// We have enough weight - register it so BastiBlocks PreInherents hook knows
+				frame_system::Pallet::<T>::register_extra_weight_unchecked(
+					combined_weight,
+					frame_support::dispatch::DispatchClass::Mandatory,
+				);
+				CanExecuteOnPoll::<T>::put(true);
+				log!(debug, "registered {:?} weight for next poll execution", combined_weight);
+			} else {
+				// Not enough weight - emit event and don't execute on_poll
 				Self::deposit_event(Event::UnexpectedPhaseTransitionHalt {
-					required: T::DbWeight::get().reads(1),
-					had: weight_meter.remaining(),
+					required: combined_weight,
+					had: remaining_weight,
 				});
+				CanExecuteOnPoll::<T>::put(false);
+				log!(
+					warn,
+					"not enough weight for phase transition, required: {:?}, had: {:?}",
+					combined_weight,
+					remaining_weight
+				);
+			}
+
+			// Return the weight consumed by on_initialize itself (storage read + conditional write)
+			T::DbWeight::get().reads_writes(1, 1)
+		}
+
+		fn on_poll(_now: BlockNumberFor<T>, weight_meter: &mut WeightMeter) {
+			// Check if on_initialize determined we have enough weight
+			if !CanExecuteOnPoll::<T>::take() {
+				// Not enough weight was available - on_initialize already emitted an event
 				return;
 			}
 
-			// if so, consume and prepare the next phase.
 			let current_phase = Self::current_phase();
-			weight_meter.consume(T::DbWeight::get().reads(1));
-
 			let (self_weight, self_exec) = Self::per_block_exec(current_phase);
 			let (verifier_weight, verifier_exc) = T::Verifier::per_block_exec();
 
 			// The following will combine `Self::per_block_exec` and `T::Verifier::per_block_exec`
 			// into a single tuple of `(Weight, Box<_>)`. Can be moved into a reusable combinator
 			// function if we have this pattern in more places.
-			let (combined_weight, combined_exec) = (
-				// pre-exec weight is simply addition.
-				self_weight.saturating_add(verifier_weight),
-				// our new exec is..
-				Box::new(move |meter: &mut WeightMeter| {
-					self_exec(meter);
-					verifier_exc(meter);
-				}),
-			);
+			let combined_weight = self_weight.saturating_add(verifier_weight);
+			let combined_exec = Box::new(move |meter: &mut WeightMeter| {
+				self_exec(meter);
+				verifier_exc(meter);
+			});
 
 			log!(
 				trace,
-				"worst-case required weight for transition from {:?} to {:?} is {:?}, has {:?}",
+				"executing phase transition from {:?} to {:?} with weight {:?}, meter has {:?}",
 				current_phase,
 				current_phase.next(),
 				combined_weight,
 				weight_meter.remaining()
 			);
-			if weight_meter.can_consume(combined_weight) {
-				combined_exec(weight_meter);
-			} else {
-				Self::deposit_event(Event::UnexpectedPhaseTransitionOutOfWeight {
-					from: current_phase,
-					to: current_phase.next(),
-					required: combined_weight,
-					had: weight_meter.remaining(),
-				});
-			}
+
+			// NOTE: Weight was pre-registered in on_initialize for BastiBlocks support.
+			// We execute the work here but don't consume from weight_meter to avoid
+			// double-counting.
+			combined_exec(&mut WeightMeter::new());
 
 			// NOTE: why in here? because it is more accessible, for example `roll_to_with_ocw`.
 			#[cfg(test)]
@@ -865,7 +899,7 @@ pub mod pallet {
 			required: Weight,
 			had: Weight,
 		},
-		/// Phase transition could not even begin becaseu of being out of weight.
+		/// Phase transition could not even begin because of being out of weight.
 		UnexpectedPhaseTransitionHalt { required: Weight, had: Weight },
 	}
 
@@ -914,6 +948,11 @@ pub mod pallet {
 	#[pallet::getter(fn current_phase)]
 	pub type CurrentPhase<T: Config> = StorageValue<_, Phase<T>, ValueQuery>;
 
+	/// Flag to indicate if on_poll should execute this block.
+	/// Set in on_initialize if there's enough weight available.
+	#[pallet::storage]
+	pub(crate) type CanExecuteOnPoll<T: Config> = StorageValue<_, bool, ValueQuery>;
+
 	/// Wrapper struct for working with snapshots.
 	///
 	/// It manages the following storage items:

substrate/frame/staking-async/runtimes/papi-tests/zn-m.toml

@@ -36,5 +36,5 @@ name = "charlie"
 rpc_port = 9946
 args = [
 	"--authoring=slot-based",
-	"-lruntime::system=debug,runtime::multiblock-election=trace,runtime::staking=debug,runtime::staking::rc-client=trace,runtime::rc-client=debug,runtime::dap=debug,runtime::staking-async=info",
+	"-lruntime::system=debug,runtime::multiblock-election=trace,runtime::staking=debug,runtime::staking::rc-client=trace,runtime::rc-client=debug",
 ]

substrate/frame/staking-async/runtimes/papi-tests/zn-s.toml

@@ -7,6 +7,10 @@ name = "alice"
 validator = true
 rpc_port = 9944
 args = [
+	# fork-aware pool's background task crashes within ~30s on this fast-runtime
+	# setup, taking the chain down before session 1 begins. Pin to single-state
+	# until the upstream issue is fixed.
+	"--pool-type=single-state",
 	"-lruntime::system=debug,runtime::session=trace,runtime::staking-async::ah-client=trace,runtime::ah-client=debug,xcm=trace",
 ]
 
@@ -15,6 +19,7 @@ name = "bob"
 validator = true
 rpc_port = 9945
 args = [
+	"--pool-type=single-state",
 	"-lruntime::system=debug,runtime::session=trace,runtime::staking-async::ah-client=trace,runtime::ah-client=debug",
 ]
 
@@ -28,5 +33,6 @@ name = "charlie"
 rpc_port = 9946
 args = [
 	"--authoring=slot-based",
-	"-lruntime::system=debug,runtime::multiblock-election=trace,runtime::staking=debug,runtime::staking::rc-client=trace,runtime::rc-client=debug,runtime::dap=debug,runtime::staking-async=info,xcm=debug,parachain-system=debug,runtime=info",
+	"--pool-type=single-state",
+	"-lruntime::system=debug,runtime::multiblock-election=trace,runtime::staking=debug,runtime::staking::rc-client=trace,runtime::rc-client=debug,xcm=debug,parachain-system=debug,runtime=info",
 ]

substrate/frame/staking-async/runtimes/parachain/src/lib.rs

@@ -43,6 +43,7 @@ use assets_common::{
 	matching::{FromNetwork, FromSiblingParachain},
 	AssetIdForPoolAssets, AssetIdForPoolAssetsConvert, AssetIdForTrustBackedAssetsConvert,
 };
+use sp_core::Get;
 use codec::{Decode, DecodeWithMemTracking, Encode, MaxEncodedLen};
 use cumulus_pallet_parachain_system::{RelayNumberMonotonicallyIncreases, RelaychainDataProvider};
 use cumulus_primitives_core::{AggregateMessageOrigin, ParaId};
@@ -91,7 +92,7 @@ use sp_runtime::{
 use sp_version::NativeVersion;
 use sp_version::RuntimeVersion;
 use testnet_parachains_constants::westend::{
-	consensus::*, currency::*, fee::WeightToFee, snowbridge::EthereumNetwork, time::*,
+	currency::*, fee::WeightToFee, snowbridge::EthereumNetwork, time::*,
 };
 use weights::{BlockExecutionWeight, ExtrinsicBaseWeight, RocksDbWeight};
 use xcm::{
@@ -144,6 +145,26 @@ pub fn native_version() -> NativeVersion {
 
 type RelayChainBlockNumberProvider = RelaychainDataProvider<Runtime>;
 
+// Bundled-blocks configuration: 12 mini-blocks per 6s relay slot = 500ms blocks.
+// With a single assigned core, the slot-based collator bundles all 12 blocks into
+// a single PoV per relay slot (see `cumulus-test-runtime`'s `block-bundling`
+// feature, which uses the same velocity).
+pub const BLOCK_PROCESSING_VELOCITY: u32 = 12;
+pub const RELAY_CHAIN_SLOT_DURATION_MILLIS: u32 = 6000;
+// Must accommodate `BLOCK_PROCESSING_VELOCITY` blocks per PoV multiplied by the
+// inclusion latency. Mirrors `cumulus-test-runtime`'s formula
+// `VELOCITY * (3 + RELAY_PARENT_OFFSET)` for `RELAY_PARENT_OFFSET = 0`.
+pub const UNINCLUDED_SEGMENT_CAPACITY: u32 = BLOCK_PROCESSING_VELOCITY * 3;
+// Override westend's 24s slot duration: parachain slots run at the relay-chain
+// cadence; the collator subdivides the slot into `BLOCK_PROCESSING_VELOCITY`
+// bundled blocks.
+pub const SLOT_DURATION: u64 = RELAY_CHAIN_SLOT_DURATION_MILLIS as u64;
+
+type MaximumBlockWeight = cumulus_pallet_parachain_system::block_weight::MaxParachainBlockWeight<
+	Runtime,
+	ConstU32<BLOCK_PROCESSING_VELOCITY>,
+>;
+
 parameter_types! {
 	pub const Version: RuntimeVersion = VERSION;
 	pub RuntimeBlockLength: BlockLength = BlockLength::builder()
@@ -158,14 +179,14 @@ parameter_types! {
 			weights.base_extrinsic = ExtrinsicBaseWeight::get();
 		})
 		.for_class(DispatchClass::Normal, |weights| {
-			weights.max_total = Some(NORMAL_DISPATCH_RATIO * MAXIMUM_BLOCK_WEIGHT);
+			weights.max_total = Some(NORMAL_DISPATCH_RATIO * MaximumBlockWeight::get());
 		})
 		.for_class(DispatchClass::Operational, |weights| {
-			weights.max_total = Some(MAXIMUM_BLOCK_WEIGHT);
+			weights.max_total = Some(MaximumBlockWeight::get());
 			// Operational transactions have some extra reserved space, so that they
-			// are included even if block reached `MAXIMUM_BLOCK_WEIGHT`.
+			// are included even if block reached `MaximumBlockWeight`.
 			weights.reserved = Some(
-				MAXIMUM_BLOCK_WEIGHT - NORMAL_DISPATCH_RATIO * MAXIMUM_BLOCK_WEIGHT
+				MaximumBlockWeight::get() - NORMAL_DISPATCH_RATIO * MaximumBlockWeight::get()
 			);
 		})
 		.avg_block_initialization(AVERAGE_ON_INITIALIZE_RATIO)
@@ -193,6 +214,10 @@ impl frame_system::Config for Runtime {
 	type MaxConsumers = frame_support::traits::ConstU32<16>;
 	type MultiBlockMigrator = MultiBlockMigrations;
 	type SingleBlockMigrations = Migrations;
+	type PreInherents = cumulus_pallet_parachain_system::block_weight::DynamicMaxBlockWeightHooks<
+		Runtime,
+		ConstU32<BLOCK_PROCESSING_VELOCITY>,
+	>;
 }
 
 impl cumulus_pallet_weight_reclaim::Config for Runtime {
@@ -802,8 +827,8 @@ impl pallet_proxy::Config for Runtime {
 }
 
 parameter_types! {
-	pub const ReservedXcmpWeight: Weight = MAXIMUM_BLOCK_WEIGHT.saturating_div(4);
-	pub const ReservedDmpWeight: Weight = MAXIMUM_BLOCK_WEIGHT.saturating_div(4);
+	pub ReservedXcmpWeight: Weight = MaximumBlockWeight::get().saturating_div(4);
+	pub ReservedDmpWeight: Weight = MaximumBlockWeight::get().saturating_div(4);
 }
 
 impl cumulus_pallet_parachain_system::Config for Runtime {
@@ -1227,19 +1252,23 @@ pub type SignedBlock = generic::SignedBlock<Block>;
 /// BlockId type as expected by this runtime.
 pub type BlockId = generic::BlockId<Block>;
 /// The extension to the basic transaction logic.
-pub type TxExtension = cumulus_pallet_weight_reclaim::StorageWeightReclaim<
+pub type TxExtension = cumulus_pallet_parachain_system::block_weight::DynamicMaxBlockWeight<
 	Runtime,
-	(
-		frame_system::CheckNonZeroSender<Runtime>,
-		frame_system::CheckSpecVersion<Runtime>,
-		frame_system::CheckTxVersion<Runtime>,
-		frame_system::CheckGenesis<Runtime>,
-		frame_system::CheckEra<Runtime>,
-		frame_system::CheckNonce<Runtime>,
-		frame_system::CheckWeight<Runtime>,
-		pallet_asset_conversion_tx_payment::ChargeAssetTxPayment<Runtime>,
-		frame_metadata_hash_extension::CheckMetadataHash<Runtime>,
-	),
+	cumulus_pallet_weight_reclaim::StorageWeightReclaim<
+		Runtime,
+		(
+			frame_system::CheckNonZeroSender<Runtime>,
+			frame_system::CheckSpecVersion<Runtime>,
+			frame_system::CheckTxVersion<Runtime>,
+			frame_system::CheckGenesis<Runtime>,
+			frame_system::CheckEra<Runtime>,
+			frame_system::CheckNonce<Runtime>,
+			frame_system::CheckWeight<Runtime>,
+			pallet_asset_conversion_tx_payment::ChargeAssetTxPayment<Runtime>,
+			frame_metadata_hash_extension::CheckMetadataHash<Runtime>,
+		),
+	>,
+	ConstU32<BLOCK_PROCESSING_VELOCITY>,
 >;
 
 pub type UncheckedExtrinsic =
@@ -1383,6 +1412,13 @@ mod benches {
 }
 
 impl_runtime_apis! {
+	// enable the mighty basti-blocks.
+	impl cumulus_primitives_core::TargetBlockRate<Block> for Runtime {
+		fn target_block_rate() -> u32 {
+			BLOCK_PROCESSING_VELOCITY
+		}
+	}
+
 	impl sp_consensus_aura::AuraApi<Block, AuraId> for Runtime {
 		fn slot_duration() -> sp_consensus_aura::SlotDuration {
 			sp_consensus_aura::SlotDuration::from_millis(SLOT_DURATION)

substrate/frame/staking-async/runtimes/parachain/src/staking.rs

@@ -17,6 +17,7 @@
 
 /// ! Staking, and election related pallet configurations.
 use super::*;
+use cumulus_pallet_parachain_system::block_weight::DynamicMaxBlockWeight;
 use cumulus_primitives_core::relay_chain::SessionIndex;
 use frame_election_provider_support::{ElectionDataProvider, SequentialPhragmen};
 use frame_support::traits::{ConstU128, EitherOf};
@@ -725,17 +726,25 @@ where
 			// so the actual block number is `n`.
 			.saturating_sub(1);
 		let tip = 0;
-		let tx_ext = TxExtension::from((
-			frame_system::CheckNonZeroSender::<Runtime>::new(),
-			frame_system::CheckSpecVersion::<Runtime>::new(),
-			frame_system::CheckTxVersion::<Runtime>::new(),
-			frame_system::CheckGenesis::<Runtime>::new(),
-			frame_system::CheckEra::<Runtime>::from(generic::Era::mortal(period, current_block)),
-			frame_system::CheckNonce::<Runtime>::from(nonce),
-			frame_system::CheckWeight::<Runtime>::new(),
-			pallet_asset_conversion_tx_payment::ChargeAssetTxPayment::<Runtime>::from(tip, None),
-			frame_metadata_hash_extension::CheckMetadataHash::<Runtime>::new(true),
-		));
+		let tx_ext = DynamicMaxBlockWeight::new(
+			(
+				frame_system::CheckNonZeroSender::<Runtime>::new(),
+				frame_system::CheckSpecVersion::<Runtime>::new(),
+				frame_system::CheckTxVersion::<Runtime>::new(),
+				frame_system::CheckGenesis::<Runtime>::new(),
+				frame_system::CheckEra::<Runtime>::from(generic::Era::mortal(
+					period,
+					current_block,
+				)),
+				frame_system::CheckNonce::<Runtime>::from(nonce),
+				frame_system::CheckWeight::<Runtime>::new(),
+				pallet_asset_conversion_tx_payment::ChargeAssetTxPayment::<Runtime>::from(
+					tip, None,
+				),
+				frame_metadata_hash_extension::CheckMetadataHash::<Runtime>::new(true),
+			)
+				.into(),
+		);
 		let raw_payload = SignedPayload::new(call, tx_ext)
 			.map_err(|e| {
 				log::warn!("Unable to create signed payload: {:?}", e);
@@ -842,17 +851,19 @@ mod tests {
 
 	#[test]
 	fn signed_weight_ratios() {
-		sp_tracing::try_init_simple();
-		let block_weight = <Runtime as frame_system::Config>::BlockWeights::get().max_block;
-		let polkadot_signed_submission =
-			mb::weights::polkadot::MultiBlockSignedWeightInfo::<Runtime>::submit_page();
-		let kusama_signed_submission =
-			mb::weights::kusama::MultiBlockSignedWeightInfo::<Runtime>::submit_page();
-
-		log::info!(target: "runtime", "Polkadot:");
-		weight_diff(block_weight, polkadot_signed_submission);
-		log::info!(target: "runtime", "Kusama:");
-		weight_diff(block_weight, kusama_signed_submission);
+		sp_io::TestExternalities::default().execute_with(|| {
+			sp_tracing::try_init_simple();
+			let block_weight = <Runtime as frame_system::Config>::BlockWeights::get().max_block;
+			let polkadot_signed_submission =
+				mb::weights::polkadot::MultiBlockSignedWeightInfo::<Runtime>::submit_page();
+			let kusama_signed_submission =
+				mb::weights::kusama::MultiBlockSignedWeightInfo::<Runtime>::submit_page();
+
+			log::info!(target: "runtime", "Polkadot:");
+			weight_diff(block_weight, polkadot_signed_submission);
+			log::info!(target: "runtime", "Kusama:");
+			weight_diff(block_weight, kusama_signed_submission);
+		})
 	}
 
 	#[test]