Deprecate para_id() from CoreState in polkadot primitives

cumulus/client/consensus/aura/src/collators/lookahead.rs

@@ -49,7 +49,9 @@ use polkadot_node_subsystem::messages::{
 	CollationGenerationMessage, RuntimeApiMessage, RuntimeApiRequest,
 };
 use polkadot_overseer::Handle as OverseerHandle;
-use polkadot_primitives::{CollatorPair, CoreIndex, Id as ParaId, OccupiedCoreAssumption};
+use polkadot_primitives::{
+	CollatorPair, CoreIndex, CoreState, Id as ParaId, OccupiedCoreAssumption,
+};
 
 use futures::{channel::oneshot, prelude::*};
 use sc_client_api::{backend::AuxStore, BlockBackend, BlockOf};
@@ -495,6 +497,7 @@ async fn cores_scheduled_for_para(
 	para_id: ParaId,
 	overseer_handle: &mut OverseerHandle,
 ) -> Vec<CoreIndex> {
+	// Get `AvailabilityCores` from runtime
 	let (tx, rx) = oneshot::channel();
 	let request = RuntimeApiRequest::AvailabilityCores(tx);
 	overseer_handle
@@ -522,11 +525,49 @@ async fn cores_scheduled_for_para(
 		},
 	};
 
+	// Get `AsyncBackingParams` from runtime
+	let (tx, rx) = oneshot::channel();
+	let request = RuntimeApiRequest::AsyncBackingParams(tx);
+	overseer_handle
+		.send_msg(RuntimeApiMessage::Request(relay_parent, request), "LookaheadCollator")
+		.await;
+	let async_backing_params = match rx.await {
+		Ok(Ok(params)) => params,
+		Ok(Err(error)) => {
+			tracing::error!(
+				target: crate::LOG_TARGET,
+				?error,
+				?relay_parent,
+				"Failed to query async backing params runtime API",
+			);
+			return Vec::new()
+		},
+		Err(oneshot::Canceled) => {
+			tracing::error!(
+				target: crate::LOG_TARGET,
+				?relay_parent,
+				"Sender for async backing params runtime request dropped",
+			);
+			return Vec::new()
+		},
+	};
+
 	cores
 		.iter()
 		.enumerate()
 		.filter_map(|(index, core)| {
-			if core.para_id() == Some(para_id) {
+			let core_para_id = match core {
+				CoreState::Scheduled(scheduled_core) => Some(scheduled_core.para_id),
+				CoreState::Occupied(occupied_core)
+					if async_backing_params.max_candidate_depth >= 1 =>
+					occupied_core
+						.next_up_on_available
+						.as_ref()
+						.map(|scheduled_core| scheduled_core.para_id),
+				CoreState::Free | CoreState::Occupied(_) => None,
+			};
+
+			if core_para_id == Some(para_id) {
 				Some(CoreIndex(index as u32))
 			} else {
 				None

polkadot/node/core/prospective-parachains/src/tests.rs

@@ -1797,7 +1797,10 @@ fn persists_pending_availability_candidate() {
 	test_state.availability_cores = test_state
 		.availability_cores
 		.into_iter()
-		.filter(|core| core.para_id().map_or(false, |id| id == para_id))
+		.filter(|core| match core {
+			CoreState::Scheduled(scheduled_core) => scheduled_core.para_id == para_id,
+			_ => false,
+		})
 		.collect();
 	assert_eq!(test_state.availability_cores.len(), 1);
 
@@ -1896,7 +1899,10 @@ fn backwards_compatible() {
 	test_state.availability_cores = test_state
 		.availability_cores
 		.into_iter()
-		.filter(|core| core.para_id().map_or(false, |id| id == para_id))
+		.filter(|core| match core {
+			CoreState::Scheduled(scheduled_core) => scheduled_core.para_id == para_id,
+			_ => false,
+		})
 		.collect();
 	assert_eq!(test_state.availability_cores.len(), 1);
 

polkadot/node/core/provisioner/src/tests.rs

@@ -918,7 +918,11 @@ mod select_candidates {
 		let committed_receipts: Vec<_> = (0..mock_cores.len())
 			.map(|i| {
 				let mut descriptor = dummy_candidate_descriptor(dummy_hash());
-				descriptor.para_id = mock_cores[i].para_id().unwrap();
+				descriptor.para_id = if let Scheduled(scheduled_core) = &mock_cores[i] {
+					scheduled_core.para_id
+				} else {
+					panic!("`moc_cores` is not initialized with `Scheduled`?")
+				};
 				descriptor.persisted_validation_data_hash = empty_hash;
 				descriptor.pov_hash = Hash::from_low_u64_be(i as u64);
 				CommittedCandidateReceipt {

polkadot/node/network/statement-distribution/src/v2/mod.rs

@@ -46,7 +46,7 @@ use polkadot_node_subsystem_util::{
 	backing_implicit_view::View as ImplicitView,
 	reputation::ReputationAggregator,
 	runtime::{request_min_backing_votes, ProspectiveParachainsMode},
-	vstaging::fetch_claim_queue,
+	vstaging::{fetch_claim_queue, ClaimQueueSnapshot},
 };
 use polkadot_primitives::{
 	AuthorityDiscoveryId, CandidateHash, CompactStatement, CoreIndex, CoreState, GroupIndex,
@@ -681,14 +681,23 @@ pub(crate) async fn handle_active_leaves_update<Context>(
 				.map_err(JfyiError::FetchValidatorGroups)?
 				.1;
 
+		let maybe_claim_queue = fetch_claim_queue(ctx.sender(), new_relay_parent)
+		.await
+		.unwrap_or_else(|err| {
+				gum::debug!(target: LOG_TARGET, ?new_relay_parent, ?err, "handle_active_leaves_update: `claim_queue` API not available");
+				None
+			});
+
 		let local_validator = per_session.local_validator.and_then(|v| {
 			if let LocalValidatorIndex::Active(idx) = v {
 				find_active_validator_state(
 					idx,
 					&per_session.groups,
 					&availability_cores,
 					&group_rotation_info,
+					&maybe_claim_queue,
 					seconding_limit,
+					max_candidate_depth,
 				)
 			} else {
 				Some(LocalValidatorState { grid_tracker: GridTracker::default(), active: None })
@@ -752,26 +761,39 @@ fn find_active_validator_state(
 	groups: &Groups,
 	availability_cores: &[CoreState],
 	group_rotation_info: &GroupRotationInfo,
+	maybe_claim_queue: &Option<ClaimQueueSnapshot>,
 	seconding_limit: usize,
+	max_candidate_depth: usize,
 ) -> Option<LocalValidatorState> {
 	if groups.all().is_empty() {
 		return None
 	}
 
 	let our_group = groups.by_validator_index(validator_index)?;
 
-	// note: this won't work well for on-demand parachains because it only works
-	// when core assignments to paras are static throughout the session.
-
-	let core = group_rotation_info.core_for_group(our_group, availability_cores.len());
-	let para = availability_cores.get(core.0 as usize).and_then(|c| c.para_id());
+	let core_index = group_rotation_info.core_for_group(our_group, availability_cores.len());
+	let core_assignment = if let Some(claim_queue) = maybe_claim_queue {
+		claim_queue.get_claim_for(core_index, 0)
+	} else {
+		availability_cores
+			.get(core_index.0 as usize)
+			.map(|core_state| match core_state {
+				CoreState::Scheduled(scheduled_core) => Some(scheduled_core.para_id),
+				CoreState::Occupied(occupied_core) if max_candidate_depth >= 1 => occupied_core
+					.next_up_on_available
+					.as_ref()
+					.map(|scheduled_core| scheduled_core.para_id),
+				CoreState::Free | CoreState::Occupied(_) => None,
+			})
+			.flatten()
+	};
 	let group_validators = groups.get(our_group)?.to_owned();
 
 	Some(LocalValidatorState {
 		active: Some(ActiveValidatorState {
 			index: validator_index,
 			group: our_group,
-			assignment: para,
+			assignment: core_assignment,
 			cluster_tracker: ClusterTracker::new(group_validators, seconding_limit)
 				.expect("group is non-empty because we are in it; qed"),
 		}),

polkadot/primitives/src/v7/mod.rs

@@ -1086,15 +1086,6 @@ pub enum CoreState<H = Hash, N = BlockNumber> {
 }
 
 impl<N> CoreState<N> {
-	/// If this core state has a `para_id`, return it.
-	pub fn para_id(&self) -> Option<Id> {
-		match self {
-			Self::Occupied(ref core) => Some(core.para_id()),
-			Self::Scheduled(core) => Some(core.para_id),
-			Self::Free => None,
-		}
-	}
-
 	/// Is this core state `Self::Occupied`?
 	pub fn is_occupied(&self) -> bool {
 		matches!(self, Self::Occupied(_))