refactor unincluded segment length into a ConsensusHook (#2501)

* refactor unincluded segment length into a ConsensusHook

* add docs

* refactor bandwidth_out calculation

Co-authored-by: Chris Sosnin <48099298+slumber@users.noreply.github.com>

* test for limits from impl

* fmt

* make tests compile

* update comment

* uncomment test

* fix collator test by adding parent to state proof

* patch HRMP watermark rules for unincluded segment

* get consensus-common tests to pass, using unincluded segment

* fix unincluded segment tests

* get all tests passing

* fmt

* rustdoc CI

* aura-ext: limit the number of authored blocks per slot (#2551)

* aura_ext consensus hook

* reverse dependency

* include weight into hook

* fix tests

* remove stray println

Co-authored-by: Chris Sosnin <48099298+slumber@users.noreply.github.com>

* fix test warning

* fix doc link

---------

Co-authored-by: Chris Sosnin <48099298+slumber@users.noreply.github.com>
Co-authored-by: Chris Sosnin <chris125_@live.com>

client/collator/Cargo.toml

@@ -42,3 +42,4 @@ polkadot-node-subsystem-test-helpers = { git = "https://github.com/paritytech/po
 # Cumulus
 cumulus-test-client = { path = "../../test/client" }
 cumulus-test-runtime = { path = "../../test/runtime" }
+cumulus-test-relay-sproof-builder = { path = "../../test/relay-sproof-builder" }

client/collator/src/lib.rs

@@ -384,10 +384,12 @@ mod tests {
 		Client, ClientBlockImportExt, DefaultTestClientBuilderExt, InitBlockBuilder,
 		TestClientBuilder, TestClientBuilderExt,
 	};
+	use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
 	use cumulus_test_runtime::{Block, Header};
 	use futures::{channel::mpsc, executor::block_on, StreamExt};
 	use polkadot_node_subsystem_test_helpers::ForwardSubsystem;
 	use polkadot_overseer::{dummy::dummy_overseer_builder, HeadSupportsParachains};
+	use polkadot_primitives::HeadData;
 	use sp_consensus::BlockOrigin;
 	use sp_core::{testing::TaskExecutor, Pair};
 	use sp_runtime::traits::BlakeTwo256;
@@ -415,10 +417,14 @@ mod tests {
 			_: PHash,
 			validation_data: &PersistedValidationData,
 		) -> Option<ParachainCandidate<Block>> {
+			let mut sproof = RelayStateSproofBuilder::default();
+			sproof.included_para_head = Some(HeadData(parent.encode()));
+			sproof.para_id = cumulus_test_runtime::PARACHAIN_ID.into();
+
 			let builder = self.client.init_block_builder_at(
 				parent.hash(),
 				Some(validation_data.clone()),
-				Default::default(),
+				sproof,
 			);
 
 			let (block, _, proof) = builder.build().expect("Creates block").into_inner();

client/consensus/common/Cargo.toml

@@ -38,3 +38,4 @@ sp-tracing = { git = "https://github.com/paritytech/substrate", branch = "master
 
 # Cumulus
 cumulus-test-client = { path = "../../../test/client" }
+cumulus-test-relay-sproof-builder = { path = "../../../test/relay-sproof-builder" }

client/consensus/common/src/tests.rs

@@ -28,8 +28,10 @@ use cumulus_test_client::{
 	runtime::{Block, Hash, Header},
 	Backend, Client, InitBlockBuilder, TestClientBuilder, TestClientBuilderExt,
 };
+use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
 use futures::{channel::mpsc, executor::block_on, select, FutureExt, Stream, StreamExt};
 use futures_timer::Delay;
+use polkadot_primitives::HeadData;
 use sc_client_api::{blockchain::Backend as _, Backend as _, UsageProvider};
 use sc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
 use sp_consensus::{BlockOrigin, BlockStatus};
@@ -209,17 +211,36 @@ impl RelayChainInterface for Relaychain {
 	}
 }
 
+fn sproof_with_best_parent(client: &Client) -> RelayStateSproofBuilder {
+	let best_hash = client.chain_info().best_hash;
+	sproof_with_parent_by_hash(client, best_hash)
+}
+
+fn sproof_with_parent_by_hash(client: &Client, hash: PHash) -> RelayStateSproofBuilder {
+	let header = client.header(hash).ok().flatten().expect("No header for parent block");
+	sproof_with_parent(HeadData(header.encode()))
+}
+
+fn sproof_with_parent(parent: HeadData) -> RelayStateSproofBuilder {
+	let mut x = RelayStateSproofBuilder::default();
+	x.para_id = cumulus_test_client::runtime::PARACHAIN_ID.into();
+	x.included_para_head = Some(parent);
+
+	x
+}
+
 fn build_block<B: InitBlockBuilder>(
 	builder: &B,
+	sproof: RelayStateSproofBuilder,
 	at: Option<Hash>,
 	timestamp: Option<u64>,
 ) -> Block {
 	let builder = match at {
 		Some(at) => match timestamp {
-			Some(ts) => builder.init_block_builder_with_timestamp(at, None, Default::default(), ts),
-			None => builder.init_block_builder_at(at, None, Default::default()),
+			Some(ts) => builder.init_block_builder_with_timestamp(at, None, sproof, ts),
+			None => builder.init_block_builder_at(at, None, sproof),
 		},
-		None => builder.init_block_builder(None, Default::default()),
+		None => builder.init_block_builder(None, sproof),
 	};
 
 	let mut block = builder.build().unwrap().block;
@@ -260,15 +281,20 @@ fn import_block_sync<I: BlockImport<Block>>(
 	block_on(import_block(importer, block, origin, import_as_best));
 }
 
-fn build_and_import_block_ext<B: InitBlockBuilder, I: BlockImport<Block>>(
-	builder: &B,
+fn build_and_import_block_ext<I: BlockImport<Block>>(
+	client: &Client,
 	origin: BlockOrigin,
 	import_as_best: bool,
 	importer: &mut I,
 	at: Option<Hash>,
 	timestamp: Option<u64>,
 ) -> Block {
-	let block = build_block(builder, at, timestamp);
+	let sproof = match at {
+		None => sproof_with_best_parent(client),
+		Some(at) => sproof_with_parent_by_hash(client, at),
+	};
+
+	let block = build_block(client, sproof, at, timestamp);
 	import_block_sync(importer, block.clone(), origin, import_as_best);
 	block
 }
@@ -337,7 +363,12 @@ fn follow_new_best_with_dummy_recovery_works() {
 		Some(recovery_chan_tx),
 	);
 
-	let block = build_block(&*client.clone(), None, None);
+	let sproof = {
+		let best = client.chain_info().best_hash;
+		let header = client.header(best).ok().flatten().expect("No header for best");
+		sproof_with_parent(HeadData(header.encode()))
+	};
+	let block = build_block(&*client.clone(), sproof, None, None);
 	let block_clone = block.clone();
 	let client_clone = client.clone();
 
@@ -423,7 +454,8 @@ fn follow_finalized_does_not_stop_on_unknown_block() {
 	let block = build_and_import_block(client.clone(), false);
 
 	let unknown_block = {
-		let block_builder = client.init_block_builder_at(block.hash(), None, Default::default());
+		let sproof = sproof_with_parent_by_hash(&*client, block.hash());
+		let block_builder = client.init_block_builder_at(block.hash(), None, sproof);
 		block_builder.build().unwrap().block
 	};
 
@@ -472,7 +504,8 @@ fn follow_new_best_sets_best_after_it_is_imported() {
 	let block = build_and_import_block(client.clone(), false);
 
 	let unknown_block = {
-		let block_builder = client.init_block_builder_at(block.hash(), None, Default::default());
+		let sproof = sproof_with_parent_by_hash(&*client, block.hash());
+		let block_builder = client.init_block_builder_at(block.hash(), None, sproof);
 		block_builder.build().unwrap().block
 	};
 

pallets/aura-ext/Cargo.toml

@@ -18,6 +18,9 @@ sp-consensus-aura = { git = "https://github.com/paritytech/substrate", default-f
 sp-runtime = { git = "https://github.com/paritytech/substrate", default-features = false, branch = "master" }
 sp-std = { git = "https://github.com/paritytech/substrate", default-features = false, branch = "master" }
 
+# Cumulus
+cumulus-pallet-parachain-system = { path = "../parachain-system", default-features = false }
+
 [dev-dependencies]
 
 # Cumulus
@@ -35,5 +38,6 @@ std = [
 	"sp-consensus-aura/std",
 	"sp-runtime/std",
 	"sp-std/std",
+	"cumulus-pallet-parachain-system/std",
 ]
 try-runtime = ["frame-support/try-runtime"]

pallets/aura-ext/src/consensus_hook.rs

@@ -0,0 +1,56 @@
+// Copyright 2023 Parity Technologies (UK) Ltd.
+// This file is part of Cumulus.
+
+// Cumulus is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Cumulus is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
+
+//! The definition of a [`FixedVelocityConsensusHook`] for consensus logic to manage
+//! block velocity.
+//!
+//! The velocity `V` refers to the rate of block processing by the relay chain.
+
+use super::pallet;
+use cumulus_pallet_parachain_system::{
+	consensus_hook::{ConsensusHook, UnincludedSegmentCapacity},
+	relay_state_snapshot::RelayChainStateProof,
+};
+use frame_support::pallet_prelude::*;
+use sp_std::{marker::PhantomData, num::NonZeroU32};
+
+/// A consensus hook for a fixed block processing velocity and unincluded segment capacity.
+pub struct FixedVelocityConsensusHook<T, const V: u32, const C: u32>(PhantomData<T>);
+
+impl<T: pallet::Config, const V: u32, const C: u32> ConsensusHook
+	for FixedVelocityConsensusHook<T, V, C>
+{
+	// Validates the number of authored blocks within the slot with respect to the `V + 1` limit.
+	fn on_state_proof(_state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		// Ensure velocity is non-zero.
+		let velocity = V.max(1);
+
+		let authored = pallet::Pallet::<T>::slot_info()
+			.map(|(_slot, authored)| authored)
+			.expect("slot info is inserted on block initialization");
+		if authored > velocity + 1 {
+			panic!("authored blocks limit is reached for the slot")
+		}
+		let weight = T::DbWeight::get().reads(1);
+
+		(
+			weight,
+			NonZeroU32::new(sp_std::cmp::max(C, 1))
+				.expect("1 is the minimum value and non-zero; qed")
+				.into(),
+		)
+	}
+}

pallets/aura-ext/src/lib.rs

@@ -37,9 +37,12 @@
 
 use frame_support::traits::{ExecuteBlock, FindAuthor};
 use sp_application_crypto::RuntimeAppPublic;
-use sp_consensus_aura::digests::CompatibleDigestItem;
+use sp_consensus_aura::{digests::CompatibleDigestItem, Slot};
 use sp_runtime::traits::{Block as BlockT, Header as HeaderT};
 
+pub mod consensus_hook;
+pub use consensus_hook::FixedVelocityConsensusHook;
+
 type Aura<T> = pallet_aura::Pallet<T>;
 
 pub use pallet::*;
@@ -68,6 +71,19 @@ pub mod pallet {
 			// Fetch the authorities once to get them into the storage proof of the PoV.
 			Authorities::<T>::get();
 
+			let new_slot = Aura::<T>::current_slot();
+
+			let (new_slot, authored) = match SlotInfo::<T>::get() {
+				Some((slot, authored)) if slot == new_slot => (slot, authored + 1),
+				Some((slot, _)) if slot < new_slot => (new_slot, 1),
+				Some(..) => {
+					panic!("slot moved backwards")
+				},
+				None => (new_slot, 1),
+			};
+
+			SlotInfo::<T>::put((new_slot, authored));
+
 			T::DbWeight::get().reads_writes(2, 1)
 		}
 	}
@@ -84,6 +100,13 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
+	/// Current slot paired with a number of authored blocks.
+	///
+	/// Updated on each block initialization.
+	#[pallet::storage]
+	#[pallet::getter(fn slot_info)]
+	pub(crate) type SlotInfo<T: Config> = StorageValue<_, (Slot, u32), OptionQuery>;
+
 	#[pallet::genesis_config]
 	#[derive(Default)]
 	pub struct GenesisConfig;

pallets/parachain-system/src/consensus_hook.rs

@@ -0,0 +1,109 @@
+// Copyright 2023 Parity Technologies (UK) Ltd.
+// This file is part of Cumulus.
+
+// Cumulus is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Cumulus is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
+
+//! The definition of a [`ConsensusHook`] trait for consensus logic to manage the backlog
+//! of parachain blocks ready to submit to the relay chain, as well as some basic implementations.
+
+use super::relay_state_snapshot::RelayChainStateProof;
+use frame_support::weights::Weight;
+use sp_std::num::NonZeroU32;
+
+/// The possible capacity of the unincluded segment.
+#[derive(Clone)]
+pub struct UnincludedSegmentCapacity(UnincludedSegmentCapacityInner);
+
+impl UnincludedSegmentCapacity {
+	pub(crate) fn get(&self) -> u32 {
+		match self.0 {
+			UnincludedSegmentCapacityInner::ExpectParentIncluded => 1,
+			UnincludedSegmentCapacityInner::Value(v) => v.get(),
+		}
+	}
+
+	pub(crate) fn is_expecting_included_parent(&self) -> bool {
+		match self.0 {
+			UnincludedSegmentCapacityInner::ExpectParentIncluded => true,
+			UnincludedSegmentCapacityInner::Value(_) => false,
+		}
+	}
+}
+
+#[derive(Clone)]
+pub(crate) enum UnincludedSegmentCapacityInner {
+	ExpectParentIncluded,
+	Value(NonZeroU32),
+}
+
+impl From<NonZeroU32> for UnincludedSegmentCapacity {
+	fn from(value: NonZeroU32) -> Self {
+		UnincludedSegmentCapacity(UnincludedSegmentCapacityInner::Value(value))
+	}
+}
+
+/// The consensus hook for dealing with the unincluded segment.
+///
+/// Higher-level and user-configurable consensus logic is more informed about the
+/// desired unincluded segment length, as well as any rules for adapting it dynamically
+/// according to the relay-chain state.
+pub trait ConsensusHook {
+	/// This hook is called partway through the `set_validation_data` inherent in parachain-system.
+	///
+	/// The hook is allowed to panic if customized consensus rules aren't met and is required
+	/// to return a maximum capacity for the unincluded segment with weight consumed.
+	fn on_state_proof(state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity);
+}
+
+/// A special consensus hook for handling the migration to asynchronous backing gracefully,
+/// even if collators haven't been updated to provide the last included parent in the state
+/// proof yet.
+///
+/// This behaves as though the parent is included, even if the relay chain state proof doesn't contain
+/// the included para head. If the para head is present in the state proof, this does ensure the
+/// parent is included.
+pub struct ExpectParentIncluded;
+
+impl ConsensusHook for ExpectParentIncluded {
+	fn on_state_proof(_state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		(
+			Weight::zero(),
+			UnincludedSegmentCapacity(UnincludedSegmentCapacityInner::ExpectParentIncluded),
+		)
+	}
+}
+
+/// A consensus hook for a fixed unincluded segment length. This hook does nothing but
+/// set the capacity of the unincluded segment to the constant N.
+///
+/// Since it is illegal to provide an unincluded segment length of 0, this sets a minimum of
+/// 1.
+pub struct FixedCapacityUnincludedSegment<const N: u32>;
+
+impl<const N: u32> ConsensusHook for FixedCapacityUnincludedSegment<N> {
+	fn on_state_proof(_state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		(
+			Weight::zero(),
+			NonZeroU32::new(sp_std::cmp::max(N, 1))
+				.expect("1 is the minimum value and non-zero; qed")
+				.into(),
+		)
+	}
+}
+
+/// A fixed-capacity unincluded segment hook, which requires that the parent block is
+/// included prior to the current block being authored.
+///
+/// This is a simple type alias around a fixed-capacity unincluded segment with a size of 1.
+pub type RequireParentIncluded = FixedCapacityUnincludedSegment<1>;