Zombienet: add new test collator and integration test (#4797)

* WIP: Wasm compilation perf Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Fix Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Undying collator WIP Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Fix build Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * more fixes Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * update test with undying Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Correctly compute post hash Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * update helper Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * squash bugs Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Add --pov-size cli parameter Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * refactor Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix strings Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Add pov-size param to export genesis state Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Fix graveyard size Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * docs + fmt Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Fix PVF bug and switch to u8 graves Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Update tests Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Build/publish undying collator as colander img Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * add undying-collator to colander image * Fix transaction overflow Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * warn fix Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * add parachain id for export genesis cli Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix the build Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * scale test up Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * default parachain id Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Add PVF complexity param Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Hash on each iteration Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Update pvf metric histogram buckets Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Ladi attempt to fix tests * Fix test * Fix typos * Fix pvf typo * CI integration Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * cargo lock missing Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix clap merge damage Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * add zombienet image back Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Use collator image from env Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * bump zombienet version * update test to check pvf prep/exec Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * delete file Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * use default bootnode in upgrade test * FIx tests Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * add some stress - pvf exec times up to 2s Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix name Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Dial down pvf execution time < 2s Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * 1100 Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * bump Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * try again Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Add connectivity check Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Fixes and refactor folder struct Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * change toml name to match Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * bump zombienet Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * impl Feedback for Review 😎 Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * spell check Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Zombienet: add disputes test (#4859) * Zombienet disputes test Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * CI: add zombienet-parachain-disputes Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Bump zombienet Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * reduce duration Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Update zombienet_tests/functional/0002-parachains-disputes.feature fix test * Update zombienet_tests/functional/0002-parachains-disputes.feature fix fix * more logs and set collator image Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * spellcheck Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix cargo lock damage Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * merge fixes Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * bump zombienet image Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * cargo lock Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * debugging CI run - scale down test Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * fix cargo merge damage Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * default command fix Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> * Revert "debugging CI run - scale down test" This reverts commit eec2ba7e65ede3f929e2f0a8fe44930df2119450. * more review feedback Signed-off-by: Andrei Sandu <andrei-mihail@parity.io> Co-authored-by: Javier Viola <javier@parity.io> Co-authored-by: Lldenaurois <Ljdenaurois@gmail.com>
2026-05-30 16:21:02 +00:00 · 2022-02-28 12:27:01 +02:00
parent 00ce69aae8
commit 0fa91e7043
23 changed files with 1588 additions and 8 deletions
@@ -0,0 +1,135 @@
+// Copyright 2022 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot 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.
+
+// Polkadot 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 Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Polkadot CLI library.
+
+use clap::Parser;
+use sc_cli::{RuntimeVersion, SubstrateCli};
+
+/// Sub-commands supported by the collator.
+#[derive(Debug, Parser)]
+pub enum Subcommand {
+	/// Export the genesis state of the parachain.
+	#[clap(name = "export-genesis-state")]
+	ExportGenesisState(ExportGenesisStateCommand),
+
+	/// Export the genesis wasm of the parachain.
+	#[clap(name = "export-genesis-wasm")]
+	ExportGenesisWasm(ExportGenesisWasmCommand),
+}
+
+/// Command for exporting the genesis state of the parachain
+#[derive(Debug, Parser)]
+pub struct ExportGenesisStateCommand {
+	/// Id of the parachain this collator collates for.
+	#[clap(long, default_value = "100")]
+	pub parachain_id: u32,
+
+	/// The target raw PoV size in bytes. Minimum value is 64.
+	#[clap(long, default_value = "1024")]
+	pub pov_size: usize,
+
+	/// The PVF execution complexity. Actually specifies how  many iterations/signatures
+	/// we compute per block.
+	#[clap(long, default_value = "1")]
+	pub pvf_complexity: u32,
+}
+
+/// Command for exporting the genesis wasm file.
+#[derive(Debug, Parser)]
+pub struct ExportGenesisWasmCommand {}
+
+#[allow(missing_docs)]
+#[derive(Debug, Parser)]
+pub struct RunCmd {
+	#[allow(missing_docs)]
+	#[clap(flatten)]
+	pub base: sc_cli::RunCmd,
+
+	/// Id of the parachain this collator collates for.
+	#[clap(long, default_value = "2000")]
+	pub parachain_id: u32,
+
+	/// The target raw PoV size in bytes. Minimum value is 64.
+	#[clap(long, default_value = "1024")]
+	pub pov_size: usize,
+
+	/// The PVF execution complexity. Actually specifies how many iterations/signatures
+	/// we compute per block.
+	#[clap(long, default_value = "1")]
+	pub pvf_complexity: u32,
+}
+
+#[allow(missing_docs)]
+#[derive(Debug, Parser)]
+pub struct Cli {
+	#[clap(subcommand)]
+	pub subcommand: Option<Subcommand>,
+
+	#[clap(flatten)]
+	pub run: RunCmd,
+}
+
+impl SubstrateCli for Cli {
+	fn impl_name() -> String {
+		"Parity Zombienet/Undying".into()
+	}
+
+	fn impl_version() -> String {
+		env!("CARGO_PKG_VERSION").into()
+	}
+
+	fn description() -> String {
+		env!("CARGO_PKG_DESCRIPTION").into()
+	}
+
+	fn author() -> String {
+		env!("CARGO_PKG_AUTHORS").into()
+	}
+
+	fn support_url() -> String {
+		"https://github.com/paritytech/polkadot/issues/new".into()
+	}
+
+	fn copyright_start_year() -> i32 {
+		2022
+	}
+
+	fn executable_name() -> String {
+		"undying-collator".into()
+	}
+
+	fn load_spec(&self, id: &str) -> std::result::Result<Box<dyn sc_service::ChainSpec>, String> {
+		let id = if id.is_empty() { "rococo" } else { id };
+		Ok(match id {
+			"rococo-staging" =>
+				Box::new(polkadot_service::chain_spec::rococo_staging_testnet_config()?),
+			"rococo-local" =>
+				Box::new(polkadot_service::chain_spec::rococo_local_testnet_config()?),
+			"rococo" => Box::new(polkadot_service::chain_spec::rococo_config()?),
+			path => {
+				let path = std::path::PathBuf::from(path);
+				Box::new(polkadot_service::RococoChainSpec::from_json_file(path)?)
+			},
+		})
+	}
+
+	fn native_runtime_version(
+		_spec: &Box<dyn polkadot_service::ChainSpec>,
+	) -> &'static RuntimeVersion {
+		&polkadot_service::rococo_runtime::VERSION
+	}
+}
@@ -0,0 +1,419 @@
+// Copyright 2022 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot 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.
+
+// Polkadot 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 Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Collator for the `Undying` test parachain.
+
+use futures::channel::oneshot;
+use futures_timer::Delay;
+use parity_scale_codec::{Decode, Encode};
+use polkadot_node_primitives::{
+	maybe_compress_pov, Collation, CollationResult, CollationSecondedSignal, CollatorFn,
+	MaybeCompressedPoV, PoV, Statement,
+};
+use polkadot_primitives::v1::{CollatorId, CollatorPair, Hash};
+use sp_core::{traits::SpawnNamed, Pair};
+use std::{
+	collections::HashMap,
+	sync::{
+		atomic::{AtomicU32, Ordering},
+		Arc, Mutex,
+	},
+	time::Duration,
+};
+use test_parachain_undying::{execute, hash_state, BlockData, GraveyardState, HeadData};
+
+/// Default PoV size which also drives state size.
+const DEFAULT_POV_SIZE: usize = 1000;
+/// Default PVF time complexity - 1 signature per block.
+const DEFAULT_PVF_COMPLEXITY: u32 = 1;
+
+/// Calculates the head and state for the block with the given `number`.
+fn calculate_head_and_state_for_number(
+	number: u64,
+	graveyard_size: usize,
+	pvf_complexity: u32,
+) -> (HeadData, GraveyardState) {
+	let index = 0u64;
+	let mut graveyard = vec![0u8; graveyard_size * graveyard_size];
+	let zombies = 0;
+	let seal = [0u8; 32];
+
+	// Ensure a larger compressed PoV.
+	graveyard.iter_mut().enumerate().for_each(|(i, grave)| {
+		*grave = i as u8;
+	});
+
+	let mut state = GraveyardState { index, graveyard, zombies, seal };
+	let mut head =
+		HeadData { number: 0, parent_hash: Hash::default().into(), post_state: hash_state(&state) };
+
+	while head.number < number {
+		let block = BlockData { state, tombstones: 1_000, iterations: pvf_complexity };
+		let (new_head, new_state) =
+			execute(head.hash(), head.clone(), block).expect("Produces valid block");
+		head = new_head;
+		state = new_state;
+	}
+
+	(head, state)
+}
+
+/// The state of the undying parachain.
+struct State {
+	// We need to keep these around until the including relay chain blocks are finalized.
+	// This is because disputes can trigger reverts up to last finalized block, so we
+	// want that state to collate on older relay chain heads.
+	head_to_state: HashMap<Arc<HeadData>, GraveyardState>,
+	number_to_head: HashMap<u64, Arc<HeadData>>,
+	/// Block number of the best block.
+	best_block: u64,
+	/// PVF time complexity.
+	pvf_complexity: u32,
+	/// Defines the state size (Vec<u8>). Our PoV includes the entire state so this value will
+	/// drive the PoV size.
+	/// Important note: block execution heavily clones this state, so something like 300.000 is
+	/// the max value here, otherwise we'll get OOM during wasm execution.
+	/// TODO: Implement a static state, and use `ballast` to inflate the PoV size. This way
+	/// we can just discard the `ballast` before processing the block.
+	graveyard_size: usize,
+}
+
+impl State {
+	/// Init the genesis state.
+	fn genesis(graveyard_size: usize, pvf_complexity: u32) -> Self {
+		let index = 0u64;
+		let mut graveyard = vec![0u8; graveyard_size * graveyard_size];
+		let zombies = 0;
+		let seal = [0u8; 32];
+
+		// Ensure a larger compressed PoV.
+		graveyard.iter_mut().enumerate().for_each(|(i, grave)| {
+			*grave = i as u8;
+		});
+
+		let state = GraveyardState { index, graveyard, zombies, seal };
+
+		let head_data =
+			HeadData { number: 0, parent_hash: Default::default(), post_state: hash_state(&state) };
+		let head_data = Arc::new(head_data);
+
+		Self {
+			head_to_state: vec![(head_data.clone(), state.clone())].into_iter().collect(),
+			number_to_head: vec![(0, head_data)].into_iter().collect(),
+			best_block: 0,
+			pvf_complexity,
+			graveyard_size,
+		}
+	}
+
+	/// Advance the state and produce a new block based on the given `parent_head`.
+	///
+	/// Returns the new [`BlockData`] and the new [`HeadData`].
+	fn advance(&mut self, parent_head: HeadData) -> (BlockData, HeadData) {
+		self.best_block = parent_head.number;
+
+		let state = if let Some(head_data) = self.number_to_head.get(&self.best_block) {
+			self.head_to_state.get(head_data).cloned().unwrap_or_else(|| {
+				calculate_head_and_state_for_number(
+					parent_head.number,
+					self.graveyard_size,
+					self.pvf_complexity,
+				)
+				.1
+			})
+		} else {
+			let (_, state) = calculate_head_and_state_for_number(
+				parent_head.number,
+				self.graveyard_size,
+				self.pvf_complexity,
+			);
+			state
+		};
+
+		// Start with prev state and transaction to execute (place 1000 tombstones).
+		let block = BlockData { state, tombstones: 1000, iterations: self.pvf_complexity };
+
+		let (new_head, new_state) =
+			execute(parent_head.hash(), parent_head, block.clone()).expect("Produces valid block");
+
+		let new_head_arc = Arc::new(new_head.clone());
+
+		self.head_to_state.insert(new_head_arc.clone(), new_state);
+		self.number_to_head.insert(new_head.number, new_head_arc);
+
+		(block, new_head)
+	}
+}
+
+/// The collator of the undying parachain.
+pub struct Collator {
+	state: Arc<Mutex<State>>,
+	key: CollatorPair,
+	seconded_collations: Arc<AtomicU32>,
+}
+
+impl Default for Collator {
+	fn default() -> Self {
+		Self::new(DEFAULT_POV_SIZE, DEFAULT_PVF_COMPLEXITY)
+	}
+}
+
+impl Collator {
+	/// Create a new collator instance with the state initialized from genesis and `pov_size`
+	/// parameter. The same parameter needs to be passed when exporting the genesis state.
+	pub fn new(pov_size: usize, pvf_complexity: u32) -> Self {
+		let graveyard_size = ((pov_size / std::mem::size_of::<u8>()) as f64).sqrt().ceil() as usize;
+
+		log::info!(
+			"PoV target size: {} bytes. Graveyard size: ({} x {})",
+			pov_size,
+			graveyard_size,
+			graveyard_size
+		);
+
+		log::info!("PVF time complexity: {}", pvf_complexity);
+
+		Self {
+			state: Arc::new(Mutex::new(State::genesis(graveyard_size, pvf_complexity))),
+			key: CollatorPair::generate().0,
+			seconded_collations: Arc::new(AtomicU32::new(0)),
+		}
+	}
+
+	/// Get the SCALE encoded genesis head of the parachain.
+	pub fn genesis_head(&self) -> Vec<u8> {
+		self.state
+			.lock()
+			.unwrap()
+			.number_to_head
+			.get(&0)
+			.expect("Genesis header exists")
+			.encode()
+	}
+
+	/// Get the validation code of the undying parachain.
+	pub fn validation_code(&self) -> &[u8] {
+		test_parachain_undying::wasm_binary_unwrap()
+	}
+
+	/// Get the collator key.
+	pub fn collator_key(&self) -> CollatorPair {
+		self.key.clone()
+	}
+
+	/// Get the collator id.
+	pub fn collator_id(&self) -> CollatorId {
+		self.key.public()
+	}
+
+	/// Create the collation function.
+	///
+	/// This collation function can be plugged into the overseer to generate collations for the undying parachain.
+	pub fn create_collation_function(
+		&self,
+		spawner: impl SpawnNamed + Clone + 'static,
+	) -> CollatorFn {
+		use futures::FutureExt as _;
+
+		let state = self.state.clone();
+		let seconded_collations = self.seconded_collations.clone();
+
+		Box::new(move |relay_parent, validation_data| {
+			let parent = HeadData::decode(&mut &validation_data.parent_head.0[..])
+				.expect("Decodes parent head");
+
+			let (block_data, head_data) = state.lock().unwrap().advance(parent);
+
+			log::info!(
+				"created a new collation on relay-parent({}): {:?}",
+				relay_parent,
+				head_data,
+			);
+
+			// The pov is the actually the initial state and the transactions.
+			let pov = PoV { block_data: block_data.encode().into() };
+
+			let collation = Collation {
+				upward_messages: Vec::new(),
+				horizontal_messages: Vec::new(),
+				new_validation_code: None,
+				head_data: head_data.encode().into(),
+				proof_of_validity: MaybeCompressedPoV::Raw(pov.clone()),
+				processed_downward_messages: 0,
+				hrmp_watermark: validation_data.relay_parent_number,
+			};
+
+			log::info!("Raw PoV size for collation: {} bytes", pov.block_data.0.len(),);
+			let compressed_pov = maybe_compress_pov(pov);
+
+			log::info!(
+				"Compressed PoV size for collation: {} bytes",
+				compressed_pov.block_data.0.len(),
+			);
+
+			let (result_sender, recv) = oneshot::channel::<CollationSecondedSignal>();
+			let seconded_collations = seconded_collations.clone();
+			spawner.spawn(
+				"undying-collator-seconded",
+				None,
+				async move {
+					if let Ok(res) = recv.await {
+						if !matches!(
+							res.statement.payload(),
+							Statement::Seconded(s) if s.descriptor.pov_hash == compressed_pov.hash(),
+						) {
+							log::error!(
+								"Seconded statement should match our collation: {:?}",
+								res.statement.payload()
+							);
+							std::process::exit(-1);
+						}
+
+						seconded_collations.fetch_add(1, Ordering::Relaxed);
+					}
+				}
+				.boxed(),
+			);
+
+			async move { Some(CollationResult { collation, result_sender: Some(result_sender) }) }
+				.boxed()
+		})
+	}
+
+	/// Wait until `blocks` are built and enacted.
+	pub async fn wait_for_blocks(&self, blocks: u64) {
+		let start_block = self.state.lock().unwrap().best_block;
+		loop {
+			Delay::new(Duration::from_secs(1)).await;
+
+			let current_block = self.state.lock().unwrap().best_block;
+
+			if start_block + blocks <= current_block {
+				return
+			}
+		}
+	}
+
+	/// Wait until `seconded` collations of this collator are seconded by a parachain validator.
+	///
+	/// The internal counter isn't de-duplicating the collations when counting the number of seconded collations. This
+	/// means when one collation is seconded by X validators, we record X seconded messages.
+	pub async fn wait_for_seconded_collations(&self, seconded: u32) {
+		let seconded_collations = self.seconded_collations.clone();
+		loop {
+			Delay::new(Duration::from_secs(1)).await;
+
+			if seconded <= seconded_collations.load(Ordering::Relaxed) {
+				return
+			}
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	use futures::executor::block_on;
+	use polkadot_parachain::primitives::{ValidationParams, ValidationResult};
+	use polkadot_primitives::v1::{Hash, PersistedValidationData};
+
+	#[test]
+	fn collator_works() {
+		let spawner = sp_core::testing::TaskExecutor::new();
+		let collator = Collator::new(1_000, 1);
+		let collation_function = collator.create_collation_function(spawner);
+
+		for i in 0..5 {
+			let parent_head =
+				collator.state.lock().unwrap().number_to_head.get(&i).unwrap().clone();
+
+			let validation_data = PersistedValidationData {
+				parent_head: parent_head.encode().into(),
+				..Default::default()
+			};
+
+			let collation =
+				block_on(collation_function(Default::default(), &validation_data)).unwrap();
+			validate_collation(&collator, (*parent_head).clone(), collation.collation);
+		}
+	}
+
+	fn validate_collation(collator: &Collator, parent_head: HeadData, collation: Collation) {
+		use polkadot_node_core_pvf::testing::validate_candidate;
+
+		let block_data = match collation.proof_of_validity {
+			MaybeCompressedPoV::Raw(pov) => pov.block_data,
+			MaybeCompressedPoV::Compressed(_) => panic!("Only works with uncompressed povs"),
+		};
+
+		let ret_buf = validate_candidate(
+			collator.validation_code(),
+			&ValidationParams {
+				parent_head: parent_head.encode().into(),
+				block_data,
+				relay_parent_number: 1,
+				relay_parent_storage_root: Hash::zero(),
+			}
+			.encode(),
+		)
+		.unwrap();
+		let ret = ValidationResult::decode(&mut &ret_buf[..]).unwrap();
+
+		let new_head = HeadData::decode(&mut &ret.head_data.0[..]).unwrap();
+		assert_eq!(
+			**collator
+				.state
+				.lock()
+				.unwrap()
+				.number_to_head
+				.get(&(parent_head.number + 1))
+				.unwrap(),
+			new_head
+		);
+	}
+
+	#[test]
+	fn advance_to_state_when_parent_head_is_missing() {
+		let collator = Collator::new(1_000, 1);
+		let graveyard_size = collator.state.lock().unwrap().graveyard_size;
+
+		let mut head = calculate_head_and_state_for_number(10, graveyard_size, 1).0;
+
+		for i in 1..10 {
+			head = collator.state.lock().unwrap().advance(head).1;
+			assert_eq!(10 + i, head.number);
+		}
+
+		let collator = Collator::new(1_000, 1);
+		let mut second_head = collator
+			.state
+			.lock()
+			.unwrap()
+			.number_to_head
+			.get(&0)
+			.cloned()
+			.unwrap()
+			.as_ref()
+			.clone();
+
+		for _ in 1..20 {
+			second_head = collator.state.lock().unwrap().advance(second_head.clone()).1;
+		}
+
+		assert_eq!(second_head, head);
+	}
+}
@@ -0,0 +1,111 @@
+// Copyright 2022 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot 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.
+
+// Polkadot 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 Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Collator for the `Undying` test parachain.
+
+use polkadot_cli::{Error, Result};
+use polkadot_node_primitives::CollationGenerationConfig;
+use polkadot_node_subsystem::messages::{CollationGenerationMessage, CollatorProtocolMessage};
+use polkadot_primitives::v1::Id as ParaId;
+use sc_cli::{Error as SubstrateCliError, Role, SubstrateCli};
+use sp_core::hexdisplay::HexDisplay;
+use test_parachain_undying_collator::Collator;
+
+mod cli;
+use cli::Cli;
+
+fn main() -> Result<()> {
+	let cli = Cli::from_args();
+
+	match cli.subcommand {
+		Some(cli::Subcommand::ExportGenesisState(params)) => {
+			// `pov_size` and `pvf_complexity` need to match the ones that we start the collator
+			// with.
+			let collator = Collator::new(params.pov_size, params.pvf_complexity);
+			println!("0x{:?}", HexDisplay::from(&collator.genesis_head()));
+
+			Ok::<_, Error>(())
+		},
+		Some(cli::Subcommand::ExportGenesisWasm(_params)) => {
+			// We pass some dummy values for `pov_size` and `pvf_complexity` as these don't
+			// matter for `wasm` export.
+			println!("0x{:?}", HexDisplay::from(&Collator::default().validation_code()));
+
+			Ok(())
+		},
+		None => {
+			let runner = cli.create_runner(&cli.run.base).map_err(|e| {
+				SubstrateCliError::Application(
+					Box::new(e) as Box<(dyn 'static + Send + Sync + std::error::Error)>
+				)
+			})?;
+
+			runner.run_node_until_exit(|config| async move {
+				let role = config.role.clone();
+
+				match role {
+					Role::Light => Err("Light client not supported".into()),
+					_ => {
+						let collator = Collator::new(cli.run.pov_size, cli.run.pvf_complexity);
+
+						let full_node = polkadot_service::build_full(
+							config,
+							polkadot_service::IsCollator::Yes(collator.collator_key()),
+							None,
+							true,
+							None,
+							None,
+							false,
+							polkadot_service::RealOverseerGen,
+						)
+						.map_err(|e| e.to_string())?;
+						let mut overseer_handle = full_node
+							.overseer_handle
+							.expect("Overseer handle should be initialized for collators");
+
+						let genesis_head_hex =
+							format!("0x{:?}", HexDisplay::from(&collator.genesis_head()));
+						let validation_code_hex =
+							format!("0x{:?}", HexDisplay::from(&collator.validation_code()));
+
+						let para_id = ParaId::from(cli.run.parachain_id);
+
+						log::info!("Running `Undying` collator for parachain id: {}", para_id);
+						log::info!("Genesis state: {}", genesis_head_hex);
+						log::info!("Validation code: {}", validation_code_hex);
+
+						let config = CollationGenerationConfig {
+							key: collator.collator_key(),
+							collator: collator
+								.create_collation_function(full_node.task_manager.spawn_handle()),
+							para_id,
+						};
+						overseer_handle
+							.send_msg(CollationGenerationMessage::Initialize(config), "Collator")
+							.await;
+
+						overseer_handle
+							.send_msg(CollatorProtocolMessage::CollateOn(para_id), "Collator")
+							.await;
+
+						Ok(full_node.task_manager)
+					},
+				}
+			})
+		},
+	}?;
+	Ok(())
+}