Requeue UMP queue items if weight exhausted

runtime/parachains/src/ump.rs

@@ -14,14 +14,15 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
+
 use crate::{
 	configuration::{self, HostConfiguration},
 	initializer,
 };
 use frame_support::pallet_prelude::*;
 use primitives::v1::{Id as ParaId, UpwardMessage};
 use sp_std::{
-	collections::{btree_map::BTreeMap, vec_deque::VecDeque},
+	collections::{btree_map::{BTreeMap, Entry}, vec_deque::VecDeque},
 	convert::TryFrom,
 	fmt,
 	marker::PhantomData,
@@ -417,9 +418,10 @@ impl<T: Config> Pallet<T> {
 				) {
 					Ok(used) => weight_used += used,
 					Err((id, required)) => {
-						// we process messages in order and don't drop them if we run out of weight, so need to break
-						// here.
+						// we process messages in order and don't drop them if we run out of weight,
+						// so need to break here and requeue the message we took off the queue.
 						Self::deposit_event(Event::WeightExhausted(id, max_weight, required));
+						queue_cache.requeue::<T>(dispatchee, upward_message);
 						break
 					},
 				}
@@ -460,6 +462,7 @@ impl<T: Config> Pallet<T> {
 /// This struct is not supposed to be dropped but rather to be consumed by [`flush`].
 struct QueueCache(BTreeMap<ParaId, QueueCacheEntry>);
 
+#[derive(Default)]
 struct QueueCacheEntry {
 	queue: VecDeque<UpwardMessage>,
 	count: u32,
@@ -493,6 +496,22 @@ impl QueueCache {
 		(upward_message, became_empty)
 	}
 
+	/// Places a message back on to a given parachain's UMP queue.
+	/// 
+	/// - `para`: The ID of the parachain.
+	/// - `message`: The message to put back on the front of the queue.
+	fn requeue<T: Config>(&mut self, para: ParaId, message: UpwardMessage) {
+		match self.0.entry(para) {
+			Entry::Occupied(entry) => {
+				let cache_entry = entry.into_mut();
+				cache_entry.count += 1;
+				cache_entry.total_size += message.len() as u32;
+				cache_entry.queue.push_front(message);
+			},
+			_ => (),
+		}
+	}
+
 	/// Flushes the updated queues into the storage.
 	fn flush<T: Config>(self) {
 		// NOTE we use an explicit method here instead of Drop impl because it has unwanted semantics