MQ pallet: add discard_overweight call

bin/node/runtime/src/lib.rs

@@ -36,9 +36,10 @@ use frame_support::{
 	traits::{
 		fungible::ItemOf,
 		tokens::{nonfungibles_v2::Inspect, GetSalary, PayFromAccount},
-		AsEnsureOriginWithArg, ConstBool, ConstU128, ConstU16, ConstU32, Currency, EitherOfDiverse,
-		EqualPrivilegeOnly, Everything, Imbalance, InstanceFilter, KeyOwnerProofSystem,
-		LockIdentifier, Nothing, OnUnbalanced, U128CurrencyToVote, WithdrawReasons,
+		AsEnsureOriginWithArg, AsEnsureOriginWithContainsPair, ConstBool, ConstU128, ConstU16,
+		ConstU32, Currency, EitherOfDiverse, EqualPrivilegeOnly, Everything, Imbalance,
+		InstanceFilter, KeyOwnerProofSystem, LockIdentifier, Nothing, OnUnbalanced,
+		U128CurrencyToVote, WithdrawReasons,
 	},
 	weights::{
 		constants::{
@@ -1178,6 +1179,8 @@ impl pallet_message_queue::Config for Runtime {
 	type WeightInfo = ();
 	/// NOTE: Always set this to `NoopMessageProcessor` for benchmarking.
 	type MessageProcessor = pallet_message_queue::mock_helpers::NoopMessageProcessor<u32>;
+	type DiscardOverweightOrigin =
+		AsEnsureOriginWithContainsPair<EnsureRoot<AccountId>, Everything>;
 	type Size = u32;
 	type QueueChangeHandler = ();
 	type HeapSize = ConstU32<{ 64 * 1024 }>;

frame/message-queue/src/benchmarking.rs

@@ -234,7 +234,7 @@ mod benchmarks {
 			}
 			.into(),
 		);
-		assert!(!Pages::<T>::contains_key(&origin, 0), "Page must be removed");
+		assert!(Pages::<T>::iter_keys().count().is_zero(), "Page must be removed");
 	}
 
 	// Worst case for `execute_overweight` where the page is updated.
@@ -274,6 +274,72 @@ mod benchmarks {
 		assert!(Pages::<T>::contains_key(&origin, 0), "Page must be updated");
 	}
 
+	#[benchmark]
+	fn discard_overweight_page_removed() -> Result<(), BenchmarkError> {
+		let queue: MessageOriginOf<T> = 0.into();
+		let origin = T::DiscardOverweightOrigin::try_successful_origin(&queue)
+			.map_err(|_| BenchmarkError::Weightless)?;
+		let (mut page, msgs) = full_page::<T>();
+		// Skip all messages.
+		for _ in 1..msgs {
+			page.skip_first(true);
+		}
+		page.skip_first(false);
+		let book = book_for::<T>(&page);
+		Pages::<T>::insert(&queue, 0, &page);
+		BookStateFor::<T>::insert(&queue, &book);
+		// Remove the last message since `peek_index` has linear complexity.
+		let message_index = ((msgs - 1) as u32).into();
+
+		#[extrinsic_call]
+		discard_overweight(origin as T::RuntimeOrigin, 0u32.into(), 0u32, message_index);
+
+		assert_last_event::<T>(
+			Event::OverweightDiscarded {
+				id: sp_io::hashing::blake2_256(&((msgs - 1) as u32).encode()),
+				origin: 0.into(),
+				page_index: 0,
+				message_index,
+			}
+			.into(),
+		);
+
+		assert!(Pages::<T>::iter_keys().count().is_zero(), "Page must be removed");
+		Ok(())
+	}
+
+	#[benchmark]
+	fn discard_overweight_page_updated() -> Result<(), BenchmarkError> {
+		let queue: MessageOriginOf<T> = 0.into();
+		let origin = T::DiscardOverweightOrigin::try_successful_origin(&queue)
+			.map_err(|_| BenchmarkError::Weightless)?;
+		let (mut page, msgs) = full_page::<T>();
+		// Skip all messages.
+		for _ in 0..msgs {
+			page.skip_first(false);
+		}
+		let book = book_for::<T>(&page);
+		Pages::<T>::insert(&queue, 0, &page);
+		BookStateFor::<T>::insert(&queue, &book);
+		let message_index = ((msgs - 1) as u32).into();
+
+		#[extrinsic_call]
+		discard_overweight(origin as T::RuntimeOrigin, 0u32.into(), 0u32, message_index);
+
+		assert_last_event::<T>(
+			Event::OverweightDiscarded {
+				id: sp_io::hashing::blake2_256(&((msgs - 1) as u32).encode()),
+				origin: 0.into(),
+				page_index: 0,
+				message_index,
+			}
+			.into(),
+		);
+
+		assert!(Pages::<T>::contains_key(&queue, 0), "Page must be updated");
+		Ok(())
+	}
+
 	impl_benchmark_test_suite! {
 		MessageQueue,
 		crate::mock::new_test_ext::<crate::integration_test::Test>(),

frame/message-queue/src/integration_test.rs

@@ -32,8 +32,9 @@ use crate::{
 use crate as pallet_message_queue;
 use frame_support::{
 	parameter_types,
-	traits::{ConstU32, ConstU64},
+	traits::{AsEnsureOriginWithArg, ConstU32, ConstU64},
 };
+use frame_system::EnsureRoot;
 use rand::{rngs::StdRng, Rng, SeedableRng};
 use rand_distr::Pareto;
 use sp_core::H256;
@@ -94,6 +95,7 @@ impl Config for Test {
 	type RuntimeEvent = RuntimeEvent;
 	type WeightInfo = MockedWeightInfo;
 	type MessageProcessor = CountingMessageProcessor;
+	type DiscardOverweightOrigin = AsEnsureOriginWithArg<EnsureRoot<Self::AccountId>>;
 	type Size = u32;
 	type QueueChangeHandler = ();
 	type HeapSize = HeapSize;
@@ -128,10 +130,11 @@ fn stress_test_enqueue_and_service() {
 	let max_queues = 10_000;
 	let max_messages_per_queue = 10_000;
 	let max_msg_len = MaxMessageLenOf::<Test>::get();
-	let mut rng = StdRng::seed_from_u64(42);
 
 	new_test_ext::<Test>().execute_with(|| {
+		let mut rng = StdRng::seed_from_u64(gen_seed());
 		let mut msgs_remaining = 0;
+
 		for _ in 0..blocks {
 			// Start by enqueuing a large number of messages.
 			let enqueued =
@@ -176,9 +179,9 @@ fn stress_test_queue_suspension() {
 	let max_messages_per_queue = 10_000;
 	let (max_suspend_per_block, max_resume_per_block) = (100, 50);
 	let max_msg_len = MaxMessageLenOf::<Test>::get();
-	let mut rng = StdRng::seed_from_u64(41);
 
 	new_test_ext::<Test>().execute_with(|| {
+		let mut rng = StdRng::seed_from_u64(gen_seed());
 		let mut suspended = BTreeSet::<u32>::new();
 		let mut msgs_remaining = 0;
 
@@ -334,3 +337,14 @@ fn post_conditions() {
 	assert_eq!(MessageQueue::service_queues(Weight::MAX), Weight::zero(), "Nothing left");
 	next_block();
 }
+
+/// Pull a seed from env `TEST_SEED` or generate a random one. Logged in both cases.
+fn gen_seed() -> u64 {
+	let s = std::env::var("TEST_SEED")
+		.ok()
+		.and_then(|s| s.parse().ok())
+		.unwrap_or_else(|| StdRng::from_entropy().gen());
+
+	log::info!("Using TEST_SEED={}", s);
+	s
+}