Emit UMP signal from parachain runtime

Author: serban300

cumulus/pallets/parachain-system/src/lib.rs

@@ -33,9 +33,10 @@ use alloc::{collections::btree_map::BTreeMap, vec, vec::Vec};
 use codec::{Decode, DecodeLimit, Encode};
 use core::cmp;
 use cumulus_primitives_core::{
-	relay_chain, AbridgedHostConfiguration, ChannelInfo, ChannelStatus, CollationInfo,
-	CumulusDigestItem, GetChannelInfo, ListChannelInfos, MessageSendError, OutboundHrmpMessage,
-	ParaId, PersistedValidationData, UpwardMessage, UpwardMessageSender, XcmpMessageHandler,
+	relay_chain::{self, UMPSignal, UMP_SEPARATOR},
+	AbridgedHostConfiguration, ChannelInfo, ChannelStatus, CollationInfo, CumulusDigestItem,
+	GetChannelInfo, ListChannelInfos, MessageSendError, OutboundHrmpMessage, ParaId,
+	PersistedValidationData, UpwardMessage, UpwardMessageSender, XcmpMessageHandler,
 	XcmpMessageSource,
 };
 use cumulus_primitives_parachain_inherent::{v0, MessageQueueChain, ParachainInherentData};
@@ -361,8 +362,8 @@ pub mod pallet {
 				UpwardMessages::<T>::put(&up[..num as usize]);
 				*up = up.split_off(num as usize);
 
-				// Send the core selector UMP signal.
-				Self::send_ump_signal();
+				// Send the pending UMP signals + the core selector UMP signal.
+				Self::send_ump_signals();
 
 				// If the total size of the pending messages is less than the threshold,
 				// we decrease the fee factor, since the queue is less congested.
@@ -585,7 +586,7 @@ pub mod pallet {
 				validation_data: vfp,
 				relay_chain_state,
 				relay_parent_descendants,
-				collator_peer_id: _,
+				collator_peer_id,
 			} = data;
 
 			// Check that the associated relay chain block number is as expected.
@@ -693,6 +694,12 @@ pub mod pallet {
 
 			<T::OnSystemEvent as OnSystemEvent>::on_validation_data(&vfp);
 
+			if let Some(collator_peer_id) = collator_peer_id {
+				PendingUpwardSignals::<T>::mutate(|signals| {
+					signals.push(UMPSignal::ApprovedPeer(collator_peer_id).encode());
+				});
+			}
+
 			total_weight.saturating_accrue(Self::enqueue_inbound_downward_messages(
 				relevant_messaging_state.dmq_mqc_head,
 				inbound_messages_data.downward_messages,
@@ -905,14 +912,20 @@ pub mod pallet {
 
 	/// Upward messages that were sent in a block.
 	///
-	/// This will be cleared in `on_initialize` of each new block.
+	/// This will be cleared in `on_initialize` for each new block.
 	#[pallet::storage]
 	pub type UpwardMessages<T: Config> = StorageValue<_, Vec<UpwardMessage>, ValueQuery>;
 
-	/// Upward messages that are still pending and not yet send to the relay chain.
+	/// Upward messages that are still pending and not yet sent to the relay chain.
 	#[pallet::storage]
 	pub type PendingUpwardMessages<T: Config> = StorageValue<_, Vec<UpwardMessage>, ValueQuery>;
 
+	/// Upward signals that are still pending and not yet sent to the relay chain.
+	///
+	/// This will be cleared in `on_finalize` for each block.
+	#[pallet::storage]
+	pub type PendingUpwardSignals<T: Config> = StorageValue<_, Vec<UpwardMessage>, ValueQuery>;
+
 	/// The factor to multiply the base delivery fee by for UMP.
 	#[pallet::storage]
 	pub type UpwardDeliveryFeeFactor<T: Config> =
@@ -1508,22 +1521,25 @@ impl<T: Config> Pallet<T> {
 	}
 
 	/// Send the ump signals
-	fn send_ump_signal() {
-		use cumulus_primitives_core::relay_chain::{UMPSignal, UMP_SEPARATOR};
+	fn send_ump_signals() {
+		// Take the pending UMP signals.
+		let mut ump_signals = PendingUpwardSignals::<T>::take();
+		// Append the core selector signal.
+		if let Some(core_info) =
+			CumulusDigestItem::find_core_info(&frame_system::Pallet::<T>::digest())
+		{
+			ump_signals.push(
+				UMPSignal::SelectCore(core_info.selector, core_info.claim_queue_offset).encode(),
+			);
+		}
 
-		UpwardMessages::<T>::mutate(|up| {
-			if let Some(core_info) =
-				CumulusDigestItem::find_core_info(&frame_system::Pallet::<T>::digest())
-			{
+		// Send all the UMP signals.
+		if !ump_signals.is_empty() {
+			UpwardMessages::<T>::mutate(|up| {
 				up.push(UMP_SEPARATOR);
-
-				// Send the core selector signal.
-				up.push(
-					UMPSignal::SelectCore(core_info.selector, core_info.claim_queue_offset)
-						.encode(),
-				);
-			}
-		});
+				up.append(&mut ump_signals);
+			});
+		}
 	}
 
 	/// Open HRMP channel for using it in benchmarks or tests.

cumulus/pallets/parachain-system/src/tests.rs

@@ -19,8 +19,12 @@
 use super::*;
 use crate::mock::*;
 
+use alloc::collections::BTreeMap;
 use core::num::NonZeroU32;
-use cumulus_primitives_core::{AbridgedHrmpChannel, InboundDownwardMessage, InboundHrmpMessage};
+use cumulus_primitives_core::{
+	relay_chain::ApprovedPeerId, AbridgedHrmpChannel, ClaimQueueOffset, CoreInfo, CoreSelector,
+	InboundDownwardMessage, InboundHrmpMessage, CUMULUS_CONSENSUS_ID,
+};
 use cumulus_primitives_parachain_inherent::{
 	v0, INHERENT_IDENTIFIER, PARACHAIN_INHERENT_IDENTIFIER_V0,
 };
@@ -31,6 +35,7 @@ use rand::Rng;
 use relay_chain::HrmpChannelId;
 use sp_core::H256;
 use sp_inherents::InherentDataProvider;
+use sp_runtime::DigestItem;
 use sp_trie::StorageProof;
 
 #[test]
@@ -1655,3 +1660,74 @@ fn ump_fee_factor_increases_and_decreases() {
 			},
 		);
 }
+
+#[test]
+fn ump_signals_are_sent_correctly() {
+	let core_info = CoreInfo {
+		selector: CoreSelector(1),
+		claim_queue_offset: ClaimQueueOffset(1),
+		number_of_cores: codec::Compact(1),
+	};
+
+	// Test cases list with the following format:
+	// `((expect_approved_peer, expect_select_core), expected_upward_messages)`
+	let test_cases = BTreeMap::from([
+		((false, false), vec![b"Test".to_vec()]),
+		(
+			(true, false),
+			vec![
+				b"Test".to_vec(),
+				UMP_SEPARATOR,
+				UMPSignal::ApprovedPeer(ApprovedPeerId::try_from(b"12345".to_vec()).unwrap())
+					.encode(),
+			],
+		),
+		(
+			(false, true),
+			vec![
+				b"Test".to_vec(),
+				UMP_SEPARATOR,
+				UMPSignal::SelectCore(core_info.selector, core_info.claim_queue_offset).encode(),
+			],
+		),
+		(
+			(true, true),
+			vec![
+				b"Test".to_vec(),
+				UMP_SEPARATOR,
+				UMPSignal::ApprovedPeer(ApprovedPeerId::try_from(b"12345".to_vec()).unwrap())
+					.encode(),
+				UMPSignal::SelectCore(core_info.selector, core_info.claim_queue_offset).encode(),
+			],
+		),
+	]);
+
+	for ((expect_approved_peer, expect_select_core), expected_upward_messages) in test_cases {
+		let core_info_digest = CumulusDigestItem::CoreInfo(core_info.clone()).encode();
+
+		BlockTests::new()
+			.with_inherent_data(move |_, _, data| {
+				if expect_approved_peer {
+					data.collator_peer_id =
+						Some(ApprovedPeerId::try_from(b"12345".to_vec()).unwrap());
+				}
+			})
+			.add_with_post_test(
+				1,
+				move || {
+					ParachainSystem::send_upward_message(b"Test".to_vec()).unwrap();
+
+					if expect_select_core {
+						System::deposit_log(DigestItem::PreRuntime(
+							CUMULUS_CONSENSUS_ID,
+							core_info_digest.clone(),
+						));
+					}
+				},
+				move || {
+					assert_eq!(PendingUpwardSignals::<Test>::get(), Vec::<Vec<u8>>::new());
+					assert_eq!(UpwardMessages::<Test>::get(), expected_upward_messages);
+				},
+			);
+	}
+}