1use crate::collator::SlotClaim;
24use codec::Codec;
25use cumulus_client_consensus_common::{self as consensus_common, ParentSearchParams};
26use cumulus_primitives_aura::{AuraUnincludedSegmentApi, Slot};
27use cumulus_primitives_core::{
28 relay_chain::Header as RelayHeader, BlockT, KeyToIncludeInRelayProof, RelayProofRequest,
29};
30use cumulus_relay_chain_interface::{OverseerHandle, RelayChainInterface};
31use polkadot_node_subsystem::messages::CollatorProtocolMessage;
32use polkadot_node_subsystem_util::runtime::ClaimQueueSnapshot;
33use polkadot_primitives::{
34 Hash as RelayHash, Id as ParaId, OccupiedCoreAssumption, ValidationCodeHash,
35};
36use sc_consensus_aura::{standalone as aura_internal, AuraApi};
37use sp_api::{ApiExt, ProvideRuntimeApi};
38use sp_core::Pair;
39use sp_keystore::KeystorePtr;
40use sp_timestamp::Timestamp;
41
42pub mod lookahead;
43pub mod slot_based;
44
45struct BackingGroupConnectionHelper {
48 keystore: sp_keystore::KeystorePtr,
49 overseer_handle: OverseerHandle,
50 our_slot: Option<Slot>,
51}
52
53impl BackingGroupConnectionHelper {
54 pub fn new(keystore: sp_keystore::KeystorePtr, overseer_handle: OverseerHandle) -> Self {
55 Self { keystore, overseer_handle, our_slot: None }
56 }
57
58 async fn send_subsystem_message(&mut self, message: CollatorProtocolMessage) {
59 self.overseer_handle.send_msg(message, "BackingGroupConnectionHelper").await;
60 }
61
62 pub async fn update<P>(&mut self, current_slot: Slot, authorities: &[P::Public])
64 where
65 P: sp_core::Pair + Send + Sync,
66 P::Public: Codec,
67 {
68 if Some(current_slot) <= self.our_slot {
69 return;
72 }
73
74 let next_slot = current_slot + 1;
75 let next_slot_is_ours =
76 aura_internal::claim_slot::<P>(next_slot, authorities, &self.keystore)
77 .await
78 .is_some();
79
80 if next_slot_is_ours {
81 if self.our_slot.is_none() {
84 tracing::debug!(target: crate::LOG_TARGET, "Our slot {} is next, connecting to backing groups", next_slot);
86 self.send_subsystem_message(CollatorProtocolMessage::ConnectToBackingGroups)
87 .await;
88 }
89 self.our_slot = Some(next_slot);
90 } else if self.our_slot.take().is_some() {
91 tracing::debug!(target: crate::LOG_TARGET, "Current slot = {}, disconnecting from backing groups", current_slot);
93 self.send_subsystem_message(CollatorProtocolMessage::DisconnectFromBackingGroups)
94 .await;
95 }
96 }
97}
98
99async fn check_validation_code_or_log(
104 local_validation_code_hash: &ValidationCodeHash,
105 para_id: ParaId,
106 relay_client: &impl RelayChainInterface,
107 relay_parent: RelayHash,
108) {
109 let state_validation_code_hash = match relay_client
110 .validation_code_hash(relay_parent, para_id, OccupiedCoreAssumption::Included)
111 .await
112 {
113 Ok(hash) => hash,
114 Err(error) => {
115 tracing::debug!(
116 target: super::LOG_TARGET,
117 %error,
118 ?relay_parent,
119 %para_id,
120 "Failed to fetch validation code hash",
121 );
122 return;
123 },
124 };
125
126 match state_validation_code_hash {
127 Some(state) => {
128 if state != *local_validation_code_hash {
129 tracing::warn!(
130 target: super::LOG_TARGET,
131 %para_id,
132 ?relay_parent,
133 ?local_validation_code_hash,
134 relay_validation_code_hash = ?state,
135 "Parachain code doesn't match validation code stored in the relay chain state.",
136 );
137 }
138 },
139 None => {
140 tracing::warn!(
141 target: super::LOG_TARGET,
142 %para_id,
143 ?relay_parent,
144 "Could not find validation code for parachain in the relay chain state.",
145 );
146 },
147 }
148}
149
150async fn claim_queue_at(
152 relay_parent: RelayHash,
153 relay_client: &impl RelayChainInterface,
154) -> ClaimQueueSnapshot {
155 match relay_client.claim_queue(relay_parent).await {
157 Ok(claim_queue) => claim_queue.into(),
158 Err(error) => {
159 tracing::error!(
160 target: crate::LOG_TARGET,
161 ?error,
162 ?relay_parent,
163 "Failed to query claim queue runtime API",
164 );
165 Default::default()
166 },
167 }
168}
169
170async fn claim_slot<Block: BlockT, Client, P>(
172 para_slot: Slot,
173 timestamp: Timestamp,
174 parent_hash: Block::Hash,
175 client: &Client,
176 keystore: &KeystorePtr,
177) -> Option<SlotClaim<P::Public>>
178where
179 Client: ProvideRuntimeApi<Block>,
180 Client::Api: AuraApi<Block, P::Public> + ApiExt<Block>,
181 P: Pair,
182 P::Public: Codec,
183 P::Signature: Codec,
184{
185 let mut runtime_api = client.runtime_api();
186 runtime_api.set_call_context(sp_core::traits::CallContext::Onchain { import: false });
187 let authorities = runtime_api.authorities(parent_hash).ok()?;
188 let author_pub = aura_internal::claim_slot::<P>(para_slot, &authorities, keystore).await?;
189 Some(SlotClaim::unchecked::<P>(author_pub, para_slot, timestamp))
190}
191
192async fn can_build_upon<Block: BlockT, Client>(
194 parent_hash: Block::Hash,
195 included_block: Block::Hash,
196 relay_slot: Slot,
197 para_slot: Slot,
198 client: &Client,
199) -> bool
200where
201 Client: ProvideRuntimeApi<Block>,
202 Client::Api: AuraUnincludedSegmentApi<Block> + ApiExt<Block>,
203{
204 if parent_hash == included_block {
210 return true;
211 }
212
213 let runtime_api = client.runtime_api();
214 let Some(api_version) = runtime_api
215 .api_version::<dyn AuraUnincludedSegmentApi<Block>>(parent_hash)
216 .ok()
217 .flatten()
218 else {
219 return false;
220 };
221
222 let slot = if api_version > 1 { relay_slot } else { para_slot };
223
224 runtime_api
225 .can_build_upon(parent_hash, included_block, slot)
226 .ok()
227 .unwrap_or(false)
228}
229
230async fn find_parent<Block>(
236 relay_client: &impl RelayChainInterface,
237 para_backend: &impl sc_client_api::Backend<Block>,
238 para_id: ParaId,
239 params: ParentSearchParams,
240 filter_parent: impl Fn(&Block::Header) -> bool,
241) -> Option<consensus_common::ParentSearchResult<Block>>
242where
243 Block: BlockT,
244{
245 let result = match cumulus_client_consensus_common::find_parent_for_building::<Block>(
249 relay_client,
250 para_backend,
251 para_id,
252 params.clone(),
253 filter_parent,
254 )
255 .await
256 {
257 Ok(Some(result)) => result,
258 Ok(None) => {
259 tracing::warn!(
260 target: crate::LOG_TARGET,
261 ?params,
262 "Could not find parent to build upon.",
263 );
264 return None;
265 },
266 Err(e) => {
267 tracing::error!(
268 target: crate::LOG_TARGET,
269 ?params,
270 err = ?e,
271 "Could not find parent to build upon"
272 );
273 return None;
274 },
275 };
276
277 Some(result)
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283 use crate::collators::BackingGroupConnectionHelper;
284 use codec::Encode;
285 use cumulus_primitives_aura::Slot;
286 use cumulus_primitives_core::BlockT;
287 use cumulus_relay_chain_interface::PHash;
288 use cumulus_test_client::{
289 runtime::{Block, Hash},
290 BuildBlockBuilder, Client, DefaultTestClientBuilderExt, TestClientBuilder,
291 TestClientBuilderExt,
292 };
293 use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
294 use futures::StreamExt;
295 use polkadot_overseer::{Event, Handle};
296 use polkadot_primitives::HeadData;
297 use sc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
298 use sp_consensus::BlockOrigin;
299 use sp_keystore::{Keystore, KeystorePtr};
300 use std::sync::{Arc, Mutex};
301
302 async fn import_block<I: BlockImport<Block>>(
303 importer: &I,
304 block: Block,
305 origin: BlockOrigin,
306 import_as_best: bool,
307 ) {
308 let (header, body) = block.deconstruct();
309
310 let mut block_import_params = BlockImportParams::new(origin, header);
311 block_import_params.fork_choice = Some(ForkChoiceStrategy::Custom(import_as_best));
312 block_import_params.body = Some(body);
313 importer.import_block(block_import_params).await.unwrap();
314 }
315
316 fn sproof_with_parent_by_hash(client: &Client, hash: PHash) -> RelayStateSproofBuilder {
317 let header = client.header(hash).ok().flatten().expect("No header for parent block");
318 let included = HeadData(header.encode());
319 let mut builder = RelayStateSproofBuilder::default();
320 builder.para_id = cumulus_test_client::runtime::PARACHAIN_ID.into();
321 builder.included_para_head = Some(included);
322
323 builder
324 }
325 async fn build_and_import_block(client: &Client, included: Hash) -> Block {
326 let sproof = sproof_with_parent_by_hash(client, included);
327
328 let block_builder = client
329 .init_block_builder_builder()
330 .with_relay_sproof_builder(sproof)
331 .build()
332 .block_builder;
333
334 let block = block_builder.build().unwrap().block;
335
336 let origin = BlockOrigin::NetworkInitialSync;
337 import_block(client, block.clone(), origin, true).await;
338 block
339 }
340
341 fn set_up_components(num_authorities: usize) -> (Arc<Client>, KeystorePtr) {
342 let keystore = Arc::new(sp_keystore::testing::MemoryKeystore::new()) as Arc<_>;
343 for key in sp_keyring::Sr25519Keyring::iter().take(num_authorities) {
344 Keystore::sr25519_generate_new(
345 &*keystore,
346 sp_application_crypto::key_types::AURA,
347 Some(&key.to_seed()),
348 )
349 .expect("Can insert key into MemoryKeyStore");
350 }
351 (Arc::new(TestClientBuilder::new().build()), keystore)
352 }
353
354 #[tokio::test]
360 async fn test_can_build_upon() {
361 sp_tracing::try_init_simple();
362
363 let (client, _keystore) = set_up_components(6);
364
365 let genesis_hash = client.chain_info().genesis_hash;
366 let mut last_hash = genesis_hash;
367
368 while can_build_upon::<_, _>(
370 last_hash,
371 genesis_hash,
372 Slot::from(u64::MAX),
373 Slot::from(u64::MAX),
374 &*client,
375 )
376 .await
377 {
378 let block = build_and_import_block(&client, genesis_hash).await;
379 last_hash = block.header().hash();
380 }
381
382 let result = can_build_upon::<_, _>(
385 last_hash,
386 last_hash,
387 Slot::from(u64::MAX),
388 Slot::from(u64::MAX),
389 &*client,
390 )
391 .await;
392 assert!(result);
393 }
394
395 fn create_overseer_handle() -> (OverseerHandle, Arc<Mutex<Vec<CollatorProtocolMessage>>>) {
397 let messages = Arc::new(Mutex::new(Vec::new()));
398 let messages_clone = messages.clone();
399
400 let (tx, mut rx) = polkadot_node_subsystem_util::metered::channel(100);
401
402 tokio::spawn(async move {
404 while let Some(event) = rx.next().await {
405 if let Event::MsgToSubsystem { msg, .. } = event {
406 if let polkadot_node_subsystem::AllMessages::CollatorProtocol(cp_msg) = msg {
407 messages_clone.lock().unwrap().push(cp_msg);
408 }
409 }
410 }
411 });
412
413 (Handle::new(tx), messages)
414 }
415
416 #[tokio::test]
417 async fn preconnect_when_next_slot_is_ours() {
418 let (client, keystore) = set_up_components(1);
419 let genesis_hash = client.chain_info().genesis_hash;
420 let (overseer_handle, messages_recorder) = create_overseer_handle();
421
422 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
423
424 let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
426
427 helper
429 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
430 .await;
431
432 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
434
435 let messages = messages_recorder.lock().unwrap();
436 assert_eq!(messages.len(), 1);
437 assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
438 assert_eq!(helper.our_slot, Some(Slot::from(6)));
439 }
440
441 #[tokio::test]
442 async fn preconnect_no_duplicate_connect_message() {
443 let (client, keystore) = set_up_components(1);
444 let genesis_hash = client.chain_info().genesis_hash;
445 let (overseer_handle, messages_recorder) = create_overseer_handle();
446
447 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
448
449 let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
451
452 helper
454 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
455 .await;
456
457 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
459 assert_eq!(messages_recorder.lock().unwrap().len(), 1);
460 messages_recorder.lock().unwrap().clear();
461
462 helper
464 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
465 .await;
466 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
467 assert_eq!(messages_recorder.lock().unwrap().len(), 0);
468
469 helper
471 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(6), &authorities)
472 .await;
473 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
474 assert_eq!(messages_recorder.lock().unwrap().len(), 0);
475 }
476
477 #[tokio::test]
478 async fn preconnect_disconnect_when_slot_passes() {
479 let (client, keystore) = set_up_components(1);
480 let genesis_hash = client.chain_info().genesis_hash;
481 let (overseer_handle, messages_recorder) = create_overseer_handle();
482
483 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
484
485 let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
487
488 helper
493 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
494 .await;
495 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
496 assert_eq!(helper.our_slot, Some(Slot::from(6)));
497 messages_recorder.lock().unwrap().clear();
498
499 helper
501 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(8), &authorities)
502 .await;
503 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
504
505 {
506 let messages = messages_recorder.lock().unwrap();
507 assert_eq!(messages.len(), 1, "Expected exactly one disconnect message");
508 assert!(matches!(messages[0], CollatorProtocolMessage::DisconnectFromBackingGroups));
509 assert_eq!(helper.our_slot, None);
510 }
511
512 messages_recorder.lock().unwrap().clear();
513
514 helper
517 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(8), &authorities)
518 .await;
519 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
520
521 let messages = messages_recorder.lock().unwrap();
522 assert_eq!(messages.len(), 0, "Expected no messages");
523 assert_eq!(helper.our_slot, None);
524 }
525
526 #[tokio::test]
527 async fn preconnect_no_disconnect_without_previous_connection() {
528 let (client, keystore) = set_up_components(1);
529 let genesis_hash = client.chain_info().genesis_hash;
530 let (overseer_handle, messages_recorder) = create_overseer_handle();
531
532 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
533
534 let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
536
537 helper
543 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(1), &authorities)
544 .await;
545
546 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
547 assert_eq!(messages_recorder.lock().unwrap().len(), 0);
549 assert_eq!(helper.our_slot, None);
550 }
551
552 #[tokio::test]
553 async fn preconnect_multiple_cycles() {
554 let (client, keystore) = set_up_components(1);
555 let genesis_hash = client.chain_info().genesis_hash;
556 let (overseer_handle, messages_recorder) = create_overseer_handle();
557
558 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
559
560 let authorities = client.runtime_api().authorities(genesis_hash).unwrap();
562
563 helper
568 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(5), &authorities)
569 .await;
570 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
571 {
572 let messages = messages_recorder.lock().unwrap();
573 assert_eq!(messages.len(), 1);
574 assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
575 }
576 assert_eq!(helper.our_slot, Some(Slot::from(6)));
577 messages_recorder.lock().unwrap().clear();
578
579 helper
581 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(7), &authorities)
582 .await;
583 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
584 {
585 let messages = messages_recorder.lock().unwrap();
586 assert_eq!(messages.len(), 1);
587 assert!(matches!(messages[0], CollatorProtocolMessage::DisconnectFromBackingGroups));
588 }
589 assert_eq!(helper.our_slot, None);
590 messages_recorder.lock().unwrap().clear();
591
592 helper
594 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(11), &authorities)
595 .await;
596 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
597 {
598 let messages = messages_recorder.lock().unwrap();
599 assert_eq!(messages.len(), 1);
600 assert!(matches!(messages[0], CollatorProtocolMessage::ConnectToBackingGroups));
601 }
602 assert_eq!(helper.our_slot, Some(Slot::from(12)));
603 }
604
605 #[tokio::test]
606 async fn preconnect_handles_empty_authorities() {
607 let keystore = Arc::new(sp_keystore::testing::MemoryKeystore::new()) as Arc<_>;
608 let (overseer_handle, messages_recorder) = create_overseer_handle();
609
610 let mut helper = BackingGroupConnectionHelper::new(keystore, overseer_handle);
611
612 let authorities = vec![];
614 helper
615 .update::<sp_consensus_aura::sr25519::AuthorityPair>(Slot::from(0), &authorities)
616 .await;
617
618 tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
619 assert_eq!(messages_recorder.lock().unwrap().len(), 0);
621 }
622}
623
624pub(crate) fn get_relay_proof_request<Block, Client>(
631 client: &Client,
632 parent_hash: Block::Hash,
633) -> RelayProofRequest
634where
635 Block: BlockT,
636 Client: ProvideRuntimeApi<Block>,
637 Client::Api: KeyToIncludeInRelayProof<Block>,
638{
639 client.runtime_api().keys_to_prove(parent_hash).unwrap_or_else(|e| {
640 tracing::debug!(
641 target: crate::LOG_TARGET,
642 error = ?e,
643 "Failed to fetch relay proof requests from runtime, using empty request"
644 );
645 Default::default()
646 })
647}
648
649#[derive(Clone)]
651pub struct RelayParentData {
652 relay_parent: RelayHeader,
654 descendants: Vec<RelayHeader>,
656}
657
658impl RelayParentData {
659 pub fn new(relay_parent: RelayHeader) -> Self {
661 Self { relay_parent, descendants: Default::default() }
662 }
663
664 pub fn new_with_descendants(relay_parent: RelayHeader, descendants: Vec<RelayHeader>) -> Self {
666 Self { relay_parent, descendants }
667 }
668
669 pub fn relay_parent(&self) -> &RelayHeader {
671 &self.relay_parent
672 }
673
674 pub fn take_descendants(&mut self) -> Vec<RelayHeader> {
678 std::mem::take(&mut self.descendants)
679 }
680
681 #[cfg(test)]
683 pub fn descendants_len(&self) -> usize {
684 self.descendants.len()
685 }
686
687 pub fn into_inherent_descendant_list(self) -> Vec<RelayHeader> {
691 let Self { relay_parent, mut descendants } = self;
692
693 if descendants.is_empty() {
694 return Default::default();
695 }
696
697 let mut result = vec![relay_parent];
698 result.append(&mut descendants);
699 result
700 }
701}