feat: Rebrand Polkadot/Substrate references to PezkuwiChain

This commit systematically rebrands various references from Parity Technologies'
Polkadot/Substrate ecosystem to PezkuwiChain within the kurdistan-sdk.

Key changes include:
- Updated external repository URLs (zombienet-sdk, parity-db, parity-scale-codec, wasm-instrument) to point to pezkuwichain forks.
- Modified internal documentation and code comments to reflect PezkuwiChain naming and structure.
- Replaced direct references to  with  or specific paths within the  for XCM, Pezkuwi, and other modules.
- Cleaned up deprecated  issue and PR references in various  and  files, particularly in  and  modules.
- Adjusted image and logo URLs in documentation to point to PezkuwiChain assets.
- Removed or rephrased comments related to external Polkadot/Substrate PRs and issues.

This is a significant step towards fully customizing the SDK for the PezkuwiChain ecosystem.

pezcumulus/client/consensus/aura/Cargo.toml

@@ -0,0 +1,104 @@
+[package]
+name = "pezcumulus-client-consensus-aura"
+description = "AURA consensus algorithm for teyrchains"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+async-trait = { workspace = true }
+codec = { features = ["derive"], workspace = true, default-features = true }
+futures = { workspace = true }
+parking_lot = { workspace = true }
+schnellru = { workspace = true }
+tokio = { workspace = true, features = ["macros"] }
+tracing = { workspace = true, default-features = true }
+
+# Bizinikiwi
+prometheus-endpoint = { workspace = true, default-features = true }
+pezsc-client-api = { workspace = true, default-features = true }
+pezsc-consensus = { workspace = true, default-features = true }
+pezsc-consensus-aura = { workspace = true, default-features = true }
+pezsc-consensus-babe = { workspace = true, default-features = true }
+pezsc-consensus-slots = { workspace = true, default-features = true }
+pezsc-network-types = { workspace = true, default-features = true }
+pezsc-telemetry = { workspace = true, default-features = true }
+pezsc-utils = { workspace = true, default-features = true }
+pezsp-api = { workspace = true, default-features = true }
+pezsp-application-crypto = { workspace = true, default-features = true }
+pezsp-block-builder = { workspace = true, default-features = true }
+pezsp-blockchain = { workspace = true, default-features = true }
+pezsp-consensus = { workspace = true, default-features = true }
+pezsp-consensus-aura = { workspace = true, default-features = true }
+pezsp-core = { workspace = true, default-features = true }
+pezsp-inherents = { workspace = true, default-features = true }
+pezsp-keystore = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+pezsp-state-machine = { workspace = true, default-features = true }
+pezsp-timestamp = { workspace = true, default-features = true }
+pezsp-trie = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-client-collator = { workspace = true, default-features = true }
+pezcumulus-client-consensus-common = { workspace = true, default-features = true }
+pezcumulus-client-consensus-proposer = { workspace = true, default-features = true }
+pezcumulus-client-teyrchain-inherent = { workspace = true, default-features = true }
+pezcumulus-primitives-aura = { workspace = true, default-features = true }
+pezcumulus-primitives-core = { workspace = true, default-features = true }
+pezcumulus-relay-chain-interface = { workspace = true, default-features = true }
+
+# Pezkuwi
+pezkuwi-node-primitives = { workspace = true, default-features = true }
+pezkuwi-node-subsystem = { workspace = true, default-features = true }
+pezkuwi-node-subsystem-util = { workspace = true, default-features = true }
+pezkuwi-overseer = { workspace = true, default-features = true }
+pezkuwi-primitives = { workspace = true, default-features = true }
+
+[dev-dependencies]
+pezcumulus-test-client = { workspace = true }
+pezcumulus-test-relay-sproof-builder = { workspace = true }
+rstest = { workspace = true }
+pezsp-keyring = { workspace = true }
+pezsp-tracing = { workspace = true }
+pezsp-version = { workspace = true }
+
+[features]
+runtime-benchmarks = [
+	"pezcumulus-client-collator/runtime-benchmarks",
+	"pezcumulus-client-consensus-common/runtime-benchmarks",
+	"pezcumulus-client-consensus-proposer/runtime-benchmarks",
+	"pezcumulus-client-teyrchain-inherent/runtime-benchmarks",
+	"pezcumulus-primitives-aura/runtime-benchmarks",
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezcumulus-relay-chain-interface/runtime-benchmarks",
+	"pezcumulus-test-client/runtime-benchmarks",
+	"pezcumulus-test-relay-sproof-builder/runtime-benchmarks",
+	"pezkuwi-node-primitives/runtime-benchmarks",
+	"pezkuwi-node-subsystem-util/runtime-benchmarks",
+	"pezkuwi-node-subsystem/runtime-benchmarks",
+	"pezkuwi-overseer/runtime-benchmarks",
+	"pezkuwi-primitives/runtime-benchmarks",
+	"pezsc-client-api/runtime-benchmarks",
+	"pezsc-consensus-aura/runtime-benchmarks",
+	"pezsc-consensus-babe/runtime-benchmarks",
+	"pezsc-consensus-slots/runtime-benchmarks",
+	"pezsc-consensus/runtime-benchmarks",
+	"pezsp-api/runtime-benchmarks",
+	"pezsp-block-builder/runtime-benchmarks",
+	"pezsp-blockchain/runtime-benchmarks",
+	"pezsp-consensus-aura/runtime-benchmarks",
+	"pezsp-consensus/runtime-benchmarks",
+	"pezsp-inherents/runtime-benchmarks",
+	"pezsp-keyring/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+	"pezsp-state-machine/runtime-benchmarks",
+	"pezsp-timestamp/runtime-benchmarks",
+	"pezsp-trie/runtime-benchmarks",
+	"pezsp-version/runtime-benchmarks",
+]

pezcumulus/client/consensus/aura/src/collator.rs

@@ -0,0 +1,442 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! The core collator logic for Aura - slot claiming, block proposing, and collation
+//! packaging.
+//!
+//! The [`Collator`] struct exposed here is meant to be a component of higher-level logic
+//! which actually manages the control flow of the collator - which slots to claim, how
+//! many collations to build, when to work, etc.
+//!
+//! This module also exposes some standalone functions for common operations when building
+//! aura-based collators.
+
+use codec::Codec;
+use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_client_consensus_common::{
+	self as consensus_common, TeyrchainBlockImportMarker, TeyrchainCandidate,
+};
+use cumulus_client_consensus_proposer::ProposerInterface;
+use cumulus_client_teyrchain_inherent::{TeyrchainInherentData, TeyrchainInherentDataProvider};
+use cumulus_primitives_core::{
+	relay_chain::Hash as PHash, DigestItem, PersistedValidationData, TeyrchainBlockData,
+};
+use cumulus_relay_chain_interface::RelayChainInterface;
+
+use pezkuwi_node_primitives::{Collation, MaybeCompressedPoV};
+use pezkuwi_primitives::{Header as PHeader, Id as ParaId};
+
+use crate::collators::RelayParentData;
+use futures::prelude::*;
+use pezsc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy, StateAction};
+use pezsc_consensus_aura::standalone as aura_internal;
+use pezsc_network_types::PeerId;
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_application_crypto::AppPublic;
+use pezsp_consensus::BlockOrigin;
+use pezsp_consensus_aura::{AuraApi, Slot, SlotDuration};
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::{CreateInherentDataProviders, InherentData, InherentDataProvider};
+use pezsp_keystore::KeystorePtr;
+use pezsp_runtime::{
+	generic::Digest,
+	traits::{Block as BlockT, HashingFor, Header as HeaderT, Member},
+};
+use pezsp_state_machine::StorageChanges;
+use pezsp_timestamp::Timestamp;
+use std::{error::Error, time::Duration};
+
+/// Parameters for instantiating a [`Collator`].
+pub struct Params<BI, CIDP, RClient, Proposer, CS> {
+	/// A builder for inherent data builders.
+	pub create_inherent_data_providers: CIDP,
+	/// The block import handle.
+	pub block_import: BI,
+	/// An interface to the relay-chain client.
+	pub relay_client: RClient,
+	/// The keystore handle used for accessing teyrchain key material.
+	pub keystore: KeystorePtr,
+	/// The collator network peer id.
+	pub collator_peer_id: PeerId,
+	/// The identifier of the teyrchain within the relay-chain.
+	pub para_id: ParaId,
+	/// The block proposer used for building blocks.
+	pub proposer: Proposer,
+	/// The collator service used for bundling proposals into collations and announcing
+	/// to the network.
+	pub collator_service: CS,
+}
+
+/// A utility struct for writing collation logic that makes use of Aura entirely
+/// or in part. See module docs for more details.
+pub struct Collator<Block, P, BI, CIDP, RClient, Proposer, CS> {
+	create_inherent_data_providers: CIDP,
+	block_import: BI,
+	relay_client: RClient,
+	keystore: KeystorePtr,
+	para_id: ParaId,
+	proposer: Proposer,
+	collator_service: CS,
+	_marker: std::marker::PhantomData<(Block, Box<dyn Fn(P) + Send + Sync + 'static>)>,
+}
+
+impl<Block, P, BI, CIDP, RClient, Proposer, CS> Collator<Block, P, BI, CIDP, RClient, Proposer, CS>
+where
+	Block: BlockT,
+	RClient: RelayChainInterface,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block>,
+	CS: CollatorServiceInterface<Block>,
+	P: Pair,
+	P::Public: AppPublic + Member,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	/// Instantiate a new instance of the `Aura` manager.
+	pub fn new(params: Params<BI, CIDP, RClient, Proposer, CS>) -> Self {
+		Collator {
+			create_inherent_data_providers: params.create_inherent_data_providers,
+			block_import: params.block_import,
+			relay_client: params.relay_client,
+			keystore: params.keystore,
+			para_id: params.para_id,
+			proposer: params.proposer,
+			collator_service: params.collator_service,
+			_marker: std::marker::PhantomData,
+		}
+	}
+
+	/// Explicitly creates the inherent data for teyrchain block authoring and overrides
+	/// the timestamp inherent data with the one provided, if any. Additionally allows to specify
+	/// relay parent descendants that can be used to prevent authoring at the tip of the relay
+	/// chain.
+	pub async fn create_inherent_data_with_rp_offset(
+		&self,
+		relay_parent: PHash,
+		validation_data: &PersistedValidationData,
+		parent_hash: Block::Hash,
+		timestamp: impl Into<Option<Timestamp>>,
+		relay_parent_descendants: Option<RelayParentData>,
+		collator_peer_id: PeerId,
+	) -> Result<(TeyrchainInherentData, InherentData), Box<dyn Error + Send + Sync + 'static>> {
+		let paras_inherent_data = TeyrchainInherentDataProvider::create_at(
+			relay_parent,
+			&self.relay_client,
+			validation_data,
+			self.para_id,
+			relay_parent_descendants
+				.map(RelayParentData::into_inherent_descendant_list)
+				.unwrap_or_default(),
+			Vec::new(),
+			collator_peer_id,
+		)
+		.await;
+
+		let paras_inherent_data = match paras_inherent_data {
+			Some(p) => p,
+			None =>
+				return Err(
+					format!("Could not create paras inherent data at {:?}", relay_parent).into()
+				),
+		};
+
+		let mut other_inherent_data = self
+			.create_inherent_data_providers
+			.create_inherent_data_providers(parent_hash, ())
+			.map_err(|e| e as Box<dyn Error + Send + Sync + 'static>)
+			.await?
+			.create_inherent_data()
+			.await
+			.map_err(Box::new)?;
+
+		if let Some(timestamp) = timestamp.into() {
+			other_inherent_data.replace_data(pezsp_timestamp::INHERENT_IDENTIFIER, &timestamp);
+		}
+
+		Ok((paras_inherent_data, other_inherent_data))
+	}
+
+	/// Explicitly creates the inherent data for teyrchain block authoring and overrides
+	/// the timestamp inherent data with the one provided, if any.
+	pub async fn create_inherent_data(
+		&self,
+		relay_parent: PHash,
+		validation_data: &PersistedValidationData,
+		parent_hash: Block::Hash,
+		timestamp: impl Into<Option<Timestamp>>,
+		collator_peer_id: PeerId,
+	) -> Result<(TeyrchainInherentData, InherentData), Box<dyn Error + Send + Sync + 'static>> {
+		self.create_inherent_data_with_rp_offset(
+			relay_parent,
+			validation_data,
+			parent_hash,
+			timestamp,
+			None,
+			collator_peer_id,
+		)
+		.await
+	}
+
+	/// Build and import a teyrchain block on the given parent header, using the given slot claim.
+	pub async fn build_block_and_import(
+		&mut self,
+		parent_header: &Block::Header,
+		slot_claim: &SlotClaim<P::Public>,
+		additional_pre_digest: impl Into<Option<Vec<DigestItem>>>,
+		inherent_data: (TeyrchainInherentData, InherentData),
+		proposal_duration: Duration,
+		max_pov_size: usize,
+	) -> Result<Option<TeyrchainCandidate<Block>>, Box<dyn Error + Send + 'static>> {
+		let mut digest = additional_pre_digest.into().unwrap_or_default();
+		digest.push(slot_claim.pre_digest.clone());
+
+		let maybe_proposal = self
+			.proposer
+			.propose(
+				&parent_header,
+				&inherent_data.0,
+				inherent_data.1,
+				Digest { logs: digest },
+				proposal_duration,
+				Some(max_pov_size),
+			)
+			.await
+			.map_err(|e| Box::new(e) as Box<dyn Error + Send>)?;
+
+		let proposal = match maybe_proposal {
+			None => return Ok(None),
+			Some(p) => p,
+		};
+
+		let sealed_importable = seal::<_, P>(
+			proposal.block,
+			proposal.storage_changes,
+			&slot_claim.author_pub,
+			&self.keystore,
+		)
+		.map_err(|e| e as Box<dyn Error + Send>)?;
+
+		let block = Block::new(
+			sealed_importable.post_header(),
+			sealed_importable
+				.body
+				.as_ref()
+				.expect("body always created with this `propose` fn; qed")
+				.clone(),
+		);
+
+		self.block_import
+			.import_block(sealed_importable)
+			.map_err(|e| Box::new(e) as Box<dyn Error + Send>)
+			.await?;
+
+		Ok(Some(TeyrchainCandidate { block, proof: proposal.proof }))
+	}
+
+	/// Propose, seal, import a block and packaging it into a collation.
+	///
+	/// Provide the slot to build at as well as any other necessary pre-digest logs,
+	/// the inherent data, and the proposal duration and PoV size limits.
+	///
+	/// The Aura pre-digest should not be explicitly provided and is set internally.
+	///
+	/// This does not announce the collation to the teyrchain network or the relay chain.
+	pub async fn collate(
+		&mut self,
+		parent_header: &Block::Header,
+		slot_claim: &SlotClaim<P::Public>,
+		additional_pre_digest: impl Into<Option<Vec<DigestItem>>>,
+		inherent_data: (TeyrchainInherentData, InherentData),
+		proposal_duration: Duration,
+		max_pov_size: usize,
+	) -> Result<Option<(Collation, TeyrchainBlockData<Block>)>, Box<dyn Error + Send + 'static>> {
+		let maybe_candidate = self
+			.build_block_and_import(
+				parent_header,
+				slot_claim,
+				additional_pre_digest,
+				inherent_data,
+				proposal_duration,
+				max_pov_size,
+			)
+			.await?;
+
+		let Some(candidate) = maybe_candidate else { return Ok(None) };
+
+		let hash = candidate.block.header().hash();
+		if let Some((collation, block_data)) =
+			self.collator_service.build_collation(parent_header, hash, candidate)
+		{
+			block_data.log_size_info();
+
+			if let MaybeCompressedPoV::Compressed(ref pov) = collation.proof_of_validity {
+				tracing::info!(
+					target: crate::LOG_TARGET,
+					"Compressed PoV size: {}kb",
+					pov.block_data.0.len() as f64 / 1024f64,
+				);
+			}
+
+			Ok(Some((collation, block_data)))
+		} else {
+			Err(Box::<dyn Error + Send + Sync>::from("Unable to produce collation"))
+		}
+	}
+
+	/// Get the underlying collator service.
+	pub fn collator_service(&self) -> &CS {
+		&self.collator_service
+	}
+}
+
+/// A claim on an Aura slot.
+pub struct SlotClaim<Pub> {
+	author_pub: Pub,
+	pre_digest: DigestItem,
+	slot: Slot,
+	timestamp: Timestamp,
+}
+
+impl<Pub> SlotClaim<Pub> {
+	/// Create a slot-claim from the given author public key, slot, and timestamp.
+	///
+	/// This does not check whether the author actually owns the slot or the timestamp
+	/// falls within the slot.
+	pub fn unchecked<P>(author_pub: Pub, slot: Slot, timestamp: Timestamp) -> Self
+	where
+		P: Pair<Public = Pub>,
+		P::Public: Codec,
+		P::Signature: Codec,
+	{
+		SlotClaim { author_pub, timestamp, pre_digest: aura_internal::pre_digest::<P>(slot), slot }
+	}
+
+	/// Get the author's public key.
+	pub fn author_pub(&self) -> &Pub {
+		&self.author_pub
+	}
+
+	/// Get the Aura pre-digest for this slot.
+	pub fn pre_digest(&self) -> &DigestItem {
+		&self.pre_digest
+	}
+
+	/// Get the slot assigned to this claim.
+	pub fn slot(&self) -> Slot {
+		self.slot
+	}
+
+	/// Get the timestamp corresponding to the relay-chain slot this claim was
+	/// generated against.
+	pub fn timestamp(&self) -> Timestamp {
+		self.timestamp
+	}
+}
+
+/// Attempt to claim a slot derived from the given relay-parent header's slot.
+pub async fn claim_slot<B, C, P>(
+	client: &C,
+	parent_hash: B::Hash,
+	relay_parent_header: &PHeader,
+	slot_duration: SlotDuration,
+	relay_chain_slot_duration: Duration,
+	keystore: &KeystorePtr,
+) -> Result<Option<SlotClaim<P::Public>>, Box<dyn Error>>
+where
+	B: BlockT,
+	C: ProvideRuntimeApi<B> + Send + Sync + 'static,
+	C::Api: AuraApi<B, P::Public>,
+	P: Pair,
+	P::Public: Codec,
+	P::Signature: Codec,
+{
+	// load authorities
+	let authorities = client.runtime_api().authorities(parent_hash).map_err(Box::new)?;
+
+	// Determine the current slot and timestamp based on the relay-parent's.
+	let (slot_now, timestamp) = match consensus_common::relay_slot_and_timestamp(
+		relay_parent_header,
+		relay_chain_slot_duration,
+	) {
+		Some((r_s, t)) => {
+			let our_slot = Slot::from_timestamp(t, slot_duration);
+			tracing::debug!(
+				target: crate::LOG_TARGET,
+				relay_slot = ?r_s,
+				para_slot = ?our_slot,
+				timestamp = ?t,
+				?slot_duration,
+				?relay_chain_slot_duration,
+				"Adjusted relay-chain slot to teyrchain slot"
+			);
+			(our_slot, t)
+		},
+		None => return Ok(None),
+	};
+
+	// Try to claim the slot locally.
+	let author_pub = {
+		let res = aura_internal::claim_slot::<P>(slot_now, &authorities, keystore).await;
+		match res {
+			Some(p) => p,
+			None => return Ok(None),
+		}
+	};
+
+	Ok(Some(SlotClaim::unchecked::<P>(author_pub, slot_now, timestamp)))
+}
+
+/// Seal a block with a signature in the header.
+pub fn seal<B: BlockT, P>(
+	pre_sealed: B,
+	storage_changes: StorageChanges<HashingFor<B>>,
+	author_pub: &P::Public,
+	keystore: &KeystorePtr,
+) -> Result<BlockImportParams<B>, Box<dyn Error + Send + Sync + 'static>>
+where
+	P: Pair,
+	P::Signature: Codec + TryFrom<Vec<u8>>,
+	P::Public: AppPublic,
+{
+	let (pre_header, body) = pre_sealed.deconstruct();
+	let pre_hash = pre_header.hash();
+	let block_number = *pre_header.number();
+
+	// seal the block.
+	let block_import_params = {
+		let seal_digest =
+			aura_internal::seal::<_, P>(&pre_hash, &author_pub, keystore).map_err(Box::new)?;
+		let mut block_import_params = BlockImportParams::new(BlockOrigin::Own, pre_header);
+		block_import_params.post_digests.push(seal_digest);
+		block_import_params.body = Some(body);
+		block_import_params.state_action =
+			StateAction::ApplyChanges(pezsc_consensus::StorageChanges::Changes(storage_changes));
+		block_import_params.fork_choice = Some(ForkChoiceStrategy::LongestChain);
+		block_import_params
+	};
+	let post_hash = block_import_params.post_hash();
+
+	tracing::info!(
+		target: crate::LOG_TARGET,
+		"🔖 Pre-sealed block for proposal at {}. Hash now {:?}, previously {:?}.",
+		block_number,
+		post_hash,
+		pre_hash,
+	);
+
+	Ok(block_import_params)
+}

pezcumulus/client/consensus/aura/src/collators/basic.rs

@@ -0,0 +1,277 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! This provides the option to run a basic relay-chain driven Aura implementation.
+//!
+//! This collator only builds on top of the most recently included block, limiting the
+//! block time to a maximum of two times the relay-chain block time, and requiring the
+//! block to be built and distributed to validators between two relay-chain blocks.
+//!
+//! For more information about AuRa, the Bizinikiwi crate should be checked.
+
+use codec::{Codec, Decode};
+use cumulus_client_collator::{
+	relay_chain_driven::CollationRequest, service::ServiceInterface as CollatorServiceInterface,
+};
+use cumulus_client_consensus_common::TeyrchainBlockImportMarker;
+use cumulus_client_consensus_proposer::ProposerInterface;
+use cumulus_primitives_core::{relay_chain::BlockId as RBlockId, CollectCollationInfo};
+use cumulus_relay_chain_interface::RelayChainInterface;
+
+use pezkuwi_node_primitives::CollationResult;
+use pezkuwi_overseer::Handle as OverseerHandle;
+use pezkuwi_primitives::{CollatorPair, Id as ParaId, ValidationCode};
+
+use futures::{channel::mpsc::Receiver, prelude::*};
+use pezsc_client_api::{backend::AuxStore, BlockBackend, BlockOf};
+use pezsc_consensus::BlockImport;
+use pezsc_network_types::PeerId;
+use pezsp_api::{CallApiAt, ProvideRuntimeApi};
+use pezsp_application_crypto::AppPublic;
+use pezsp_blockchain::HeaderBackend;
+use pezsp_consensus_aura::AuraApi;
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_keystore::KeystorePtr;
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT, Member};
+use pezsp_state_machine::Backend as _;
+use std::{sync::Arc, time::Duration};
+
+use crate::collator as collator_util;
+
+/// Parameters for [`run`].
+pub struct Params<BI, CIDP, Client, RClient, Proposer, CS> {
+	/// Inherent data providers. Only non-consensus inherent data should be provided, i.e.
+	/// the timestamp, slot, and paras inherents should be omitted, as they are set by this
+	/// collator.
+	pub create_inherent_data_providers: CIDP,
+	/// Used to actually import blocks.
+	pub block_import: BI,
+	/// The underlying para client.
+	pub para_client: Arc<Client>,
+	/// A handle to the relay-chain client.
+	pub relay_client: RClient,
+	/// The underlying keystore, which should contain Aura consensus keys.
+	pub keystore: KeystorePtr,
+	/// The collator key used to sign collations before submitting to validators.
+	pub collator_key: CollatorPair,
+	/// The collator network peer id.
+	pub collator_peer_id: PeerId,
+	/// The para's ID.
+	pub para_id: ParaId,
+	/// A handle to the relay-chain client's "Overseer" or task orchestrator.
+	pub overseer_handle: OverseerHandle,
+	/// The length of slots in the relay chain.
+	pub relay_chain_slot_duration: Duration,
+	/// The underlying block proposer this should call into.
+	pub proposer: Proposer,
+	/// The generic collator service used to plug into this consensus engine.
+	pub collator_service: CS,
+	/// The amount of time to spend authoring each block.
+	pub authoring_duration: Duration,
+	/// Receiver for collation requests. If `None`, Aura consensus will establish a new receiver.
+	/// Should be used when a chain migrates from a different consensus algorithm and was already
+	/// processing collation requests before initializing Aura.
+	pub collation_request_receiver: Option<Receiver<CollationRequest>>,
+}
+
+/// Run bare Aura consensus as a relay-chain-driven collator.
+pub fn run<Block, P, BI, CIDP, Client, RClient, Proposer, CS>(
+	params: Params<BI, CIDP, Client, RClient, Proposer, CS>,
+) -> impl Future<Output = ()> + Send + 'static
+where
+	Block: BlockT + Send,
+	Client: ProvideRuntimeApi<Block>
+		+ BlockOf
+		+ AuxStore
+		+ HeaderBackend<Block>
+		+ BlockBackend<Block>
+		+ CallApiAt<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	Client::Api: AuraApi<Block, P::Public> + CollectCollationInfo<Block>,
+	RClient: RelayChainInterface + Send + Clone + 'static,
+	CIDP: CreateInherentDataProviders<Block, ()> + Send + 'static,
+	CIDP::InherentDataProviders: Send,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block> + Send + Sync + 'static,
+	CS: CollatorServiceInterface<Block> + Send + Sync + 'static,
+	P: Pair,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	async move {
+		let mut collation_requests = match params.collation_request_receiver {
+			Some(receiver) => receiver,
+			None =>
+				cumulus_client_collator::relay_chain_driven::init(
+					params.collator_key,
+					params.para_id,
+					params.overseer_handle,
+				)
+				.await,
+		};
+
+		let mut collator = {
+			let params = collator_util::Params {
+				create_inherent_data_providers: params.create_inherent_data_providers,
+				block_import: params.block_import,
+				relay_client: params.relay_client.clone(),
+				keystore: params.keystore.clone(),
+				collator_peer_id: params.collator_peer_id,
+				para_id: params.para_id,
+				proposer: params.proposer,
+				collator_service: params.collator_service,
+			};
+
+			collator_util::Collator::<Block, P, _, _, _, _, _>::new(params)
+		};
+
+		let mut last_processed_slot = 0;
+		let mut last_relay_chain_block = Default::default();
+
+		while let Some(request) = collation_requests.next().await {
+			macro_rules! reject_with_error {
+				($err:expr) => {{
+					request.complete(None);
+					tracing::error!(target: crate::LOG_TARGET, err = ?{ $err });
+					continue;
+				}};
+			}
+
+			macro_rules! try_request {
+				($x:expr) => {{
+					match $x {
+						Ok(x) => x,
+						Err(e) => reject_with_error!(e),
+					}
+				}};
+			}
+
+			let validation_data = request.persisted_validation_data();
+
+			let parent_header =
+				try_request!(Block::Header::decode(&mut &validation_data.parent_head.0[..]));
+
+			let parent_hash = parent_header.hash();
+
+			if !collator.collator_service().check_block_status(parent_hash, &parent_header) {
+				continue;
+			}
+
+			let Ok(Some(code)) =
+				params.para_client.state_at(parent_hash).map_err(drop).and_then(|s| {
+					s.storage(&pezsp_core::storage::well_known_keys::CODE).map_err(drop)
+				})
+			else {
+				continue;
+			};
+
+			super::check_validation_code_or_log(
+				&ValidationCode::from(code).hash(),
+				params.para_id,
+				&params.relay_client,
+				*request.relay_parent(),
+			)
+			.await;
+
+			let relay_parent_header =
+				match params.relay_client.header(RBlockId::hash(*request.relay_parent())).await {
+					Err(e) => reject_with_error!(e),
+					Ok(None) => continue, // sanity: would be inconsistent to get `None` here
+					Ok(Some(h)) => h,
+				};
+
+			let slot_duration = match params.para_client.runtime_api().slot_duration(parent_hash) {
+				Ok(d) => d,
+				Err(e) => reject_with_error!(e),
+			};
+
+			let claim = match collator_util::claim_slot::<_, _, P>(
+				&*params.para_client,
+				parent_hash,
+				&relay_parent_header,
+				slot_duration,
+				params.relay_chain_slot_duration,
+				&params.keystore,
+			)
+			.await
+			{
+				Ok(None) => continue,
+				Ok(Some(c)) => c,
+				Err(e) => reject_with_error!(e),
+			};
+
+			// With async backing this function will be called every relay chain block.
+			//
+			// Most teyrchains currently run with 12 seconds slots and thus, they would try to
+			// produce multiple blocks per slot which very likely would fail on chain. Thus, we have
+			// this "hack" to only produce one block per slot per relay chain fork.
+			//
+			// With https://github.com/pezkuwichain/pezkuwi-sdk/issues/127 this implementation will be
+			// obsolete and also the underlying issue will be fixed.
+			if last_processed_slot >= *claim.slot() &&
+				last_relay_chain_block < *relay_parent_header.number()
+			{
+				continue;
+			}
+
+			let (teyrchain_inherent_data, other_inherent_data) = try_request!(
+				collator
+					.create_inherent_data(
+						*request.relay_parent(),
+						&validation_data,
+						parent_hash,
+						claim.timestamp(),
+						params.collator_peer_id,
+					)
+					.await
+			);
+
+			let allowed_pov_size = (validation_data.max_pov_size / 2) as usize;
+
+			let maybe_collation = try_request!(
+				collator
+					.collate(
+						&parent_header,
+						&claim,
+						None,
+						(teyrchain_inherent_data, other_inherent_data),
+						params.authoring_duration,
+						allowed_pov_size,
+					)
+					.await
+			);
+
+			if let Some((collation, block_data)) = maybe_collation {
+				let Some(block_hash) = block_data.blocks().first().map(|b| b.hash()) else {
+					continue;
+				};
+				let result_sender =
+					Some(collator.collator_service().announce_with_barrier(block_hash));
+				request.complete(Some(CollationResult { collation, result_sender }));
+			} else {
+				request.complete(None);
+				tracing::debug!(target: crate::LOG_TARGET, "No block proposal");
+			}
+
+			last_processed_slot = *claim.slot();
+			last_relay_chain_block = *relay_parent_header.number();
+		}
+	}
+}

pezcumulus/client/consensus/aura/src/collators/lookahead.rs

@@ -0,0 +1,511 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! A collator for Aura that looks ahead of the most recently included teyrchain block
+//! when determining what to build upon.
+//!
+//! This collator also builds additional blocks when the maximum backlog is not saturated.
+//! The size of the backlog is determined by invoking a runtime API. If that runtime API
+//! is not supported, this assumes a maximum backlog size of 1.
+//!
+//! This takes more advantage of asynchronous backing, though not complete advantage.
+//! When the backlog is not saturated, this approach lets the backlog temporarily 'catch up'
+//! with periods of higher throughput. When the backlog is saturated, we typically
+//! fall back to the limited cadence of a single teyrchain block per relay-chain block.
+//!
+//! Despite this, the fact that there is a backlog at all allows us to spend more time
+//! building the block, as there is some buffer before it can get posted to the relay-chain.
+//! The main limitation is block propagation time - i.e. the new blocks created by an author
+//! must be propagated to the next author before their turn.
+
+use codec::{Codec, Encode};
+use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_client_consensus_common::{self as consensus_common, TeyrchainBlockImportMarker};
+use cumulus_client_consensus_proposer::ProposerInterface;
+use cumulus_primitives_aura::AuraUnincludedSegmentApi;
+use cumulus_primitives_core::{CollectCollationInfo, PersistedValidationData};
+use cumulus_relay_chain_interface::RelayChainInterface;
+
+use pezkuwi_node_primitives::SubmitCollationParams;
+use pezkuwi_node_subsystem::messages::CollationGenerationMessage;
+use pezkuwi_overseer::Handle as OverseerHandle;
+use pezkuwi_primitives::{CollatorPair, Id as ParaId, OccupiedCoreAssumption};
+
+use crate::{
+	collator as collator_util,
+	collators::{claim_queue_at, BackingGroupConnectionHelper},
+	export_pov_to_path,
+};
+use futures::prelude::*;
+use pezsc_client_api::{backend::AuxStore, BlockBackend, BlockOf};
+use pezsc_consensus::BlockImport;
+use pezsc_network_types::PeerId;
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_application_crypto::AppPublic;
+use pezsp_blockchain::HeaderBackend;
+use pezsp_consensus_aura::{AuraApi, Slot};
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_keystore::KeystorePtr;
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT, Member};
+use pezsp_timestamp::Timestamp;
+use std::{path::PathBuf, sync::Arc, time::Duration};
+
+/// Parameters for [`run`].
+pub struct Params<BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS> {
+	/// Inherent data providers. Only non-consensus inherent data should be provided, i.e.
+	/// the timestamp, slot, and paras inherents should be omitted, as they are set by this
+	/// collator.
+	pub create_inherent_data_providers: CIDP,
+	/// Used to actually import blocks.
+	pub block_import: BI,
+	/// The underlying para client.
+	pub para_client: Arc<Client>,
+	/// The para client's backend, used to access the database.
+	pub para_backend: Arc<Backend>,
+	/// A handle to the relay-chain client.
+	pub relay_client: RClient,
+	/// A validation code hash provider, used to get the current validation code hash.
+	pub code_hash_provider: CHP,
+	/// The underlying keystore, which should contain Aura consensus keys.
+	pub keystore: KeystorePtr,
+	/// The collator key used to sign collations before submitting to validators.
+	pub collator_key: CollatorPair,
+	/// The collator network peer id.
+	pub collator_peer_id: PeerId,
+	/// The para's ID.
+	pub para_id: ParaId,
+	/// A handle to the relay-chain client's "Overseer" or task orchestrator.
+	pub overseer_handle: OverseerHandle,
+	/// The length of slots in the relay chain.
+	pub relay_chain_slot_duration: Duration,
+	/// The underlying block proposer this should call into.
+	pub proposer: Proposer,
+	/// The generic collator service used to plug into this consensus engine.
+	pub collator_service: CS,
+	/// The amount of time to spend authoring each block.
+	pub authoring_duration: Duration,
+	/// Whether we should reinitialize the collator config (i.e. we are transitioning to aura).
+	pub reinitialize: bool,
+	/// The maximum percentage of the maximum PoV size that the collator can use.
+	/// It will be removed once <https://github.com/pezkuwichain/pezkuwi-sdk/issues/23> is fixed.
+	pub max_pov_percentage: Option<u32>,
+}
+
+/// Get the current teyrchain slot from a given block hash.
+///
+/// Returns the teyrchain slot, relay chain slot, and timestamp.
+fn get_teyrchain_slot<Block, Client, P>(
+	para_client: &Client,
+	block_hash: Block::Hash,
+	relay_parent_header: &pezkuwi_primitives::Header,
+	relay_chain_slot_duration: Duration,
+) -> Option<(Slot, Slot, Timestamp)>
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>,
+	Client::Api: AuraApi<Block, P>,
+	P: Codec,
+{
+	let slot_duration =
+		match pezsc_consensus_aura::standalone::slot_duration_at(para_client, block_hash) {
+			Ok(sd) => sd,
+			Err(err) => {
+				tracing::error!(target: crate::LOG_TARGET, ?err, "Failed to acquire teyrchain slot duration");
+				return None;
+			},
+		};
+
+	tracing::debug!(target: crate::LOG_TARGET, ?slot_duration, ?block_hash, "Teyrchain slot duration acquired");
+
+	let (relay_slot, timestamp) =
+		consensus_common::relay_slot_and_timestamp(relay_parent_header, relay_chain_slot_duration)?;
+
+	let slot_now = Slot::from_timestamp(timestamp, slot_duration);
+
+	tracing::debug!(
+		target: crate::LOG_TARGET,
+		?relay_slot,
+		para_slot = ?slot_now,
+		?timestamp,
+		?slot_duration,
+		?relay_chain_slot_duration,
+		"Adjusted relay-chain slot to teyrchain slot"
+	);
+
+	Some((slot_now, relay_slot, timestamp))
+}
+
+/// Run async-backing-friendly Aura.
+pub fn run<Block, P, BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS>(
+	params: Params<BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS>,
+) -> impl Future<Output = ()> + Send + 'static
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>
+		+ BlockOf
+		+ AuxStore
+		+ HeaderBackend<Block>
+		+ BlockBackend<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	Client::Api:
+		AuraApi<Block, P::Public> + CollectCollationInfo<Block> + AuraUnincludedSegmentApi<Block>,
+	Backend: pezsc_client_api::Backend<Block> + 'static,
+	RClient: RelayChainInterface + Clone + 'static,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+	CIDP::InherentDataProviders: Send,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block> + Send + Sync + 'static,
+	CS: CollatorServiceInterface<Block> + Send + Sync + 'static,
+	CHP: consensus_common::ValidationCodeHashProvider<Block::Hash> + Send + 'static,
+	P: Pair + Send + Sync + 'static,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	run_with_export::<_, P, _, _, _, _, _, _, _, _>(ParamsWithExport { params, export_pov: None })
+}
+
+/// Parameters for [`run_with_export`].
+pub struct ParamsWithExport<BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS> {
+	/// The parameters.
+	pub params: Params<BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS>,
+
+	/// When set, the collator will export every produced `POV` to this folder.
+	pub export_pov: Option<PathBuf>,
+}
+
+/// Run async-backing-friendly Aura.
+///
+/// This is exactly the same as [`run`], but it supports the optional export of each produced `POV`
+/// to the file system.
+pub fn run_with_export<Block, P, BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS>(
+	ParamsWithExport { mut params, export_pov }: ParamsWithExport<
+		BI,
+		CIDP,
+		Client,
+		Backend,
+		RClient,
+		CHP,
+		Proposer,
+		CS,
+	>,
+) -> impl Future<Output = ()> + Send + 'static
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>
+		+ BlockOf
+		+ AuxStore
+		+ HeaderBackend<Block>
+		+ BlockBackend<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	Client::Api:
+		AuraApi<Block, P::Public> + CollectCollationInfo<Block> + AuraUnincludedSegmentApi<Block>,
+	Backend: pezsc_client_api::Backend<Block> + 'static,
+	RClient: RelayChainInterface + Clone + 'static,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+	CIDP::InherentDataProviders: Send,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block> + Send + Sync + 'static,
+	CS: CollatorServiceInterface<Block> + Send + Sync + 'static,
+	CHP: consensus_common::ValidationCodeHashProvider<Block::Hash> + Send + 'static,
+	P: Pair + Send + Sync + 'static,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	async move {
+		cumulus_client_collator::initialize_collator_subsystems(
+			&mut params.overseer_handle,
+			params.collator_key,
+			params.para_id,
+			params.reinitialize,
+		)
+		.await;
+
+		let mut import_notifications = match params.relay_client.import_notification_stream().await
+		{
+			Ok(s) => s,
+			Err(err) => {
+				tracing::error!(
+					target: crate::LOG_TARGET,
+					?err,
+					"Failed to initialize consensus: no relay chain import notification stream"
+				);
+
+				return;
+			},
+		};
+
+		let mut collator = {
+			let params = collator_util::Params {
+				create_inherent_data_providers: params.create_inherent_data_providers,
+				block_import: params.block_import,
+				relay_client: params.relay_client.clone(),
+				keystore: params.keystore.clone(),
+				collator_peer_id: params.collator_peer_id,
+				para_id: params.para_id,
+				proposer: params.proposer,
+				collator_service: params.collator_service,
+			};
+
+			collator_util::Collator::<Block, P, _, _, _, _, _>::new(params)
+		};
+
+		let mut connection_helper = BackingGroupConnectionHelper::new(
+			params.keystore.clone(),
+			params.overseer_handle.clone(),
+		);
+
+		while let Some(relay_parent_header) = import_notifications.next().await {
+			let relay_parent = relay_parent_header.hash();
+
+			let Some(core_index) = claim_queue_at(relay_parent, &mut params.relay_client)
+				.await
+				.iter_claims_at_depth_for_para(0, params.para_id)
+				.next()
+			else {
+				tracing::trace!(
+					target: crate::LOG_TARGET,
+					?relay_parent,
+					?params.para_id,
+					"Para is not scheduled on any core, skipping import notification",
+				);
+
+				continue;
+			};
+
+			let max_pov_size = match params
+				.relay_client
+				.persisted_validation_data(
+					relay_parent,
+					params.para_id,
+					OccupiedCoreAssumption::Included,
+				)
+				.await
+			{
+				Ok(None) => continue,
+				Ok(Some(pvd)) => pvd.max_pov_size,
+				Err(err) => {
+					tracing::error!(target: crate::LOG_TARGET, ?err, "Failed to gather information from relay-client");
+					continue;
+				},
+			};
+
+			let (included_block, initial_parent) = match crate::collators::find_parent(
+				relay_parent,
+				params.para_id,
+				&*params.para_backend,
+				&params.relay_client,
+			)
+			.await
+			{
+				Some(value) => value,
+				None => continue,
+			};
+
+			let para_client = &*params.para_client;
+			let keystore = &params.keystore;
+			let can_build_upon = |block_hash| {
+				let (slot_now, relay_slot, timestamp) = get_teyrchain_slot::<_, _, P::Public>(
+					para_client,
+					block_hash,
+					&relay_parent_header,
+					params.relay_chain_slot_duration,
+				)?;
+
+				Some(super::can_build_upon::<_, _, P>(
+					slot_now,
+					relay_slot,
+					timestamp,
+					block_hash,
+					included_block.hash(),
+					para_client,
+					&keystore,
+				))
+			};
+
+			// Build in a loop until not allowed. Note that the authorities can change
+			// at any block, so we need to re-claim our slot every time.
+			let mut parent_hash = initial_parent.hash;
+			let mut parent_header = initial_parent.header;
+			let overseer_handle = &mut params.overseer_handle;
+
+			// Do not try to build upon an unknown, pruned or bad block
+			if !collator.collator_service().check_block_status(parent_hash, &parent_header) {
+				continue;
+			}
+
+			// Trigger pre-conect to backing groups if necessary.
+			if let (Some((slot_now, _relay_slot, _timestamp)), Ok(authorities)) = (
+				get_teyrchain_slot::<_, _, P::Public>(
+					para_client,
+					parent_hash,
+					&relay_parent_header,
+					params.relay_chain_slot_duration,
+				),
+				para_client.runtime_api().authorities(parent_hash),
+			) {
+				connection_helper.update::<P>(slot_now, &authorities).await;
+			}
+
+			// This needs to change to support elastic scaling, but for continuously
+			// scheduled chains this ensures that the backlog will grow steadily.
+			for n_built in 0..2 {
+				let slot_claim = match can_build_upon(parent_hash) {
+					Some(fut) => match fut.await {
+						None => break,
+						Some(c) => c,
+					},
+					None => break,
+				};
+
+				tracing::debug!(
+					target: crate::LOG_TARGET,
+					?relay_parent,
+					unincluded_segment_len = initial_parent.depth + n_built,
+					"Slot claimed. Building"
+				);
+
+				let validation_data = PersistedValidationData {
+					parent_head: parent_header.encode().into(),
+					relay_parent_number: *relay_parent_header.number(),
+					relay_parent_storage_root: *relay_parent_header.state_root(),
+					max_pov_size,
+				};
+
+				// Build and announce collations recursively until
+				// `can_build_upon` fails or building a collation fails.
+				let (teyrchain_inherent_data, other_inherent_data) = match collator
+					.create_inherent_data(
+						relay_parent,
+						&validation_data,
+						parent_hash,
+						slot_claim.timestamp(),
+						params.collator_peer_id,
+					)
+					.await
+				{
+					Err(err) => {
+						tracing::error!(target: crate::LOG_TARGET, ?err);
+						break;
+					},
+					Ok(x) => x,
+				};
+
+				let Some(validation_code_hash) =
+					params.code_hash_provider.code_hash_at(parent_hash)
+				else {
+					tracing::error!(target: crate::LOG_TARGET, ?parent_hash, "Could not fetch validation code hash");
+					break;
+				};
+
+				super::check_validation_code_or_log(
+					&validation_code_hash,
+					params.para_id,
+					&params.relay_client,
+					relay_parent,
+				)
+				.await;
+
+				let allowed_pov_size = if let Some(max_pov_percentage) = params.max_pov_percentage {
+					validation_data.max_pov_size * max_pov_percentage / 100
+				} else {
+					// Set the block limit to 85% of the maximum PoV size.
+					//
+					// Once https://github.com/pezkuwichain/pezkuwi-sdk/issues/23 issue is
+					// fixed, the reservation should be removed.
+					validation_data.max_pov_size * 85 / 100
+				} as usize;
+
+				match collator
+					.collate(
+						&parent_header,
+						&slot_claim,
+						None,
+						(teyrchain_inherent_data, other_inherent_data),
+						params.authoring_duration,
+						allowed_pov_size,
+					)
+					.await
+				{
+					Ok(Some((collation, block_data))) => {
+						let Some(new_block_header) =
+							block_data.blocks().first().map(|b| b.header().clone())
+						else {
+							tracing::error!(target: crate::LOG_TARGET,  "Produced PoV doesn't contain any blocks");
+							break;
+						};
+
+						let new_block_hash = new_block_header.hash();
+
+						// Here we are assuming that the import logic protects against equivocations
+						// and provides sybil-resistance, as it should.
+						collator.collator_service().announce_block(new_block_hash, None);
+
+						if let Some(ref export_pov) = export_pov {
+							export_pov_to_path::<Block>(
+								export_pov.clone(),
+								collation.proof_of_validity.clone().into_compressed(),
+								new_block_hash,
+								*new_block_header.number(),
+								parent_header.clone(),
+								*relay_parent_header.state_root(),
+								*relay_parent_header.number(),
+								validation_data.max_pov_size,
+							);
+						}
+
+						// Send a submit-collation message to the collation generation subsystem,
+						// which then distributes this to validators.
+						//
+						// Here we are assuming that the leaf is imported, as we've gotten an
+						// import notification.
+						overseer_handle
+							.send_msg(
+								CollationGenerationMessage::SubmitCollation(
+									SubmitCollationParams {
+										relay_parent,
+										collation,
+										parent_head: parent_header.encode().into(),
+										validation_code_hash,
+										result_sender: None,
+										core_index,
+									},
+								),
+								"SubmitCollation",
+							)
+							.await;
+
+						parent_hash = new_block_hash;
+						parent_header = new_block_header;
+					},
+					Ok(None) => {
+						tracing::debug!(target: crate::LOG_TARGET, "No block proposal");
+						break;
+					},
+					Err(err) => {
+						tracing::error!(target: crate::LOG_TARGET, ?err);
+						break;
+					},
+				}
+			}
+		}
+	}
+}

pezcumulus/client/consensus/aura/src/collators/mod.rs

@@ -0,0 +1,709 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! Stock, pure Aura collators.
+//!
+//! This includes the [`basic`] collator, which only builds on top of the most recently
+//! included teyrchain block, as well as the [`lookahead`] collator, which prospectively
+//! builds on teyrchain blocks which have not yet been included in the relay chain.
+
+use crate::collator::SlotClaim;
+use codec::Codec;
+use cumulus_client_consensus_common::{self as consensus_common, ParentSearchParams};
+use cumulus_primitives_aura::{AuraUnincludedSegmentApi, Slot};
+use cumulus_primitives_core::{relay_chain::Header as RelayHeader, BlockT};
+use cumulus_relay_chain_interface::{OverseerHandle, RelayChainInterface};
+use pezkuwi_node_subsystem::messages::{CollatorProtocolMessage, RuntimeApiRequest};
+use pezkuwi_node_subsystem_util::runtime::ClaimQueueSnapshot;
+use pezkuwi_primitives::{
+	Hash as RelayHash, Id as ParaId, OccupiedCoreAssumption, ValidationCodeHash,
+	DEFAULT_SCHEDULING_LOOKAHEAD,
+};
+use pezsc_consensus_aura::{standalone as aura_internal, AuraApi};
+use pezsp_api::{ApiExt, ProvideRuntimeApi, RuntimeApiInfo};
+use pezsp_core::Pair;
+use pezsp_keystore::KeystorePtr;
+use pezsp_timestamp::Timestamp;
+
+pub mod basic;
+pub mod lookahead;
+pub mod slot_based;
+
+// This is an arbitrary value which is guaranteed to exceed the required depth for 500ms blocks
+// built with a relay parent offset of 1. It must be larger than the unincluded segment capacity.
+//
+// The formula we use to compute the capacity of the unincluded segment in the teyrchain runtime
+// is:
+// UNINCLUDED_SEGMENT_CAPACITY = (2 + RELAY_PARENT_OFFSET) * BLOCK_PROCESSING_VELOCITY + 1.
+//
+// Since we only search for parent blocks which have already been imported,
+// we can guarantee that all imported blocks respect the unincluded segment
+// rules specified by the teyrchain's runtime and thus will never be too deep. This is just an extra
+// sanity check.
+const PARENT_SEARCH_DEPTH: usize = 40;
+
+// Helper to pre-connect to the backing group we got assigned to and keep the connection
+// open until backing group changes or own slot ends.
+struct BackingGroupConnectionHelper {
+	keystore: pezsp_keystore::KeystorePtr,
+	overseer_handle: OverseerHandle,
+	our_slot: Option<Slot>,
+}
+
+impl BackingGroupConnectionHelper {
+	pub fn new(keystore: pezsp_keystore::KeystorePtr, overseer_handle: OverseerHandle) -> Self {
+		Self { keystore, overseer_handle, our_slot: None }
+	}
+
+	async fn send_subsystem_message(&mut self, message: CollatorProtocolMessage) {
+		self.overseer_handle.send_msg(message, "BackingGroupConnectionHelper").await;
+	}
+
+	/// Update the current slot and initiate connections to backing groups if needed.
+	pub async fn update<P>(&mut self, current_slot: Slot, authorities: &[P::Public])
+	where
+		P: pezsp_core::Pair + Send + Sync,
+		P::Public: Codec,
+	{
+		if Some(current_slot) <= self.our_slot {
+			// Current slot or next slot is ours.
+			// We already sent pre-connect message, no need to proceed further.
+			return;
+		}
+
+		let next_slot = current_slot + 1;
+		let next_slot_is_ours =
+			aura_internal::claim_slot::<P>(next_slot, authorities, &self.keystore)
+				.await
+				.is_some();
+
+		if next_slot_is_ours {
+			// Only send message if we were not connected. This avoids sending duplicate messages
+			// when running with a single collator.
+			if self.our_slot.is_none() {
+				// Next slot is ours, send connect message.
+				tracing::debug!(target: crate::LOG_TARGET, "Our slot {} is next, connecting to backing groups", next_slot);
+				self.send_subsystem_message(CollatorProtocolMessage::ConnectToBackingGroups)
+					.await;
+			}
+			self.our_slot = Some(next_slot);
+		} else if self.our_slot.take().is_some() {
+			// Next slot is not ours, send disconnect only if we had a slot before.
+			tracing::debug!(target: crate::LOG_TARGET, "Current slot = {}, disconnecting from backing groups", current_slot);
+			self.send_subsystem_message(CollatorProtocolMessage::DisconnectFromBackingGroups)
+				.await;
+		}
+	}
+}
+
+/// Check the `local_validation_code_hash` against the validation code hash in the relay chain
+/// state.
+///
+/// If the code hashes do not match, it prints a warning.
+async fn check_validation_code_or_log(
+	local_validation_code_hash: &ValidationCodeHash,
+	para_id: ParaId,
+	relay_client: &impl RelayChainInterface,
+	relay_parent: RelayHash,
+) {
+	let state_validation_code_hash = match relay_client
+		.validation_code_hash(relay_parent, para_id, OccupiedCoreAssumption::Included)
+		.await
+	{
+		Ok(hash) => hash,
+		Err(error) => {
+			tracing::debug!(
+				target: super::LOG_TARGET,
+				%error,
+				?relay_parent,
+				%para_id,
+				"Failed to fetch validation code hash",
+			);
+			return;
+		},
+	};
+
+	match state_validation_code_hash {
+		Some(state) =>
+			if state != *local_validation_code_hash {
+				tracing::warn!(
+					target: super::LOG_TARGET,
+					%para_id,
+					?relay_parent,
+					?local_validation_code_hash,
+					relay_validation_code_hash = ?state,
+					"Teyrchain code doesn't match validation code stored in the relay chain state.",
+				);
+			},
+		None => {
+			tracing::warn!(
+				target: super::LOG_TARGET,
+				%para_id,
+				?relay_parent,
+				"Could not find validation code for teyrchain in the relay chain state.",
+			);
+		},
+	}
+}
+
+/// Fetch scheduling lookahead at given relay parent.
+async fn scheduling_lookahead(
+	relay_parent: RelayHash,
+	relay_client: &impl RelayChainInterface,
+) -> Option<u32> {
+	let runtime_api_version = relay_client
+		.version(relay_parent)
+		.await
+		.map_err(|e| {
+			tracing::error!(
+				target: super::LOG_TARGET,
+				error = ?e,
+				"Failed to fetch relay chain runtime version.",
+			)
+		})
+		.ok()?;
+
+	let teyrchain_host_runtime_api_version = runtime_api_version
+		.api_version(
+			&<dyn pezkuwi_primitives::runtime_api::TeyrchainHost<pezkuwi_primitives::Block>>::ID,
+		)
+		.unwrap_or_default();
+
+	if teyrchain_host_runtime_api_version <
+		RuntimeApiRequest::SCHEDULING_LOOKAHEAD_RUNTIME_REQUIREMENT
+	{
+		return None;
+	}
+
+	match relay_client.scheduling_lookahead(relay_parent).await {
+		Ok(scheduling_lookahead) => Some(scheduling_lookahead),
+		Err(err) => {
+			tracing::error!(
+				target: crate::LOG_TARGET,
+				?err,
+				?relay_parent,
+				"Failed to fetch scheduling lookahead from relay chain",
+			);
+			None
+		},
+	}
+}
+
+// Returns the claim queue at the given relay parent.
+async fn claim_queue_at(
+	relay_parent: RelayHash,
+	relay_client: &impl RelayChainInterface,
+) -> ClaimQueueSnapshot {
+	// Get `ClaimQueue` from runtime
+	match relay_client.claim_queue(relay_parent).await {
+		Ok(claim_queue) => claim_queue.into(),
+		Err(error) => {
+			tracing::error!(
+				target: crate::LOG_TARGET,
+				?error,
+				?relay_parent,
+				"Failed to query claim queue runtime API",
+			);
+			Default::default()
+		},
+	}
+}
+
+// Checks if we own the slot at the given block and whether there
+// is space in the unincluded segment.
+async fn can_build_upon<Block: BlockT, Client, P>(
+	para_slot: Slot,
+	relay_slot: Slot,
+	timestamp: Timestamp,
+	parent_hash: Block::Hash,
+	included_block: Block::Hash,
+	client: &Client,
+	keystore: &KeystorePtr,
+) -> Option<SlotClaim<P::Public>>
+where
+	Client: ProvideRuntimeApi<Block>,
+	Client::Api: AuraApi<Block, P::Public> + AuraUnincludedSegmentApi<Block> + ApiExt<Block>,
+	P: Pair,
+	P::Public: Codec,
+	P::Signature: Codec,
+{
+	let runtime_api = client.runtime_api();
+	let authorities = runtime_api.authorities(parent_hash).ok()?;
+	let author_pub = aura_internal::claim_slot::<P>(para_slot, &authorities, keystore).await?;
+
+	// This function is typically called when we want to build block N. At that point, the
+	// unincluded segment in the runtime is unaware of the hash of block N-1. If the unincluded
+	// segment in the runtime is full, but block N-1 is the included block, the unincluded segment
+	// should have length 0 and we can build. Since the hash is not available to the runtime
+	// however, we need this extra check here.
+	if parent_hash == included_block {
+		return Some(SlotClaim::unchecked::<P>(author_pub, para_slot, timestamp));
+	}
+
+	let api_version = runtime_api
+		.api_version::<dyn AuraUnincludedSegmentApi<Block>>(parent_hash)
+		.ok()
+		.flatten()?;
+
+	let slot = if api_version > 1 { relay_slot } else { para_slot };
+
+	runtime_api
+		.can_build_upon(parent_hash, included_block, slot)
+		.ok()?
+		.then(|| SlotClaim::unchecked::<P>(author_pub, para_slot, timestamp))
+}
+
+/// Use [`cumulus_client_consensus_common::find_potential_parents`] to find teyrchain blocks that
+/// we can build on. Once a list of potential parents is retrieved, return the last one of the
+/// longest chain.
+async fn find_parent<Block>(
+	relay_parent: RelayHash,
+	para_id: ParaId,
+	para_backend: &impl pezsc_client_api::Backend<Block>,
+	relay_client: &impl RelayChainInterface,
+) -> Option<(<Block as BlockT>::Header, consensus_common::PotentialParent<Block>)>
+where
+	Block: BlockT,
+{
+	let parent_search_params = ParentSearchParams {
+		relay_parent,
+		para_id,
+		ancestry_lookback: scheduling_lookahead(relay_parent, relay_client)
+			.await
+			.unwrap_or(DEFAULT_SCHEDULING_LOOKAHEAD)
+			.saturating_sub(1) as usize,
+		max_depth: PARENT_SEARCH_DEPTH,
+		ignore_alternative_branches: true,
+	};
+
+	let potential_parents = cumulus_client_consensus_common::find_potential_parents::<Block>(
+		parent_search_params,
+		para_backend,
+		relay_client,
+	)
+	.await;
+
+	let potential_parents = match potential_parents {
+		Err(e) => {
+			tracing::error!(
+				target: crate::LOG_TARGET,
+				?relay_parent,
+				err = ?e,
+				"Could not fetch potential parents to build upon"
+			);
+
+			return None;
+		},
+		Ok(x) => x,
+	};
+
+	let included_block = potential_parents.iter().find(|x| x.depth == 0)?.header.clone();
+	potential_parents
+		.into_iter()
+		.max_by_key(|a| a.depth)
+		.map(|parent| (included_block, parent))
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::collators::{can_build_upon, BackingGroupConnectionHelper};
+	use codec::Encode;
+	use cumulus_primitives_aura::Slot;
+	use cumulus_primitives_core::BlockT;
+	use cumulus_relay_chain_interface::PHash;
+	use cumulus_test_client::{
+		runtime::{Block, Hash},
+		Client, DefaultTestClientBuilderExt, InitBlockBuilder, TestClientBuilder,
+		TestClientBuilderExt,
+	};
+	use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+	use futures::StreamExt;
+	use pezkuwi_overseer::{Event, Handle};
+	use pezkuwi_primitives::HeadData;
+	use pezsc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
+	use pezsp_consensus::BlockOrigin;
+	use pezsp_keystore::{Keystore, KeystorePtr};
+	use pezsp_timestamp::Timestamp;
+	use std::sync::{Arc, Mutex};
+
+	async fn import_block<I: BlockImport<Block>>(
+		importer: &I,
+		block: Block,
+		origin: BlockOrigin,
+		import_as_best: bool,
+	) {
+		let (header, body) = block.deconstruct();
+
+		let mut block_import_params = BlockImportParams::new(origin, header);
+		block_import_params.fork_choice = Some(ForkChoiceStrategy::Custom(import_as_best));
+		block_import_params.body = Some(body);
+		importer.import_block(block_import_params).await.unwrap();
+	}
+
+	fn sproof_with_parent_by_hash(client: &Client, hash: PHash) -> RelayStateSproofBuilder {
+		let header = client.header(hash).ok().flatten().expect("No header for parent block");
+		let included = HeadData(header.encode());
+		let mut builder = RelayStateSproofBuilder::default();
+		builder.para_id = cumulus_test_client::runtime::TEYRCHAIN_ID.into();
+		builder.included_para_head = Some(included);
+
+		builder
+	}
+	async fn build_and_import_block(client: &Client, included: Hash) -> Block {
+		let sproof = sproof_with_parent_by_hash(client, included);
+
+		let block_builder = client.init_block_builder(None, sproof).block_builder;
+
+		let block = block_builder.build().unwrap().block;
+
+		let origin = BlockOrigin::NetworkInitialSync;
+		import_block(client, block.clone(), origin, true).await;
+		block
+	}
+
+	fn set_up_components(num_authorities: usize) -> (Arc<Client>, KeystorePtr) {
+		let keystore = Arc::new(pezsp_keystore::testing::MemoryKeystore::new()) as Arc<_>;
+		for key in pezsp_keyring::Sr25519Keyring::iter().take(num_authorities) {
+			Keystore::sr25519_generate_new(
+				&*keystore,
+				pezsp_application_crypto::key_types::AURA,
+				Some(&key.to_seed()),
+			)
+			.expect("Can insert key into MemoryKeyStore");
+		}
+		(Arc::new(TestClientBuilder::new().build()), keystore)
+	}
+
+	/// This tests a special scenario where the unincluded segment in the runtime
+	/// is full. We are calling `can_build_upon`, passing the last built block as the
+	/// included one. In the runtime we will not find the hash of the included block in the
+	/// unincluded segment. The `can_build_upon` runtime API would therefore return `false`, but
+	/// we are ensuring on the node side that we are are always able to build on the included block.
+	#[tokio::test]
+	async fn test_can_build_upon() {
+		let (client, keystore) = set_up_components(6);
+
+		let genesis_hash = client.chain_info().genesis_hash;
+		let mut last_hash = genesis_hash;
+
+		// Fill up the unincluded segment tracker in the runtime.
+		while can_build_upon::<_, _, pezsp_consensus_aura::sr25519::AuthorityPair>(
+			Slot::from(u64::MAX),
+			Slot::from(u64::MAX),
+			Timestamp::default(),
+			last_hash,
+			genesis_hash,
+			&*client,
+			&keystore,
+		)
+		.await
+		.is_some()
+		{
+			let block = build_and_import_block(&client, genesis_hash).await;
+			last_hash = block.header().hash();
+		}
+
+		// Blocks were built with the genesis hash set as included block.
+		// We call `can_build_upon` with the last built block as the included block.
+		let result = can_build_upon::<_, _, pezsp_consensus_aura::sr25519::AuthorityPair>(
+			Slot::from(u64::MAX),
+			Slot::from(u64::MAX),
+			Timestamp::default(),
+			last_hash,
+			last_hash,
+			&*client,
+			&keystore,
+		)
+		.await;
+		assert!(result.is_some());
+	}
+
+	/// Helper to create a mock overseer handle and message recorder
+	fn create_overseer_handle() -> (OverseerHandle, Arc<Mutex<Vec<CollatorProtocolMessage>>>) {
+		let messages = Arc::new(Mutex::new(Vec::new()));
+		let messages_clone = messages.clone();
+
+		let (tx, mut rx) = pezkuwi_node_subsystem_util::metered::channel(100);
+
+		// Spawn a task to receive and record overseer messages
+		tokio::spawn(async move {
+			while let Some(event) = rx.next().await {
+				if let Event::MsgToSubsystem { msg, .. } = event {
+					if let pezkuwi_node_subsystem::AllMessages::CollatorProtocol(cp_msg) = msg {
+						messages_clone.lock().unwrap().push(cp_msg);
+					}
+				}
+			}
+		});
+
+		(Handle::new(tx), messages)
+	}
+
+	#[tokio::test]
+	async fn preconnect_when_next_slot_is_ours() {
+		let (client, keystore) = set_up_components(1);
+		let genesis_hash = client.chain_info().genesis_hash;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Fetch authorities for the update call
+		let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
+
+		// Update with slot 5, next slot (6) should be ours
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
+			.await;
+
+		// Give time for message to be processed
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+
+		let messages = messages_recorder.lock().unwrap();
+		assert_eq!(messages.len(), 1);
+		assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
+		assert_eq!(helper.our_slot, Some(Slot::from(6)));
+	}
+
+	#[tokio::test]
+	async fn preconnect_no_duplicate_connect_message() {
+		let (client, keystore) = set_up_components(1);
+		let genesis_hash = client.chain_info().genesis_hash;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Fetch authorities for the update calls
+		let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
+
+		// Update with slot 5, next slot (6) is ours
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
+			.await;
+
+		// Give time for message to be processed
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		assert_eq!(messages_recorder.lock().unwrap().len(), 1);
+		messages_recorder.lock().unwrap().clear();
+
+		// Update with slot 5 again - should not send another message
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		assert_eq!(messages_recorder.lock().unwrap().len(), 0);
+
+		// Update with slot 1 (our slot) - should not send another message
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(6), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		assert_eq!(messages_recorder.lock().unwrap().len(), 0);
+	}
+
+	#[tokio::test]
+	async fn preconnect_disconnect_when_slot_passes() {
+		let (client, keystore) = set_up_components(1);
+		let genesis_hash = client.chain_info().genesis_hash;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Fetch authorities for the update calls
+		let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
+
+		// Slot 0 -> Alice, Slot 1 -> Bob, Slot 2 -> Charlie, Slot 3 -> Dave, Slot 4 -> Eve,
+		// Slot 5 -> Ferdie, Slot 6 -> Alice
+
+		// Update with slot 5, next slot (6) is ours -> should connect
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		assert_eq!(helper.our_slot, Some(Slot::from(6)));
+		messages_recorder.lock().unwrap().clear();
+
+		// Update with slot 8, next slot (9) is Charlie's -> should disconnect
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(8), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+
+		{
+			let messages = messages_recorder.lock().unwrap();
+			assert_eq!(messages.len(), 1, "Expected exactly one disconnect message");
+			assert!(matches!(messages[0], CollatorProtocolMessage::DisconnectFromBackingGroups));
+			assert_eq!(helper.our_slot, None);
+		}
+
+		messages_recorder.lock().unwrap().clear();
+
+		// Update again with slot 8, next slot (9) is Charlie's -> should not send another
+		// disconnect message
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(8), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+
+		let messages = messages_recorder.lock().unwrap();
+		assert_eq!(messages.len(), 0, "Expected no messages");
+		assert_eq!(helper.our_slot, None);
+	}
+
+	#[tokio::test]
+	async fn preconnect_no_disconnect_without_previous_connection() {
+		let (client, keystore) = set_up_components(1);
+		let genesis_hash = client.chain_info().genesis_hash;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Fetch authorities for the update call
+		let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
+
+		// Slot 0 -> Alice, Slot 1 -> Bob, Slot 2 -> Charlie, Slot 3 -> Dave, Slot 4 -> Eve,
+		// Slot 5 -> Ferdie
+
+		// Update with slot 1 (Bob's slot), next slot (2) is Charlie's
+		// Since we never connected before (our_slot is None), we should not send disconnect
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(1), &authorities)
+			.await;
+
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		// Should not send any message since we never connected
+		assert_eq!(messages_recorder.lock().unwrap().len(), 0);
+		assert_eq!(helper.our_slot, None);
+	}
+
+	#[tokio::test]
+	async fn preconnect_multiple_cycles() {
+		let (client, keystore) = set_up_components(1);
+		let genesis_hash = client.chain_info().genesis_hash;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Fetch authorities for the update calls
+		let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
+
+		// Slot 0 -> Alice, Slot 1 -> Bob, Slot 2 -> Charlie, Slot 3 -> Dave, Slot 4 -> Eve,
+		// Slot 5 -> Ferdie, Slot 6 -> Alice, Slot 7 -> Bob, ...
+
+		// Cycle 1: Connect at slot 5, next slot (6) is ours
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		{
+			let messages = messages_recorder.lock().unwrap();
+			assert_eq!(messages.len(), 1);
+			assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
+		}
+		assert_eq!(helper.our_slot, Some(Slot::from(6)));
+		messages_recorder.lock().unwrap().clear();
+
+		// Cycle 1: Disconnect at slot 7, next slot (8) is Charlie's
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(7), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		{
+			let messages = messages_recorder.lock().unwrap();
+			assert_eq!(messages.len(), 1);
+			assert!(matches!(messages[0], CollatorProtocolMessage::DisconnectFromBackingGroups));
+		}
+		assert_eq!(helper.our_slot, None);
+		messages_recorder.lock().unwrap().clear();
+
+		// Cycle 2: Connect again at slot 11, next slot (12) is ours
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(11), &authorities)
+			.await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		{
+			let messages = messages_recorder.lock().unwrap();
+			assert_eq!(messages.len(), 1);
+			assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
+		}
+		assert_eq!(helper.our_slot, Some(Slot::from(12)));
+	}
+
+	#[tokio::test]
+	async fn preconnect_handles_empty_authorities() {
+		let keystore = Arc::new(pezsp_keystore::testing::MemoryKeystore::new()) as Arc<_>;
+		let (overseer_handle, messages_recorder) = create_overseer_handle();
+
+		let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
+
+		// Pass empty authorities list
+		let authorities = vec![];
+		helper
+			.update::<pezsp_consensus_aura::sr25519::AuthorityPair>(Slot::from(0), &authorities)
+			.await;
+
+		tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+		// Should not send any message if authorities list is empty
+		assert_eq!(messages_recorder.lock().unwrap().len(), 0);
+	}
+}
+
+/// Holds a relay parent and its descendants.
+pub struct RelayParentData {
+	/// The relay parent block header
+	relay_parent: RelayHeader,
+	/// Ordered collection of descendant block headers, from oldest to newest
+	descendants: Vec<RelayHeader>,
+}
+
+impl RelayParentData {
+	/// Creates a new instance with the given relay parent and no descendants.
+	pub fn new(relay_parent: RelayHeader) -> Self {
+		Self { relay_parent, descendants: Default::default() }
+	}
+
+	/// Creates a new instance with the given relay parent and descendants.
+	pub fn new_with_descendants(relay_parent: RelayHeader, descendants: Vec<RelayHeader>) -> Self {
+		Self { relay_parent, descendants }
+	}
+
+	/// Returns a reference to the relay parent header.
+	pub fn relay_parent(&self) -> &RelayHeader {
+		&self.relay_parent
+	}
+
+	/// Returns the number of descendants.
+	#[cfg(test)]
+	pub fn descendants_len(&self) -> usize {
+		self.descendants.len()
+	}
+
+	/// Consumes the structure and returns a vector containing the relay parent followed by its
+	/// descendants in chronological order. The resulting list should be provided to the teyrchain
+	/// inherent data.
+	pub fn into_inherent_descendant_list(self) -> Vec<RelayHeader> {
+		let Self { relay_parent, mut descendants } = self;
+
+		if descendants.is_empty() {
+			return Default::default();
+		}
+
+		let mut result = vec![relay_parent];
+		result.append(&mut descendants);
+		result
+	}
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/block_builder_task.rs

@@ -0,0 +1,647 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use codec::{Codec, Encode};
+
+use super::CollatorMessage;
+use crate::{
+	collator as collator_util,
+	collators::{
+		check_validation_code_or_log,
+		slot_based::{
+			relay_chain_data_cache::{RelayChainData, RelayChainDataCache},
+			slot_timer::{SlotInfo, SlotTimer},
+		},
+		BackingGroupConnectionHelper, RelayParentData,
+	},
+	LOG_TARGET,
+};
+use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_client_consensus_common::{self as consensus_common, TeyrchainBlockImportMarker};
+use cumulus_client_consensus_proposer::ProposerInterface;
+use cumulus_primitives_aura::{AuraUnincludedSegmentApi, Slot};
+use cumulus_primitives_core::{
+	extract_relay_parent, rpsr_digest, ClaimQueueOffset, CoreInfo, CoreSelector, CumulusDigestItem,
+	PersistedValidationData, RelayParentOffsetApi,
+};
+use cumulus_relay_chain_interface::RelayChainInterface;
+use futures::prelude::*;
+use pezkuwi_primitives::{
+	Block as RelayBlock, CoreIndex, Hash as RelayHash, Header as RelayHeader, Id as ParaId,
+};
+use pezsc_client_api::{backend::AuxStore, BlockBackend, BlockOf, UsageProvider};
+use pezsc_consensus::BlockImport;
+use pezsc_consensus_aura::SlotDuration;
+use pezsc_network_types::PeerId;
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_application_crypto::AppPublic;
+use pezsp_blockchain::HeaderBackend;
+use pezsp_consensus_aura::AuraApi;
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_keystore::KeystorePtr;
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT, Member, Zero};
+use std::{collections::VecDeque, sync::Arc, time::Duration};
+
+/// Parameters for [`run_block_builder`].
+pub struct BuilderTaskParams<
+	Block: BlockT,
+	BI,
+	CIDP,
+	Client,
+	Backend,
+	RelayClient,
+	CHP,
+	Proposer,
+	CS,
+> {
+	/// Inherent data providers. Only non-consensus inherent data should be provided, i.e.
+	/// the timestamp, slot, and paras inherents should be omitted, as they are set by this
+	/// collator.
+	pub create_inherent_data_providers: CIDP,
+	/// Used to actually import blocks.
+	pub block_import: BI,
+	/// The underlying para client.
+	pub para_client: Arc<Client>,
+	/// The para client's backend, used to access the database.
+	pub para_backend: Arc<Backend>,
+	/// A handle to the relay-chain client.
+	pub relay_client: RelayClient,
+	/// A validation code hash provider, used to get the current validation code hash.
+	pub code_hash_provider: CHP,
+	/// The underlying keystore, which should contain Aura consensus keys.
+	pub keystore: KeystorePtr,
+	/// The collator network peer id.
+	pub collator_peer_id: PeerId,
+	/// The para's ID.
+	pub para_id: ParaId,
+	/// The underlying block proposer this should call into.
+	pub proposer: Proposer,
+	/// The generic collator service used to plug into this consensus engine.
+	pub collator_service: CS,
+	/// The amount of time to spend authoring each block.
+	pub authoring_duration: Duration,
+	/// Channel to send built blocks to the collation task.
+	pub collator_sender: pezsc_utils::mpsc::TracingUnboundedSender<CollatorMessage<Block>>,
+	/// Slot duration of the relay chain.
+	pub relay_chain_slot_duration: Duration,
+	/// Offset all time operations by this duration.
+	///
+	/// This is a time quantity that is subtracted from the actual timestamp when computing
+	/// the time left to enter a new slot. In practice, this *left-shifts* the clock time with the
+	/// intent to keep our "clock" slightly behind the relay chain one and thus reducing the
+	/// likelihood of encountering unfavorable notification arrival timings (i.e. we don't want to
+	/// wait for relay chain notifications because we woke up too early).
+	pub slot_offset: Duration,
+	/// The maximum percentage of the maximum PoV size that the collator can use.
+	/// It will be removed once https://github.com/pezkuwichain/pezkuwi-sdk/issues/23 is fixed.
+	pub max_pov_percentage: Option<u32>,
+}
+
+/// Run block-builder.
+pub fn run_block_builder<Block, P, BI, CIDP, Client, Backend, RelayClient, CHP, Proposer, CS>(
+	params: BuilderTaskParams<Block, BI, CIDP, Client, Backend, RelayClient, CHP, Proposer, CS>,
+) -> impl Future<Output = ()> + Send + 'static
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>
+		+ UsageProvider<Block>
+		+ BlockOf
+		+ AuxStore
+		+ HeaderBackend<Block>
+		+ BlockBackend<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	Client::Api:
+		AuraApi<Block, P::Public> + RelayParentOffsetApi<Block> + AuraUnincludedSegmentApi<Block>,
+	Backend: pezsc_client_api::Backend<Block> + 'static,
+	RelayClient: RelayChainInterface + Clone + 'static,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+	CIDP::InherentDataProviders: Send,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block> + Send + Sync + 'static,
+	CS: CollatorServiceInterface<Block> + Send + Sync + 'static,
+	CHP: consensus_common::ValidationCodeHashProvider<Block::Hash> + Send + 'static,
+	P: Pair + Send + Sync + 'static,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	async move {
+		tracing::info!(target: LOG_TARGET, "Starting slot-based block-builder task.");
+		let BuilderTaskParams {
+			relay_client,
+			create_inherent_data_providers,
+			para_client,
+			keystore,
+			block_import,
+			collator_peer_id,
+			para_id,
+			proposer,
+			collator_service,
+			collator_sender,
+			code_hash_provider,
+			authoring_duration,
+			relay_chain_slot_duration,
+			para_backend,
+			slot_offset,
+			max_pov_percentage,
+		} = params;
+
+		let mut slot_timer = SlotTimer::<_, _, P>::new_with_offset(
+			para_client.clone(),
+			slot_offset,
+			relay_chain_slot_duration,
+		);
+
+		let mut collator = {
+			let params = collator_util::Params {
+				create_inherent_data_providers,
+				block_import,
+				relay_client: relay_client.clone(),
+				keystore: keystore.clone(),
+				collator_peer_id,
+				para_id,
+				proposer,
+				collator_service,
+			};
+
+			collator_util::Collator::<Block, P, _, _, _, _, _>::new(params)
+		};
+
+		let mut relay_chain_data_cache = RelayChainDataCache::new(relay_client.clone(), para_id);
+		let mut connection_helper = BackingGroupConnectionHelper::new(
+			keystore.clone(),
+			relay_client
+				.overseer_handle()
+				// Should never fail. If it fails, then providing collations to relay chain
+				// doesn't work either. So it is fine to panic here.
+				.expect("Relay chain interface must provide overseer handle."),
+		);
+
+		loop {
+			// We wait here until the next slot arrives.
+			if slot_timer.wait_until_next_slot().await.is_err() {
+				tracing::error!(target: LOG_TARGET, "Unable to wait for next slot.");
+				return;
+			};
+
+			let Ok(relay_best_hash) = relay_client.best_block_hash().await else {
+				tracing::warn!(target: crate::LOG_TARGET, "Unable to fetch latest relay chain block hash.");
+				continue;
+			};
+
+			let best_hash = para_client.info().best_hash;
+			let relay_parent_offset =
+				para_client.runtime_api().relay_parent_offset(best_hash).unwrap_or_default();
+
+			let Ok(para_slot_duration) = crate::slot_duration(&*para_client) else {
+				tracing::error!(target: LOG_TARGET, "Failed to fetch slot duration from runtime.");
+				continue;
+			};
+
+			let Ok(Some(rp_data)) = offset_relay_parent_find_descendants(
+				&mut relay_chain_data_cache,
+				relay_best_hash,
+				relay_parent_offset,
+			)
+			.await
+			else {
+				continue;
+			};
+
+			let Some(para_slot) = adjust_para_to_relay_parent_slot(
+				rp_data.relay_parent(),
+				relay_chain_slot_duration,
+				para_slot_duration,
+			) else {
+				continue;
+			};
+
+			let relay_parent = rp_data.relay_parent().hash();
+			let relay_parent_header = rp_data.relay_parent().clone();
+
+			let Some((included_header, parent)) =
+				crate::collators::find_parent(relay_parent, para_id, &*para_backend, &relay_client)
+					.await
+			else {
+				continue;
+			};
+
+			let parent_hash = parent.hash;
+			let parent_header = &parent.header;
+
+			// Retrieve the core.
+			let core = match determine_core(
+				&mut relay_chain_data_cache,
+				&relay_parent_header,
+				para_id,
+				parent_header,
+				relay_parent_offset,
+			)
+			.await
+			{
+				Err(()) => {
+					tracing::debug!(
+						target: LOG_TARGET,
+						?relay_parent,
+						"Failed to determine core"
+					);
+
+					continue;
+				},
+				Ok(Some(cores)) => {
+					tracing::debug!(
+						target: LOG_TARGET,
+						?relay_parent,
+						core_selector = ?cores.selector,
+						claim_queue_offset = ?cores.claim_queue_offset,
+						"Going to claim core",
+					);
+
+					cores
+				},
+				Ok(None) => {
+					tracing::debug!(
+						target: LOG_TARGET,
+						?relay_parent,
+						"No core scheduled"
+					);
+
+					continue;
+				},
+			};
+
+			let Ok(RelayChainData { max_pov_size, last_claimed_core_selector, .. }) =
+				relay_chain_data_cache.get_mut_relay_chain_data(relay_parent).await
+			else {
+				continue;
+			};
+
+			slot_timer.update_scheduling(core.total_cores().into());
+
+			// We mainly call this to inform users at genesis if there is a mismatch with the
+			// on-chain data.
+			collator.collator_service().check_block_status(parent_hash, parent_header);
+
+			let Ok(relay_slot) =
+				pezsc_consensus_babe::find_pre_digest::<RelayBlock>(&relay_parent_header)
+					.map(|babe_pre_digest| babe_pre_digest.slot())
+			else {
+				tracing::error!(target: crate::LOG_TARGET, "Relay chain does not contain babe slot. This should never happen.");
+				continue;
+			};
+
+			let included_header_hash = included_header.hash();
+
+			if let Ok(authorities) = para_client.runtime_api().authorities(parent_hash) {
+				connection_helper.update::<P>(para_slot.slot, &authorities).await;
+			}
+
+			let slot_claim = match crate::collators::can_build_upon::<_, _, P>(
+				para_slot.slot,
+				relay_slot,
+				para_slot.timestamp,
+				parent_hash,
+				included_header_hash,
+				&*para_client,
+				&keystore,
+			)
+			.await
+			{
+				Some(slot) => slot,
+				None => {
+					tracing::debug!(
+						target: crate::LOG_TARGET,
+						unincluded_segment_len = parent.depth,
+						relay_parent = ?relay_parent,
+						relay_parent_num = %relay_parent_header.number(),
+						included_hash = ?included_header_hash,
+						included_num = %included_header.number(),
+						parent = ?parent_hash,
+						slot = ?para_slot.slot,
+						"Not building block."
+					);
+					continue;
+				},
+			};
+
+			tracing::debug!(
+				target: crate::LOG_TARGET,
+				unincluded_segment_len = parent.depth,
+				relay_parent = %relay_parent,
+				relay_parent_num = %relay_parent_header.number(),
+				relay_parent_offset,
+				included_hash = %included_header_hash,
+				included_num = %included_header.number(),
+				parent = %parent_hash,
+				slot = ?para_slot.slot,
+				"Building block."
+			);
+
+			let validation_data = PersistedValidationData {
+				parent_head: parent_header.encode().into(),
+				relay_parent_number: *relay_parent_header.number(),
+				relay_parent_storage_root: *relay_parent_header.state_root(),
+				max_pov_size: *max_pov_size,
+			};
+
+			let (teyrchain_inherent_data, other_inherent_data) = match collator
+				.create_inherent_data_with_rp_offset(
+					relay_parent,
+					&validation_data,
+					parent_hash,
+					slot_claim.timestamp(),
+					Some(rp_data),
+					collator_peer_id,
+				)
+				.await
+			{
+				Err(err) => {
+					tracing::error!(target: crate::LOG_TARGET, ?err);
+					break;
+				},
+				Ok(x) => x,
+			};
+
+			let validation_code_hash = match code_hash_provider.code_hash_at(parent_hash) {
+				None => {
+					tracing::error!(target: crate::LOG_TARGET, ?parent_hash, "Could not fetch validation code hash");
+					break;
+				},
+				Some(v) => v,
+			};
+
+			check_validation_code_or_log(
+				&validation_code_hash,
+				para_id,
+				&relay_client,
+				relay_parent,
+			)
+			.await;
+
+			let allowed_pov_size = if let Some(max_pov_percentage) = max_pov_percentage {
+				validation_data.max_pov_size * max_pov_percentage / 100
+			} else {
+				// Set the block limit to 85% of the maximum PoV size.
+				//
+				// Once https://github.com/pezkuwichain/pezkuwi-sdk/issues/23 issue is
+				// fixed, this should be removed.
+				validation_data.max_pov_size * 85 / 100
+			} as usize;
+
+			let adjusted_authoring_duration =
+				slot_timer.adjust_authoring_duration(authoring_duration);
+			tracing::debug!(target: crate::LOG_TARGET, duration = ?adjusted_authoring_duration, "Adjusted proposal duration.");
+
+			let Some(adjusted_authoring_duration) = adjusted_authoring_duration else {
+				tracing::debug!(
+					target: crate::LOG_TARGET,
+					unincluded_segment_len = parent.depth,
+					relay_parent = ?relay_parent,
+					relay_parent_num = %relay_parent_header.number(),
+					included_hash = ?included_header_hash,
+					included_num = %included_header.number(),
+					parent = ?parent_hash,
+					slot = ?para_slot.slot,
+					"Not building block due to insufficient authoring duration."
+				);
+
+				continue;
+			};
+
+			let Ok(Some(candidate)) = collator
+				.build_block_and_import(
+					&parent_header,
+					&slot_claim,
+					Some(vec![CumulusDigestItem::CoreInfo(core.core_info()).to_digest_item()]),
+					(teyrchain_inherent_data, other_inherent_data),
+					adjusted_authoring_duration,
+					allowed_pov_size,
+				)
+				.await
+			else {
+				tracing::error!(target: crate::LOG_TARGET, "Unable to build block at slot.");
+				continue;
+			};
+
+			let new_block_hash = candidate.block.header().hash();
+
+			// Announce the newly built block to our peers.
+			collator.collator_service().announce_block(new_block_hash, None);
+
+			*last_claimed_core_selector = Some(core.core_selector());
+
+			if let Err(err) = collator_sender.unbounded_send(CollatorMessage {
+				relay_parent,
+				parent_header: parent_header.clone(),
+				teyrchain_candidate: candidate,
+				validation_code_hash,
+				core_index: core.core_index(),
+				max_pov_size: validation_data.max_pov_size,
+			}) {
+				tracing::error!(target: crate::LOG_TARGET, ?err, "Unable to send block to collation task.");
+				return;
+			}
+		}
+	}
+}
+
+/// Translate the slot of the relay parent to the slot of the teyrchain.
+fn adjust_para_to_relay_parent_slot(
+	relay_header: &RelayHeader,
+	relay_chain_slot_duration: Duration,
+	para_slot_duration: SlotDuration,
+) -> Option<SlotInfo> {
+	let relay_slot = pezsc_consensus_babe::find_pre_digest::<RelayBlock>(&relay_header)
+		.map(|babe_pre_digest| babe_pre_digest.slot())
+		.ok()?;
+	let new_slot = Slot::from_timestamp(
+		relay_slot
+			.timestamp(SlotDuration::from_millis(relay_chain_slot_duration.as_millis() as u64))?,
+		para_slot_duration,
+	);
+	let para_slot = SlotInfo { slot: new_slot, timestamp: new_slot.timestamp(para_slot_duration)? };
+	tracing::debug!(
+		target: LOG_TARGET,
+		timestamp = ?para_slot.timestamp,
+		slot = ?para_slot.slot,
+		"Teyrchain slot adjusted to relay chain.",
+	);
+	Some(para_slot)
+}
+
+/// Finds a relay chain parent block at a specified offset from the best block, collecting its
+/// descendants.
+///
+/// # Returns
+/// * `Ok(RelayParentData)` - Contains the target relay parent and its ordered list of descendants
+/// * `Err(())` - If any relay chain block header cannot be retrieved
+///
+/// The function traverses backwards from the best block until it finds the block at the specified
+/// offset, collecting all blocks in between to maintain the chain of ancestry.
+pub(crate) async fn offset_relay_parent_find_descendants<RelayClient>(
+	relay_chain_data_cache: &mut RelayChainDataCache<RelayClient>,
+	relay_best_block: RelayHash,
+	relay_parent_offset: u32,
+) -> Result<Option<RelayParentData>, ()>
+where
+	RelayClient: RelayChainInterface + Clone + 'static,
+{
+	let Ok(mut relay_header) = relay_chain_data_cache
+		.get_mut_relay_chain_data(relay_best_block)
+		.await
+		.map(|d| d.relay_parent_header.clone())
+	else {
+		tracing::error!(target: LOG_TARGET, ?relay_best_block, "Unable to fetch best relay chain block header.");
+		return Err(());
+	};
+
+	if relay_parent_offset == 0 {
+		return Ok(Some(RelayParentData::new(relay_header)));
+	}
+
+	if pezsc_consensus_babe::contains_epoch_change::<RelayBlock>(&relay_header) {
+		tracing::debug!(target: LOG_TARGET, ?relay_best_block, relay_best_block_number = relay_header.number(), "Relay parent is in previous session.");
+		return Ok(None);
+	}
+
+	let mut required_ancestors: VecDeque<RelayHeader> = Default::default();
+	required_ancestors.push_front(relay_header.clone());
+	while required_ancestors.len() < relay_parent_offset as usize {
+		let next_header = relay_chain_data_cache
+			.get_mut_relay_chain_data(*relay_header.parent_hash())
+			.await?
+			.relay_parent_header
+			.clone();
+		if pezsc_consensus_babe::contains_epoch_change::<RelayBlock>(&next_header) {
+			tracing::debug!(target: LOG_TARGET, ?relay_best_block, ancestor = %next_header.hash(), ancestor_block_number = next_header.number(), "Ancestor of best block is in previous session.");
+			return Ok(None);
+		}
+		required_ancestors.push_front(next_header.clone());
+		relay_header = next_header;
+	}
+
+	let relay_parent = relay_chain_data_cache
+		.get_mut_relay_chain_data(*relay_header.parent_hash())
+		.await?
+		.relay_parent_header
+		.clone();
+
+	tracing::debug!(
+		target: LOG_TARGET,
+		relay_parent_hash = %relay_parent.hash(),
+		relay_parent_num = relay_parent.number(),
+		num_descendants = required_ancestors.len(),
+		"Relay parent descendants."
+	);
+
+	Ok(Some(RelayParentData::new_with_descendants(relay_parent, required_ancestors.into())))
+}
+
+/// Return value of [`determine_core`].
+pub(crate) struct Core {
+	selector: CoreSelector,
+	claim_queue_offset: ClaimQueueOffset,
+	core_index: CoreIndex,
+	number_of_cores: u16,
+}
+
+impl Core {
+	/// Returns the current [`CoreInfo`].
+	fn core_info(&self) -> CoreInfo {
+		CoreInfo {
+			selector: self.selector,
+			claim_queue_offset: self.claim_queue_offset,
+			number_of_cores: self.number_of_cores.into(),
+		}
+	}
+
+	/// Returns the current [`CoreSelector`].
+	pub(crate) fn core_selector(&self) -> CoreSelector {
+		self.selector
+	}
+
+	/// Returns the current [`CoreIndex`].
+	pub(crate) fn core_index(&self) -> CoreIndex {
+		self.core_index
+	}
+
+	/// Returns the total number of cores.
+	pub(crate) fn total_cores(&self) -> u16 {
+		self.number_of_cores
+	}
+}
+
+/// Determine the core for the given `para_id`.
+pub(crate) async fn determine_core<H: HeaderT, RI: RelayChainInterface + 'static>(
+	relay_chain_data_cache: &mut RelayChainDataCache<RI>,
+	relay_parent: &RelayHeader,
+	para_id: ParaId,
+	para_parent: &H,
+	relay_parent_offset: u32,
+) -> Result<Option<Core>, ()> {
+	let cores_at_offset = &relay_chain_data_cache
+		.get_mut_relay_chain_data(relay_parent.hash())
+		.await?
+		.claim_queue
+		.iter_claims_at_depth_for_para(relay_parent_offset as usize, para_id)
+		.collect::<Vec<_>>();
+
+	let is_new_relay_parent = if para_parent.number().is_zero() {
+		true
+	} else {
+		match extract_relay_parent(para_parent.digest()) {
+			Some(last_relay_parent) => last_relay_parent != relay_parent.hash(),
+			None =>
+				rpsr_digest::extract_relay_parent_storage_root(para_parent.digest())
+					.ok_or(())?
+					.0 != *relay_parent.state_root(),
+		}
+	};
+
+	let core_info = CumulusDigestItem::find_core_info(para_parent.digest());
+
+	// If we are using a new relay parent, we can start over from the start.
+	let (selector, core_index) = if is_new_relay_parent {
+		let Some(core_index) = cores_at_offset.get(0) else { return Ok(None) };
+
+		(0, *core_index)
+	} else if let Some(core_info) = core_info {
+		let selector = core_info.selector.0 as usize + 1;
+		let Some(core_index) = cores_at_offset.get(selector) else { return Ok(None) };
+
+		(selector, *core_index)
+	} else {
+		let last_claimed_core_selector = relay_chain_data_cache
+			.get_mut_relay_chain_data(relay_parent.hash())
+			.await?
+			.last_claimed_core_selector;
+
+		let selector = last_claimed_core_selector.map_or(0, |cs| cs.0 as usize) + 1;
+		let Some(core_index) = cores_at_offset.get(selector) else { return Ok(None) };
+
+		(selector, *core_index)
+	};
+
+	Ok(Some(Core {
+		selector: CoreSelector(selector as u8),
+		core_index,
+		claim_queue_offset: ClaimQueueOffset(relay_parent_offset as u8),
+		number_of_cores: cores_at_offset.len() as u16,
+	}))
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/block_import.rs

@@ -0,0 +1,145 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use futures::{stream::FusedStream, StreamExt};
+use pezsc_consensus::{BlockImport, StateAction};
+use pezsc_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
+use pezsp_api::{ApiExt, CallApiAt, CallContext, Core, ProvideRuntimeApi, StorageProof};
+use pezsp_runtime::traits::{Block as BlockT, Header as _};
+use pezsp_trie::proof_size_extension::ProofSizeExt;
+use std::sync::Arc;
+
+/// Handle for receiving the block and the storage proof from the [`SlotBasedBlockImport`].
+///
+/// This handle should be passed to [`Params`](super::Params) or can also be dropped if the node is
+/// not running as collator.
+pub struct SlotBasedBlockImportHandle<Block> {
+	receiver: TracingUnboundedReceiver<(Block, StorageProof)>,
+}
+
+impl<Block> SlotBasedBlockImportHandle<Block> {
+	/// Returns the next item.
+	///
+	/// The future will never return when the internal channel is closed.
+	pub async fn next(&mut self) -> (Block, StorageProof) {
+		loop {
+			if self.receiver.is_terminated() {
+				futures::pending!()
+			} else if let Some(res) = self.receiver.next().await {
+				return res;
+			}
+		}
+	}
+}
+
+/// Special block import for the slot based collator.
+pub struct SlotBasedBlockImport<Block, BI, Client> {
+	inner: BI,
+	client: Arc<Client>,
+	sender: TracingUnboundedSender<(Block, StorageProof)>,
+}
+
+impl<Block, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
+	/// Create a new instance.
+	///
+	/// The returned [`SlotBasedBlockImportHandle`] needs to be passed to the
+	/// [`Params`](super::Params), so that this block import instance can communicate with the
+	/// collation task. If the node is not running as a collator, just dropping the handle is fine.
+	pub fn new(inner: BI, client: Arc<Client>) -> (Self, SlotBasedBlockImportHandle<Block>) {
+		let (sender, receiver) = tracing_unbounded("SlotBasedBlockImportChannel", 1000);
+
+		(Self { sender, client, inner }, SlotBasedBlockImportHandle { receiver })
+	}
+}
+
+impl<Block, BI: Clone, Client> Clone for SlotBasedBlockImport<Block, BI, Client> {
+	fn clone(&self) -> Self {
+		Self { inner: self.inner.clone(), client: self.client.clone(), sender: self.sender.clone() }
+	}
+}
+
+#[async_trait::async_trait]
+impl<Block, BI, Client> BlockImport<Block> for SlotBasedBlockImport<Block, BI, Client>
+where
+	Block: BlockT,
+	BI: BlockImport<Block> + Send + Sync,
+	BI::Error: Into<pezsp_consensus::Error>,
+	Client: ProvideRuntimeApi<Block> + CallApiAt<Block> + Send + Sync,
+	Client::StateBackend: Send,
+	Client::Api: Core<Block>,
+{
+	type Error = pezsp_consensus::Error;
+
+	async fn check_block(
+		&self,
+		block: pezsc_consensus::BlockCheckParams<Block>,
+	) -> Result<pezsc_consensus::ImportResult, Self::Error> {
+		self.inner.check_block(block).await.map_err(Into::into)
+	}
+
+	async fn import_block(
+		&self,
+		mut params: pezsc_consensus::BlockImportParams<Block>,
+	) -> Result<pezsc_consensus::ImportResult, Self::Error> {
+		// If the channel exists and it is required to execute the block, we will execute the block
+		// here. This is done to collect the storage proof and to prevent re-execution, we push
+		// downwards the state changes. `StateAction::ApplyChanges` is ignored, because it either
+		// means that the node produced the block itself or the block was imported via state sync.
+		if !self.sender.is_closed() && !matches!(params.state_action, StateAction::ApplyChanges(_))
+		{
+			let mut runtime_api = self.client.runtime_api();
+
+			runtime_api.set_call_context(CallContext::Onchain);
+
+			runtime_api.record_proof();
+			let recorder = runtime_api
+				.proof_recorder()
+				.expect("Proof recording is enabled in the line above; qed.");
+			runtime_api.register_extension(ProofSizeExt::new(recorder));
+
+			let parent_hash = *params.header.parent_hash();
+
+			let block = Block::new(params.header.clone(), params.body.clone().unwrap_or_default());
+
+			runtime_api
+				.execute_block(parent_hash, block.clone().into())
+				.map_err(|e| Box::new(e) as Box<_>)?;
+
+			let storage_proof =
+				runtime_api.extract_proof().expect("Proof recording was enabled above; qed");
+
+			let state = self.client.state_at(parent_hash).map_err(|e| Box::new(e) as Box<_>)?;
+			let gen_storage_changes = runtime_api
+				.into_storage_changes(&state, parent_hash)
+				.map_err(pezsp_consensus::Error::ChainLookup)?;
+
+			if params.header.state_root() != &gen_storage_changes.transaction_storage_root {
+				return Err(pezsp_consensus::Error::Other(Box::new(
+					pezsp_blockchain::Error::InvalidStateRoot,
+				)));
+			}
+
+			params.state_action = StateAction::ApplyChanges(pezsc_consensus::StorageChanges::Changes(
+				gen_storage_changes,
+			));
+
+			let _ = self.sender.unbounded_send((block, storage_proof));
+		}
+
+		self.inner.import_block(params).await.map_err(Into::into)
+	}
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/collation_task.rs

@@ -0,0 +1,189 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use codec::Encode;
+use std::path::PathBuf;
+
+use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_relay_chain_interface::RelayChainInterface;
+
+use pezkuwi_node_primitives::{MaybeCompressedPoV, SubmitCollationParams};
+use pezkuwi_node_subsystem::messages::CollationGenerationMessage;
+use pezkuwi_overseer::Handle as OverseerHandle;
+use pezkuwi_primitives::{CollatorPair, Id as ParaId};
+
+use cumulus_primitives_core::relay_chain::BlockId;
+use futures::prelude::*;
+
+use crate::export_pov_to_path;
+use pezsc_utils::mpsc::TracingUnboundedReceiver;
+use pezsp_runtime::traits::{Block as BlockT, Header};
+
+use super::CollatorMessage;
+
+const LOG_TARGET: &str = "aura::pezcumulus::collation_task";
+
+/// Parameters for the collation task.
+pub struct Params<Block: BlockT, RClient, CS> {
+	/// A handle to the relay-chain client.
+	pub relay_client: RClient,
+	/// The collator key used to sign collations before submitting to validators.
+	pub collator_key: CollatorPair,
+	/// The para's ID.
+	pub para_id: ParaId,
+	/// Whether we should reinitialize the collator config (i.e. we are transitioning to aura).
+	pub reinitialize: bool,
+	/// Collator service interface
+	pub collator_service: CS,
+	/// Receiver channel for communication with the block builder task.
+	pub collator_receiver: TracingUnboundedReceiver<CollatorMessage<Block>>,
+	/// The handle from the special slot based block import.
+	pub block_import_handle: super::SlotBasedBlockImportHandle<Block>,
+	/// When set, the collator will export every produced `POV` to this folder.
+	pub export_pov: Option<PathBuf>,
+}
+
+/// Asynchronously executes the collation task for a teyrchain.
+///
+/// This function initializes the collator subsystems necessary for producing and submitting
+/// collations to the relay chain. It listens for new best relay chain block notifications and
+/// handles collator messages. If our teyrchain is scheduled on a core and we have a candidate,
+/// the task will build a collation and send it to the relay chain.
+pub async fn run_collation_task<Block, RClient, CS>(
+	Params {
+		relay_client,
+		collator_key,
+		para_id,
+		reinitialize,
+		collator_service,
+		mut collator_receiver,
+		mut block_import_handle,
+		export_pov,
+	}: Params<Block, RClient, CS>,
+) where
+	Block: BlockT,
+	CS: CollatorServiceInterface<Block> + Send + Sync + 'static,
+	RClient: RelayChainInterface + Clone + 'static,
+{
+	let Ok(mut overseer_handle) = relay_client.overseer_handle() else {
+		tracing::error!(target: LOG_TARGET, "Failed to get overseer handle.");
+		return;
+	};
+
+	cumulus_client_collator::initialize_collator_subsystems(
+		&mut overseer_handle,
+		collator_key,
+		para_id,
+		reinitialize,
+	)
+	.await;
+
+	loop {
+		futures::select! {
+			collator_message = collator_receiver.next() => {
+				let Some(message) = collator_message else {
+					return;
+				};
+
+				handle_collation_message(message, &collator_service, &mut overseer_handle,relay_client.clone(),export_pov.clone()).await;
+			},
+			block_import_msg = block_import_handle.next().fuse() => {
+				// TODO: Implement me.
+				// Issue: https://github.com/pezkuwichain/pezkuwi-sdk/issues/24
+				let _ = block_import_msg;
+			}
+		}
+	}
+}
+
+/// Handle an incoming collation message from the block builder task.
+/// This builds the collation from the [`CollatorMessage`] and submits it to
+/// the collation-generation subsystem of the relay chain.
+async fn handle_collation_message<Block: BlockT, RClient: RelayChainInterface + Clone + 'static>(
+	message: CollatorMessage<Block>,
+	collator_service: &impl CollatorServiceInterface<Block>,
+	overseer_handle: &mut OverseerHandle,
+	relay_client: RClient,
+	export_pov: Option<PathBuf>,
+) {
+	let CollatorMessage {
+		parent_header,
+		teyrchain_candidate,
+		validation_code_hash,
+		relay_parent,
+		core_index,
+		max_pov_size,
+	} = message;
+
+	let hash = teyrchain_candidate.block.header().hash();
+	let number = *teyrchain_candidate.block.header().number();
+	let (collation, block_data) =
+		match collator_service.build_collation(&parent_header, hash, teyrchain_candidate) {
+			Some(collation) => collation,
+			None => {
+				tracing::warn!(target: LOG_TARGET, %hash, ?number, ?core_index, "Unable to build collation.");
+				return;
+			},
+		};
+
+	block_data.log_size_info();
+
+	if let MaybeCompressedPoV::Compressed(ref pov) = collation.proof_of_validity {
+		if let Some(pov_path) = export_pov {
+			if let Ok(Some(relay_parent_header)) =
+				relay_client.header(BlockId::Hash(relay_parent)).await
+			{
+				if let Some(header) = block_data.blocks().first().map(|b| b.header()) {
+					export_pov_to_path::<Block>(
+						pov_path.clone(),
+						pov.clone(),
+						header.hash(),
+						*header.number(),
+						parent_header.clone(),
+						relay_parent_header.state_root,
+						relay_parent_header.number,
+						max_pov_size,
+					);
+				}
+			} else {
+				tracing::error!(target: LOG_TARGET, "Failed to get relay parent header from hash: {relay_parent:?}");
+			}
+		}
+
+		tracing::info!(
+			target: LOG_TARGET,
+			"Compressed PoV size: {}kb",
+			pov.block_data.0.len() as f64 / 1024f64,
+		);
+	}
+
+	tracing::debug!(target: LOG_TARGET, ?core_index, ?hash, %number, "Submitting collation for core.");
+
+	overseer_handle
+		.send_msg(
+			CollationGenerationMessage::SubmitCollation(SubmitCollationParams {
+				relay_parent,
+				collation,
+				parent_head: parent_header.encode().into(),
+				validation_code_hash,
+				core_index,
+				result_sender: None,
+			}),
+			"SubmitCollation",
+		)
+		.await;
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/mod.rs

@@ -0,0 +1,270 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! # Architecture Overview
+//!
+//! The block building mechanism operates through two coordinated tasks:
+//!
+//! 1. **Block Builder Task**: Orchestrates the timing and execution of teyrchain block production
+//! 2. **Collator Task**: Processes built blocks into collations for relay chain submission
+//!
+//! # Block Builder Task Details
+//!
+//! The block builder task manages block production timing and execution through an iterative
+//! process:
+//!
+//! 1. Awaits the next production signal from the internal timer
+//! 2. Retrieves the current best relay chain block and identifies a valid parent block (see
+//!    [find_potential_parents][cumulus_client_consensus_common::find_potential_parents] for parent
+//!    selection criteria)
+//! 3. Validates that:
+//!    - The teyrchain has an assigned core on the relay chain
+//!    - No block has been previously built on the target core
+//! 4. Executes block building and import operations
+//! 5. Transmits the completed block to the collator task
+//!
+//! # Block Production Timing
+//!
+//! When a block is produced is determined by the following parameters:
+//!
+//! - Teyrchain slot duration
+//! - Number of assigned teyrchain cores
+//! - Teyrchain runtime configuration
+//!
+//! ## Timing Examples
+//!
+//! The following table demonstrates various timing configurations and their effects. The "AURA
+//! Slot" column shows which author is responsible for the block.
+//!
+//! | Slot Duration (ms) | Cores | Production Attempts (ms) | AURA Slot  |
+//! |-------------------|--------|-------------------------|------------|
+//! | 2000              | 3      | 0, 2000, 4000, 6000    | 0, 1, 2, 3 |
+//! | 6000              | 1      | 0, 6000, 12000, 18000  | 0, 1, 2, 3 |
+//! | 6000              | 3      | 0, 2000, 4000, 6000    | 0, 0, 0, 1 |
+//! | 12000             | 1      | 0, 6000, 12000, 18000  | 0, 0, 1, 1 |
+//! | 12000             | 3      | 0, 2000, 4000, 6000    | 0, 0, 0, 0 |
+//!
+//! # Collator Task Details
+//!
+//! The collator task receives built blocks from the block builder task and performs two primary
+//! functions:
+//!
+//! 1. Block compression
+//! 2. Submission to the collation-generation subsystem
+
+use self::{block_builder_task::run_block_builder, collation_task::run_collation_task};
+pub use block_import::{SlotBasedBlockImport, SlotBasedBlockImportHandle};
+use codec::Codec;
+use consensus_common::TeyrchainCandidate;
+use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_client_consensus_common::{self as consensus_common, TeyrchainBlockImportMarker};
+use cumulus_client_consensus_proposer::ProposerInterface;
+use cumulus_primitives_aura::AuraUnincludedSegmentApi;
+use cumulus_primitives_core::RelayParentOffsetApi;
+use cumulus_relay_chain_interface::RelayChainInterface;
+use futures::FutureExt;
+use pezkuwi_primitives::{
+	CollatorPair, CoreIndex, Hash as RelayHash, Id as ParaId, ValidationCodeHash,
+};
+use pezsc_client_api::{backend::AuxStore, BlockBackend, BlockOf, UsageProvider};
+use pezsc_consensus::BlockImport;
+use pezsc_network_types::PeerId;
+use pezsc_utils::mpsc::tracing_unbounded;
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_application_crypto::AppPublic;
+use pezsp_blockchain::HeaderBackend;
+use pezsp_consensus_aura::AuraApi;
+use pezsp_core::{crypto::Pair, traits::SpawnEssentialNamed};
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_keystore::KeystorePtr;
+use pezsp_runtime::traits::{Block as BlockT, Member};
+use std::{path::PathBuf, sync::Arc, time::Duration};
+
+mod block_builder_task;
+mod block_import;
+mod collation_task;
+mod relay_chain_data_cache;
+mod slot_timer;
+
+#[cfg(test)]
+mod tests;
+
+/// Parameters for [`run`].
+pub struct Params<Block, BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS, Spawner> {
+	/// Inherent data providers. Only non-consensus inherent data should be provided, i.e.
+	/// the timestamp, slot, and paras inherents should be omitted, as they are set by this
+	/// collator.
+	pub create_inherent_data_providers: CIDP,
+	/// Used to actually import blocks.
+	pub block_import: BI,
+	/// The underlying para client.
+	pub para_client: Arc<Client>,
+	/// The para client's backend, used to access the database.
+	pub para_backend: Arc<Backend>,
+	/// A handle to the relay-chain client.
+	pub relay_client: RClient,
+	/// A validation code hash provider, used to get the current validation code hash.
+	pub code_hash_provider: CHP,
+	/// The underlying keystore, which should contain Aura consensus keys.
+	pub keystore: KeystorePtr,
+	/// The collator key used to sign collations before submitting to validators.
+	pub collator_key: CollatorPair,
+	/// The collator network peer id.
+	pub collator_peer_id: PeerId,
+	/// The para's ID.
+	pub para_id: ParaId,
+	/// The underlying block proposer this should call into.
+	pub proposer: Proposer,
+	/// The generic collator service used to plug into this consensus engine.
+	pub collator_service: CS,
+	/// The amount of time to spend authoring each block.
+	pub authoring_duration: Duration,
+	/// Whether we should reinitialize the collator config (i.e. we are transitioning to aura).
+	pub reinitialize: bool,
+	/// Offset slots by a fixed duration. This can be used to create more preferrable authoring
+	/// timings.
+	pub slot_offset: Duration,
+	/// The handle returned by [`SlotBasedBlockImport`].
+	pub block_import_handle: SlotBasedBlockImportHandle<Block>,
+	/// Spawner for spawning futures.
+	pub spawner: Spawner,
+	/// Slot duration of the relay chain
+	pub relay_chain_slot_duration: Duration,
+	/// When set, the collator will export every produced `POV` to this folder.
+	pub export_pov: Option<PathBuf>,
+	/// The maximum percentage of the maximum PoV size that the collator can use.
+	/// It will be removed once <https://github.com/pezkuwichain/pezkuwi-sdk/issues/23> is fixed.
+	pub max_pov_percentage: Option<u32>,
+}
+
+/// Run aura-based block building and collation task.
+pub fn run<Block, P, BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS, Spawner>(
+	params: Params<Block, BI, CIDP, Client, Backend, RClient, CHP, Proposer, CS, Spawner>,
+) where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block>
+		+ BlockOf
+		+ AuxStore
+		+ HeaderBackend<Block>
+		+ BlockBackend<Block>
+		+ UsageProvider<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	Client::Api:
+		AuraApi<Block, P::Public> + AuraUnincludedSegmentApi<Block> + RelayParentOffsetApi<Block>,
+	Backend: pezsc_client_api::Backend<Block> + 'static,
+	RClient: RelayChainInterface + Clone + 'static,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+	CIDP::InherentDataProviders: Send,
+	BI: BlockImport<Block> + TeyrchainBlockImportMarker + Send + Sync + 'static,
+	Proposer: ProposerInterface<Block> + Send + Sync + 'static,
+	CS: CollatorServiceInterface<Block> + Send + Sync + Clone + 'static,
+	CHP: consensus_common::ValidationCodeHashProvider<Block::Hash> + Send + 'static,
+	P: Pair + Send + Sync + 'static,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+	Spawner: SpawnEssentialNamed + Clone + 'static,
+{
+	let Params {
+		create_inherent_data_providers,
+		block_import,
+		para_client,
+		para_backend,
+		relay_client,
+		code_hash_provider,
+		keystore,
+		collator_key,
+		collator_peer_id,
+		para_id,
+		proposer,
+		collator_service,
+		authoring_duration,
+		reinitialize,
+		slot_offset,
+		block_import_handle,
+		spawner,
+		export_pov,
+		relay_chain_slot_duration,
+		max_pov_percentage,
+	} = params;
+
+	let (tx, rx) = tracing_unbounded("mpsc_builder_to_collator", 100);
+	let collator_task_params = collation_task::Params {
+		relay_client: relay_client.clone(),
+		collator_key,
+		para_id,
+		reinitialize,
+		collator_service: collator_service.clone(),
+		collator_receiver: rx,
+		block_import_handle,
+		export_pov,
+	};
+
+	let collation_task_fut = run_collation_task::<Block, _, _>(collator_task_params);
+
+	let block_builder_params = block_builder_task::BuilderTaskParams {
+		create_inherent_data_providers,
+		block_import,
+		para_client,
+		para_backend,
+		relay_client,
+		code_hash_provider,
+		keystore,
+		collator_peer_id,
+		para_id,
+		proposer,
+		collator_service,
+		authoring_duration,
+		collator_sender: tx,
+		relay_chain_slot_duration,
+		slot_offset,
+		max_pov_percentage,
+	};
+
+	let block_builder_fut =
+		run_block_builder::<Block, P, _, _, _, _, _, _, _, _>(block_builder_params);
+
+	spawner.spawn_essential_blocking(
+		"slot-based-block-builder",
+		Some("slot-based-collator"),
+		block_builder_fut.boxed(),
+	);
+	spawner.spawn_essential_blocking(
+		"slot-based-collation",
+		Some("slot-based-collator"),
+		collation_task_fut.boxed(),
+	);
+}
+
+/// Message to be sent from the block builder to the collation task.
+///
+/// Contains all data necessary to submit a collation to the relay chain.
+struct CollatorMessage<Block: BlockT> {
+	/// The hash of the relay chain block that provides the context for the teyrchain block.
+	pub relay_parent: RelayHash,
+	/// The header of the parent block.
+	pub parent_header: Block::Header,
+	/// The teyrchain block candidate.
+	pub teyrchain_candidate: TeyrchainCandidate<Block>,
+	/// The validation code hash at the parent block.
+	pub validation_code_hash: ValidationCodeHash,
+	/// Core index that this block should be submitted on
+	pub core_index: CoreIndex,
+	/// Maximum pov size. Currently needed only for exporting PoV.
+	pub max_pov_size: u32,
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/relay_chain_data_cache.rs

@@ -0,0 +1,122 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! Utility for caching [`RelayChainData`] for different relay blocks.
+
+use crate::collators::claim_queue_at;
+use cumulus_primitives_core::CoreSelector;
+use cumulus_relay_chain_interface::RelayChainInterface;
+use pezkuwi_node_subsystem_util::runtime::ClaimQueueSnapshot;
+use pezkuwi_primitives::{
+	Hash as RelayHash, Header as RelayHeader, Id as ParaId, OccupiedCoreAssumption,
+};
+use pezsp_runtime::generic::BlockId;
+
+/// Contains relay chain data necessary for teyrchain block building.
+#[derive(Clone, Debug)]
+pub struct RelayChainData {
+	/// Current relay chain parent header.
+	pub relay_parent_header: RelayHeader,
+	/// The claim queue at the relay parent.
+	pub claim_queue: ClaimQueueSnapshot,
+	/// Maximum configured PoV size on the relay chain.
+	pub max_pov_size: u32,
+	/// The last [`CoreSelector`] we used.
+	pub last_claimed_core_selector: Option<CoreSelector>,
+}
+
+/// Simple helper to fetch relay chain data and cache it based on the current relay chain best block
+/// hash.
+pub struct RelayChainDataCache<RI> {
+	relay_client: RI,
+	para_id: ParaId,
+	cached_data: schnellru::LruMap<RelayHash, RelayChainData>,
+}
+
+impl<RI> RelayChainDataCache<RI>
+where
+	RI: RelayChainInterface + 'static,
+{
+	pub fn new(relay_client: RI, para_id: ParaId) -> Self {
+		Self {
+			relay_client,
+			para_id,
+			// 50 cached relay chain blocks should be more than enough.
+			cached_data: schnellru::LruMap::new(schnellru::ByLength::new(50)),
+		}
+	}
+
+	/// Fetch required [`RelayChainData`] from the relay chain.
+	/// If this data has been fetched in the past for the incoming hash, it will reuse
+	/// cached data.
+	pub async fn get_mut_relay_chain_data(
+		&mut self,
+		relay_parent: RelayHash,
+	) -> Result<&mut RelayChainData, ()> {
+		let insert_data = if self.cached_data.peek(&relay_parent).is_some() {
+			tracing::trace!(target: crate::LOG_TARGET, %relay_parent, "Using cached data for relay parent.");
+			None
+		} else {
+			tracing::trace!(target: crate::LOG_TARGET, %relay_parent, "Relay chain best block changed, fetching new data from relay chain.");
+			Some(self.update_for_relay_parent(relay_parent).await?)
+		};
+
+		Ok(self
+			.cached_data
+			.get_or_insert(relay_parent, || {
+				insert_data.expect("`insert_data` exists if not cached yet; qed")
+			})
+			.expect("There is space for at least one element; qed"))
+	}
+
+	/// Fetch fresh data from the relay chain for the given relay parent hash.
+	async fn update_for_relay_parent(&self, relay_parent: RelayHash) -> Result<RelayChainData, ()> {
+		let claim_queue = claim_queue_at(relay_parent, &self.relay_client).await;
+
+		let Ok(Some(relay_parent_header)) =
+			self.relay_client.header(BlockId::Hash(relay_parent)).await
+		else {
+			tracing::warn!(target: crate::LOG_TARGET, "Unable to fetch latest relay chain block header.");
+			return Err(());
+		};
+
+		let max_pov_size = match self
+			.relay_client
+			.persisted_validation_data(relay_parent, self.para_id, OccupiedCoreAssumption::Included)
+			.await
+		{
+			Ok(None) => return Err(()),
+			Ok(Some(pvd)) => pvd.max_pov_size,
+			Err(err) => {
+				tracing::error!(target: crate::LOG_TARGET, ?err, "Failed to gather information from relay-client");
+				return Err(());
+			},
+		};
+
+		Ok(RelayChainData {
+			relay_parent_header,
+			claim_queue,
+			max_pov_size,
+			last_claimed_core_selector: None,
+		})
+	}
+
+	#[cfg(test)]
+	pub(crate) fn insert_test_data(&mut self, relay_parent: RelayHash, data: RelayChainData) {
+		self.cached_data.insert(relay_parent, data);
+	}
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/slot_timer.rs

@@ -0,0 +1,647 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use crate::LOG_TARGET;
+use codec::Codec;
+use cumulus_primitives_aura::Slot;
+use cumulus_primitives_core::BlockT;
+use pezsc_client_api::UsageProvider;
+use pezsc_consensus_aura::SlotDuration;
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_application_crypto::AppPublic;
+use pezsp_consensus_aura::AuraApi;
+use pezsp_core::Pair;
+use pezsp_runtime::traits::Member;
+use pezsp_timestamp::Timestamp;
+use std::{
+	cmp::{max, min},
+	sync::Arc,
+	time::Duration,
+};
+
+/// Lower limits of allowed block production interval.
+/// Defensive mechanism, corresponds to 12 cores at 6 second block time.
+const BLOCK_PRODUCTION_MINIMUM_INTERVAL_MS: Duration = Duration::from_millis(500);
+
+/// Theoretically, the block production is capped at `BLOCK_PRODUCTION_MINIMUM_INTERVAL_MS`.
+/// In practice, there might be slight deviations due to timing inaccuracies and delays.
+///
+/// This constant is taken into account while adjusting the authoring duration to fit into the slot.
+/// Therefore, it will only reduce the authoring duration if we are within the
+/// `BLOCK_PRODUCTION_ADJUSTMENT_MS` threshold of the next slot.
+///
+/// ### 12 cores 500ms blocks
+///
+/// For example, for 12 cores 500ms blocks: the next slot is scheduled in 490ms due to delays.
+/// In that case, we still want to attempt producing the block, as missing the slot would be worse
+/// than producing slightly too fast.
+const BLOCK_PRODUCTION_THRESHOLD_MS: Duration = Duration::from_millis(100);
+
+/// The amount of time the authoring duration of the last block production attempt
+/// should be reduced by to fit into the slot timing.
+const BLOCK_PRODUCTION_ADJUSTMENT_MS: Duration = Duration::from_millis(1000);
+
+#[derive(Debug)]
+pub(crate) struct SlotInfo {
+	pub timestamp: Timestamp,
+	pub slot: Slot,
+}
+
+/// Manages block-production timings based on chain parameters and assigned cores.
+#[derive(Debug)]
+pub(crate) struct SlotTimer<Block, Client, P> {
+	/// Teyrchain client that is used for runtime calls
+	client: Arc<Client>,
+	/// Offset the current time by this duration.
+	time_offset: Duration,
+	/// Last reported core count.
+	last_reported_core_num: Option<u32>,
+	/// Slot duration of the relay chain. This is used to compute how man block-production
+	/// attempts we should trigger per relay chain block.
+	relay_slot_duration: Duration,
+	/// Stores the latest slot that was reported by [`Self::wait_until_next_slot`].
+	last_reported_slot: Option<Slot>,
+	_marker: std::marker::PhantomData<(Block, Box<dyn Fn(P) + Send + Sync + 'static>)>,
+}
+
+/// Compute when to try block-authoring next.
+/// The exact time point is determined by the slot duration of relay- and teyrchain as
+/// well as the last observed core count. If more cores are available, we attempt to author blocks
+/// for them.
+///
+/// Returns a tuple with:
+/// - `Duration`: How long to wait until the next slot.
+/// - `Slot`: The AURA slot used for authoring
+fn compute_next_wake_up_time(
+	para_slot_duration: SlotDuration,
+	relay_slot_duration: Duration,
+	core_count: Option<u32>,
+	time_now: Duration,
+	time_offset: Duration,
+) -> (Duration, Slot) {
+	let para_slots_per_relay_block =
+		(relay_slot_duration.as_millis() / para_slot_duration.as_millis() as u128) as u32;
+	let assigned_core_num = core_count.unwrap_or(1);
+
+	// Trigger at least once per relay block, if we have for example 12 second slot duration,
+	// we should still produce two blocks if we are scheduled on every relay block.
+	let mut block_production_interval = min(para_slot_duration.as_duration(), relay_slot_duration);
+
+	if assigned_core_num > para_slots_per_relay_block &&
+		para_slot_duration.as_duration() >= relay_slot_duration
+	{
+		block_production_interval =
+			max(relay_slot_duration / assigned_core_num, BLOCK_PRODUCTION_MINIMUM_INTERVAL_MS);
+		tracing::debug!(
+			target: LOG_TARGET,
+			?block_production_interval,
+			"Expected to produce for {assigned_core_num} cores but only have {para_slots_per_relay_block} slots. Attempting to produce multiple blocks per slot."
+		);
+	}
+
+	let (duration, timestamp) =
+		time_until_next_attempt(time_now, block_production_interval, time_offset);
+	let aura_slot = Slot::from_timestamp(timestamp, para_slot_duration);
+	(duration, aura_slot)
+}
+
+/// Compute the time until the next slot changes.
+///
+/// Returns None if the next slot cannot be computed.
+fn compute_time_until_next_slot_change(
+	para_slot_duration: SlotDuration,
+	time_now: Duration,
+	time_offset: Duration,
+	last_reported_slot: Slot,
+) -> Option<(Duration, Slot)> {
+	let now = time_now.saturating_sub(time_offset);
+	let next_slot = last_reported_slot + Slot::from(1);
+
+	let Some(next_slot_timestamp) = next_slot.timestamp(para_slot_duration) else {
+		return None;
+	};
+	let remaining_time = next_slot_timestamp.as_duration().saturating_sub(now);
+
+	Some((remaining_time, next_slot))
+}
+
+/// Returns current duration since Unix epoch.
+fn duration_now() -> Duration {
+	use std::time::SystemTime;
+	let now = SystemTime::now();
+	now.duration_since(SystemTime::UNIX_EPOCH).unwrap_or_else(|e| {
+		panic!("Current time {:?} is before Unix epoch. Something is wrong: {:?}", now, e)
+	})
+}
+
+/// Adjust the authoring duration.
+fn adjust_authoring_duration(
+	mut authoring_duration: Duration,
+	next_block: (Duration, Slot),
+	next_slot_change: (Duration, Slot),
+	different_authors: bool,
+) -> Option<Duration> {
+	let (duration, next_block_slot) = next_block;
+	let (duration_until_next_slot, next_slot) = next_slot_change;
+
+	// The authoring of blocks must stop 1 second before the slot ends.
+	let duration_until_deadline =
+		duration_until_next_slot.saturating_sub(BLOCK_PRODUCTION_ADJUSTMENT_MS);
+	tracing::debug!(
+		target: LOG_TARGET,
+		?authoring_duration,
+		?duration,
+		?next_block_slot,
+		?duration_until_next_slot,
+		?next_slot,
+		?duration_until_deadline,
+		?different_authors,
+		"Adjusting authoring duration for slot.",
+	);
+
+	// Ensure no blocks are produced in the last second of the slot,
+	// regardless of authoring duration.
+	if duration_until_deadline == Duration::ZERO {
+		if different_authors {
+			tracing::warn!(
+				target: LOG_TARGET,
+				?duration_until_next_slot,
+				?next_slot,
+				"Not enough time left in the slot to adjust authoring duration. Skipping block production for the slot."
+			);
+
+			return None;
+		}
+
+		// If authors are the same, we can still attempt producing the block
+		// considering the next block duration.
+		return Some(authoring_duration.min(duration));
+	}
+
+	// Clamp the authoring duration to fit into the slot deadline only if authors are different.
+	// For most cases, the deadline is farther in the future than the authoring duration.
+	if different_authors && authoring_duration >= duration_until_deadline {
+		authoring_duration = duration_until_deadline;
+
+		// Ensure we are not going below the minimum interval within a reasonable threshold.
+		// For 12 cores, we might have a scenario where the last 3 blocks are skipped:
+		// - Block 10: next slot change in 1.493s:
+		// 	 - After adjusting the deadline: 1.493s - 1s = 0.493s the block could be produced
+		//     without issues.
+		// - Block 11: next slot change in 0.993s - skipped by the deadline
+		// - Block 12: next slot change in 0.493s - skipped by the deadline
+		if authoring_duration <
+			BLOCK_PRODUCTION_MINIMUM_INTERVAL_MS.saturating_sub(BLOCK_PRODUCTION_THRESHOLD_MS)
+		{
+			tracing::debug!(
+				target: LOG_TARGET,
+				?authoring_duration,
+				?next_slot,
+				"Authoring duration is below minimum. Skipping block production for the slot."
+			);
+			return None;
+		}
+	}
+
+	// The `duration` intends to slightly adjust when then block production
+	// attempt happens. This goes slightly below the `BLOCK_PRODUCTION_MINIMUM_INTERVAL_MS`
+	// threshold.
+	Some(authoring_duration.min(duration))
+}
+
+/// Returns the duration until the next block production should be attempted.
+/// Returns:
+/// - Duration: The duration until the next attempt.
+fn time_until_next_attempt(
+	now: Duration,
+	block_production_interval: Duration,
+	offset: Duration,
+) -> (Duration, Timestamp) {
+	let now = now.as_millis().saturating_sub(offset.as_millis());
+
+	let next_slot_time = ((now + block_production_interval.as_millis()) /
+		block_production_interval.as_millis()) *
+		block_production_interval.as_millis();
+	let remaining_millis = next_slot_time - now;
+	(Duration::from_millis(remaining_millis as u64), Timestamp::from(next_slot_time as u64))
+}
+
+impl<Block, Client, P> SlotTimer<Block, Client, P>
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block> + UsageProvider<Block> + Send + Sync + 'static,
+	Client::Api: AuraApi<Block, P::Public>,
+	P: Pair,
+	P::Public: AppPublic + Member + Codec,
+	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
+{
+	/// Create a new slot timer.
+	pub fn new_with_offset(
+		client: Arc<Client>,
+		time_offset: Duration,
+		relay_slot_duration: Duration,
+	) -> Self {
+		Self {
+			client,
+			time_offset,
+			last_reported_core_num: None,
+			relay_slot_duration,
+			last_reported_slot: Default::default(),
+			_marker: Default::default(),
+		}
+	}
+
+	/// Inform the slot timer about the last seen number of cores.
+	pub fn update_scheduling(&mut self, num_cores_next_block: u32) {
+		self.last_reported_core_num = Some(num_cores_next_block);
+	}
+
+	/// Returns the slot and how much time left until the next block production attempt.
+	pub fn time_until_next_block(&mut self, slot_duration: SlotDuration) -> (Duration, Slot) {
+		compute_next_wake_up_time(
+			slot_duration,
+			self.relay_slot_duration,
+			self.last_reported_core_num,
+			duration_now(),
+			self.time_offset,
+		)
+	}
+
+	/// Compute the time until the next slot changes.
+	fn time_until_next_slot_change(
+		&mut self,
+		slot_duration: SlotDuration,
+	) -> Option<(Duration, Slot)> {
+		compute_time_until_next_slot_change(
+			slot_duration,
+			duration_now(),
+			self.time_offset,
+			self.last_reported_slot.unwrap_or_default(),
+		)
+	}
+
+	/// Check if two slots have different authors based on AURA round-robin algorithm.
+	///
+	/// Returns true if the authors for the two slots are different.
+	fn check_different_slot_authors(&self, slot: Slot, next_slot: Slot) -> bool {
+		let best_hash = self.client.usage_info().chain.best_hash;
+
+		let Ok(authorities) = self.client.runtime_api().authorities(best_hash) else {
+			tracing::warn!(target: LOG_TARGET, "Failed to fetch authorities for slot author comparison");
+			// Presume they are different, this will adjust the slot authoring duration more
+			// conservatively.
+			return true;
+		};
+
+		let authorities_len = authorities.len() as u64;
+		if authorities_len <= 1 {
+			return false;
+		}
+
+		let author1_idx = *slot % authorities_len;
+		let author2_idx = *next_slot % authorities_len;
+
+		author1_idx != author2_idx
+	}
+
+	/// Adjust the authoring duration to fit into the slot timing.
+	///
+	/// Returns the adjusted authoring duration and the slot that it corresponds to.
+	pub fn adjust_authoring_duration(&mut self, authoring_duration: Duration) -> Option<Duration> {
+		let Ok(slot_duration) = crate::slot_duration(&*self.client) else {
+			tracing::error!(target: LOG_TARGET, "Failed to fetch slot duration from runtime.");
+			return None;
+		};
+
+		let next_block = self.time_until_next_block(slot_duration);
+		let Some(next_slot_change) = self.time_until_next_slot_change(slot_duration) else {
+			tracing::error!(
+				target: LOG_TARGET,
+				"Failed to compute time until next slot change. Using unadjusted authoring duration."
+			);
+			return Some(authoring_duration);
+		};
+
+		// Check if authors at current and next slots are different
+		let current_slot = self.last_reported_slot.unwrap_or(next_block.1);
+		let different_authors = self.check_different_slot_authors(current_slot, next_slot_change.1);
+
+		adjust_authoring_duration(
+			authoring_duration,
+			next_block,
+			next_slot_change,
+			different_authors,
+		)
+	}
+
+	/// Returns a future that resolves when the next block production should be attempted.
+	pub async fn wait_until_next_slot(&mut self) -> Result<(), ()> {
+		let slot_duration = match crate::slot_duration(&*self.client) {
+			Ok(d) => d,
+			Err(error) => {
+				tracing::error!(target: LOG_TARGET, %error, "Failed to fetch slot duration from runtime.");
+				return Err(());
+			},
+		};
+
+		let (time_until_next_attempt, mut next_aura_slot) =
+			self.time_until_next_block(slot_duration);
+
+		tracing::trace!(
+			target: LOG_TARGET,
+			?time_until_next_attempt,
+			aura_slot = ?next_aura_slot,
+			last_reported = ?self.last_reported_slot,
+			"Determined next block production opportunity."
+		);
+
+		match self.last_reported_slot {
+			// If we already reported a slot, we don't want to skip a slot. But we also don't want
+			// to go through all the slots if a node was halted for some reason.
+			Some(ls) if ls + 1 < next_aura_slot && next_aura_slot <= ls + 3 => {
+				next_aura_slot = ls + 1u64;
+			},
+			None | Some(_) => {
+				tracing::trace!(target: LOG_TARGET, ?time_until_next_attempt, "Sleeping until the next slot.");
+				tokio::time::sleep(time_until_next_attempt).await;
+			},
+		}
+
+		tracing::debug!(
+			target: LOG_TARGET,
+			?slot_duration,
+			aura_slot = ?next_aura_slot,
+			"New block production opportunity."
+		);
+
+		self.last_reported_slot = Some(next_aura_slot);
+
+		Ok(())
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use rstest::rstest;
+	use pezsc_consensus_aura::SlotDuration;
+	const RELAY_CHAIN_SLOT_DURATION: u64 = 6000;
+
+	#[rstest]
+	// Test that different now timestamps have correct impact
+	//                    ||||
+	#[case(6000, Some(1), 1000, 0, 5000)]
+	#[case(6000, Some(1), 0, 0, 6000)]
+	#[case(6000, Some(1), 6000, 0, 6000)]
+	#[case(6000, Some(0), 6000, 0, 6000)]
+	// Test that `None` core defaults to 1
+	//           ||||
+	#[case(6000, None, 1000, 0, 5000)]
+	#[case(6000, None, 0, 0, 6000)]
+	#[case(6000, None, 6000, 0, 6000)]
+	// Test that offset affects the current time correctly
+	//                          ||||
+	#[case(6000, Some(1), 1000, 1000, 6000)]
+	#[case(6000, Some(1), 12000, 2000, 2000)]
+	#[case(6000, Some(1), 12000, 6000, 6000)]
+	#[case(6000, Some(1), 12000, 7000, 1000)]
+	// Test that number of cores affects the block production interval
+	//           |||||||
+	#[case(6000, Some(3), 12000, 0, 2000)]
+	#[case(6000, Some(2), 12000, 0, 3000)]
+	#[case(6000, Some(3), 11999, 0, 1)]
+	// High core count
+	//           ||||||||
+	#[case(6000, Some(12), 0, 0, 500)]
+	/// Test that the minimum block interval is respected
+	/// at high core counts.
+	///          |||||||||
+	#[case(6000, Some(100), 0, 0, 500)]
+	// Test that slot_duration works correctly
+	//     ||||
+	#[case(2000, Some(1), 1000, 0, 1000)]
+	#[case(2000, Some(1), 3000, 0, 1000)]
+	#[case(2000, Some(1), 10000, 0, 2000)]
+	#[case(2000, Some(2), 1000, 0, 1000)]
+	// Cores are ignored if relay_slot_duration != para_slot_duration
+	//           |||||||
+	#[case(2000, Some(3), 3000, 0, 1000)]
+	// For long slot durations, we should still check
+	// every relay chain block for the slot.
+	//     |||||
+	#[case(12000, None, 0, 0, 6000)]
+	#[case(12000, None, 6100, 0, 5900)]
+	#[case(12000, None, 6000, 2000, 2000)]
+	#[case(12000, Some(2), 6000, 0, 3000)]
+	#[case(12000, Some(3), 6000, 0, 2000)]
+	#[case(12000, Some(3), 8100, 0, 1900)]
+	fn test_get_next_slot(
+		#[case] para_slot_millis: u64,
+		#[case] core_count: Option<u32>,
+		#[case] time_now: u64,
+		#[case] offset_millis: u64,
+		#[case] expected_wait_duration: u128,
+	) {
+		let para_slot_duration = SlotDuration::from_millis(para_slot_millis); // 6 second slots
+		let relay_slot_duration = Duration::from_millis(RELAY_CHAIN_SLOT_DURATION);
+		let time_now = Duration::from_millis(time_now); // 1 second passed
+		let offset = Duration::from_millis(offset_millis);
+
+		let (wait_duration, _) = compute_next_wake_up_time(
+			para_slot_duration,
+			relay_slot_duration,
+			core_count,
+			time_now,
+			offset,
+		);
+
+		assert_eq!(wait_duration.as_millis(), expected_wait_duration, "Wait time mismatch.");
+		// Should wait 5 seconds
+	}
+
+	#[rstest]
+	// Basic slot change scenarios
+	#[case(6000, 0, 0, Slot::from(0), 6000, Slot::from(1))]
+	#[case(6000, 1000, 0, Slot::from(0), 5000, Slot::from(1))]
+	#[case(6000, 6000, 0, Slot::from(1), 6000, Slot::from(2))]
+	#[case(6000, 12000, 0, Slot::from(2), 6000, Slot::from(3))]
+	// Test with offset
+	#[case(6000, 1000, 1000, Slot::from(0), 6000, Slot::from(1))]
+	#[case(6000, 2000, 1000, Slot::from(0), 5000, Slot::from(1))]
+	#[case(6000, 6000, 3000, Slot::from(0), 3000, Slot::from(1))]
+	// Different slot durations
+	#[case(3000, 1000, 0, Slot::from(0), 2000, Slot::from(1))]
+	#[case(3000, 3000, 0, Slot::from(1), 3000, Slot::from(2))]
+	#[case(12000, 6000, 0, Slot::from(0), 6000, Slot::from(1))]
+	#[case(12000, 12000, 0, Slot::from(1), 12000, Slot::from(2))]
+	// Edge cases - at slot boundary
+	#[case(6000, 5999, 0, Slot::from(0), 1, Slot::from(1))]
+	#[case(6000, 11999, 0, Slot::from(1), 1, Slot::from(2))]
+	fn test_compute_time_until_next_slot_change(
+		#[case] para_slot_millis: u64,
+		#[case] time_now: u64,
+		#[case] offset_millis: u64,
+		#[case] last_reported_slot: Slot,
+		#[case] expected_duration: u128,
+		#[case] expected_next_slot: Slot,
+	) {
+		let para_slot_duration = SlotDuration::from_millis(para_slot_millis);
+		let time_now = Duration::from_millis(time_now);
+		let offset = Duration::from_millis(offset_millis);
+
+		let result = compute_time_until_next_slot_change(
+			para_slot_duration,
+			time_now,
+			offset,
+			last_reported_slot,
+		);
+
+		assert!(result.is_some(), "Expected result to be Some");
+		let (duration, next_slot) = result.unwrap();
+		assert_eq!(duration.as_millis(), expected_duration, "Duration mismatch");
+		assert_eq!(next_slot, expected_next_slot, "Next slot mismatch");
+	}
+
+	#[rstest]
+	// Various scenarios for 2s block production adjustment.
+	#[case::blocks_2s_fits_next_block(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(2000), Slot::from(1)), // Next block
+		(Duration::from_millis(4000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(2000)), // Expected
+	)]
+	#[case::blocks_2s_closer_next_slot(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1950), Slot::from(1)), // Next block
+		(Duration::from_millis(4000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(1950)), // Expected
+	)]
+	#[case::blocks_2s_closer_next_slot_bigger(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1500), Slot::from(1)), // Next block
+		(Duration::from_millis(4000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(1500)), // Expected
+	)]
+	#[case::blocks_2s_reduce_by_1s(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(2000), Slot::from(1)), // Next block
+		(Duration::from_millis(2000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(1000)), // Expected
+	)]
+	#[case::blocks_2s_reduce_by_1s_plus_offset(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1950), Slot::from(1)), // Next block
+		(Duration::from_millis(1950), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(950)), // Expected
+	)]
+	#[case::blocks_2s_reduce_to_minimum(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1400), Slot::from(1)), // Next block
+		(Duration::from_millis(1400), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(400)), // Expected
+	)]
+	#[case::blocks_2s_reduce_below_minimum(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1300), Slot::from(1)), // Next block
+		(Duration::from_millis(1300), Slot::from(2)), // Next slot change
+		true, // Different authors
+		None, // Expected to reduce below minimum
+	)]
+	#[case::blocks_2s_same_author(
+		Duration::from_millis(2000), // Authoring duration
+		(Duration::from_millis(1400), Slot::from(1)), // Next block
+		(Duration::from_millis(1400), Slot::from(2)), // Next slot change
+		false, // Different authors
+		Some(Duration::from_millis(1400)), // Expected no adjustment for last second.
+	)]
+	// Various scenarios for 500ms block production adjustment.
+	#[case::blocks_500ms_fits_next_block(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(500), Slot::from(1)), // Next block
+		(Duration::from_millis(2000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(500)), // Expected
+	)]
+	#[case::blocks_500ms_closer_next_slot(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(450), Slot::from(1)), // Next block
+		(Duration::from_millis(2000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(450)), // Expected
+	)]
+	#[case::blocks_500ms_closer_next_slot_bigger(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(400), Slot::from(1)), // Next block
+		(Duration::from_millis(1500), Slot::from(2)), // Next slot change
+		true, // Different authors
+		Some(Duration::from_millis(400)), // Expected
+	)]
+	#[case::blocks_500ms_reduce_by_1s(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(500), Slot::from(1)), // Next block
+		(Duration::from_millis(1000), Slot::from(2)), // Next slot change
+		true, // Different authors
+		None, // Expected
+	)]
+	#[case::blocks_500ms_reduce_by_1s_closer(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(500), Slot::from(1)), // Next block
+		(Duration::from_millis(500), Slot::from(2)), // Next slot change
+		true, // Different authors
+		None, // Expected
+	)]
+	// If we are producing with 1 collator for 500ms authoring duration,
+	// we must produce the last two slots and ignore the 1s adjustment.
+	#[case::blocks_500ms_same_author(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(410), Slot::from(1)), // Next block
+		(Duration::from_millis(1000), Slot::from(2)), // Next slot change
+		false, // Different authors
+		Some(Duration::from_millis(410)), // Expected no adjustment for last second.
+	)]
+	#[case::blocks_500ms_same_author_closer(
+		Duration::from_millis(500), // Authoring duration
+		(Duration::from_millis(400), Slot::from(1)), // Next block
+		(Duration::from_millis(400), Slot::from(2)), // Next slot change
+		false, // Different authors
+		Some(Duration::from_millis(400)), // Expected no adjustment for last second.
+	)]
+	fn test_adjust_authoring_duration(
+		#[case] authoring_duration: Duration,
+		#[case] next_block: (Duration, Slot),
+		#[case] next_slot_change: (Duration, Slot),
+		#[case] different_authors: bool,
+		#[case] expected: Option<Duration>,
+	) {
+		pezsp_tracing::init_for_tests();
+
+		let result = adjust_authoring_duration(
+			authoring_duration,
+			next_block,
+			next_slot_change,
+			different_authors,
+		);
+		tracing::debug!("Adjusted authoring duration: {:?}", result);
+		assert_eq!(result, expected);
+	}
+}

pezcumulus/client/consensus/aura/src/collators/slot_based/tests.rs

@@ -0,0 +1,730 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use super::{
+	block_builder_task::{determine_core, offset_relay_parent_find_descendants},
+	relay_chain_data_cache::{RelayChainData, RelayChainDataCache},
+};
+use async_trait::async_trait;
+use codec::Encode;
+use cumulus_primitives_core::{ClaimQueueOffset, CoreInfo, CoreSelector, CumulusDigestItem};
+use cumulus_relay_chain_interface::*;
+use futures::Stream;
+use pezkuwi_node_subsystem_util::runtime::ClaimQueueSnapshot;
+use pezkuwi_primitives::{
+	CandidateEvent, CommittedCandidateReceiptV2, CoreIndex, Hash as RelayHash,
+	Header as RelayHeader, Id as ParaId,
+};
+use rstest::rstest;
+use pezsc_consensus_babe::{
+	AuthorityId, ConsensusLog as BabeConsensusLog, NextEpochDescriptor, BABE_ENGINE_ID,
+};
+use pezsp_core::sr25519;
+use pezsp_runtime::{generic::BlockId, testing::Header as TestHeader, traits::Header};
+use pezsp_version::RuntimeVersion;
+use std::{
+	collections::{BTreeMap, HashMap, VecDeque},
+	pin::Pin,
+};
+
+#[tokio::test]
+async fn offset_test_zero_offset() {
+	let (headers, best_hash) = create_header_chain();
+
+	let client = TestRelayClient::new(headers);
+
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, best_hash, 0).await;
+	assert!(result.is_ok());
+	let data = result.unwrap().unwrap();
+	assert_eq!(data.descendants_len(), 0);
+	assert_eq!(data.relay_parent().hash(), best_hash);
+	assert!(data.into_inherent_descendant_list().is_empty());
+}
+
+#[tokio::test]
+async fn offset_test_two_offset() {
+	let (headers, best_hash) = create_header_chain();
+
+	let client = TestRelayClient::new(headers);
+
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, best_hash, 2).await;
+	assert!(result.is_ok());
+	let data = result.unwrap().unwrap();
+	assert_eq!(data.descendants_len(), 2);
+	assert_eq!(*data.relay_parent().number(), 98);
+	let descendant_list = data.into_inherent_descendant_list();
+	assert_eq!(descendant_list.len(), 3);
+	assert_eq!(*descendant_list.first().unwrap().number(), 98);
+	assert_eq!(*descendant_list.last().unwrap().number(), 100);
+}
+
+#[tokio::test]
+async fn offset_test_five_offset() {
+	let (headers, best_hash) = create_header_chain();
+
+	let client = TestRelayClient::new(headers);
+
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, best_hash, 5).await;
+	assert!(result.is_ok());
+	let data = result.unwrap().unwrap();
+	assert_eq!(data.descendants_len(), 5);
+	assert_eq!(*data.relay_parent().number(), 95);
+	let descendant_list = data.into_inherent_descendant_list();
+	assert_eq!(descendant_list.len(), 6);
+	assert_eq!(*descendant_list.first().unwrap().number(), 95);
+	assert_eq!(*descendant_list.last().unwrap().number(), 100);
+}
+
+#[tokio::test]
+async fn offset_test_too_long() {
+	let (headers, _best_hash) = create_header_chain();
+
+	let client = TestRelayClient::new(headers);
+
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, _best_hash, 200).await;
+	assert!(result.is_err());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, _best_hash, 101).await;
+	assert!(result.is_err());
+}
+
+#[derive(PartialEq)]
+enum HasEpochChange {
+	Yes,
+	No,
+}
+
+#[rstest]
+#[case::in_best(
+	&[HasEpochChange::No, HasEpochChange::No, HasEpochChange::Yes],
+)]
+#[case::in_first_ancestor(
+	&[HasEpochChange::No, HasEpochChange::Yes, HasEpochChange::No],
+)]
+#[case::in_second_ancestor(
+	&[HasEpochChange::Yes, HasEpochChange::No, HasEpochChange::No],
+)]
+#[tokio::test]
+async fn offset_returns_none_when_epoch_change_encountered(#[case] flags: &[HasEpochChange]) {
+	let (headers, best_hash) = build_headers_with_epoch_flags(flags);
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	let result = offset_relay_parent_find_descendants(&mut cache, best_hash, 3).await;
+	assert!(result.is_ok());
+	assert!(result.unwrap().is_none());
+}
+
+#[tokio::test]
+async fn determine_core_new_relay_parent() {
+	let (headers, _best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: Default::default(),
+		number: 100,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a test para parent header at block 0 (genesis)
+	let para_parent = TestHeader::new_from_number(0);
+
+	// Setup claim queue data for the cache
+	cache.set_test_data(relay_parent.clone(), vec![CoreIndex(0), CoreIndex(1)]);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	let core = result.unwrap();
+	let core = core.unwrap();
+	assert_eq!(core.core_selector(), CoreSelector(0));
+	assert_eq!(core.core_index(), CoreIndex(0));
+	assert_eq!(core.total_cores(), 2);
+}
+
+#[tokio::test]
+async fn determine_core_with_core_info() {
+	let (headers, best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: best_hash,
+		number: 101,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a para parent header with core info in digest
+	let core_info = CoreInfo {
+		selector: CoreSelector(0),
+		claim_queue_offset: ClaimQueueOffset(0),
+		number_of_cores: 3.into(),
+	};
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	digest.push(CumulusDigestItem::CoreInfo(core_info).to_digest_item());
+	// Add relay parent storage root to make it a non-new relay parent
+	digest.push(cumulus_primitives_core::rpsr_digest::relay_parent_storage_root_item(
+		*relay_parent.state_root(),
+		*relay_parent.number(),
+	));
+
+	let para_parent = TestHeader {
+		parent_hash: best_hash.into(),
+		number: 1,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest,
+	};
+
+	// Setup claim queue data for the cache
+	cache.set_test_data(relay_parent.clone(), vec![CoreIndex(0), CoreIndex(1), CoreIndex(2)]);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	match result {
+		Ok(Some(core)) => {
+			assert_eq!(core.core_selector(), CoreSelector(1)); // Should be next selector (0 + 1)
+			assert_eq!(core.core_index(), CoreIndex(1));
+			assert_eq!(core.total_cores(), 3);
+		},
+		Ok(None) => panic!("Expected Some core, got None"),
+		Err(()) => panic!("determine_core returned error"),
+	}
+}
+
+#[tokio::test]
+async fn determine_core_no_cores_available() {
+	let (headers, _best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: Default::default(),
+		number: 100,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a test para parent header at block 0 (genesis)
+	let para_parent = TestHeader::new_from_number(0);
+
+	// Setup empty claim queue
+	cache.set_test_data(relay_parent.clone(), vec![]);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	let core = result.unwrap();
+	assert!(core.is_none());
+}
+
+#[tokio::test]
+async fn determine_core_selector_overflow() {
+	let (headers, best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: best_hash,
+		number: 101,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	let core_info = CoreInfo {
+		selector: CoreSelector(1),
+		claim_queue_offset: ClaimQueueOffset(0),
+		number_of_cores: 2.into(),
+	};
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	digest.push(CumulusDigestItem::CoreInfo(core_info).to_digest_item());
+	// Add relay parent storage root to make it a non-new relay parent
+	digest.push(cumulus_primitives_core::rpsr_digest::relay_parent_storage_root_item(
+		*relay_parent.state_root(),
+		*relay_parent.number(),
+	));
+
+	let para_parent = TestHeader {
+		parent_hash: best_hash.into(),
+		number: 1,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest,
+	};
+
+	// Setup claim queue with only 2 cores
+	cache.set_test_data(relay_parent.clone(), vec![CoreIndex(0), CoreIndex(1)]);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	let core = result.unwrap();
+	assert!(core.is_none()); // Should return None when selector overflows
+}
+
+#[tokio::test]
+async fn determine_core_uses_last_claimed_core_selector() {
+	let (headers, best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: best_hash,
+		number: 101,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a para parent header without core info in digest (non-genesis)
+	// Need to add relay parent storage root to digest to make it a non-new relay parent
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	digest.push(cumulus_primitives_core::rpsr_digest::relay_parent_storage_root_item(
+		*relay_parent.state_root(),
+		*relay_parent.number(),
+	));
+
+	let para_parent = TestHeader {
+		parent_hash: best_hash.into(),
+		number: 1,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest,
+	};
+
+	// Setup claim queue data with last_claimed_core_selector set to 1
+	cache.set_test_data_with_last_selector(
+		relay_parent.clone(),
+		vec![CoreIndex(0), CoreIndex(1), CoreIndex(2)],
+		Some(CoreSelector(1)),
+	);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	match result {
+		Ok(Some(core)) => {
+			// Should use last_claimed_core_selector (1) + 1 = 2
+			assert_eq!(core.core_selector(), CoreSelector(2));
+			assert_eq!(core.core_index(), CoreIndex(2));
+			assert_eq!(core.total_cores(), 3);
+		},
+		Ok(None) => panic!("Expected Some core, got None"),
+		Err(()) => panic!("determine_core returned error"),
+	}
+}
+
+#[tokio::test]
+async fn determine_core_uses_last_claimed_core_selector_wraps_around() {
+	let (headers, best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: best_hash,
+		number: 101,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a para parent header without core info in digest (non-genesis)
+	// Need to add relay parent storage root to digest to make it a non-new relay parent
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	digest.push(cumulus_primitives_core::rpsr_digest::relay_parent_storage_root_item(
+		*relay_parent.state_root(),
+		*relay_parent.number(),
+	));
+
+	let para_parent = TestHeader {
+		parent_hash: best_hash.into(),
+		number: 1,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest,
+	};
+
+	// Setup claim queue data with last_claimed_core_selector set to 2 (last index)
+	// Next selector should wrap around to out of bounds and return None
+	cache.set_test_data_with_last_selector(
+		relay_parent.clone(),
+		vec![CoreIndex(0), CoreIndex(1), CoreIndex(2)],
+		Some(CoreSelector(2)),
+	);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	match result {
+		Ok(Some(_)) => panic!("Expected None due to selector overflow"),
+		Ok(None) => {
+			// This is expected - selector 2 + 1 = 3, but only cores 0,1,2 available
+		},
+		Err(()) => panic!("determine_core returned error"),
+	}
+}
+
+#[tokio::test]
+async fn determine_core_no_last_claimed_core_selector() {
+	let (headers, best_hash) = create_header_chain();
+	let client = TestRelayClient::new(headers);
+	let mut cache = RelayChainDataCache::new(client, 1.into());
+
+	// Create a test relay parent header
+	let relay_parent = RelayHeader {
+		parent_hash: best_hash,
+		number: 101,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest: Default::default(),
+	};
+
+	// Create a para parent header without core info in digest (non-genesis)
+	// Need to add relay parent storage root to digest to make it a non-new relay parent
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	digest.push(cumulus_primitives_core::rpsr_digest::relay_parent_storage_root_item(
+		*relay_parent.state_root(),
+		*relay_parent.number(),
+	));
+
+	let para_parent = TestHeader {
+		parent_hash: best_hash.into(),
+		number: 1,
+		state_root: Default::default(),
+		extrinsics_root: Default::default(),
+		digest,
+	};
+
+	// Setup claim queue data with no last_claimed_core_selector (None)
+	cache.set_test_data_with_last_selector(
+		relay_parent.clone(),
+		vec![CoreIndex(0), CoreIndex(1), CoreIndex(2)],
+		None,
+	);
+
+	let result = determine_core(&mut cache, &relay_parent, 1.into(), &para_parent, 0).await;
+
+	match result {
+		Ok(Some(core)) => {
+			// Should start from selector 0 + 1 = 1 when no last selector
+			assert_eq!(core.core_selector(), CoreSelector(1));
+			assert_eq!(core.core_index(), CoreIndex(1));
+			assert_eq!(core.total_cores(), 3);
+		},
+		Ok(None) => panic!("Expected Some core, got None"),
+		Err(()) => panic!("determine_core returned error"),
+	}
+}
+
+#[derive(Clone)]
+struct TestRelayClient {
+	headers: HashMap<RelayHash, RelayHeader>,
+}
+
+impl TestRelayClient {
+	fn new(headers: HashMap<RelayHash, RelayHeader>) -> Self {
+		Self { headers }
+	}
+}
+
+#[async_trait]
+impl RelayChainInterface for TestRelayClient {
+	async fn validators(&self, _: RelayHash) -> RelayChainResult<Vec<ValidatorId>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn best_block_hash(&self) -> RelayChainResult<RelayHash> {
+		unimplemented!("Not needed for test")
+	}
+	async fn finalized_block_hash(&self) -> RelayChainResult<RelayHash> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn retrieve_dmq_contents(
+		&self,
+		_: ParaId,
+		_: RelayHash,
+	) -> RelayChainResult<Vec<InboundDownwardMessage>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn retrieve_all_inbound_hrmp_channel_contents(
+		&self,
+		_: ParaId,
+		_: RelayHash,
+	) -> RelayChainResult<BTreeMap<ParaId, Vec<InboundHrmpMessage>>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn persisted_validation_data(
+		&self,
+		_: RelayHash,
+		_: ParaId,
+		_: OccupiedCoreAssumption,
+	) -> RelayChainResult<Option<PersistedValidationData>> {
+		use cumulus_primitives_core::PersistedValidationData;
+		Ok(Some(PersistedValidationData {
+			parent_head: Default::default(),
+			relay_parent_number: 100,
+			relay_parent_storage_root: Default::default(),
+			max_pov_size: 1024 * 1024,
+		}))
+	}
+
+	async fn validation_code_hash(
+		&self,
+		_: RelayHash,
+		_: ParaId,
+		_: OccupiedCoreAssumption,
+	) -> RelayChainResult<Option<ValidationCodeHash>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn candidate_pending_availability(
+		&self,
+		_: RelayHash,
+		_: ParaId,
+	) -> RelayChainResult<Option<CommittedCandidateReceiptV2>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn candidates_pending_availability(
+		&self,
+		_: RelayHash,
+		_: ParaId,
+	) -> RelayChainResult<Vec<CommittedCandidateReceiptV2>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn session_index_for_child(&self, _: RelayHash) -> RelayChainResult<SessionIndex> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn import_notification_stream(
+		&self,
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn finality_notification_stream(
+		&self,
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn is_major_syncing(&self) -> RelayChainResult<bool> {
+		unimplemented!("Not needed for test")
+	}
+
+	fn overseer_handle(&self) -> RelayChainResult<OverseerHandle> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn get_storage_by_key(
+		&self,
+		_: RelayHash,
+		_: &[u8],
+	) -> RelayChainResult<Option<StorageValue>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn prove_read(
+		&self,
+		_: RelayHash,
+		_: &Vec<Vec<u8>>,
+	) -> RelayChainResult<pezsc_client_api::StorageProof> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn wait_for_block(&self, _: RelayHash) -> RelayChainResult<()> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn new_best_notification_stream(
+		&self,
+	) -> RelayChainResult<Pin<Box<dyn Stream<Item = PHeader> + Send>>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn header(
+		&self,
+		block_id: BlockId<pezkuwi_primitives::Block>,
+	) -> RelayChainResult<Option<PHeader>> {
+		let hash = match block_id {
+			BlockId::Hash(hash) => hash,
+			BlockId::Number(_) => unimplemented!("Not needed for test"),
+		};
+		let header = self.headers.get(&hash);
+
+		Ok(header.cloned())
+	}
+
+	async fn availability_cores(
+		&self,
+		_relay_parent: RelayHash,
+	) -> RelayChainResult<Vec<CoreState<RelayHash, BlockNumber>>> {
+		unimplemented!("Not needed for test");
+	}
+
+	async fn version(&self, _: RelayHash) -> RelayChainResult<RuntimeVersion> {
+		unimplemented!("Not needed for test");
+	}
+
+	async fn claim_queue(
+		&self,
+		_: RelayHash,
+	) -> RelayChainResult<BTreeMap<CoreIndex, VecDeque<ParaId>>> {
+		// Return empty claim queue for offset tests
+		Ok(BTreeMap::new())
+	}
+
+	async fn call_runtime_api(
+		&self,
+		_method_name: &'static str,
+		_hash: RelayHash,
+		_payload: &[u8],
+	) -> RelayChainResult<Vec<u8>> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn scheduling_lookahead(&self, _: RelayHash) -> RelayChainResult<u32> {
+		unimplemented!("Not needed for test")
+	}
+
+	async fn candidate_events(&self, _: RelayHash) -> RelayChainResult<Vec<CandidateEvent>> {
+		unimplemented!("Not needed for test")
+	}
+}
+
+/// Build a consecutive set of relay headers whose digest entries optionally carry a BABE
+/// epoch-change marker, returning the underlying map and the hash of the last header.
+fn build_headers_with_epoch_flags(
+	flags: &[HasEpochChange],
+) -> (HashMap<RelayHash, RelayHeader>, RelayHash) {
+	let mut headers = HashMap::new();
+	let mut parent_hash = RelayHash::default();
+	let mut last_hash = RelayHash::default();
+
+	for (index, has_epoch_change) in flags.iter().enumerate() {
+		let digest = if *has_epoch_change == HasEpochChange::Yes {
+			babe_epoch_change_digest()
+		} else {
+			Default::default()
+		};
+
+		let header = RelayHeader {
+			parent_hash,
+			number: (index as u32 + 1),
+			state_root: Default::default(),
+			extrinsics_root: Default::default(),
+			digest,
+		};
+
+		let hash = header.hash();
+		headers.insert(hash, header);
+		parent_hash = hash;
+		last_hash = hash;
+	}
+
+	(headers, last_hash)
+}
+
+/// Create a digest containing a single BABE `NextEpochData` item for use in tests.
+fn babe_epoch_change_digest() -> pezsp_runtime::generic::Digest {
+	let mut digest = pezsp_runtime::generic::Digest::default();
+	let authority_id = AuthorityId::from(sr25519::Public::from_raw([1u8; 32]));
+	let next_epoch =
+		NextEpochDescriptor { authorities: vec![(authority_id, 1u64)], randomness: [0u8; 32] };
+	let log = BabeConsensusLog::NextEpochData(next_epoch);
+	digest.push(pezsp_runtime::generic::DigestItem::Consensus(BABE_ENGINE_ID, log.encode()));
+	digest
+}
+
+fn create_header_chain() -> (HashMap<RelayHash, RelayHeader>, RelayHash) {
+	let mut headers = HashMap::new();
+	let mut current_parent = None;
+	let mut header_hash = RelayHash::repeat_byte(0x1);
+
+	// Create chain from highest to lowest number
+	for number in 1..=100 {
+		let mut header = RelayHeader {
+			parent_hash: Default::default(),
+			number,
+			state_root: Default::default(),
+			extrinsics_root: Default::default(),
+			digest: Default::default(),
+		};
+		if let Some(hash) = current_parent {
+			header.parent_hash = hash;
+		}
+
+		header_hash = header.hash();
+		// Store header and update parent for next iteration
+		headers.insert(header_hash, header.clone());
+		current_parent = Some(header_hash);
+	}
+
+	(headers, header_hash)
+}
+
+// Test extension for RelayChainDataCache
+impl RelayChainDataCache<TestRelayClient> {
+	fn set_test_data(&mut self, relay_parent_header: RelayHeader, cores: Vec<CoreIndex>) {
+		self.set_test_data_with_last_selector(relay_parent_header, cores, None);
+	}
+
+	fn set_test_data_with_last_selector(
+		&mut self,
+		relay_parent_header: RelayHeader,
+		cores: Vec<CoreIndex>,
+		last_claimed_core_selector: Option<CoreSelector>,
+	) {
+		let relay_parent_hash = relay_parent_header.hash();
+
+		let mut claim_queue = BTreeMap::new();
+		for core_index in cores {
+			claim_queue.insert(core_index, [ParaId::from(1)].into());
+		}
+
+		let claim_queue_snapshot = ClaimQueueSnapshot::from(claim_queue);
+
+		let data = RelayChainData {
+			relay_parent_header,
+			claim_queue: claim_queue_snapshot,
+			max_pov_size: 1024 * 1024,
+			last_claimed_core_selector,
+		};
+
+		self.insert_test_data(relay_parent_hash, data);
+	}
+}

pezcumulus/client/consensus/aura/src/equivocation_import_queue.rs

@@ -0,0 +1,392 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+/// An import queue which provides some equivocation resistance with lenient trait bounds.
+///
+/// Equivocation resistance in general is a hard problem, as different nodes in the network
+/// may see equivocations in a different order, and therefore may not agree on which blocks
+/// should be thrown out and which ones should be kept.
+use codec::Codec;
+use cumulus_client_consensus_common::TeyrchainBlockImportMarker;
+use cumulus_primitives_core::{CumulusDigestItem, RelayBlockIdentifier};
+use parking_lot::Mutex;
+use pezkuwi_primitives::Hash as RHash;
+use pezsc_consensus::{
+	import_queue::{BasicQueue, Verifier as VerifierT},
+	BlockImport, BlockImportParams, ForkChoiceStrategy,
+};
+use pezsc_consensus_aura::{standalone as aura_internal, AuthoritiesTracker};
+use pezsc_telemetry::{telemetry, TelemetryHandle, CONSENSUS_DEBUG, CONSENSUS_TRACE};
+use schnellru::{ByLength, LruMap};
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_block_builder::BlockBuilder as BlockBuilderApi;
+use pezsp_blockchain::{HeaderBackend, HeaderMetadata};
+use pezsp_consensus::{error::Error as ConsensusError, BlockOrigin};
+use pezsp_consensus_aura::{AuraApi, Slot, SlotDuration};
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT, NumberFor};
+use std::{fmt::Debug, sync::Arc};
+
+const LRU_WINDOW: u32 = 512;
+const EQUIVOCATION_LIMIT: usize = 16;
+
+struct NaiveEquivocationDefender<N> {
+	/// We distinguish blocks by `(Slot, BlockNumber, RelayParent)`.
+	cache: LruMap<(u64, N, RHash), usize>,
+}
+
+impl<N: std::hash::Hash + PartialEq> Default for NaiveEquivocationDefender<N> {
+	fn default() -> Self {
+		NaiveEquivocationDefender { cache: LruMap::new(ByLength::new(LRU_WINDOW)) }
+	}
+}
+
+impl<N: std::hash::Hash + PartialEq> NaiveEquivocationDefender<N> {
+	// Returns `true` if equivocation is beyond the limit.
+	fn insert_and_check(&mut self, slot: Slot, block_number: N, relay_chain_parent: RHash) -> bool {
+		let val = self
+			.cache
+			.get_or_insert((*slot, block_number, relay_chain_parent), || 0)
+			.expect("insertion with ByLength limiter always succeeds; qed");
+
+		if *val == EQUIVOCATION_LIMIT {
+			true
+		} else {
+			*val += 1;
+			false
+		}
+	}
+}
+
+/// A teyrchain block import verifier that checks for equivocation limits within each slot.
+pub struct Verifier<P: Pair, Client, Block: BlockT, CIDP> {
+	client: Arc<Client>,
+	create_inherent_data_providers: CIDP,
+	defender: Mutex<NaiveEquivocationDefender<NumberFor<Block>>>,
+	telemetry: Option<TelemetryHandle>,
+	// Unused for now. Will be plugged in with a later PR.
+	_authorities_tracker: AuthoritiesTracker<P, Block, Client>,
+}
+
+impl<P, Client, Block, CIDP> Verifier<P, Client, Block, CIDP>
+where
+	P: Pair,
+	P::Signature: Codec,
+	P::Public: Codec + Debug,
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block> + Send + Sync,
+	<Client as ProvideRuntimeApi<Block>>::Api: BlockBuilderApi<Block> + AuraApi<Block, P::Public>,
+
+	CIDP: CreateInherentDataProviders<Block, ()>,
+{
+	/// Creates a new Verifier instance for handling teyrchain block import verification in Aura
+	/// consensus.
+	pub fn new(
+		client: Arc<Client>,
+		inherent_data_provider: CIDP,
+		telemetry: Option<TelemetryHandle>,
+	) -> Self {
+		Self {
+			client: client.clone(),
+			create_inherent_data_providers: inherent_data_provider,
+			defender: Mutex::new(NaiveEquivocationDefender::default()),
+			telemetry,
+			_authorities_tracker: AuthoritiesTracker::new(client),
+		}
+	}
+}
+
+#[async_trait::async_trait]
+impl<P, Client, Block, CIDP> VerifierT<Block> for Verifier<P, Client, Block, CIDP>
+where
+	P: Pair,
+	P::Signature: Codec,
+	P::Public: Codec + Debug,
+	Block: BlockT,
+	Client: HeaderBackend<Block>
+		+ HeaderMetadata<Block, Error = pezsp_blockchain::Error>
+		+ ProvideRuntimeApi<Block>
+		+ Send
+		+ Sync,
+	<Client as ProvideRuntimeApi<Block>>::Api: BlockBuilderApi<Block> + AuraApi<Block, P::Public>,
+
+	CIDP: CreateInherentDataProviders<Block, ()>,
+{
+	async fn verify(
+		&self,
+		mut block_params: BlockImportParams<Block>,
+	) -> Result<BlockImportParams<Block>, String> {
+		// Skip checks that include execution, if being told so, or when importing only state.
+		//
+		// This is done for example when gap syncing and it is expected that the block after the gap
+		// was checked/chosen properly, e.g. by warp syncing to this block using a finality proof.
+		if block_params.state_action.skip_execution_checks() || block_params.with_state() {
+			block_params.fork_choice = Some(ForkChoiceStrategy::Custom(block_params.with_state()));
+			return Ok(block_params);
+		}
+
+		let post_hash = block_params.header.hash();
+		let parent_hash = *block_params.header.parent_hash();
+
+		// check seal and update pre-hash/post-hash
+		{
+			let authorities = aura_internal::fetch_authorities(self.client.as_ref(), parent_hash)
+				.map_err(|e| {
+				format!("Could not fetch authorities at {:?}: {}", parent_hash, e)
+			})?;
+
+			let slot_duration = self
+				.client
+				.runtime_api()
+				.slot_duration(parent_hash)
+				.map_err(|e| e.to_string())?;
+
+			let slot_now = slot_now(slot_duration);
+			let res = aura_internal::check_header_slot_and_seal::<Block, P>(
+				slot_now,
+				block_params.header,
+				&authorities,
+			);
+
+			match res {
+				Ok((pre_header, slot, seal_digest)) => {
+					telemetry!(
+						self.telemetry;
+						CONSENSUS_TRACE;
+						"aura.checked_and_importing";
+						"pre_header" => ?pre_header,
+					);
+
+					// We need some kind of identifier for the relay parent, in the worst case we
+					// take the all `0` hash.
+					let relay_parent =
+						match CumulusDigestItem::find_relay_block_identifier(pre_header.digest()) {
+							None => Default::default(),
+							Some(RelayBlockIdentifier::ByHash(h)) |
+							Some(RelayBlockIdentifier::ByStorageRoot {
+								storage_root: h, ..
+							}) => h,
+						};
+
+					block_params.header = pre_header;
+					block_params.post_digests.push(seal_digest);
+					block_params.fork_choice = Some(ForkChoiceStrategy::LongestChain);
+					block_params.post_hash = Some(post_hash);
+
+					// Check for and reject egregious amounts of equivocations.
+					//
+					// If the `origin` is `ConsensusBroadcast`, we ignore the result of the
+					// equivocation check. This `origin` is for example used by pov-recovery.
+					if self.defender.lock().insert_and_check(
+						slot,
+						*block_params.header.number(),
+						relay_parent,
+					) && !matches!(block_params.origin, BlockOrigin::ConsensusBroadcast)
+					{
+						return Err(format!(
+							"Rejecting block {:?} due to excessive equivocations at slot",
+							post_hash,
+						));
+					}
+				},
+				Err(aura_internal::SealVerificationError::Deferred(hdr, slot)) => {
+					telemetry!(
+						self.telemetry;
+						CONSENSUS_DEBUG;
+						"aura.header_too_far_in_future";
+						"hash" => ?post_hash,
+						"a" => ?hdr,
+						"b" => ?slot,
+					);
+
+					return Err(format!(
+						"Rejecting block ({:?}) from future slot {:?}",
+						post_hash, slot
+					));
+				},
+				Err(e) =>
+					return Err(format!(
+						"Rejecting block ({:?}) with invalid seal ({:?})",
+						post_hash, e
+					)),
+			}
+		}
+
+		// Check inherents.
+		if let Some(body) = block_params.body.clone() {
+			let block = Block::new(block_params.header.clone(), body);
+			let create_inherent_data_providers = self
+				.create_inherent_data_providers
+				.create_inherent_data_providers(parent_hash, ())
+				.await
+				.map_err(|e| format!("Could not create inherent data {:?}", e))?;
+
+			pezsp_block_builder::check_inherents(
+				self.client.clone(),
+				parent_hash,
+				block,
+				&create_inherent_data_providers,
+			)
+			.await
+			.map_err(|e| format!("Error checking block inherents {:?}", e))?;
+		}
+
+		Ok(block_params)
+	}
+}
+
+fn slot_now(slot_duration: SlotDuration) -> Slot {
+	let timestamp = pezsp_timestamp::InherentDataProvider::from_system_time().timestamp();
+	Slot::from_timestamp(timestamp, slot_duration)
+}
+
+/// Start an import queue for a Pezcumulus node which checks blocks' seals and inherent data.
+///
+/// Pass in only inherent data providers which don't include aura or teyrchain consensus inherents,
+/// e.g. things like timestamp and custom inherents for the runtime.
+///
+/// The others are generated explicitly internally.
+///
+/// This should only be used for runtimes where the runtime does not check all inherents and
+/// seals in `execute_block` (see <https://github.com/pezkuwichain/kurdistan-sdk/issues/91>)
+pub fn fully_verifying_import_queue<P, Client, Block: BlockT, I, CIDP>(
+	client: Arc<Client>,
+	block_import: I,
+	create_inherent_data_providers: CIDP,
+	spawner: &impl pezsp_core::traits::SpawnEssentialNamed,
+	registry: Option<&prometheus_endpoint::Registry>,
+	telemetry: Option<TelemetryHandle>,
+) -> BasicQueue<Block>
+where
+	P: Pair + 'static,
+	P::Signature: Codec,
+	P::Public: Codec + Debug,
+	I: BlockImport<Block, Error = ConsensusError>
+		+ TeyrchainBlockImportMarker
+		+ Send
+		+ Sync
+		+ 'static,
+	Client: HeaderBackend<Block>
+		+ HeaderMetadata<Block, Error = pezsp_blockchain::Error>
+		+ ProvideRuntimeApi<Block>
+		+ Send
+		+ Sync
+		+ 'static,
+	<Client as ProvideRuntimeApi<Block>>::Api: BlockBuilderApi<Block> + AuraApi<Block, P::Public>,
+	CIDP: CreateInherentDataProviders<Block, ()> + 'static,
+{
+	let verifier = Verifier::<P, _, _, _> {
+		client: client.clone(),
+		create_inherent_data_providers,
+		defender: Mutex::new(NaiveEquivocationDefender::default()),
+		telemetry,
+		_authorities_tracker: AuthoritiesTracker::new(client.clone()),
+	};
+
+	BasicQueue::new(verifier, Box::new(block_import), None, spawner, registry)
+}
+
+#[cfg(test)]
+mod test {
+	use super::*;
+	use codec::Encode;
+	use cumulus_test_client::{
+		runtime::Block, seal_block, Client, InitBlockBuilder, TestClientBuilder,
+		TestClientBuilderExt,
+	};
+	use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+	use futures::FutureExt;
+	use pezkuwi_primitives::{HeadData, PersistedValidationData};
+	use pezsc_client_api::HeaderBackend;
+	use pezsp_consensus_aura::sr25519;
+	use pezsp_tracing::try_init_simple;
+	use std::{collections::HashSet, sync::Arc};
+
+	#[test]
+	fn import_equivocated_blocks_from_recovery() {
+		try_init_simple();
+
+		let client = Arc::new(TestClientBuilder::default().build());
+
+		let verifier = Verifier::<sr25519::AuthorityPair, Client, Block, _> {
+			client: client.clone(),
+			create_inherent_data_providers: |_, _| async move {
+				Ok(pezsp_timestamp::InherentDataProvider::from_system_time())
+			},
+			defender: Mutex::new(NaiveEquivocationDefender::default()),
+			telemetry: None,
+			_authorities_tracker: AuthoritiesTracker::new(client.clone()),
+		};
+
+		let genesis = client.info().best_hash;
+		let mut sproof = RelayStateSproofBuilder::default();
+		sproof.included_para_head = Some(HeadData(client.header(genesis).unwrap().encode()));
+		sproof.para_id = cumulus_test_client::runtime::TEYRCHAIN_ID.into();
+
+		let validation_data = PersistedValidationData {
+			relay_parent_number: 1,
+			parent_head: client.header(genesis).unwrap().encode().into(),
+			..Default::default()
+		};
+
+		let block_builder = client.init_block_builder(Some(validation_data), sproof);
+		let block = block_builder.block_builder.build().unwrap();
+
+		let mut blocks = Vec::new();
+		for _ in 0..EQUIVOCATION_LIMIT + 1 {
+			blocks.push(seal_block(block.block.clone(), &client))
+		}
+
+		// sr25519 should generate a different signature every time you sign something and thus, all
+		// blocks get a different hash (even if they are the same block).
+		assert_eq!(blocks.iter().map(|b| b.hash()).collect::<HashSet<_>>().len(), blocks.len());
+
+		blocks.iter().take(EQUIVOCATION_LIMIT).for_each(|block| {
+			let mut params =
+				BlockImportParams::new(BlockOrigin::NetworkBroadcast, block.header().clone());
+			params.body = Some(block.extrinsics().to_vec());
+			verifier.verify(params).now_or_never().unwrap().unwrap();
+		});
+
+		// Now let's try some previously verified block and a block we have not verified yet.
+		//
+		// Verify should fail, because we are above the limit. However, when we change the origin to
+		// `ConsensusBroadcast`, it should work.
+		let extra_blocks =
+			vec![blocks[EQUIVOCATION_LIMIT / 2].clone(), blocks.last().unwrap().clone()];
+
+		extra_blocks.into_iter().for_each(|block| {
+			let mut params =
+				BlockImportParams::new(BlockOrigin::NetworkBroadcast, block.header().clone());
+			params.body = Some(block.extrinsics().to_vec());
+			assert!(verifier
+				.verify(params)
+				.now_or_never()
+				.unwrap()
+				.map(drop)
+				.unwrap_err()
+				.contains("excessive equivocations at slot"));
+
+			// When it comes from `pov-recovery`, we will accept it
+			let mut params =
+				BlockImportParams::new(BlockOrigin::ConsensusBroadcast, block.header().clone());
+			params.body = Some(block.extrinsics().to_vec());
+			assert!(verifier.verify(params).now_or_never().unwrap().is_ok());
+		});
+	}
+}

pezcumulus/client/consensus/aura/src/import_queue.rs

@@ -0,0 +1,126 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! Teyrchain specific wrapper for the AuRa import queue.
+
+use codec::Codec;
+use cumulus_client_consensus_common::TeyrchainBlockImportMarker;
+use prometheus_endpoint::Registry;
+use pezsc_client_api::{backend::AuxStore, BlockOf, UsageProvider};
+use pezsc_consensus::{import_queue::DefaultImportQueue, BlockImport};
+use pezsc_consensus_aura::{AuraVerifier, CompatibilityMode};
+use pezsc_consensus_slots::InherentDataProviderExt;
+use pezsc_telemetry::TelemetryHandle;
+use pezsp_api::{ApiExt, ProvideRuntimeApi};
+use pezsp_block_builder::BlockBuilder as BlockBuilderApi;
+use pezsp_blockchain::{HeaderBackend, HeaderMetadata};
+use pezsp_consensus::Error as ConsensusError;
+use pezsp_consensus_aura::AuraApi;
+use pezsp_core::crypto::Pair;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_runtime::traits::Block as BlockT;
+use std::{fmt::Debug, sync::Arc};
+
+/// Parameters for [`import_queue`].
+pub struct ImportQueueParams<'a, I, C, CIDP, S> {
+	/// The block import to use.
+	pub block_import: I,
+	/// The client to interact with the chain.
+	pub client: Arc<C>,
+	/// The inherent data providers, to create the inherent data.
+	pub create_inherent_data_providers: CIDP,
+	/// The spawner to spawn background tasks.
+	pub spawner: &'a S,
+	/// The prometheus registry.
+	pub registry: Option<&'a Registry>,
+	/// The telemetry handle.
+	pub telemetry: Option<TelemetryHandle>,
+}
+
+/// Start an import queue for the Aura consensus algorithm.
+pub fn import_queue<P, Block, I, C, S, CIDP>(
+	ImportQueueParams {
+		block_import,
+		client,
+		create_inherent_data_providers,
+		spawner,
+		registry,
+		telemetry,
+	}: ImportQueueParams<'_, I, C, CIDP, S>,
+) -> Result<DefaultImportQueue<Block>, pezsp_consensus::Error>
+where
+	Block: BlockT,
+	C::Api: BlockBuilderApi<Block> + AuraApi<Block, P::Public> + ApiExt<Block>,
+	C: 'static
+		+ ProvideRuntimeApi<Block>
+		+ BlockOf
+		+ Send
+		+ Sync
+		+ AuxStore
+		+ UsageProvider<Block>
+		+ HeaderBackend<Block>
+		+ HeaderMetadata<Block, Error = pezsp_blockchain::Error>,
+	I: BlockImport<Block, Error = ConsensusError>
+		+ TeyrchainBlockImportMarker
+		+ Send
+		+ Sync
+		+ 'static,
+	P: Pair + 'static,
+	P::Public: Debug + Codec,
+	P::Signature: Codec,
+	S: pezsp_core::traits::SpawnEssentialNamed,
+	CIDP: CreateInherentDataProviders<Block, ()> + Sync + Send + 'static,
+	CIDP::InherentDataProviders: InherentDataProviderExt + Send + Sync,
+{
+	pezsc_consensus_aura::import_queue::<P, _, _, _, _, _>(pezsc_consensus_aura::ImportQueueParams {
+		block_import,
+		justification_import: None,
+		client,
+		create_inherent_data_providers,
+		spawner,
+		registry,
+		check_for_equivocation: pezsc_consensus_aura::CheckForEquivocation::No,
+		telemetry,
+		compatibility_mode: CompatibilityMode::None,
+	})
+}
+
+/// Parameters of [`build_verifier`].
+pub struct BuildVerifierParams<C, CIDP> {
+	/// The client to interact with the chain.
+	pub client: Arc<C>,
+	/// The inherent data providers, to create the inherent data.
+	pub create_inherent_data_providers: CIDP,
+	/// The telemetry handle.
+	pub telemetry: Option<TelemetryHandle>,
+}
+
+/// Build the [`AuraVerifier`].
+pub fn build_verifier<P: Pair, C, CIDP, B: BlockT>(
+	BuildVerifierParams { client, create_inherent_data_providers, telemetry }: BuildVerifierParams<
+		C,
+		CIDP,
+	>,
+) -> AuraVerifier<C, P, CIDP, B> {
+	pezsc_consensus_aura::build_verifier(pezsc_consensus_aura::BuildVerifierParams {
+		client,
+		create_inherent_data_providers,
+		telemetry,
+		check_for_equivocation: pezsc_consensus_aura::CheckForEquivocation::No,
+		compatibility_mode: CompatibilityMode::None,
+	})
+}

pezcumulus/client/consensus/aura/src/lib.rs

@@ -0,0 +1,86 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! The AuRa consensus algorithm for teyrchains.
+//!
+//! This extends the Bizinikiwi provided AuRa consensus implementation to make it compatible for
+//! teyrchains.
+//!
+//! For more information about AuRa, the Bizinikiwi crate should be checked.
+
+use codec::Encode;
+use cumulus_primitives_core::PersistedValidationData;
+
+use cumulus_primitives_core::relay_chain::HeadData;
+use pezkuwi_primitives::{BlockNumber as RBlockNumber, Hash as RHash};
+use pezsp_runtime::traits::{Block as BlockT, NumberFor};
+use std::{fs, fs::File, path::PathBuf};
+
+mod import_queue;
+
+pub use import_queue::{build_verifier, import_queue, BuildVerifierParams, ImportQueueParams};
+use pezkuwi_node_primitives::PoV;
+pub use pezsc_consensus_aura::{
+	slot_duration, standalone::slot_duration_at, AuraVerifier, BuildAuraWorkerParams,
+	SlotProportion,
+};
+pub use pezsc_consensus_slots::InherentDataProviderExt;
+
+pub mod collator;
+pub mod collators;
+pub mod equivocation_import_queue;
+
+const LOG_TARGET: &str = "aura::pezcumulus";
+
+/// Export the given `pov` to the file system at `path`.
+///
+/// The file will be named `block_hash_block_number.pov`.
+///
+/// The `parent_header`, `relay_parent_storage_root` and `relay_parent_number` will also be
+/// stored in the file alongside the `pov`. This enables stateless validation of the `pov`.
+pub(crate) fn export_pov_to_path<Block: BlockT>(
+	path: PathBuf,
+	pov: PoV,
+	block_hash: Block::Hash,
+	block_number: NumberFor<Block>,
+	parent_header: Block::Header,
+	relay_parent_storage_root: RHash,
+	relay_parent_number: RBlockNumber,
+	max_pov_size: u32,
+) {
+	if let Err(error) = fs::create_dir_all(&path) {
+		tracing::error!(target: LOG_TARGET, %error, path = %path.display(), "Failed to create PoV export directory");
+		return;
+	}
+
+	let mut file = match File::create(path.join(format!("{block_hash:?}_{block_number}.pov"))) {
+		Ok(f) => f,
+		Err(error) => {
+			tracing::error!(target: LOG_TARGET, %error, "Failed to export PoV.");
+			return;
+		},
+	};
+
+	pov.encode_to(&mut file);
+	PersistedValidationData {
+		parent_head: HeadData(parent_header.encode()),
+		relay_parent_number,
+		relay_parent_storage_root,
+		max_pov_size,
+	}
+	.encode_to(&mut file);
+}

pezcumulus/client/consensus/common/Cargo.toml

@@ -0,0 +1,78 @@
+[package]
+name = "pezcumulus-client-consensus-common"
+description = "Pezcumulus specific common consensus implementations"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+async-trait = { workspace = true }
+codec = { features = ["derive"], workspace = true, default-features = true }
+dyn-clone = { workspace = true }
+futures = { workspace = true }
+log = { workspace = true, default-features = true }
+tracing = { workspace = true, default-features = true }
+
+# Bizinikiwi
+prometheus-endpoint = { workspace = true, default-features = true }
+pezsc-client-api = { workspace = true, default-features = true }
+pezsc-consensus = { workspace = true, default-features = true }
+pezsc-consensus-babe = { workspace = true, default-features = true }
+pezsc-network = { workspace = true, default-features = true }
+pezsp-blockchain = { workspace = true, default-features = true }
+pezsp-consensus = { workspace = true, default-features = true }
+pezsp-consensus-slots = { workspace = true, default-features = true }
+pezsp-core = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+pezsp-timestamp = { workspace = true, default-features = true }
+pezsp-trie = { workspace = true, default-features = true }
+pezsp-version = { workspace = true, default-features = true }
+
+# Pezkuwi
+pezkuwi-primitives = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-client-pov-recovery = { workspace = true, default-features = true }
+pezcumulus-primitives-core = { workspace = true, default-features = true }
+pezcumulus-relay-chain-interface = { workspace = true, default-features = true }
+pezcumulus-relay-chain-streams = { workspace = true, default-features = true }
+schnellru = { workspace = true }
+
+[dev-dependencies]
+futures-timer = { workspace = true }
+tokio = { features = ["macros"], workspace = true }
+
+# Bizinikiwi
+pezsp-tracing = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-test-client = { workspace = true }
+pezcumulus-test-relay-sproof-builder = { workspace = true, default-features = true }
+
+[features]
+runtime-benchmarks = [
+	"pezcumulus-client-pov-recovery/runtime-benchmarks",
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezcumulus-relay-chain-interface/runtime-benchmarks",
+	"pezcumulus-relay-chain-streams/runtime-benchmarks",
+	"pezcumulus-test-client/runtime-benchmarks",
+	"pezcumulus-test-relay-sproof-builder/runtime-benchmarks",
+	"pezkuwi-primitives/runtime-benchmarks",
+	"pezsc-client-api/runtime-benchmarks",
+	"pezsc-consensus-babe/runtime-benchmarks",
+	"pezsc-consensus/runtime-benchmarks",
+	"pezsc-network/runtime-benchmarks",
+	"pezsp-blockchain/runtime-benchmarks",
+	"pezsp-consensus-slots/runtime-benchmarks",
+	"pezsp-consensus/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+	"pezsp-timestamp/runtime-benchmarks",
+	"pezsp-trie/runtime-benchmarks",
+	"pezsp-version/runtime-benchmarks",
+]

pezcumulus/client/consensus/common/src/import_queue.rs

@@ -0,0 +1,77 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! (unstable) Composable utilities for constructing import queues for teyrchains.
+//!
+//! Unlike standalone chains, teyrchains have the requirement that all consensus logic
+//! must be checked within the runtime. This property means that work which is normally
+//! done in the import queue per-block, such as checking signatures, quorums, and whether
+//! inherent extrinsics were constructed faithfully do not need to be done, per se.
+//!
+//! It may seem that it would be beneficial for the client to do these checks regardless,
+//! but in practice this means that clients would just reject blocks which are _valid_ according
+//! to their Teyrchain Validation Function, which is the ultimate source of consensus truth.
+//!
+//! However, teyrchain runtimes expose two different access points for executing blocks
+//! in full nodes versus executing those blocks in the teyrchain validation environment.
+//! At the time of writing, the inherent and consensus checks in most Pezcumulus runtimes
+//! are only performed during teyrchain validation, not full node block execution.
+//!
+//! See <https://github.com/pezkuwichain/kurdistan-sdk/issues/91> for details.
+
+use pezsp_consensus::error::Error as ConsensusError;
+use pezsp_runtime::traits::Block as BlockT;
+
+use pezsc_consensus::{
+	block_import::{BlockImport, BlockImportParams},
+	import_queue::{BasicQueue, Verifier},
+};
+
+use crate::TeyrchainBlockImportMarker;
+
+/// A [`Verifier`] for blocks which verifies absolutely nothing.
+///
+/// This should only be used when the runtime is responsible for checking block seals and inherents.
+pub struct VerifyNothing;
+
+#[async_trait::async_trait]
+impl<Block: BlockT> Verifier<Block> for VerifyNothing {
+	async fn verify(
+		&self,
+		params: BlockImportParams<Block>,
+	) -> Result<BlockImportParams<Block>, String> {
+		Ok(params)
+	}
+}
+
+/// An import queue which does no verification.
+///
+/// This should only be used when the runtime is responsible for checking block seals and inherents.
+pub fn verify_nothing_import_queue<Block: BlockT, I>(
+	block_import: I,
+	spawner: &impl pezsp_core::traits::SpawnEssentialNamed,
+	registry: Option<&prometheus_endpoint::Registry>,
+) -> BasicQueue<Block>
+where
+	I: BlockImport<Block, Error = ConsensusError>
+		+ TeyrchainBlockImportMarker
+		+ Send
+		+ Sync
+		+ 'static,
+{
+	BasicQueue::new(VerifyNothing, Box::new(block_import), None, spawner, registry)
+}

pezcumulus/client/consensus/common/src/level_monitor.rs

@@ -0,0 +1,407 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use pezsc_client_api::{blockchain::Backend as _, Backend, HeaderBackend as _};
+use pezsp_blockchain::{HashAndNumber, HeaderMetadata, TreeRoute};
+use pezsp_runtime::traits::{Block as BlockT, NumberFor, One, Saturating, UniqueSaturatedInto, Zero};
+use std::{
+	collections::{HashMap, HashSet},
+	sync::Arc,
+};
+
+const LOG_TARGET: &str = "level-monitor";
+
+/// Value good enough to be used with teyrchains using the current backend implementation
+/// that ships with Bizinikiwi. This value may change in the future.
+pub const MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT: usize = 32;
+
+// Counter threshold after which we are going to eventually cleanup our internal data.
+const CLEANUP_THRESHOLD: u32 = 32;
+
+/// Upper bound to the number of leaves allowed for each level of the blockchain.
+///
+/// If the limit is set and more leaves are detected on block import, then the older ones are
+/// dropped to make space for the fresh blocks.
+///
+/// In environments where blocks confirmations from the relay chain may be "slow", then
+/// setting an upper bound helps keeping the chain health by dropping old (presumably) stale
+/// leaves and prevents discarding new blocks because we've reached the backend max value.
+pub enum LevelLimit {
+	/// Limit set to [`MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT`].
+	Default,
+	/// No explicit limit, however a limit may be implicitly imposed by the backend implementation.
+	None,
+	/// Custom value.
+	Some(usize),
+}
+
+/// Support structure to constrain the number of leaves at each level.
+pub struct LevelMonitor<Block: BlockT, BE> {
+	/// Max number of leaves for each level.
+	level_limit: usize,
+	/// Monotonic counter used to keep track of block freshness.
+	pub(crate) import_counter: NumberFor<Block>,
+	/// Map between blocks hashes and freshness.
+	pub(crate) freshness: HashMap<Block::Hash, NumberFor<Block>>,
+	/// Blockchain levels cache.
+	pub(crate) levels: HashMap<NumberFor<Block>, HashSet<Block::Hash>>,
+	/// Lower level number stored by the levels map.
+	lowest_level: NumberFor<Block>,
+	/// Backend reference to remove blocks on level saturation.
+	backend: Arc<BE>,
+}
+
+/// Contains information about the target scheduled for removal.
+struct TargetInfo<Block: BlockT> {
+	/// Index of freshest leaf in the leaves array.
+	freshest_leaf_idx: usize,
+	/// Route from target to its freshest leaf.
+	freshest_route: TreeRoute<Block>,
+}
+
+impl<Block, BE> LevelMonitor<Block, BE>
+where
+	Block: BlockT,
+	BE: Backend<Block>,
+{
+	/// Instance a new monitor structure.
+	pub fn new(level_limit: usize, backend: Arc<BE>) -> Self {
+		let mut monitor = LevelMonitor {
+			level_limit,
+			import_counter: Zero::zero(),
+			freshness: HashMap::new(),
+			levels: HashMap::new(),
+			lowest_level: Zero::zero(),
+			backend,
+		};
+		monitor.restore();
+		monitor
+	}
+
+	/// Restore the structure using the backend.
+	///
+	/// Blocks freshness values are inferred from the height and not from the effective import
+	/// moment. This is a not accurate but "good-enough" best effort solution.
+	///
+	/// Level limits are not enforced during this phase.
+	fn restore(&mut self) {
+		let info = self.backend.blockchain().info();
+
+		log::debug!(
+			target: LOG_TARGET,
+			"Restoring chain level monitor from last finalized block: {} {}",
+			info.finalized_number,
+			info.finalized_hash
+		);
+
+		self.lowest_level = info.finalized_number;
+		self.import_counter = info.finalized_number;
+
+		for leaf in self.backend.blockchain().leaves().unwrap_or_default() {
+			let Ok(mut meta) = self.backend.blockchain().header_metadata(leaf) else {
+				log::debug!(
+					target: LOG_TARGET,
+					"Could not fetch header metadata for leaf: {leaf:?}",
+				);
+
+				continue;
+			};
+
+			self.import_counter = self.import_counter.max(meta.number);
+
+			// Populate the monitor until we don't hit an already imported branch
+			while !self.freshness.contains_key(&meta.hash) {
+				self.freshness.insert(meta.hash, meta.number);
+				self.levels.entry(meta.number).or_default().insert(meta.hash);
+				if meta.number <= self.lowest_level {
+					break;
+				}
+
+				meta = match self.backend.blockchain().header_metadata(meta.parent) {
+					Ok(m) => m,
+					Err(_) => {
+						// This can happen after we have warp synced a node.
+						log::debug!(
+							target: LOG_TARGET,
+							"Could not fetch header metadata for parent: {:?}",
+							meta.parent,
+						);
+						break;
+					},
+				}
+			}
+		}
+
+		log::debug!(
+			target: LOG_TARGET,
+			"Restored chain level monitor up to height {}",
+			self.import_counter
+		);
+	}
+
+	/// Check and enforce the limit bound at the given height.
+	///
+	/// In practice this will enforce the given height in having a number of blocks less than
+	/// the limit passed to the constructor.
+	///
+	/// If the given level is found to have a number of blocks greater than or equal the limit
+	/// then the limit is enforced by choosing one (or more) blocks to remove.
+	///
+	/// The removal strategy is driven by the block freshness.
+	///
+	/// A block freshness is determined by the most recent leaf freshness descending from the block
+	/// itself. In other words its freshness is equal to its more "fresh" descendant.
+	///
+	/// The least "fresh" blocks are eventually removed.
+	pub fn enforce_limit(&mut self, number: NumberFor<Block>) {
+		let level_len = self.levels.get(&number).map(|l| l.len()).unwrap_or_default();
+		if level_len < self.level_limit {
+			return;
+		}
+
+		// Sort leaves by freshness only once (less fresh first) and keep track of
+		// leaves that were invalidated on removal.
+		let mut leaves = self.backend.blockchain().leaves().unwrap_or_default();
+		leaves.sort_unstable_by(|a, b| self.freshness.get(a).cmp(&self.freshness.get(b)));
+		let mut invalidated_leaves = HashSet::new();
+
+		// This may not be the most efficient way to remove **multiple** entries, but is the easy
+		// one :-). Should be considered that in "normal" conditions the number of blocks to remove
+		// is 0 or 1, it is not worth to complicate the code too much. One condition that may
+		// trigger multiple removals (2+) is if we restart the node using an existing db and a
+		// smaller limit wrt the one previously used.
+		let remove_count = level_len - self.level_limit + 1;
+
+		log::debug!(
+			target: LOG_TARGET,
+			"Detected leaves overflow at height {number}, removing {remove_count} obsolete blocks",
+		);
+
+		(0..remove_count).all(|_| {
+			self.find_target(number, &leaves, &invalidated_leaves).map_or(false, |target| {
+				self.remove_target(target, number, &leaves, &mut invalidated_leaves);
+				true
+			})
+		});
+	}
+
+	// Helper function to find the best candidate to be removed.
+	//
+	// Given a set of blocks with height equal to `number` (potential candidates)
+	// 1. For each candidate fetch all the leaves that are descending from it.
+	// 2. Set the candidate freshness equal to the fresher of its descending leaves.
+	// 3. The target is set as the candidate that is less fresh.
+	//
+	// Input `leaves` are assumed to be already ordered by "freshness" (less fresh first).
+	//
+	// Returns the index of the target fresher leaf within `leaves` and the route from target to
+	// such leaf.
+	fn find_target(
+		&self,
+		number: NumberFor<Block>,
+		leaves: &[Block::Hash],
+		invalidated_leaves: &HashSet<usize>,
+	) -> Option<TargetInfo<Block>> {
+		let mut target_info: Option<TargetInfo<Block>> = None;
+		let blockchain = self.backend.blockchain();
+		let best_hash = blockchain.info().best_hash;
+
+		// Leaves that where already assigned to some node and thus can be skipped
+		// during the search.
+		let mut assigned_leaves = HashSet::new();
+
+		let level = self.levels.get(&number)?;
+
+		for blk_hash in level.iter().filter(|hash| **hash != best_hash) {
+			// Search for the fresher leaf information for this block
+			let candidate_info = leaves
+				.iter()
+				.enumerate()
+				.filter(|(leaf_idx, _)| {
+					!assigned_leaves.contains(leaf_idx) && !invalidated_leaves.contains(leaf_idx)
+				})
+				.rev()
+				.find_map(|(leaf_idx, leaf_hash)| {
+					if blk_hash == leaf_hash {
+						let entry = HashAndNumber { number, hash: *blk_hash };
+						TreeRoute::new(vec![entry], 0).ok().map(|freshest_route| TargetInfo {
+							freshest_leaf_idx: leaf_idx,
+							freshest_route,
+						})
+					} else {
+						match pezsp_blockchain::tree_route(blockchain, *blk_hash, *leaf_hash) {
+							Ok(route) if route.retracted().is_empty() => Some(TargetInfo {
+								freshest_leaf_idx: leaf_idx,
+								freshest_route: route,
+							}),
+							Err(err) => {
+								log::warn!(
+									target: LOG_TARGET,
+									"(Lookup) Unable getting route from {:?} to {:?}: {}",
+									blk_hash,
+									leaf_hash,
+									err,
+								);
+								None
+							},
+							_ => None,
+						}
+					}
+				});
+
+			let candidate_info = match candidate_info {
+				Some(candidate_info) => {
+					assigned_leaves.insert(candidate_info.freshest_leaf_idx);
+					candidate_info
+				},
+				None => {
+					// This should never happen
+					log::error!(
+						target: LOG_TARGET,
+						"Unable getting route to any leaf from {:?} (this is a bug)",
+						blk_hash,
+					);
+					continue;
+				},
+			};
+
+			// Found fresher leaf for this candidate.
+			// This candidate is set as the new target if:
+			// 1. its fresher leaf is less fresh than the previous target fresher leaf AND
+			// 2. best block is not in its route
+
+			let is_less_fresh = || {
+				target_info
+					.as_ref()
+					.map(|ti| candidate_info.freshest_leaf_idx < ti.freshest_leaf_idx)
+					.unwrap_or(true)
+			};
+			let not_contains_best = || {
+				candidate_info
+					.freshest_route
+					.enacted()
+					.iter()
+					.all(|entry| entry.hash != best_hash)
+			};
+
+			if is_less_fresh() && not_contains_best() {
+				let early_stop = candidate_info.freshest_leaf_idx == 0;
+				target_info = Some(candidate_info);
+				if early_stop {
+					// We will never find a candidate with an worst freshest leaf than this.
+					break;
+				}
+			}
+		}
+
+		target_info
+	}
+
+	// Remove the target block and all its descendants.
+	//
+	// Leaves should have already been ordered by "freshness" (less fresh first).
+	fn remove_target(
+		&mut self,
+		target: TargetInfo<Block>,
+		number: NumberFor<Block>,
+		leaves: &[Block::Hash],
+		invalidated_leaves: &mut HashSet<usize>,
+	) {
+		let mut remove_leaf = |number, hash| {
+			log::debug!(target: LOG_TARGET, "Removing block (@{}) {:?}", number, hash);
+			if let Err(err) = self.backend.remove_leaf_block(hash) {
+				log::debug!(target: LOG_TARGET, "Remove not possible for {}: {}", hash, err);
+				return false;
+			}
+			self.levels.get_mut(&number).map(|level| level.remove(&hash));
+			self.freshness.remove(&hash);
+			true
+		};
+
+		invalidated_leaves.insert(target.freshest_leaf_idx);
+
+		// Takes care of route removal. Starts from the leaf and stops as soon as an error is
+		// encountered. In this case an error is interpreted as the block being not a leaf
+		// and it will be removed while removing another route from the same block but to a
+		// different leaf.
+		let mut remove_route = |route: TreeRoute<Block>| {
+			route.enacted().iter().rev().all(|elem| remove_leaf(elem.number, elem.hash));
+		};
+
+		let target_hash = target.freshest_route.common_block().hash;
+		debug_assert_eq!(
+			target.freshest_route.common_block().number,
+			number,
+			"This is a bug in LevelMonitor::find_target() or the Backend is corrupted"
+		);
+
+		// Remove freshest (cached) route first.
+		remove_route(target.freshest_route);
+
+		// Don't bother trying with leaves we already found to not be our descendants.
+		let to_skip = leaves.len() - target.freshest_leaf_idx;
+		leaves.iter().enumerate().rev().skip(to_skip).for_each(|(leaf_idx, leaf_hash)| {
+			if invalidated_leaves.contains(&leaf_idx) {
+				return;
+			}
+			match pezsp_blockchain::tree_route(self.backend.blockchain(), target_hash, *leaf_hash) {
+				Ok(route) if route.retracted().is_empty() => {
+					invalidated_leaves.insert(leaf_idx);
+					remove_route(route);
+				},
+				Err(err) => {
+					log::warn!(
+						target: LOG_TARGET,
+						"(Removal) unable getting route from {:?} to {:?}: {}",
+						target_hash,
+						leaf_hash,
+						err,
+					);
+				},
+				_ => (),
+			};
+		});
+
+		remove_leaf(number, target_hash);
+	}
+
+	/// Add a new imported block information to the monitor.
+	pub fn block_imported(&mut self, number: NumberFor<Block>, hash: Block::Hash) {
+		let finalized_num = self.backend.blockchain().info().finalized_number;
+
+		if number > finalized_num {
+			// Only blocks above the last finalized block should be added to the monitor
+			self.import_counter += One::one();
+			self.freshness.insert(hash, self.import_counter);
+			self.levels.entry(number).or_default().insert(hash);
+		}
+
+		let delta: u32 = finalized_num.saturating_sub(self.lowest_level).unique_saturated_into();
+		if delta >= CLEANUP_THRESHOLD {
+			// Do cleanup once in a while, we are allowed to have some obsolete information.
+			for i in 0..delta {
+				let number = self.lowest_level + i.unique_saturated_into();
+				self.levels.remove(&number).map(|level| {
+					level.iter().for_each(|hash| {
+						self.freshness.remove(hash);
+					})
+				});
+			}
+			self.lowest_level = finalized_num;
+		}
+	}
+}

pezcumulus/client/consensus/common/src/lib.rs

@@ -0,0 +1,220 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use codec::Decode;
+use pezkuwi_primitives::{Block as PBlock, Hash as PHash, Header as PHeader, ValidationCodeHash};
+
+use cumulus_primitives_core::{relay_chain, AbridgedHostConfiguration};
+use cumulus_relay_chain_interface::{RelayChainError, RelayChainInterface};
+
+use pezsc_client_api::Backend;
+use pezsc_consensus::{shared_data::SharedData, BlockImport, ImportResult};
+use pezsp_consensus_slots::Slot;
+
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT};
+use pezsp_timestamp::Timestamp;
+
+use std::{sync::Arc, time::Duration};
+
+mod level_monitor;
+mod parent_search;
+#[cfg(test)]
+mod tests;
+mod teyrchain_consensus;
+
+pub use parent_search::*;
+
+pub use cumulus_relay_chain_streams::finalized_heads;
+pub use teyrchain_consensus::spawn_teyrchain_consensus_tasks;
+
+use level_monitor::LevelMonitor;
+pub use level_monitor::{LevelLimit, MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT};
+
+pub mod import_queue;
+
+/// Provides the hash of validation code used for authoring/execution of blocks at a given
+/// hash.
+pub trait ValidationCodeHashProvider<Hash> {
+	fn code_hash_at(&self, at: Hash) -> Option<ValidationCodeHash>;
+}
+
+impl<F, Hash> ValidationCodeHashProvider<Hash> for F
+where
+	F: Fn(Hash) -> Option<ValidationCodeHash>,
+{
+	fn code_hash_at(&self, at: Hash) -> Option<ValidationCodeHash> {
+		(self)(at)
+	}
+}
+
+/// The result from building a collation.
+pub struct TeyrchainCandidate<B> {
+	/// The block that was built for this candidate.
+	pub block: B,
+	/// The proof that was recorded while building the block.
+	pub proof: pezsp_trie::StorageProof,
+}
+
+/// Teyrchain specific block import.
+///
+/// Specialized block import for teyrchains. It supports to delay setting the best block until the
+/// relay chain has included a candidate in its best block. By default the delayed best block
+/// setting is disabled. The block import also monitors the imported blocks and prunes by default if
+/// there are too many blocks at the same height. Too many blocks at the same height can for example
+/// happen if the relay chain is rejecting the teyrchain blocks in the validation.
+pub struct TeyrchainBlockImport<Block: BlockT, BI, BE> {
+	inner: BI,
+	monitor: Option<SharedData<LevelMonitor<Block, BE>>>,
+	delayed_best_block: bool,
+}
+
+impl<Block: BlockT, BI, BE: Backend<Block>> TeyrchainBlockImport<Block, BI, BE> {
+	/// Create a new instance.
+	///
+	/// The number of leaves per level limit is set to `LevelLimit::Default`.
+	pub fn new(inner: BI, backend: Arc<BE>) -> Self {
+		Self::new_with_limit(inner, backend, LevelLimit::Default)
+	}
+
+	/// Create a new instance with an explicit limit to the number of leaves per level.
+	///
+	/// This function alone doesn't enforce the limit on levels for old imported blocks,
+	/// the limit is eventually enforced only when new blocks are imported.
+	pub fn new_with_limit(inner: BI, backend: Arc<BE>, level_leaves_max: LevelLimit) -> Self {
+		let level_limit = match level_leaves_max {
+			LevelLimit::None => None,
+			LevelLimit::Some(limit) => Some(limit),
+			LevelLimit::Default => Some(MAX_LEAVES_PER_LEVEL_SENSIBLE_DEFAULT),
+		};
+
+		let monitor =
+			level_limit.map(|level_limit| SharedData::new(LevelMonitor::new(level_limit, backend)));
+
+		Self { inner, monitor, delayed_best_block: false }
+	}
+
+	/// Create a new instance which delays setting the best block.
+	///
+	/// The number of leaves per level limit is set to `LevelLimit::Default`.
+	pub fn new_with_delayed_best_block(inner: BI, backend: Arc<BE>) -> Self {
+		Self {
+			delayed_best_block: true,
+			..Self::new_with_limit(inner, backend, LevelLimit::Default)
+		}
+	}
+}
+
+impl<Block: BlockT, I: Clone, BE> Clone for TeyrchainBlockImport<Block, I, BE> {
+	fn clone(&self) -> Self {
+		TeyrchainBlockImport {
+			inner: self.inner.clone(),
+			monitor: self.monitor.clone(),
+			delayed_best_block: self.delayed_best_block,
+		}
+	}
+}
+
+#[async_trait::async_trait]
+impl<Block, BI, BE> BlockImport<Block> for TeyrchainBlockImport<Block, BI, BE>
+where
+	Block: BlockT,
+	BI: BlockImport<Block> + Send + Sync,
+	BE: Backend<Block>,
+{
+	type Error = BI::Error;
+
+	async fn check_block(
+		&self,
+		block: pezsc_consensus::BlockCheckParams<Block>,
+	) -> Result<pezsc_consensus::ImportResult, Self::Error> {
+		self.inner.check_block(block).await
+	}
+
+	async fn import_block(
+		&self,
+		mut params: pezsc_consensus::BlockImportParams<Block>,
+	) -> Result<pezsc_consensus::ImportResult, Self::Error> {
+		// Blocks are stored within the backend by using POST hash.
+		let hash = params.post_hash();
+		let number = *params.header.number();
+
+		if params.with_state() {
+			// Force imported state finality.
+			// Required for warp sync. We assume that preconditions have been
+			// checked properly and we are importing a finalized block with state.
+			params.finalized = true;
+		}
+
+		if self.delayed_best_block {
+			// Best block is determined by the relay chain, or if we are doing the initial sync
+			// we import all blocks as new best.
+			params.fork_choice = Some(pezsc_consensus::ForkChoiceStrategy::Custom(
+				params.origin == pezsp_consensus::BlockOrigin::NetworkInitialSync,
+			));
+		}
+
+		let maybe_lock = self.monitor.as_ref().map(|monitor_lock| {
+			let mut monitor = monitor_lock.shared_data_locked();
+			monitor.enforce_limit(number);
+			monitor.release_mutex()
+		});
+
+		let res = self.inner.import_block(params).await?;
+
+		if let (Some(mut monitor_lock), ImportResult::Imported(_)) = (maybe_lock, &res) {
+			let mut monitor = monitor_lock.upgrade();
+			monitor.block_imported(number, hash);
+		}
+
+		Ok(res)
+	}
+}
+
+/// Marker trait denoting a block import type that fits the teyrchain requirements.
+pub trait TeyrchainBlockImportMarker {}
+
+impl<B: BlockT, BI, BE> TeyrchainBlockImportMarker for TeyrchainBlockImport<B, BI, BE> {}
+
+/// Get the relay-parent slot and timestamp from a header.
+pub fn relay_slot_and_timestamp(
+	relay_parent_header: &PHeader,
+	relay_chain_slot_duration: Duration,
+) -> Option<(Slot, Timestamp)> {
+	pezsc_consensus_babe::find_pre_digest::<PBlock>(relay_parent_header)
+		.map(|babe_pre_digest| {
+			let slot = babe_pre_digest.slot();
+			let t = Timestamp::new(relay_chain_slot_duration.as_millis() as u64 * *slot);
+
+			(slot, t)
+		})
+		.ok()
+}
+
+/// Reads abridged host configuration from the relay chain storage at the given relay parent.
+pub async fn load_abridged_host_configuration(
+	relay_parent: PHash,
+	relay_client: &impl RelayChainInterface,
+) -> Result<Option<AbridgedHostConfiguration>, RelayChainError> {
+	relay_client
+		.get_storage_by_key(relay_parent, relay_chain::well_known_keys::ACTIVE_CONFIG)
+		.await?
+		.map(|bytes| {
+			AbridgedHostConfiguration::decode(&mut &bytes[..])
+				.map_err(RelayChainError::DeserializationError)
+		})
+		.transpose()
+}

pezcumulus/client/consensus/common/src/parent_search.rs

@@ -0,0 +1,419 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use codec::Decode;
+use pezkuwi_primitives::Hash as RelayHash;
+
+use cumulus_primitives_core::{
+	relay_chain::{BlockId as RBlockId, OccupiedCoreAssumption},
+	ParaId,
+};
+use cumulus_relay_chain_interface::{RelayChainError, RelayChainInterface};
+
+use pezsc_client_api::{Backend, HeaderBackend};
+
+use pezsp_blockchain::{Backend as BlockchainBackend, TreeRoute};
+
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT};
+
+const PARENT_SEARCH_LOG_TARGET: &str = "consensus::common::find_potential_parents";
+
+/// Parameters when searching for suitable parents to build on top of.
+#[derive(Debug)]
+pub struct ParentSearchParams {
+	/// The relay-parent that is intended to be used.
+	pub relay_parent: RelayHash,
+	/// The ID of the teyrchain.
+	pub para_id: ParaId,
+	/// A limitation on the age of relay parents for teyrchain blocks that are being
+	/// considered. This is relative to the `relay_parent` number.
+	pub ancestry_lookback: usize,
+	/// How "deep" parents can be relative to the included teyrchain block at the relay-parent.
+	/// The included block has depth 0.
+	pub max_depth: usize,
+	/// Whether to only ignore "alternative" branches, i.e. branches of the chain
+	/// which do not contain the block pending availability.
+	pub ignore_alternative_branches: bool,
+}
+
+/// A potential parent block returned from [`find_potential_parents`]
+#[derive(PartialEq)]
+pub struct PotentialParent<B: BlockT> {
+	/// The hash of the block.
+	pub hash: B::Hash,
+	/// The header of the block.
+	pub header: B::Header,
+	/// The depth of the block with respect to the included block.
+	pub depth: usize,
+	/// Whether the block is the included block, is itself pending on-chain, or descends
+	/// from the block pending availability.
+	pub aligned_with_pending: bool,
+}
+
+impl<B: BlockT> std::fmt::Debug for PotentialParent<B> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("PotentialParent")
+			.field("hash", &self.hash)
+			.field("depth", &self.depth)
+			.field("aligned_with_pending", &self.aligned_with_pending)
+			.field("number", &self.header.number())
+			.finish()
+	}
+}
+
+/// Perform a recursive search through blocks to find potential
+/// parent blocks for a new block.
+///
+/// This accepts a relay-chain block to be used as an anchor and a maximum search depth,
+/// along with some arguments for filtering teyrchain blocks and performs a recursive search
+/// for teyrchain blocks. The search begins at the last included teyrchain block and returns
+/// a set of [`PotentialParent`]s which could be potential parents of a new block with this
+/// relay-parent according to the search parameters.
+///
+/// A teyrchain block is a potential parent if it is either the last included teyrchain block, the
+/// pending teyrchain block (when `max_depth` >= 1), or all of the following hold:
+///   * its parent is a potential parent
+///   * its relay-parent is within `ancestry_lookback` of the targeted relay-parent.
+///   * its relay-parent is within the same session as the targeted relay-parent.
+///   * the block number is within `max_depth` blocks of the included block
+pub async fn find_potential_parents<B: BlockT>(
+	params: ParentSearchParams,
+	backend: &impl Backend<B>,
+	relay_client: &impl RelayChainInterface,
+) -> Result<Vec<PotentialParent<B>>, RelayChainError> {
+	tracing::trace!("Parent search parameters: {params:?}");
+	// Get the included block.
+	let Some((included_header, included_hash)) =
+		fetch_included_from_relay_chain(relay_client, backend, params.para_id, params.relay_parent)
+			.await?
+	else {
+		return Ok(Default::default());
+	};
+
+	let only_included = vec![PotentialParent {
+		hash: included_hash,
+		header: included_header.clone(),
+		depth: 0,
+		aligned_with_pending: true,
+	}];
+
+	if params.max_depth == 0 {
+		return Ok(only_included);
+	};
+
+	// Pending header and hash.
+	let maybe_pending = {
+		// Fetch the most recent pending header from the relay chain. We use
+		// `OccupiedCoreAssumption::Included` so the candidate pending availability gets enacted
+		// before being returned to us.
+		let pending_header = relay_client
+			.persisted_validation_data(
+				params.relay_parent,
+				params.para_id,
+				OccupiedCoreAssumption::Included,
+			)
+			.await?
+			.and_then(|p| B::Header::decode(&mut &p.parent_head.0[..]).ok())
+			.filter(|x| x.hash() != included_hash);
+
+		// If the pending block is not locally known, we can't do anything.
+		if let Some(header) = pending_header {
+			let pending_hash = header.hash();
+			match backend.blockchain().header(pending_hash) {
+				// We are supposed to ignore branches that don't contain the pending block, but we
+				// do not know the pending block locally.
+				Ok(None) | Err(_) if params.ignore_alternative_branches => {
+					tracing::warn!(
+						target: PARENT_SEARCH_LOG_TARGET,
+						%pending_hash,
+						"Failed to get header for pending block.",
+					);
+					return Ok(Default::default());
+				},
+				Ok(Some(_)) => Some((header, pending_hash)),
+				_ => None,
+			}
+		} else {
+			None
+		}
+	};
+
+	let maybe_route_to_last_pending = maybe_pending
+		.as_ref()
+		.map(|(_, pending)| {
+			pezsp_blockchain::tree_route(backend.blockchain(), included_hash, *pending)
+		})
+		.transpose()?;
+
+	// If we want to ignore alternative branches there is no reason to start
+	// the parent search at the included block. We can add the included block and
+	// the path to the pending block to the potential parents directly (limited by max_depth).
+	let (frontier, potential_parents) = match (
+		&maybe_pending,
+		params.ignore_alternative_branches,
+		&maybe_route_to_last_pending,
+	) {
+		(Some((pending_header, pending_hash)), true, Some(ref route_to_pending)) => {
+			let mut potential_parents = only_included;
+
+			// This is a defensive check, should never happen.
+			if !route_to_pending.retracted().is_empty() {
+				tracing::warn!(target: PARENT_SEARCH_LOG_TARGET, "Included block not an ancestor of pending block. This should not happen.");
+				return Ok(Default::default());
+			}
+
+			// Add all items on the path included -> pending - 1 to the potential parents, but
+			// not more than `max_depth`.
+			let num_parents_on_path =
+				route_to_pending.enacted().len().saturating_sub(1).min(params.max_depth);
+			for (num, block) in
+				route_to_pending.enacted().iter().take(num_parents_on_path).enumerate()
+			{
+				let Ok(Some(header)) = backend.blockchain().header(block.hash) else { continue };
+
+				potential_parents.push(PotentialParent {
+					hash: block.hash,
+					header,
+					depth: 1 + num,
+					aligned_with_pending: true,
+				});
+			}
+
+			// The search for additional potential parents should now start at the children of
+			// the pending block.
+			(
+				vec![PotentialParent {
+					hash: *pending_hash,
+					header: pending_header.clone(),
+					depth: route_to_pending.enacted().len(),
+					aligned_with_pending: true,
+				}],
+				potential_parents,
+			)
+		},
+		_ => (only_included, Default::default()),
+	};
+
+	if potential_parents.len() > params.max_depth {
+		return Ok(potential_parents);
+	}
+
+	// Build up the ancestry record of the relay chain to compare against.
+	let rp_ancestry =
+		build_relay_parent_ancestry(params.ancestry_lookback, params.relay_parent, relay_client)
+			.await?;
+
+	Ok(search_child_branches_for_parents(
+		frontier,
+		maybe_route_to_last_pending,
+		included_header,
+		maybe_pending.map(|(_, hash)| hash),
+		backend,
+		params.max_depth,
+		params.ignore_alternative_branches,
+		rp_ancestry,
+		potential_parents,
+	))
+}
+
+/// Fetch the included block from the relay chain.
+async fn fetch_included_from_relay_chain<B: BlockT>(
+	relay_client: &impl RelayChainInterface,
+	backend: &impl Backend<B>,
+	para_id: ParaId,
+	relay_parent: RelayHash,
+) -> Result<Option<(B::Header, B::Hash)>, RelayChainError> {
+	// Fetch the pending header from the relay chain. We use `OccupiedCoreAssumption::TimedOut`
+	// so that even if there is a pending candidate, we assume it is timed out and we get the
+	// included head.
+	let included_header = relay_client
+		.persisted_validation_data(relay_parent, para_id, OccupiedCoreAssumption::TimedOut)
+		.await?;
+	let included_header = match included_header {
+		Some(pvd) => pvd.parent_head,
+		None => return Ok(None), // this implies the para doesn't exist.
+	};
+
+	let included_header = match B::Header::decode(&mut &included_header.0[..]).ok() {
+		None => return Ok(None),
+		Some(x) => x,
+	};
+
+	let included_hash = included_header.hash();
+	// If the included block is not locally known, we can't do anything.
+	match backend.blockchain().header(included_hash) {
+		Ok(None) => {
+			tracing::warn!(
+				target: PARENT_SEARCH_LOG_TARGET,
+				%included_hash,
+				"Failed to get header for included block.",
+			);
+			return Ok(None);
+		},
+		Err(e) => {
+			tracing::warn!(
+				target: PARENT_SEARCH_LOG_TARGET,
+				%included_hash,
+				%e,
+				"Failed to get header for included block.",
+			);
+			return Ok(None);
+		},
+		_ => {},
+	};
+
+	Ok(Some((included_header, included_hash)))
+}
+
+/// Build an ancestry of relay parents that are acceptable.
+///
+/// An acceptable relay parent is one that is no more than `ancestry_lookback` + 1 blocks below the
+/// relay parent we want to build on. Teyrchain blocks anchored on relay parents older than that can
+/// not be considered potential parents for block building. They have no chance of still getting
+/// included, so our newly build teyrchain block would also not get included.
+///
+/// On success, returns a vector of `(header_hash, state_root)` of the relevant relay chain
+/// ancestry blocks.
+async fn build_relay_parent_ancestry(
+	ancestry_lookback: usize,
+	relay_parent: RelayHash,
+	relay_client: &impl RelayChainInterface,
+) -> Result<Vec<(RelayHash, RelayHash)>, RelayChainError> {
+	let mut ancestry = Vec::with_capacity(ancestry_lookback + 1);
+	let mut current_rp = relay_parent;
+	let mut required_session = None;
+	while ancestry.len() <= ancestry_lookback {
+		let Some(header) = relay_client.header(RBlockId::hash(current_rp)).await? else { break };
+
+		let session = relay_client.session_index_for_child(current_rp).await?;
+		if required_session.get_or_insert(session) != &session {
+			// Respect the relay-chain rule not to cross session boundaries.
+			break;
+		}
+
+		ancestry.push((current_rp, *header.state_root()));
+		current_rp = *header.parent_hash();
+
+		// don't iterate back into the genesis block.
+		if header.number == 1 {
+			break;
+		}
+	}
+	Ok(ancestry)
+}
+
+/// Start search for child blocks that can be used as parents.
+pub fn search_child_branches_for_parents<Block: BlockT>(
+	mut frontier: Vec<PotentialParent<Block>>,
+	maybe_route_to_last_pending: Option<TreeRoute<Block>>,
+	included_header: Block::Header,
+	pending_hash: Option<Block::Hash>,
+	backend: &impl Backend<Block>,
+	max_depth: usize,
+	ignore_alternative_branches: bool,
+	rp_ancestry: Vec<(RelayHash, RelayHash)>,
+	mut potential_parents: Vec<PotentialParent<Block>>,
+) -> Vec<PotentialParent<Block>> {
+	let included_hash = included_header.hash();
+	let is_hash_in_ancestry = |hash| rp_ancestry.iter().any(|x| x.0 == hash);
+	let is_root_in_ancestry = |root| rp_ancestry.iter().any(|x| x.1 == root);
+
+	// The distance between pending and included block. Is later used to check if a child
+	// is aligned with pending when it is between pending and included block.
+	let pending_distance = maybe_route_to_last_pending.as_ref().map(|route| route.enacted().len());
+
+	// If a block is on the path included -> pending, we consider it `aligned_with_pending`.
+	let is_child_pending = |hash| {
+		maybe_route_to_last_pending
+			.as_ref()
+			.map_or(true, |route| route.enacted().iter().any(|x| x.hash == hash))
+	};
+
+	tracing::trace!(
+		target: PARENT_SEARCH_LOG_TARGET,
+		?included_hash,
+		included_num = ?included_header.number(),
+		?pending_hash ,
+		?rp_ancestry,
+		"Searching relay chain ancestry."
+	);
+	while let Some(entry) = frontier.pop() {
+		let is_pending = pending_hash.as_ref().map_or(false, |h| &entry.hash == h);
+		let is_included = included_hash == entry.hash;
+
+		// note: even if the pending block or included block have a relay parent
+		// outside of the expected part of the relay chain, they are always allowed
+		// because they have already been posted on chain.
+		let is_potential = is_pending || is_included || {
+			let digest = entry.header.digest();
+			let is_hash_in_ancestry_check = cumulus_primitives_core::extract_relay_parent(digest)
+				.map_or(false, is_hash_in_ancestry);
+			let is_root_in_ancestry_check =
+				cumulus_primitives_core::rpsr_digest::extract_relay_parent_storage_root(digest)
+					.map(|(r, _n)| r)
+					.map_or(false, is_root_in_ancestry);
+
+			is_hash_in_ancestry_check || is_root_in_ancestry_check
+		};
+
+		let parent_aligned_with_pending = entry.aligned_with_pending;
+		let child_depth = entry.depth + 1;
+		let hash = entry.hash;
+
+		tracing::trace!(
+			target: PARENT_SEARCH_LOG_TARGET,
+			?hash,
+			is_potential,
+			is_pending,
+			is_included,
+			"Checking potential parent."
+		);
+
+		if is_potential {
+			potential_parents.push(entry);
+		}
+
+		if !is_potential || child_depth > max_depth {
+			continue;
+		}
+
+		// push children onto search frontier.
+		for child in backend.blockchain().children(hash).ok().into_iter().flatten() {
+			tracing::trace!(target: PARENT_SEARCH_LOG_TARGET, ?child, child_depth, ?pending_distance, "Looking at child.");
+
+			let aligned_with_pending = parent_aligned_with_pending &&
+				(pending_distance.map_or(true, |dist| child_depth > dist) ||
+					is_child_pending(child));
+
+			if ignore_alternative_branches && !aligned_with_pending {
+				tracing::trace!(target: PARENT_SEARCH_LOG_TARGET, ?child, "Child is not aligned with pending block.");
+				continue;
+			}
+
+			let Ok(Some(header)) = backend.blockchain().header(child) else { continue };
+
+			frontier.push(PotentialParent {
+				hash: child,
+				header,
+				depth: child_depth,
+				aligned_with_pending,
+			});
+		}
+	}
+
+	potential_parents
+}

pezcumulus/client/consensus/common/src/teyrchain_consensus.rs

@@ -0,0 +1,544 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use cumulus_relay_chain_streams::{finalized_heads, new_best_heads};
+use pezsc_client_api::{
+	Backend, BlockBackend, BlockImportNotification, BlockchainEvents, Finalizer, UsageProvider,
+};
+use pezsc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
+use schnellru::{ByLength, LruMap};
+use pezsp_blockchain::Error as ClientError;
+use pezsp_consensus::{BlockOrigin, BlockStatus};
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT};
+
+use cumulus_client_pov_recovery::{RecoveryKind, RecoveryRequest};
+use cumulus_relay_chain_interface::RelayChainInterface;
+
+use pezkuwi_primitives::Id as ParaId;
+
+use codec::Decode;
+use futures::{
+	channel::mpsc::{Sender, UnboundedSender},
+	pin_mut, select, FutureExt, SinkExt, Stream, StreamExt,
+};
+use pezsp_core::traits::SpawnEssentialNamed;
+
+use std::sync::Arc;
+
+const LOG_TARGET: &str = "pezcumulus-consensus";
+const FINALIZATION_CACHE_SIZE: u32 = 40;
+
+fn handle_new_finalized_head<P, Block, B>(
+	teyrchain: &Arc<P>,
+	header: Block::Header,
+	last_seen_finalized_hashes: &mut LruMap<Block::Hash, ()>,
+) where
+	Block: BlockT,
+	B: Backend<Block>,
+	P: Finalizer<Block, B> + UsageProvider<Block> + BlockchainEvents<Block>,
+{
+	let hash = header.hash();
+
+	last_seen_finalized_hashes.insert(hash, ());
+
+	// Only finalize if we are below the incoming finalized teyrchain head
+	if teyrchain.usage_info().chain.finalized_number < *header.number() {
+		tracing::debug!(
+			target: LOG_TARGET,
+			block_hash = ?hash,
+			"Attempting to finalize header.",
+		);
+		if let Err(e) = teyrchain.finalize_block(hash, None, true) {
+			match e {
+				ClientError::UnknownBlock(_) => tracing::debug!(
+					target: LOG_TARGET,
+					block_hash = ?hash,
+					"Could not finalize block because it is unknown.",
+				),
+				_ => tracing::warn!(
+					target: LOG_TARGET,
+					error = ?e,
+					block_hash = ?hash,
+					"Failed to finalize block",
+				),
+			}
+		}
+	}
+}
+
+/// Streams finalized teyrchain heads from the relay chain.
+///
+/// This worker continuously monitors the relay chain for finalized blocks and extracts
+/// the corresponding teyrchain head data for the given `para_id`. The extracted head
+/// data is sent through the provided channel for consumption by the consensus system.
+///
+/// This is necessary because finalization of blocks can take a long
+/// time. During this blocking operation, we should not keep references to finality notifications,
+/// because that prevents the corresponding blocks from getting pruned.
+pub async fn finalized_head_stream_worker<R: RelayChainInterface + Clone, Block: BlockT>(
+	mut tx: UnboundedSender<Block::Header>,
+	para_id: ParaId,
+	relay_chain: R,
+) {
+	let finalized_heads = match finalized_heads(relay_chain.clone(), para_id).await {
+		Ok(finalized_heads_stream) => finalized_heads_stream.fuse(),
+		Err(err) => {
+			tracing::error!(target: LOG_TARGET, error = ?err, "Unable to retrieve finalized heads stream.");
+			return;
+		},
+	};
+
+	pin_mut!(finalized_heads);
+	loop {
+		if let Some((head_data, _)) = finalized_heads.next().await {
+			let header = match Block::Header::decode(&mut &head_data[..]) {
+				Ok(header) => header,
+				Err(err) => {
+					tracing::debug!(
+						target: LOG_TARGET,
+						error = ?err,
+						"Could not decode teyrchain header while following finalized heads.",
+					);
+					continue;
+				},
+			};
+			if let Err(e) = tx.send(header).await {
+				tracing::error!(target: LOG_TARGET, ?e, "Error while sending finalized head.");
+				return;
+			};
+		}
+	}
+}
+
+/// Follow the finalized head of the given teyrchain.
+///
+/// For every finalized block of the relay chain, it will get the included teyrchain header
+/// corresponding to `para_id` and will finalize it in the teyrchain.
+async fn follow_finalized_head<P, Block, B>(
+	teyrchain: Arc<P>,
+	finalized_head_stream: Box<impl Stream<Item = Block::Header> + Unpin + Send>,
+) where
+	Block: BlockT,
+	P: Finalizer<Block, B> + UsageProvider<Block> + BlockchainEvents<Block>,
+	B: Backend<Block>,
+{
+	let mut imported_blocks = teyrchain.import_notification_stream().fuse();
+	let mut finalized_head_stream = finalized_head_stream.fuse();
+
+	// We use this cache to finalize blocks that are imported late.
+	// For example, a block that has been recovered via PoV-Recovery
+	// on a full node can have several minutes delay. With this cache
+	// we have some "memory" of recently finalized blocks.
+	let mut last_seen_finalized_hashes = LruMap::new(ByLength::new(FINALIZATION_CACHE_SIZE));
+
+	loop {
+		select! {
+			fin = finalized_head_stream.next() => {
+				match fin {
+					Some(finalized_head) =>
+						handle_new_finalized_head(&teyrchain, finalized_head, &mut last_seen_finalized_hashes),
+					None => {
+						tracing::debug!(target: LOG_TARGET, "Stopping following finalized head.");
+						return
+					}
+				}
+			},
+			imported = imported_blocks.next() => {
+				match imported {
+					Some(imported_block) => {
+						// When we see a block import that is already finalized, we immediately finalize it.
+						if last_seen_finalized_hashes.peek(&imported_block.hash).is_some() {
+							tracing::debug!(
+								target: LOG_TARGET,
+								block_hash = ?imported_block.hash,
+								"Setting newly imported block as finalized.",
+							);
+
+							if let Err(e) = teyrchain.finalize_block(imported_block.hash, None, true) {
+								match e {
+									ClientError::UnknownBlock(_) => tracing::debug!(
+										target: LOG_TARGET,
+										block_hash = ?imported_block.hash,
+										"Could not finalize block because it is unknown.",
+									),
+									_ => tracing::warn!(
+										target: LOG_TARGET,
+										error = ?e,
+										block_hash = ?imported_block.hash,
+										"Failed to finalize block",
+									),
+								}
+							}
+						}
+					},
+					None => {
+						tracing::debug!(
+							target: LOG_TARGET,
+							"Stopping following imported blocks.",
+						);
+						return
+					}
+				}
+			}
+		}
+	}
+}
+
+/// Spawns the essential finalization tasks for teyrchain consensus.
+///
+/// This function creates and spawns two critical background tasks:
+/// 1. A finalized head stream worker that monitors relay chain finality and extracts included
+///    headers
+/// 2. The main teyrchain consensus task that handles finalization and best block updates
+pub fn spawn_teyrchain_consensus_tasks<P, R, Block, B, S>(
+	para_id: ParaId,
+	teyrchain: Arc<P>,
+	relay_chain: R,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
+	recovery_chan_tx: Option<Sender<RecoveryRequest<Block>>>,
+	spawn_handle: S,
+) where
+	Block: BlockT,
+	P: Finalizer<Block, B>
+		+ UsageProvider<Block>
+		+ Send
+		+ Sync
+		+ BlockBackend<Block>
+		+ BlockchainEvents<Block>
+		+ 'static,
+	for<'a> &'a P: BlockImport<Block>,
+	R: RelayChainInterface + Clone + 'static,
+	S: SpawnEssentialNamed + 'static,
+	B: Backend<Block> + 'static,
+{
+	let (tx, rx) = futures::channel::mpsc::unbounded();
+	let worker = finalized_head_stream_worker::<_, Block>(tx, para_id, relay_chain.clone());
+	let consensus = run_teyrchain_consensus(
+		para_id,
+		teyrchain,
+		relay_chain,
+		announce_block,
+		Box::new(rx),
+		recovery_chan_tx,
+	);
+
+	spawn_handle.spawn_essential_blocking("pezcumulus-consensus", None, Box::pin(consensus));
+	spawn_handle.spawn_essential_blocking(
+		"pezcumulus-consensus-finality-stream",
+		None,
+		Box::pin(worker),
+	);
+}
+
+/// Run the teyrchain consensus.
+///
+/// This will follow the given `relay_chain` to act as consensus for the teyrchain that corresponds
+/// to the given `para_id`. It will set the new best block of the teyrchain as it gets aware of it.
+/// The same happens for the finalized block.
+///
+/// # Note
+///
+/// This will access the backend of the teyrchain and thus, this future should be spawned as
+/// blocking task.
+pub async fn run_teyrchain_consensus<P, R, Block, B>(
+	para_id: ParaId,
+	teyrchain: Arc<P>,
+	relay_chain: R,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
+	finalized_head_stream: Box<impl Stream<Item = Block::Header> + Unpin + Send>,
+	recovery_chan_tx: Option<Sender<RecoveryRequest<Block>>>,
+) where
+	Block: BlockT,
+	P: Finalizer<Block, B>
+		+ UsageProvider<Block>
+		+ Send
+		+ Sync
+		+ BlockBackend<Block>
+		+ BlockchainEvents<Block>,
+	for<'a> &'a P: BlockImport<Block>,
+	R: RelayChainInterface + Clone,
+	B: Backend<Block>,
+{
+	let follow_new_best = follow_new_best(
+		para_id,
+		teyrchain.clone(),
+		relay_chain.clone(),
+		announce_block,
+		recovery_chan_tx,
+	);
+	let follow_finalized_head = follow_finalized_head(teyrchain, finalized_head_stream);
+	select! {
+		_ = follow_new_best.fuse() => {},
+		_ = follow_finalized_head.fuse() => {},
+	}
+}
+
+/// Follow the relay chain new best head, to update the Teyrchain new best head.
+async fn follow_new_best<P, R, Block, B>(
+	para_id: ParaId,
+	teyrchain: Arc<P>,
+	relay_chain: R,
+	announce_block: Arc<dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync>,
+	mut recovery_chan_tx: Option<Sender<RecoveryRequest<Block>>>,
+) where
+	Block: BlockT,
+	P: Finalizer<Block, B>
+		+ UsageProvider<Block>
+		+ Send
+		+ Sync
+		+ BlockBackend<Block>
+		+ BlockchainEvents<Block>,
+	for<'a> &'a P: BlockImport<Block>,
+	R: RelayChainInterface + Clone,
+	B: Backend<Block>,
+{
+	let new_best_heads = match new_best_heads(relay_chain, para_id).await {
+		Ok(best_heads_stream) => best_heads_stream.fuse(),
+		Err(err) => {
+			tracing::error!(target: LOG_TARGET, error = ?err, "Unable to retrieve best heads stream.");
+			return;
+		},
+	};
+
+	pin_mut!(new_best_heads);
+
+	let mut imported_blocks = teyrchain.import_notification_stream().fuse();
+	// The unset best header of the teyrchain. Will be `Some(_)` when we have imported a relay chain
+	// block before the associated teyrchain block. In this case we need to wait for this block to
+	// be imported to set it as new best.
+	let mut unset_best_header = None;
+
+	loop {
+		select! {
+			h = new_best_heads.next() => {
+				match h {
+					Some(h) => handle_new_best_teyrchain_head(
+						h,
+						&*teyrchain,
+						&mut unset_best_header,
+						recovery_chan_tx.as_mut(),
+					).await,
+					None => {
+						tracing::debug!(
+							target: LOG_TARGET,
+							"Stopping following new best.",
+						);
+						return
+					}
+				}
+			},
+			i = imported_blocks.next() => {
+				match i {
+					Some(i) => handle_new_block_imported(
+						i,
+						&mut unset_best_header,
+						&*teyrchain,
+						&*announce_block,
+					).await,
+					None => {
+						tracing::debug!(
+							target: LOG_TARGET,
+							"Stopping following imported blocks.",
+						);
+						return
+					}
+				}
+			},
+		}
+	}
+}
+
+/// Handle a new import block of the teyrchain.
+async fn handle_new_block_imported<Block, P>(
+	notification: BlockImportNotification<Block>,
+	unset_best_header_opt: &mut Option<Block::Header>,
+	teyrchain: &P,
+	announce_block: &(dyn Fn(Block::Hash, Option<Vec<u8>>) + Send + Sync),
+) where
+	Block: BlockT,
+	P: UsageProvider<Block> + Send + Sync + BlockBackend<Block>,
+	for<'a> &'a P: BlockImport<Block>,
+{
+	// HACK
+	//
+	// Remove after https://github.com/pezkuwichain/kurdistan-sdk/issues/76 or similar is merged
+	if notification.origin != BlockOrigin::Own {
+		announce_block(notification.hash, None);
+	}
+
+	let unset_best_header = match (notification.is_new_best, &unset_best_header_opt) {
+		// If this is the new best block or we don't have any unset block, we can end it here.
+		(true, _) | (_, None) => return,
+		(false, Some(ref u)) => u,
+	};
+
+	let unset_hash = if notification.header.number() < unset_best_header.number() {
+		return;
+	} else if notification.header.number() == unset_best_header.number() {
+		let unset_hash = unset_best_header.hash();
+
+		if unset_hash != notification.hash {
+			return;
+		} else {
+			unset_hash
+		}
+	} else {
+		unset_best_header.hash()
+	};
+
+	match teyrchain.block_status(unset_hash) {
+		Ok(BlockStatus::InChainWithState) => {
+			let unset_best_header = unset_best_header_opt
+				.take()
+				.expect("We checked above that the value is set; qed");
+			tracing::debug!(
+				target: LOG_TARGET,
+				?unset_hash,
+				"Importing block as new best for teyrchain.",
+			);
+			import_block_as_new_best(unset_hash, unset_best_header, teyrchain).await;
+		},
+		state => tracing::debug!(
+			target: LOG_TARGET,
+			?unset_best_header,
+			?notification.header,
+			?state,
+			"Unexpected state for unset best header.",
+		),
+	}
+}
+
+/// Handle the new best teyrchain head as extracted from the new best relay chain.
+async fn handle_new_best_teyrchain_head<Block, P>(
+	head: Vec<u8>,
+	teyrchain: &P,
+	unset_best_header: &mut Option<Block::Header>,
+	mut recovery_chan_tx: Option<&mut Sender<RecoveryRequest<Block>>>,
+) where
+	Block: BlockT,
+	P: UsageProvider<Block> + Send + Sync + BlockBackend<Block>,
+	for<'a> &'a P: BlockImport<Block>,
+{
+	let teyrchain_head = match <<Block as BlockT>::Header>::decode(&mut &head[..]) {
+		Ok(header) => header,
+		Err(err) => {
+			tracing::debug!(
+				target: LOG_TARGET,
+				error = ?err,
+				"Could not decode Teyrchain header while following best heads.",
+			);
+			return;
+		},
+	};
+
+	let hash = teyrchain_head.hash();
+
+	if teyrchain.usage_info().chain.best_hash == hash {
+		tracing::debug!(
+			target: LOG_TARGET,
+			block_hash = ?hash,
+			"Skipping set new best block, because block is already the best.",
+		);
+		return;
+	}
+
+	// Make sure the block is already known or otherwise we skip setting new best.
+	match teyrchain.block_status(hash) {
+		Ok(BlockStatus::InChainWithState) => {
+			unset_best_header.take();
+			tracing::debug!(
+				target: LOG_TARGET,
+				included = ?hash,
+				"Importing block as new best for teyrchain.",
+			);
+			import_block_as_new_best(hash, teyrchain_head, teyrchain).await;
+		},
+		Ok(BlockStatus::InChainPruned) => {
+			tracing::error!(
+				target: LOG_TARGET,
+				block_hash = ?hash,
+				"Trying to set pruned block as new best!",
+			);
+		},
+		Ok(BlockStatus::Unknown) => {
+			*unset_best_header = Some(teyrchain_head);
+
+			tracing::debug!(
+				target: LOG_TARGET,
+				block_hash = ?hash,
+				"Teyrchain block not yet imported, waiting for import to enact as best block.",
+			);
+
+			if let Some(ref mut recovery_chan_tx) = recovery_chan_tx {
+				// Best effort channel to actively encourage block recovery.
+				// An error here is not fatal; the relay chain continuously re-announces
+				// the best block, thus we will have other opportunities to retry.
+				let req = RecoveryRequest { hash, kind: RecoveryKind::Full };
+				if let Err(err) = recovery_chan_tx.try_send(req) {
+					tracing::warn!(
+						target: LOG_TARGET,
+						block_hash = ?hash,
+						error = ?err,
+						"Unable to notify block recovery subsystem"
+					)
+				}
+			}
+		},
+		Err(e) => {
+			tracing::error!(
+				target: LOG_TARGET,
+				block_hash = ?hash,
+				error = ?e,
+				"Failed to get block status of block.",
+			);
+		},
+		_ => {},
+	}
+}
+
+async fn import_block_as_new_best<Block, P>(hash: Block::Hash, header: Block::Header, teyrchain: &P)
+where
+	Block: BlockT,
+	P: UsageProvider<Block> + Send + Sync + BlockBackend<Block>,
+	for<'a> &'a P: BlockImport<Block>,
+{
+	let best_number = teyrchain.usage_info().chain.best_number;
+	if *header.number() < best_number {
+		tracing::debug!(
+			target: LOG_TARGET,
+			%best_number,
+			block_number = %header.number(),
+			"Skipping importing block as new best block, because there already exists a \
+			 best block with an higher number",
+		);
+		return;
+	}
+
+	// Make it the new best block
+	let mut block_import_params = BlockImportParams::new(BlockOrigin::ConsensusBroadcast, header);
+	block_import_params.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+	block_import_params.import_existing = true;
+
+	if let Err(err) = teyrchain.import_block(block_import_params).await {
+		tracing::warn!(
+			target: LOG_TARGET,
+			block_hash = ?hash,
+			error = ?err,
+			"Failed to set new best block.",
+		);
+	}
+}

pezcumulus/client/consensus/proposer/Cargo.toml

@@ -0,0 +1,45 @@
+[package]
+name = "pezcumulus-client-consensus-proposer"
+description = "A Bizinikiwi `Proposer` for building teyrchain blocks"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+anyhow = { workspace = true, default-features = true }
+async-trait = { workspace = true }
+thiserror = { workspace = true }
+
+# Bizinikiwi
+pezsc-basic-authorship = { workspace = true }
+pezsc-block-builder = { workspace = true }
+pezsc-transaction-pool-api = { workspace = true }
+pezsp-api = { workspace = true, default-features = true }
+pezsp-blockchain = { workspace = true, default-features = true }
+pezsp-consensus = { workspace = true, default-features = true }
+pezsp-inherents = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+pezsp-state-machine = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-primitives-teyrchain-inherent = { workspace = true, default-features = true }
+
+[features]
+runtime-benchmarks = [
+	"pezcumulus-primitives-teyrchain-inherent/runtime-benchmarks",
+	"pezsc-basic-authorship/runtime-benchmarks",
+	"pezsc-block-builder/runtime-benchmarks",
+	"pezsc-transaction-pool-api/runtime-benchmarks",
+	"pezsp-api/runtime-benchmarks",
+	"pezsp-blockchain/runtime-benchmarks",
+	"pezsp-consensus/runtime-benchmarks",
+	"pezsp-inherents/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+	"pezsp-state-machine/runtime-benchmarks",
+]

pezcumulus/client/consensus/proposer/src/lib.rs

@@ -0,0 +1,126 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! The Pezcumulus [`ProposerInterface`] is an extension of the Bizinikiwi [`ProposerFactory`]
+//! for creating new teyrchain blocks.
+//!
+//! This utility is designed to be composed within any collator consensus algorithm.
+
+use async_trait::async_trait;
+use cumulus_primitives_teyrchain_inherent::TeyrchainInherentData;
+use pezsc_basic_authorship::{ProposeArgs, ProposerFactory};
+use pezsc_block_builder::BlockBuilderApi;
+use pezsc_transaction_pool_api::TransactionPool;
+use pezsp_api::{ApiExt, CallApiAt, ProvideRuntimeApi};
+use pezsp_blockchain::HeaderBackend;
+use pezsp_consensus::{EnableProofRecording, Environment, Proposal};
+use pezsp_inherents::{InherentData, InherentDataProvider};
+use pezsp_runtime::{traits::Block as BlockT, Digest};
+use pezsp_state_machine::StorageProof;
+use std::{fmt::Debug, time::Duration};
+
+/// Errors that can occur when proposing a teyrchain block.
+#[derive(thiserror::Error, Debug)]
+#[error(transparent)]
+pub struct Error {
+	inner: anyhow::Error,
+}
+
+impl Error {
+	/// Create an error tied to the creation of a proposer.
+	pub fn proposer_creation(err: impl Into<anyhow::Error>) -> Self {
+		Error { inner: err.into().context("Proposer Creation") }
+	}
+
+	/// Create an error tied to the proposing logic itself.
+	pub fn proposing(err: impl Into<anyhow::Error>) -> Self {
+		Error { inner: err.into().context("Proposing") }
+	}
+}
+
+/// A type alias for easily referring to the type of a proposal produced by a specific
+/// [`ProposerInterface`].
+pub type ProposalOf<B> = Proposal<B, StorageProof>;
+
+/// An interface for proposers.
+#[async_trait]
+pub trait ProposerInterface<Block: BlockT> {
+	/// Propose a collation using the supplied `InherentData` and the provided
+	/// `TeyrchainInherentData`.
+	///
+	/// Also specify any required inherent digests, the maximum proposal duration,
+	/// and the block size limit in bytes. See the documentation on
+	/// [`pezsp_consensus::Proposer::propose`] for more details on how to interpret these parameters.
+	///
+	/// The `InherentData` and `Digest` are left deliberately general in order to accommodate
+	/// all possible collator selection algorithms or inherent creation mechanisms,
+	/// while the `TeyrchainInherentData` is made explicit so it will be constructed appropriately.
+	///
+	/// If the `InherentData` passed into this function already has a `TeyrchainInherentData`,
+	/// this should throw an error.
+	async fn propose(
+		&mut self,
+		parent_header: &Block::Header,
+		paras_inherent_data: &TeyrchainInherentData,
+		other_inherent_data: InherentData,
+		inherent_digests: Digest,
+		max_duration: Duration,
+		block_size_limit: Option<usize>,
+	) -> Result<Option<Proposal<Block, StorageProof>>, Error>;
+}
+
+#[async_trait]
+impl<Block, A, C> ProposerInterface<Block> for ProposerFactory<A, C, EnableProofRecording>
+where
+	A: TransactionPool<Block = Block> + 'static,
+	C: HeaderBackend<Block> + ProvideRuntimeApi<Block> + CallApiAt<Block> + Send + Sync + 'static,
+	C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
+	Block: pezsp_runtime::traits::Block,
+{
+	async fn propose(
+		&mut self,
+		parent_header: &Block::Header,
+		paras_inherent_data: &TeyrchainInherentData,
+		other_inherent_data: InherentData,
+		inherent_digests: Digest,
+		max_duration: Duration,
+		block_size_limit: Option<usize>,
+	) -> Result<Option<Proposal<Block, StorageProof>>, Error> {
+		let proposer = self
+			.init(parent_header)
+			.await
+			.map_err(|e| Error::proposer_creation(anyhow::Error::new(e)))?;
+
+		let mut inherent_data = other_inherent_data;
+		paras_inherent_data
+			.provide_inherent_data(&mut inherent_data)
+			.await
+			.map_err(|e| Error::proposing(anyhow::Error::new(e)))?;
+
+		proposer
+			.propose_block(ProposeArgs {
+				inherent_data,
+				inherent_digests,
+				max_duration,
+				block_size_limit,
+				ignored_nodes_by_proof_recording: None,
+			})
+			.await
+			.map(Some)
+			.map_err(|e| Error::proposing(anyhow::Error::new(e)).into())
+	}
+}

pezcumulus/client/consensus/relay-chain/Cargo.toml

@@ -0,0 +1,48 @@
+[package]
+name = "pezcumulus-client-consensus-relay-chain"
+description = "The relay-chain provided consensus algorithm"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+async-trait = { workspace = true }
+futures = { workspace = true }
+parking_lot = { workspace = true, default-features = true }
+tracing = { workspace = true, default-features = true }
+
+# Bizinikiwi
+prometheus-endpoint = { workspace = true, default-features = true }
+pezsc-consensus = { workspace = true, default-features = true }
+pezsp-api = { workspace = true, default-features = true }
+pezsp-block-builder = { workspace = true, default-features = true }
+pezsp-blockchain = { workspace = true, default-features = true }
+pezsp-consensus = { workspace = true, default-features = true }
+pezsp-core = { workspace = true, default-features = true }
+pezsp-inherents = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-client-consensus-common = { workspace = true, default-features = true }
+pezcumulus-primitives-core = { workspace = true, default-features = true }
+pezcumulus-relay-chain-interface = { workspace = true, default-features = true }
+
+[features]
+runtime-benchmarks = [
+	"pezcumulus-client-consensus-common/runtime-benchmarks",
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezcumulus-relay-chain-interface/runtime-benchmarks",
+	"pezsc-consensus/runtime-benchmarks",
+	"pezsp-api/runtime-benchmarks",
+	"pezsp-block-builder/runtime-benchmarks",
+	"pezsp-blockchain/runtime-benchmarks",
+	"pezsp-consensus/runtime-benchmarks",
+	"pezsp-inherents/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+]

pezcumulus/client/consensus/relay-chain/src/import_queue.rs

@@ -0,0 +1,88 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+use std::{marker::PhantomData, sync::Arc};
+
+use pezsc_consensus::{import_queue::Verifier as VerifierT, BlockImportParams};
+use pezsp_api::ProvideRuntimeApi;
+use pezsp_block_builder::BlockBuilder as BlockBuilderApi;
+use pezsp_inherents::CreateInherentDataProviders;
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT};
+
+/// A verifier that just checks the inherents.
+pub struct Verifier<Client, Block, CIDP> {
+	client: Arc<Client>,
+	create_inherent_data_providers: CIDP,
+	_marker: PhantomData<Block>,
+}
+
+impl<Client, Block, CIDP> Verifier<Client, Block, CIDP> {
+	/// Create a new instance.
+	pub fn new(client: Arc<Client>, create_inherent_data_providers: CIDP) -> Self {
+		Self { client, create_inherent_data_providers, _marker: PhantomData }
+	}
+}
+
+#[async_trait::async_trait]
+impl<Client, Block, CIDP> VerifierT<Block> for Verifier<Client, Block, CIDP>
+where
+	Block: BlockT,
+	Client: ProvideRuntimeApi<Block> + Send + Sync,
+	<Client as ProvideRuntimeApi<Block>>::Api: BlockBuilderApi<Block>,
+	CIDP: CreateInherentDataProviders<Block, ()>,
+{
+	async fn verify(
+		&self,
+		mut block_params: BlockImportParams<Block>,
+	) -> Result<BlockImportParams<Block>, String> {
+		block_params.fork_choice = Some(pezsc_consensus::ForkChoiceStrategy::Custom(
+			block_params.origin == pezsp_consensus::BlockOrigin::NetworkInitialSync,
+		));
+
+		// Skip checks that include execution, if being told so, or when importing only state.
+		//
+		// This is done for example when gap syncing and it is expected that the block after the gap
+		// was checked/chosen properly, e.g. by warp syncing to this block using a finality proof.
+		if block_params.state_action.skip_execution_checks() || block_params.with_state() {
+			return Ok(block_params);
+		}
+
+		if let Some(inner_body) = block_params.body.take() {
+			let inherent_data_providers = self
+				.create_inherent_data_providers
+				.create_inherent_data_providers(*block_params.header.parent_hash(), ())
+				.await
+				.map_err(|e| e.to_string())?;
+
+			let block = Block::new(block_params.header.clone(), inner_body);
+			pezsp_block_builder::check_inherents(
+				self.client.clone(),
+				*block.header().parent_hash(),
+				block.clone(),
+				&inherent_data_providers,
+			)
+			.await
+			.map_err(|e| format!("Error checking block inherents {:?}", e))?;
+
+			let (_, inner_body) = block.deconstruct();
+			block_params.body = Some(inner_body);
+		}
+
+		block_params.post_hash = Some(block_params.header.hash());
+		Ok(block_params)
+	}
+}

pezcumulus/client/consensus/relay-chain/src/lib.rs

@@ -0,0 +1,39 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// Pezcumulus 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.
+
+// Pezcumulus 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 Pezcumulus. If not, see <https://www.gnu.org/licenses/>.
+
+//! The relay-chain provided consensus algorithm for teyrchains.
+//!
+//! This is the simplest consensus algorithm you can use when developing a teyrchain. It is a
+//! permission-less consensus algorithm that doesn't require any staking or similar to join as a
+//! collator. In this algorithm the consensus is provided by the relay-chain. This works in the
+//! following way.
+//!
+//! 1. Each node that sees itself as a collator is free to build a teyrchain candidate.
+//!
+//! 2. This teyrchain candidate is send to the teyrchain validators that are part of the relay
+//! chain.
+//!
+//! 3. The teyrchain validators validate at most X different teyrchain candidates, where X is the
+//! total number of teyrchain validators.
+//!
+//! 4. The teyrchain candidate that is backed by the most validators is chosen by the relay-chain
+//! block producer to be added as backed candidate on chain.
+//!
+//! 5. After the teyrchain candidate got backed and included, all collators start at 1.
+
+mod import_queue;
+pub use import_queue::Verifier;