Add try_state check to Pallet MessageQueue

frame/message-queue/src/integration_test.rs

@@ -22,8 +22,9 @@
 
 use crate::{
 	mock::{
-		new_test_ext, CountingMessageProcessor, IntoWeight, MockedWeightInfo, NumMessagesProcessed,
-		SuspendedQueues,
+		build_and_execute, new_test_ext, CountingMessageProcessor, IntoWeight, MockedWeightInfo,
+		NumMessagesProcessed, SuspendedQueues,
+
 	},
 	mock_helpers::MessageOrigin,
 	*,
@@ -130,7 +131,7 @@ fn stress_test_enqueue_and_service() {
 	let max_msg_len = MaxMessageLenOf::<Test>::get();
 	let mut rng = StdRng::seed_from_u64(42);
 
-	new_test_ext::<Test>().execute_with(|| {
+	build_and_execute::<Test>(|| {
 		let mut msgs_remaining = 0;
 		for _ in 0..blocks {
 			// Start by enqueuing a large number of messages.
@@ -178,7 +179,7 @@ fn stress_test_queue_suspension() {
 	let max_msg_len = MaxMessageLenOf::<Test>::get();
 	let mut rng = StdRng::seed_from_u64(41);
 
-	new_test_ext::<Test>().execute_with(|| {
+	build_and_execute::<Test>(|| {
 		let mut suspended = BTreeSet::<u32>::new();
 		let mut msgs_remaining = 0;
 

frame/message-queue/src/lib.rs

@@ -571,6 +571,11 @@ pub mod pallet {
 			}
 		}
 
+		#[cfg(feature = "try-runtime")]
+		fn try_state(_n: BlockNumberFor<T>) -> Result<(), &'static str> {
+			Self::do_try_state()
+		}
+
 		/// Check all assumptions about [`crate::Config`].
 		fn integrity_test() {
 			assert!(!MaxMessageLenOf::<T>::get().is_zero(), "HeapSize too low");
@@ -1085,6 +1090,111 @@ impl<T: Config> Pallet<T> {
 		ItemExecutionStatus::Executed(is_processed)
 	}
 
+	/// Ensure the correctness of state of this pallet.
+	///
+	/// # Assumptions-
+	///
+	/// If `serviceHead` points to a ready Queue, then BookState of that Queue has:
+	///
+	/// * `message_count` > 0
+	/// * `size` > 0
+	/// * `end` > `begin`
+	/// * Some(ready_neighbours)
+	/// * If `ready_neighbours.next` == self.origin, then `ready_neighbours.prev` == self.origin
+	///   (only queue in ring)
+	///
+	/// For Pages(begin to end-1) in BookState:
+	///
+	/// * `remaining` > 0
+	/// * `remaining_size` > 0
+	/// * `first` <= `last`
+	/// * Every page can be decoded into peek_* functions
+
+	pub fn do_try_state() -> Result<(), &'static str> {
+		// Checking memory corruption for BookStateFor
+		assert_eq!(
+			BookStateFor::<T>::iter_keys().count(),
+			BookStateFor::<T>::iter_values().count(),
+			"Memory Corruption in BookStateFor"
+		);
+
+		// Checking memory corruption for Pages
+		assert_eq!(
+			Pages::<T>::iter_keys().count(),
+			Pages::<T>::iter_values().count(),
+			"Memory Corruption in Pages"
+		);
+
+		// No state to check
+		if ServiceHead::<T>::get().is_none() {
+			return Ok(())
+		}
+
+		//loop around this origin
+		let starting_origin = ServiceHead::<T>::get().unwrap();
+
+		while let Some(head) = Self::bump_service_head(&mut WeightMeter::max_limit()) {
+			ensure!(
+				BookStateFor::<T>::contains_key(&head),
+				"Service head must point to an existing book"
+			);
+
+			let head_book_state = BookStateFor::<T>::get(&head);
+			ensure!(
+				head_book_state.message_count > 0,
+				"There must be some messages if in ReadyRing"
+			);
+			ensure!(head_book_state.size > 0, "There must be some messages if in ReadyRing");
+			ensure!(
+				head_book_state.end > head_book_state.begin,
+				"End > Begin if unprocessed messages exists"
+			);
+			ensure!(
+				head_book_state.ready_neighbours.is_some(),
+				"There must be neighbours if in ReadyRing"
+			);
+
+			if head_book_state.ready_neighbours.as_ref().unwrap().next == head {
+				assert_eq!(
+					head_book_state.ready_neighbours.as_ref().unwrap().prev,
+					head,
+					"Can only happen if only queue in ReadyRing"
+				);
+			}
+
+			for page_index in head_book_state.begin..head_book_state.end {
+				let page = Pages::<T>::get(&head, page_index).unwrap();
+				let remaining_messages = page.remaining;
+				let mut counted_remaining_messages = 0;
+				assert!(
+					remaining_messages > 0.into(),
+					"These must be some messages that have not been processed yet!"
+				);
+
+				for i in 0..u32::MAX {
+					if let Some((_, processed, _)) = page.peek_index(i as usize) {
+						if !processed {
+							counted_remaining_messages += 1;
+						}
+					} else {
+						break
+					}
+				}
+
+				assert_eq!(
+					remaining_messages,
+					counted_remaining_messages.into(),
+					"Memory Corruption"
+				);
+			}
+
+			if head_book_state.ready_neighbours.as_ref().unwrap().next == starting_origin {
+				break
+			}
+		}
+		Ok(())
+	}
+
 	/// Print the pages in each queue and the messages in each page.
 	///
 	/// Processed messages are prefixed with a `*` and the current `begin`ning page with a `>`.

frame/message-queue/src/mock.rs

@@ -285,6 +285,17 @@ where
 	ext
 }
 
+/// Run the function pointer inside externalities and asserts the try_state hook at the end.
+pub fn build_and_execute<T: Config>(test: impl FnOnce() -> ())
+where
+	<T as frame_system::Config>::BlockNumber: From<u32>,
+{
+	new_test_ext::<T>().execute_with(|| {
+		test();
+		MessageQueue::do_try_state().unwrap();
+	})
+}
+
 /// Set the weight of a specific weight function.
 pub fn set_weight(name: &str, w: Weight) {
 	MockedWeightInfo::set_weight::<Test>(name, w);