pezkuwi-sdk/pezcumulus/client/consensus/common/src/teyrchain_consensus.rs

// 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.",
		);
	}
}