From fc627930d54785acaaa30fe003190244221ca36a Mon Sep 17 00:00:00 2001
From: "paritytech-cmd-bot-polkadot-sdk[bot]"
 <179002856+paritytech-cmd-bot-polkadot-sdk[bot]@users.noreply.github.com>
Date: Thu, 17 Oct 2024 12:59:13 +0200
Subject: [PATCH] [stable2407] Backport #4803 (#6097)

Backport #4803 into `stable2407` from gotnoshoeson.

See the
[documentation](https://github.com/paritytech/polkadot-sdk/blob/master/docs/BACKPORT.md)
on how to use this bot.

<!--
  # To be used by other automation, do not modify:
  original-pr-number: #${pull_number}
-->

Co-authored-by: Miles Patterson <miles_patterson@icloud.com>
---
 prdoc/pr_4803.prdoc                        |  14 +++
 substrate/frame/message-queue/src/lib.rs   | 109 ++++++++++++---------
 substrate/frame/message-queue/src/tests.rs |   2 +-
 3 files changed, 80 insertions(+), 45 deletions(-)
 create mode 100644 prdoc/pr_4803.prdoc

diff --git a/prdoc/pr_4803.prdoc b/prdoc/pr_4803.prdoc
new file mode 100644
index 00000000000..0d2ad08d610
--- /dev/null
+++ b/prdoc/pr_4803.prdoc
@@ -0,0 +1,14 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Fix for issue #4762
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      When the status of the queue is on_initialize, throw a defensive message and return weight of 0, 
+      however when status is on_idle, do not throw a defensive message, only return weight of 0
+
+crates:
+  - name: pallet-message-queue
+    bump: patch
diff --git a/substrate/frame/message-queue/src/lib.rs b/substrate/frame/message-queue/src/lib.rs
index 2dbffef7e5a..736c7745ea5 100644
--- a/substrate/frame/message-queue/src/lib.rs
+++ b/substrate/frame/message-queue/src/lib.rs
@@ -649,7 +649,7 @@ pub mod pallet {
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(_n: BlockNumberFor<T>) -> Weight {
 			if let Some(weight_limit) = T::ServiceWeight::get() {
-				Self::service_queues(weight_limit)
+				Self::service_queues_impl(weight_limit, ServiceQueuesContext::OnInitialize)
 			} else {
 				Weight::zero()
 			}
@@ -658,7 +658,10 @@ pub mod pallet {
 		fn on_idle(_n: BlockNumberFor<T>, remaining_weight: Weight) -> Weight {
 			if let Some(weight_limit) = T::IdleMaxServiceWeight::get() {
 				// Make use of the remaining weight to process enqueued messages.
-				Self::service_queues(weight_limit.min(remaining_weight))
+				Self::service_queues_impl(
+					weight_limit.min(remaining_weight),
+					ServiceQueuesContext::OnIdle,
+				)
 			} else {
 				Weight::zero()
 			}
@@ -777,6 +780,18 @@ enum MessageExecutionStatus {
 	StackLimitReached,
 }
 
+/// The context to pass to [`Pallet::service_queues_impl`] through on_idle and on_initialize hooks
+/// We don't want to throw the defensive message if called from on_idle hook
+#[derive(PartialEq)]
+enum ServiceQueuesContext {
+	/// Context of on_idle hook.
+	OnIdle,
+	/// Context of on_initialize hook.
+	OnInitialize,
+	/// Context `service_queues` trait function.
+	ServiceQueues,
+}
+
 impl<T: Config> Pallet<T> {
 	/// Knit `origin` into the ready ring right at the end.
 	///
@@ -1489,6 +1504,53 @@ impl<T: Config> Pallet<T> {
 			},
 		}
 	}
+
+	fn service_queues_impl(weight_limit: Weight, context: ServiceQueuesContext) -> Weight {
+		let mut weight = WeightMeter::with_limit(weight_limit);
+
+		// Get the maximum weight that processing a single message may take:
+		let max_weight = Self::max_message_weight(weight_limit).unwrap_or_else(|| {
+			if matches!(context, ServiceQueuesContext::OnInitialize) {
+				defensive!("Not enough weight to service a single message.");
+			}
+			Weight::zero()
+		});
+
+		match with_service_mutex(|| {
+			let mut next = match Self::bump_service_head(&mut weight) {
+				Some(h) => h,
+				None => return weight.consumed(),
+			};
+			// The last queue that did not make any progress.
+			// The loop aborts as soon as it arrives at this queue again without making any progress
+			// on other queues in between.
+			let mut last_no_progress = None;
+
+			loop {
+				let (progressed, n) = Self::service_queue(next.clone(), &mut weight, max_weight);
+				next = match n {
+					Some(n) =>
+						if !progressed {
+							if last_no_progress == Some(n.clone()) {
+								break
+							}
+							if last_no_progress.is_none() {
+								last_no_progress = Some(next.clone())
+							}
+							n
+						} else {
+							last_no_progress = None;
+							n
+						},
+					None => break,
+				}
+			}
+			weight.consumed()
+		}) {
+			Err(()) => weight.consumed(),
+			Ok(w) => w,
+		}
+	}
 }
 
 /// Run a closure that errors on re-entrance. Meant to be used by anything that services queues.
@@ -1558,48 +1620,7 @@ impl<T: Config> ServiceQueues for Pallet<T> {
 	type OverweightMessageAddress = (MessageOriginOf<T>, PageIndex, T::Size);
 
 	fn service_queues(weight_limit: Weight) -> Weight {
-		let mut weight = WeightMeter::with_limit(weight_limit);
-
-		// Get the maximum weight that processing a single message may take:
-		let max_weight = Self::max_message_weight(weight_limit).unwrap_or_else(|| {
-			defensive!("Not enough weight to service a single message.");
-			Weight::zero()
-		});
-
-		match with_service_mutex(|| {
-			let mut next = match Self::bump_service_head(&mut weight) {
-				Some(h) => h,
-				None => return weight.consumed(),
-			};
-			// The last queue that did not make any progress.
-			// The loop aborts as soon as it arrives at this queue again without making any progress
-			// on other queues in between.
-			let mut last_no_progress = None;
-
-			loop {
-				let (progressed, n) = Self::service_queue(next.clone(), &mut weight, max_weight);
-				next = match n {
-					Some(n) =>
-						if !progressed {
-							if last_no_progress == Some(n.clone()) {
-								break
-							}
-							if last_no_progress.is_none() {
-								last_no_progress = Some(next.clone())
-							}
-							n
-						} else {
-							last_no_progress = None;
-							n
-						},
-					None => break,
-				}
-			}
-			weight.consumed()
-		}) {
-			Err(()) => weight.consumed(),
-			Ok(w) => w,
-		}
+		Self::service_queues_impl(weight_limit, ServiceQueuesContext::ServiceQueues)
 	}
 
 	/// Execute a single overweight message.
diff --git a/substrate/frame/message-queue/src/tests.rs b/substrate/frame/message-queue/src/tests.rs
index e89fdb8b320..e232cd07536 100644
--- a/substrate/frame/message-queue/src/tests.rs
+++ b/substrate/frame/message-queue/src/tests.rs
@@ -279,7 +279,7 @@ fn service_queues_low_weight_defensive() {
 		assert!(MessageQueue::do_integrity_test().is_err());
 
 		MessageQueue::enqueue_message(msg("weight=0"), Here);
-		MessageQueue::service_queues(104.into_weight());
+		MessageQueue::service_queues_impl(104.into_weight(), ServiceQueuesContext::OnInitialize);
 	});
 }
 
-- 
GitLab