Allow providing a custom block announce validator

cli/src/command.rs

@@ -193,6 +193,7 @@ where
 			authority_discovery_enabled,
 			6000,
 			grandpa_pause,
+			None,
 		).map(|(s, _)| s),
 		D::native_version().runtime_version,
 	)

collator/src/lib.rs

@@ -63,15 +63,17 @@ use polkadot_primitives::{
 	}
 };
 use polkadot_cli::{
-	ProvideRuntimeApi, AbstractService, ParachainHost, IdentifyVariant,
-	service::{self, Role}
+	ProvideRuntimeApi, AbstractService, ParachainHost, service,
 };
 pub use polkadot_cli::service::Configuration;
 pub use polkadot_cli::Cli;
 pub use polkadot_validation::SignedStatement;
 pub use polkadot_primitives::parachain::CollatorId;
 pub use sc_network::PeerId;
 pub use service::RuntimeApiCollection;
+use service::{
+	BlockAnnounceValidator, PolkadotExecutor, KusamaExecutor, polkadot_runtime, kusama_runtime, westend_runtime, Role,
+};
 pub use sc_cli::SubstrateCli;
 
 const COLLATION_TIMEOUT: Duration = Duration::from_secs(30);
@@ -339,39 +341,103 @@ fn build_collator_service<S, P, Extrinsic>(
 	Ok(service)
 }
 
-/// Async function that will run the collator node with the given `RelayChainContext` and `ParachainContext`
+macro_rules! start_collator {
+	(
+		$new_full:path, $config:expr, $key:expr, $para_id:expr, $build_parachain_context:expr,
+		$block_announce_validator:expr $(,)?
+	) => {
+		if matches!($config.role, Role::Light) {
+			Err(service::Error::Other("light nodes are unsupported as collator".into()))
+		} else {
+			$new_full($config, Some(($key.public(), $para_id)), None, false, 6000, None, $block_announce_validator)
+				.and_then(|service|
+					$crate::build_collator_service(
+						service,
+						$para_id,
+						$key,
+						$build_parachain_context,
+					)
+				)
+		}
+	};
+}
+
+/// Async function that will run the polkadot collator node with the given `RelayChainContext` and `ParachainContext`
 /// built by the given `BuildParachainContext` and arguments to the underlying polkadot node.
-pub async fn start_collator<P>(
+pub async fn start_collator_polkadot<P>(
 	build_parachain_context: P,
 	para_id: ParaId,
 	key: Arc<CollatorPair>,
 	config: Configuration,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, polkadot_runtime::RuntimeApi,
+			PolkadotExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
 ) -> Result<(), polkadot_service::Error>
 where
 	P: BuildParachainContext,
 	P::ParachainContext: Send + 'static,
 	<P::ParachainContext as ParachainContext>::ProduceCandidate: Send,
 {
-	let is_kusama = config.chain_spec.is_kusama();
-	match (is_kusama, &config.role) {
-		(_, Role::Light) => return Err(
-			polkadot_service::Error::Other("light nodes are unsupported as collator".into())
-		).into(),
-		(true, _) =>
-			build_collator_service(
-				service::kusama_new_full(config, Some((key.public(), para_id)), None, false, 6000, None)?,
-				para_id,
-				key,
-				build_parachain_context,
-			)?.await,
-		(false, _) =>
-			build_collator_service(
-				service::polkadot_new_full(config, Some((key.public(), para_id)), None, false, 6000, None)?,
-				para_id,
-				key,
-				build_parachain_context,
-			)?.await,
-	}
+	start_collator!(
+		service::polkadot_new_full,
+		config,
+		key,
+		para_id,
+		build_parachain_context,
+		block_announce_validator_builder,
+	)?.await
+}
+
+/// Async function that will run the kusama collator node with the given `RelayChainContext` and `ParachainContext`
+/// built by the given `BuildParachainContext` and arguments to the underlying polkadot node.
+pub async fn start_collator_kusama<P>(
+	build_parachain_context: P,
+	para_id: ParaId,
+	key: Arc<CollatorPair>,
+	config: Configuration,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, kusama_runtime::RuntimeApi,
+			KusamaExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
+) -> Result<(), polkadot_service::Error>
+where
+	P: BuildParachainContext,
+	P::ParachainContext: Send + 'static,
+	<P::ParachainContext as ParachainContext>::ProduceCandidate: Send,
+{
+	start_collator!(
+		service::kusama_new_full,
+		config,
+		key,
+		para_id,
+		build_parachain_context,
+		block_announce_validator_builder,
+	)?.await
+}
+
+/// Async function that will run the westend collator node with the given `RelayChainContext` and `ParachainContext`
+/// built by the given `BuildParachainContext` and arguments to the underlying polkadot node.
+pub async fn start_collator_westend<P>(
+	build_parachain_context: P,
+	para_id: ParaId,
+	key: Arc<CollatorPair>,
+	config: Configuration,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, westend_runtime::RuntimeApi,
+			KusamaExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
+) -> Result<(), polkadot_service::Error>
+where
+	P: BuildParachainContext,
+	P::ParachainContext: Send + 'static,
+	<P::ParachainContext as ParachainContext>::ProduceCandidate: Send,
+{
+	start_collator!(
+		service::westend_new_full,
+		config,
+		key,
+		para_id,
+		build_parachain_context,
+		block_announce_validator_builder,
+	)?.await
 }
 
 fn compute_targets(para_id: ParaId, session_keys: &[ValidatorId], roster: DutyRoster) -> HashSet<ValidatorId> {
@@ -423,7 +489,7 @@ mod tests {
 		}
 	}
 
-	// Make sure that the future returned by `start_collator` implementes `Send`.
+	// Make sure that the future returned by `start_collator` implements `Send`.
 	#[test]
 	fn start_collator_is_send() {
 		fn check_send<T: Send>(_: T) {}
@@ -432,11 +498,12 @@ mod tests {
 		let task_executor = Arc::new(|_| unimplemented!());
 		let config = cli.create_configuration(&cli.run.base, task_executor).unwrap();
 
-		check_send(start_collator(
+		check_send(start_collator_polkadot(
 			BuildDummyParachainContext,
 			0.into(),
 			Arc::new(CollatorPair::generate().0),
 			config,
+			None,
 		));
 	}
 }

parachain/test-parachains/adder/collator/src/main.rs

@@ -136,11 +136,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 	let cli = Cli::from_iter(&["-dev"]);
 	let runner = cli.create_runner(&cli.run.base)?;
 	runner.async_run(|config| {
-		collator::start_collator(
+		collator::start_collator_polkadot(
 			context,
 			id,
 			key,
 			config,
+			None,
 		).map_err(|e| e.into())
 	})?;
 

service/src/lib.rs

@@ -41,7 +41,7 @@ pub use sc_client::{ExecutionStrategy, CallExecutor, Client};
 pub use sc_client_api::backend::Backend;
 pub use sp_api::{Core as CoreApi, ConstructRuntimeApi, ProvideRuntimeApi, StateBackend};
 pub use sp_runtime::traits::{HashFor, NumberFor};
-pub use consensus_common::SelectChain;
+pub use consensus_common::{block_validation::BlockAnnounceValidator, SelectChain};
 pub use polkadot_primitives::parachain::{CollatorId, ParachainHost};
 pub use polkadot_primitives::Block;
 pub use sp_runtime::traits::{Block as BlockT, self as runtime_traits, BlakeTwo256};
@@ -233,6 +233,9 @@ pub fn polkadot_new_full(
 	authority_discovery_enabled: bool,
 	slot_duration: u64,
 	grandpa_pause: Option<(u32, u32)>,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, polkadot_runtime::RuntimeApi,
+			PolkadotExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
 )
 	-> Result<(
 		impl AbstractService<
@@ -252,6 +255,7 @@ pub fn polkadot_new_full(
 		authority_discovery_enabled,
 		slot_duration,
 		grandpa_pause,
+		block_announce_validator_builder,
 	)
 }
 
@@ -264,6 +268,9 @@ pub fn kusama_new_full(
 	authority_discovery_enabled: bool,
 	slot_duration: u64,
 	grandpa_pause: Option<(u32, u32)>,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, kusama_runtime::RuntimeApi,
+			KusamaExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
 )
 	-> Result<(
 		impl AbstractService<
@@ -283,6 +290,7 @@ pub fn kusama_new_full(
 		authority_discovery_enabled,
 		slot_duration,
 		grandpa_pause,
+		block_announce_validator_builder,
 	)
 }
 
@@ -295,6 +303,9 @@ pub fn westend_new_full(
 	authority_discovery_enabled: bool,
 	slot_duration: u64,
 	grandpa_pause: Option<(u32, u32)>,
+	block_announce_validator_builder:
+		Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, westend_runtime::RuntimeApi,
+			KusamaExecutor>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
 )
 	-> Result<(
 		impl AbstractService<
@@ -314,6 +325,7 @@ pub fn westend_new_full(
 		authority_discovery_enabled,
 		slot_duration,
 		grandpa_pause,
+		block_announce_validator_builder,
 	)
 }
 
@@ -337,6 +349,8 @@ pub fn new_full<Runtime, Dispatch, Extrinsic>(
 	authority_discovery_enabled: bool,
 	slot_duration: u64,
 	grandpa_pause: Option<(u32, u32)>,
+	block_announce_validator_builder: Option<Box<dyn FnOnce(Arc<service::TFullClient<Block, Runtime,
+		Dispatch>>) -> Box<dyn BlockAnnounceValidator<Block> + Send> + Send + 'static>>,
 )
 	-> Result<(
 		impl AbstractService<
@@ -375,10 +389,14 @@ pub fn new_full<Runtime, Dispatch, Extrinsic>(
 	let authority_discovery_enabled = authority_discovery_enabled;
 	let slot_duration = slot_duration;
 
-	let (builder, mut import_setup, inherent_data_providers) = new_full_start!(config, Runtime, Dispatch);
+	let (mut builder, mut import_setup, inherent_data_providers) = new_full_start!(config, Runtime, Dispatch);
 
 	let backend = builder.backend().clone();
 
+	if let Some(f) = block_announce_validator_builder {
+		builder = builder.with_block_announce_validator(f)?;
+	}
+
 	let service = builder
 		.with_finality_proof_provider(|client, backend| {
 			let provider = client as Arc<dyn grandpa::StorageAndProofProvider<_, _>>;