cumulus: resubmission store

Cargo.toml

@@ -77,6 +77,7 @@ members = [
 	"cumulus/client/relay-chain-minimal-node",
 	"cumulus/client/relay-chain-rpc-interface",
 	"cumulus/client/relay-chain-streams",
+	"cumulus/client/resubmission-store",
 	"cumulus/client/service",
 	"cumulus/pallets/ah-ops",
 	"cumulus/pallets/aura-ext",
@@ -777,6 +778,7 @@ cumulus-client-network = { path = "cumulus/client/network", default-features = f
 cumulus-client-parachain-inherent = { path = "cumulus/client/parachain-inherent", default-features = false }
 cumulus-client-pov-recovery = { path = "cumulus/client/pov-recovery", default-features = false }
 cumulus-client-proof-size-recording = { path = "cumulus/client/proof-size-recording", default-features = false }
+cumulus-client-resubmission-store = { path = "cumulus/client/resubmission-store", default-features = false }
 cumulus-client-service = { path = "cumulus/client/service", default-features = false }
 cumulus-pallet-aura-ext = { path = "cumulus/pallets/aura-ext", default-features = false }
 cumulus-pallet-parachain-system = { path = "cumulus/pallets/parachain-system", default-features = false }

cumulus/client/consensus/aura/Cargo.toml

@@ -17,6 +17,7 @@ codec = { features = ["derive"], workspace = true, default-features = true }
 futures = { workspace = true }
 parking_lot = { workspace = true }
 schnellru = { workspace = true }
+thiserror = { workspace = true }
 tokio = { workspace = true, features = ["macros"] }
 tracing = { workspace = true, default-features = true }
 
@@ -50,6 +51,7 @@ cumulus-client-collator = { workspace = true, default-features = true }
 cumulus-client-consensus-common = { workspace = true, default-features = true }
 cumulus-client-parachain-inherent = { workspace = true, default-features = true }
 cumulus-client-proof-size-recording = { workspace = true, default-features = true }
+cumulus-client-resubmission-store = { workspace = true, default-features = true }
 cumulus-primitives-aura = { workspace = true, default-features = true }
 cumulus-primitives-core = { workspace = true, default-features = true }
 cumulus-relay-chain-interface = { workspace = true, default-features = true }

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

@@ -15,7 +15,10 @@
 // You should have received a copy of the GNU General Public License
 // along with Cumulus. If not, see <https://www.gnu.org/licenses/>.
 
-use super::CollatorMessage;
+use super::{
+	resubmission::{resolve_pvd, resolve_session},
+	CollatorMessage,
+};
 use crate::{
 	collator::{self as collator_util, BuildBlockAndImportParams, Collator, SlotClaim},
 	collators::{
@@ -36,6 +39,7 @@ use cumulus_client_consensus_common::{
 	ParachainBlockImportMarker, ParentSearchParams,
 };
 use cumulus_client_proof_size_recording::prepare_proof_size_recording_aux_data;
+use cumulus_client_resubmission_store::{now_unix_ms, prepare_resubmission_aux_data};
 use cumulus_primitives_aura::{AuraUnincludedSegmentApi, Slot};
 use cumulus_primitives_core::{
 	BlockBundleInfo, ClaimQueueOffset, CoreInfo, CoreSelector, CumulusDigestItem,
@@ -731,6 +735,29 @@ where
 	check_validation_code_or_log(&validation_code_hash, para_id, relay_client, relay_parent_hash)
 		.await;
 
+	let session = resolve_session(relay_client, relay_parent_hash)
+		.await
+		.inspect_err(|err| {
+			tracing::warn!(
+				target: LOG_TARGET,
+				?relay_parent_hash,
+				?err,
+				"Could not resolve relay-parent session; resubmission entry will be skipped.",
+			);
+		})
+		.ok();
+	let pvd = resolve_pvd(relay_client, relay_parent_hash, para_id)
+		.await
+		.inspect_err(|err| {
+			tracing::warn!(
+				target: LOG_TARGET,
+				?relay_parent_hash,
+				?err,
+				"Could not resolve persisted validation data; resubmission entry will be skipped.",
+			);
+		})
+		.ok();
+
 	let mut blocks = Vec::new();
 	let mut proofs = Vec::new();
 	let mut ignored_nodes = IgnoredNodes::default();
@@ -875,6 +902,23 @@ where
 			);
 		}
 
+		let time_ms = now_unix_ms();
+		let proof = Arc::new(built_block.proof);
+		if let (Some(relay_parent_session), Some(pvd)) = (session, pvd.clone()) {
+			prepare_resubmission_aux_data::<Block>(
+				parent_hash,
+				time_ms,
+				proof.clone(),
+				relay_parent_header.clone(),
+				relay_parent_session,
+				pvd,
+				core_info.selector,
+			)
+			.for_each(|(k, v)| {
+				import_block.auxiliary.push((k, Some(v)));
+			});
+		}
+
 		if let Err(error) = collator.import_block(import_block).await {
 			tracing::error!(target: LOG_TARGET, ?error, "Failed to import built block.");
 			return Ok(None);
@@ -884,7 +928,7 @@ where
 		collator.collator_service().announce_block(parent_hash, None);
 
 		blocks.push(built_block.block);
-		proofs.push(built_block.proof);
+		proofs.push(Arc::unwrap_or_clone(proof));
 
 		let full_core_digest = CumulusDigestItem::contains_use_full_core(parent_header.digest());
 		let runtime_upgrade_digest = parent_header

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

@@ -17,6 +17,7 @@
 
 use crate::LOG_TARGET;
 use codec::{Decode, Encode};
+use cumulus_client_consensus_common::old_finalized_hash;
 use cumulus_client_proof_size_recording::prepare_proof_size_recording_aux_data;
 use cumulus_primitives_core::{BlockBundleInfo, CoreInfo, CumulusDigestItem, RelayBlockIdentifier};
 use futures::{stream::FusedStream, StreamExt};
@@ -73,14 +74,14 @@ fn load_ignored_nodes<Block: BlockT, B: AuxStore>(
 /// 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)>,
+	receiver: TracingUnboundedReceiver<(Block, Arc<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) {
+	pub async fn next(&mut self) -> (Block, Arc<StorageProof>) {
 		loop {
 			if self.receiver.is_terminated() {
 				futures::pending!()
@@ -100,14 +101,11 @@ where
 {
 	let client_for_closure = client.clone();
 	let on_finality = move |notification: &FinalityNotification<Block>| -> AuxDataOperations {
-		// The old finalized block is the parent of the first block in the tree route,
-		// or the parent of the finalized block if the tree route is empty.
-		let old_finalized_hash = notification
-			.tree_route
-			.first()
-			.and_then(|hash| client_for_closure.header(*hash).ok().flatten())
-			.map(|h| *h.parent_hash())
-			.unwrap_or_else(|| *notification.header.parent_hash());
+		let old_finalized = old_finalized_hash::<_, Block>(
+			&*client_for_closure,
+			&notification.tree_route,
+			*notification.header.parent_hash(),
+		);
 
 		notification
 			.stale_blocks
@@ -125,7 +123,7 @@ where
 					.copied()
 					.map(|hash| (ignored_nodes_key(hash), None)),
 			)
-			.chain(std::iter::once((ignored_nodes_key(old_finalized_hash), None)))
+			.chain(std::iter::once((ignored_nodes_key(old_finalized), None)))
 			.collect()
 	};
 
@@ -136,7 +134,7 @@ where
 pub struct SlotBasedBlockImport<Block: BlockT, BI, Client> {
 	inner: BI,
 	client: Arc<Client>,
-	sender: TracingUnboundedSender<(Block, StorageProof)>,
+	sender: TracingUnboundedSender<(Block, Arc<StorageProof>)>,
 }
 
 impl<Block: BlockT, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
@@ -209,10 +207,11 @@ impl<Block: BlockT, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
 	/// The proof must be recorded in exactly the same manner as during block building, because the
 	/// proof size is tracked via `ProofSizeExt` and affects runtime state. Without identical proof
 	/// recording, the computed state root would differ and block import would fail.
-	fn execute_block_and_collect_storage_proof(
+	async fn execute_block_and_collect_storage_proof(
 		&self,
 		params: &mut sc_consensus::BlockImportParams<Block>,
-	) -> Result<(), sp_consensus::Error>
+		collect_for_resubmission: bool,
+	) -> Result<Option<(Block, Arc<StorageProof>)>, sp_consensus::Error>
 	where
 		Client: ProvideRuntimeApi<Block>
 			+ CallApiAt<Block>
@@ -236,7 +235,7 @@ impl<Block: BlockT, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
 				number = ?params.header.number(),
 				"no bundle digests, skipping execute_block_and_collect_storage_proof",
 			);
-			return Ok(());
+			return Ok(None);
 		};
 
 		let parent_hash = *params.header.parent_hash();
@@ -256,7 +255,9 @@ impl<Block: BlockT, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
 		runtime_api.record_proof_with_recorder(recorder.clone());
 		runtime_api.register_extension(ProofSizeExt::new(proof_size_recorder.clone()));
 
-		let block = Block::new(params.header.clone(), params.body.clone().unwrap_or_default());
+		let body = params.body.clone().unwrap_or_default();
+		let resubmission_body = collect_for_resubmission.then(|| body.clone());
+		let block = Block::new(params.header.clone(), body);
 
 		tracing::debug!(
 			target: LOG_TARGET,
@@ -318,7 +319,10 @@ impl<Block: BlockT, BI, Client> SlotBasedBlockImport<Block, BI, Client> {
 		params.state_action =
 			StateAction::ApplyChanges(sc_consensus::StorageChanges::Changes(gen_storage_changes));
 
-		Ok(())
+		let Some(resubmission_body) = resubmission_body else { return Ok(None) };
+
+		let block = Block::new(params.header.clone(), resubmission_body);
+		Ok(Some((block, Arc::new(storage_proof))))
 	}
 }
 
@@ -352,13 +356,31 @@ where
 		&self,
 		mut params: sc_consensus::BlockImportParams<Block>,
 	) -> Result<sc_consensus::ImportResult, Self::Error> {
-		if !(params.origin == BlockOrigin::Own ||
+		let origin = params.origin;
+
+		// Only blocks imported at the tip (built by other collators) are forwarded to the
+		// resubmission backfill task. Sync/gap/warp imports are historical and never part of our
+		// unincluded segment, so recording resubmission data for them would be wasted work (and
+		// would buffer their potentially large storage proofs).
+		let collect_for_resubmission =
+			matches!(origin, BlockOrigin::NetworkBroadcast | BlockOrigin::ConsensusBroadcast);
+
+		let imported = if !(origin == BlockOrigin::Own ||
 			params.with_state() ||
 			params.state_action.skip_execution_checks())
 		{
-			self.execute_block_and_collect_storage_proof(&mut params)?;
+			self.execute_block_and_collect_storage_proof(&mut params, collect_for_resubmission)
+				.await?
+		} else {
+			None
+		};
+
+		let result = self.inner.import_block(params).await.map_err(Into::into)?;
+
+		if let Some((block, proof)) = imported {
+			let _ = self.sender.unbounded_send((block, proof));
 		}
 
-		self.inner.import_block(params).await.map_err(Into::into)
+		Ok(result)
 	}
 }

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

@@ -50,8 +50,6 @@ pub struct Params<Block: BlockT, RClient, CS> {
 	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>,
 }
@@ -70,7 +68,6 @@ pub async fn run_collation_task<Block, RClient, CS>(
 		reinitialize,
 		collator_service,
 		mut collator_receiver,
-		mut block_import_handle,
 		export_pov,
 	}: Params<Block, RClient, CS>,
 ) where
@@ -91,21 +88,15 @@ pub async fn run_collation_task<Block, RClient, CS>(
 	)
 	.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/paritytech/polkadot-sdk/issues/6495
-				let _ = block_import_msg;
-			}
-		}
+	while let Some(message) = collator_receiver.next().await {
+		handle_collation_message(
+			message,
+			&collator_service,
+			&mut overseer_handle,
+			relay_client.clone(),
+			export_pov.clone(),
+		)
+		.await;
 	}
 }