collator-protocol: preconnect (v2)

node/collation-generation/Cargo.toml

@@ -18,5 +18,6 @@ thiserror = "1.0.30"
 parity-scale-codec = { version = "3.1.2", default-features = false, features = ["bit-vec", "derive"] }
 
 [dev-dependencies]
+async-trait = "0.1.53"
 polkadot-node-subsystem-test-helpers = { path = "../subsystem-test-helpers" }
 test-helpers = { package = "polkadot-primitives-test-helpers", path = "../../primitives/test-helpers" }

node/collation-generation/src/lib.rs

@@ -20,7 +20,7 @@
 
 use futures::{channel::mpsc, future::FutureExt, join, select, sink::SinkExt, stream::StreamExt};
 use parity_scale_codec::Encode;
-use polkadot_node_primitives::{AvailableData, CollationGenerationConfig, PoV};
+use polkadot_node_primitives::{AvailableData, CollationGenerationConfig, Collator, PoV};
 use polkadot_node_subsystem::{
 	messages::{AllMessages, CollationGenerationMessage, CollatorProtocolMessage},
 	overseer, ActiveLeavesUpdate, FromOverseer, OverseerSignal, SpawnedSubsystem, SubsystemContext,
@@ -205,7 +205,21 @@ async fn handle_new_activations<Context: SubsystemContext>(
 			let (scheduled_core, assumption) = match core {
 				CoreState::Scheduled(scheduled_core) =>
 					(scheduled_core, OccupiedCoreAssumption::Free),
-				CoreState::Occupied(_occupied_core) => {
+				CoreState::Occupied(occupied_core) => {
+					if Collator::is_collating_on_child(&*config.collator, relay_parent).await {
+						let _ = ctx
+							.send_message(AllMessages::CollatorProtocol(
+								CollatorProtocolMessage::PreConnectAsCollator(relay_parent),
+							))
+							.await;
+						gum::debug!(
+							target: LOG_TARGET,
+							para_id = %occupied_core.para_id(),
+							relay_parent = ?relay_parent,
+							"Sent pre-connect request",
+						);
+					}
+
 					// TODO: https://github.com/paritytech/polkadot/issues/1573
 					gum::trace!(
 						target: LOG_TARGET,
@@ -294,18 +308,23 @@ async fn handle_new_activations<Context: SubsystemContext>(
 				Box::pin(async move {
 					let persisted_validation_data_hash = validation_data.hash();
 
-					let (collation, result_sender) =
-						match (task_config.collator)(relay_parent, &validation_data).await {
-							Some(collation) => collation.into_inner(),
-							None => {
-								gum::debug!(
-									target: LOG_TARGET,
-									para_id = %scheduled_core.para_id,
-									"collator returned no collation on collate",
-								);
-								return
-							},
-						};
+					let (collation, result_sender) = match Collator::produce_collation(
+						&*task_config.collator,
+						relay_parent,
+						&validation_data,
+					)
+					.await
+					{
+						Some(collation) => collation.into_inner(),
+						None => {
+							gum::debug!(
+								target: LOG_TARGET,
+								para_id = %scheduled_core.para_id,
+								"collator returned no collation on collate",
+							);
+							return
+						},
+					};
 
 					// Apply compression to the block data.
 					let pov = {

node/collation-generation/src/tests.rs

@@ -17,13 +17,9 @@
 mod handle_new_activations {
 	use super::super::*;
 	use ::test_helpers::{dummy_hash, dummy_head_data, dummy_validator};
-	use futures::{
-		lock::Mutex,
-		task::{Context as FuturesContext, Poll},
-		Future,
-	};
+	use futures::lock::Mutex;
 	use polkadot_node_primitives::{
-		BlockData, Collation, CollationResult, MaybeCompressedPoV, PoV,
+		BlockData, Collation, CollationResult, Collator, MaybeCompressedPoV, PoV,
 	};
 	use polkadot_node_subsystem::{
 		errors::RuntimeApiError,
@@ -35,7 +31,6 @@ mod handle_new_activations {
 	use polkadot_primitives::v2::{
 		CollatorPair, Id as ParaId, PersistedValidationData, ScheduledCore, ValidationCode,
 	};
-	use std::pin::Pin;
 
 	fn test_collation() -> Collation {
 		Collation {
@@ -62,14 +57,16 @@ mod handle_new_activations {
 		persisted_validation_data
 	}
 
-	// Box<dyn Future<Output = Collation> + Unpin + Send
 	struct TestCollator;
 
-	impl Future for TestCollator {
-		type Output = Option<CollationResult>;
-
-		fn poll(self: Pin<&mut Self>, _cx: &mut FuturesContext) -> Poll<Self::Output> {
-			Poll::Ready(Some(CollationResult { collation: test_collation(), result_sender: None }))
+	#[async_trait::async_trait]
+	impl Collator for TestCollator {
+		async fn produce_collation(
+			&self,
+			_: Hash,
+			_: &PersistedValidationData,
+		) -> Option<CollationResult> {
+			Some(CollationResult { collation: test_collation(), result_sender: None })
 		}
 	}
 
@@ -78,7 +75,7 @@ mod handle_new_activations {
 	fn test_config<Id: Into<ParaId>>(para_id: Id) -> Arc<CollationGenerationConfig> {
 		Arc::new(CollationGenerationConfig {
 			key: CollatorPair::generate().0,
-			collator: Box::new(|_: Hash, _vd: &PersistedValidationData| TestCollator.boxed()),
+			collator: Box::new(TestCollator),
 			para_id: para_id.into(),
 		})
 	}