ump.rs 27.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// Copyright 2020 Parity Technologies (UK) Ltd.
// This file is part of Polkadot.

// Polkadot is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Polkadot is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.

17
18
19
20
use crate::{
	configuration::{self, HostConfiguration},
	initializer,
};
21
use sp_std::{fmt, prelude::*};
22
use sp_std::collections::{btree_map::BTreeMap, vec_deque::VecDeque};
23
use frame_support::{decl_module, decl_storage, StorageMap, StorageValue, weights::Weight, traits::Get};
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
use primitives::v1::{Id as ParaId, UpwardMessage};

/// All upward messages coming from parachains will be funneled into an implementation of this trait.
///
/// The message is opaque from the perspective of UMP. The message size can range from 0 to
/// `config.max_upward_message_size`.
///
/// It's up to the implementation of this trait to decide what to do with a message as long as it
/// returns the amount of weight consumed in the process of handling. Ignoring a message is a valid
/// strategy.
///
/// There are no guarantees on how much time it takes for the message sent by a candidate to end up
/// in the sink after the candidate was enacted. That typically depends on the UMP traffic, the sizes
/// of upward messages and the configuration of UMP.
///
/// It is possible that by the time the message is sank the origin parachain was offboarded. It is
/// up to the implementer to check that if it cares.
pub trait UmpSink {
	/// Process an incoming upward message and return the amount of weight it consumed.
	///
	/// See the trait docs for more details.
	fn process_upward_message(origin: ParaId, msg: Vec<u8>) -> Weight;
}

/// An implementation of a sink that just swallows the message without consuming any weight.
impl UmpSink for () {
	fn process_upward_message(_: ParaId, _: Vec<u8>) -> Weight {
		0
	}
}

Shawn Tabrizi's avatar
Shawn Tabrizi committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
/// A specific implementation of a UmpSink where messages are in the XCM format
/// and will be forwarded to the XCM Executor.
pub struct XcmSink<Config>(sp_std::marker::PhantomData<Config>);

impl<Config: xcm_executor::Config> UmpSink for XcmSink<Config> {
	fn process_upward_message(origin: ParaId, msg: Vec<u8>) -> Weight {
		use parity_scale_codec::Decode;
		use xcm::VersionedXcm;
		use xcm::v0::{Junction, MultiLocation, ExecuteXcm};
		use xcm_executor::XcmExecutor;

		let weight: Weight = 0;

		if let Ok(versioned_xcm_message) = VersionedXcm::decode(&mut &msg[..]) {
			match versioned_xcm_message {
				VersionedXcm::V0(xcm_message) => {
					let xcm_junction: Junction = Junction::Parachain { id: origin.into() };
					let xcm_location: MultiLocation = xcm_junction.into();
					// TODO: Do something with result.
					let _result = XcmExecutor::<Config>::execute_xcm(xcm_location, xcm_message);
				}
			}
		} else {
			frame_support::debug::error!(
				target: "xcm",
				"Failed to decode versioned XCM from upward message.",
			);
		}

		// TODO: to be sound, this implementation must ensure that returned (and thus consumed)
		// weight is limited to some small portion of the total block weight (as a ballpark, 1/4, 1/8
		// or lower).
		weight
	}
}

91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
/// An error returned by [`check_upward_messages`] that indicates a violation of one of acceptance
/// criteria rules.
pub enum AcceptanceCheckErr {
	MoreMessagesThanPermitted {
		sent: u32,
		permitted: u32,
	},
	MessageSize {
		idx: u32,
		msg_size: u32,
		max_size: u32,
	},
	CapacityExceeded {
		count: u32,
		limit: u32,
	},
	TotalSizeExceeded {
		total_size: u32,
		limit: u32,
	},
}

impl fmt::Debug for AcceptanceCheckErr {
	fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			AcceptanceCheckErr::MoreMessagesThanPermitted { sent, permitted } => write!(
				fmt,
				"more upward messages than permitted by config ({} > {})",
119
				sent, permitted,
120
121
122
123
124
125
126
127
			),
			AcceptanceCheckErr::MessageSize {
				idx,
				msg_size,
				max_size,
			} => write!(
				fmt,
				"upward message idx {} larger than permitted by config ({} > {})",
128
				idx, msg_size, max_size,
129
130
131
132
			),
			AcceptanceCheckErr::CapacityExceeded { count, limit } => write!(
				fmt,
				"the ump queue would have more items than permitted by config ({} > {})",
133
				count, limit,
134
135
136
137
			),
			AcceptanceCheckErr::TotalSizeExceeded { total_size, limit } => write!(
				fmt,
				"the ump queue would have grown past the max size permitted by config ({} > {})",
138
				total_size, limit,
139
140
141
142
			),
		}
	}
}
143

144
pub trait Config: frame_system::Config + configuration::Config {
145
146
147
148
149
	/// A place where all received upward messages are funneled.
	type UmpSink: UmpSink;
}

decl_storage! {
150
	trait Store for Module<T: Config> as Ump {
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
		/// Paras that are to be cleaned up at the end of the session.
		/// The entries are sorted ascending by the para id.
		OutgoingParas: Vec<ParaId>;

		/// The messages waiting to be handled by the relay-chain originating from a certain parachain.
		///
		/// Note that some upward messages might have been already processed by the inclusion logic. E.g.
		/// channel management messages.
		///
		/// The messages are processed in FIFO order.
		RelayDispatchQueues: map hasher(twox_64_concat) ParaId => VecDeque<UpwardMessage>;
		/// Size of the dispatch queues. Caches sizes of the queues in `RelayDispatchQueue`.
		///
		/// First item in the tuple is the count of messages and second
		/// is the total length (in bytes) of the message payloads.
		///
		/// Note that this is an auxilary mapping: it's possible to tell the byte size and the number of
		/// messages only looking at `RelayDispatchQueues`. This mapping is separate to avoid the cost of
		/// loading the whole message queue if only the total size and count are required.
		///
		/// Invariant:
		/// - The set of keys should exactly match the set of keys of `RelayDispatchQueues`.
		RelayDispatchQueueSize: map hasher(twox_64_concat) ParaId => (u32, u32);
		/// The ordered list of `ParaId`s that have a `RelayDispatchQueue` entry.
		///
		/// Invariant:
		/// - The set of items from this vector should be exactly the set of the keys in
		///   `RelayDispatchQueues` and `RelayDispatchQueueSize`.
		NeedsDispatch: Vec<ParaId>;
		/// This is the para that gets will get dispatched first during the next upward dispatchable queue
		/// execution round.
		///
		/// Invariant:
		/// - If `Some(para)`, then `para` must be present in `NeedsDispatch`.
		NextDispatchRoundStartWith: Option<ParaId>;
	}
}

decl_module! {
	/// The UMP module.
191
	pub struct Module<T: Config> for enum Call where origin: <T as frame_system::Config>::Origin {
192
193
194
	}
}

195
/// Routines related to the upward message passing.
196
impl<T: Config> Module<T> {
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
	/// Block initialization logic, called by initializer.
	pub(crate) fn initializer_initialize(_now: T::BlockNumber) -> Weight {
		0
	}

	/// Block finalization logic, called by initializer.
	pub(crate) fn initializer_finalize() {}

	/// Called by the initializer to note that a new session has started.
	pub(crate) fn initializer_on_new_session(
		_notification: &initializer::SessionChangeNotification<T::BlockNumber>,
	) {
		Self::perform_outgoing_para_cleanup();
	}

	/// Iterate over all paras that were registered for offboarding and remove all the data
	/// associated with them.
	fn perform_outgoing_para_cleanup() {
		let outgoing = OutgoingParas::take();
		for outgoing_para in outgoing {
			Self::clean_ump_after_outgoing(outgoing_para);
		}
	}

	/// Schedule a para to be cleaned up at the start of the next session.
	pub(crate) fn schedule_para_cleanup(id: ParaId) {
		OutgoingParas::mutate(|v| {
			if let Err(i) = v.binary_search(&id) {
				v.insert(i, id);
			}
		});
	}

	fn clean_ump_after_outgoing(outgoing_para: ParaId) {
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
		<Self as Store>::RelayDispatchQueueSize::remove(&outgoing_para);
		<Self as Store>::RelayDispatchQueues::remove(&outgoing_para);

		// Remove the outgoing para from the `NeedsDispatch` list and from
		// `NextDispatchRoundStartWith`.
		//
		// That's needed for maintaining invariant that `NextDispatchRoundStartWith` points to an
		// existing item in `NeedsDispatch`.
		<Self as Store>::NeedsDispatch::mutate(|v| {
			if let Ok(i) = v.binary_search(&outgoing_para) {
				v.remove(i);
			}
		});
		<Self as Store>::NextDispatchRoundStartWith::mutate(|v| {
			*v = v.filter(|p| *p == outgoing_para)
		});
	}

	/// Check that all the upward messages sent by a candidate pass the acceptance criteria. Returns
	/// false, if any of the messages doesn't pass.
	pub(crate) fn check_upward_messages(
		config: &HostConfiguration<T::BlockNumber>,
		para: ParaId,
		upward_messages: &[UpwardMessage],
255
	) -> Result<(), AcceptanceCheckErr> {
256
		if upward_messages.len() as u32 > config.max_upward_message_num_per_candidate {
257
258
259
260
			return Err(AcceptanceCheckErr::MoreMessagesThanPermitted {
				sent: upward_messages.len() as u32,
				permitted: config.max_upward_message_num_per_candidate,
			});
261
262
263
264
265
		}

		let (mut para_queue_count, mut para_queue_size) =
			<Self as Store>::RelayDispatchQueueSize::get(&para);

266
		for (idx, msg) in upward_messages.into_iter().enumerate() {
267
268
			let msg_size = msg.len() as u32;
			if msg_size > config.max_upward_message_size {
269
270
				return Err(AcceptanceCheckErr::MessageSize {
					idx: idx as u32,
271
					msg_size,
272
273
					max_size: config.max_upward_message_size,
				});
274
275
276
277
278
279
280
			}
			para_queue_count += 1;
			para_queue_size += msg_size;
		}

		// make sure that the queue is not overfilled.
		// we do it here only once since returning false invalidates the whole relay-chain block.
281
		if para_queue_count > config.max_upward_queue_count {
282
283
284
285
			return Err(AcceptanceCheckErr::CapacityExceeded {
				count: para_queue_count,
				limit: config.max_upward_queue_count,
			});
286
287
		}
		if para_queue_size > config.max_upward_queue_size {
288
289
290
291
			return Err(AcceptanceCheckErr::TotalSizeExceeded {
				total_size: para_queue_size,
				limit: config.max_upward_queue_size,
			});
292
		}
293
294

		Ok(())
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
	}

	/// Enacts all the upward messages sent by a candidate.
	pub(crate) fn enact_upward_messages(
		para: ParaId,
		upward_messages: Vec<UpwardMessage>,
	) -> Weight {
		let mut weight = 0;

		if !upward_messages.is_empty() {
			let (extra_cnt, extra_size) = upward_messages
				.iter()
				.fold((0, 0), |(cnt, size), d| (cnt + 1, size + d.len() as u32));

			<Self as Store>::RelayDispatchQueues::mutate(&para, |v| {
				v.extend(upward_messages.into_iter())
			});

313
314
315
316
			<Self as Store>::RelayDispatchQueueSize::mutate(&para, |(ref mut cnt, ref mut size)| {
				*cnt += extra_cnt;
				*size += extra_size;
			});
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403

			<Self as Store>::NeedsDispatch::mutate(|v| {
				if let Err(i) = v.binary_search(&para) {
					v.insert(i, para);
				}
			});

			weight += T::DbWeight::get().reads_writes(3, 3);
		}

		weight
	}

	/// Devote some time into dispatching pending upward messages.
	pub(crate) fn process_pending_upward_messages() {
		let mut used_weight_so_far = 0;

		let config = <configuration::Module<T>>::config();
		let mut cursor = NeedsDispatchCursor::new::<T>();
		let mut queue_cache = QueueCache::new();

		while let Some(dispatchee) = cursor.peek() {
			if used_weight_so_far >= config.preferred_dispatchable_upward_messages_step_weight {
				// Then check whether we've reached or overshoot the
				// preferred weight for the dispatching stage.
				//
				// if so - bail.
				break;
			}

			// dequeue the next message from the queue of the dispatchee
			let (upward_message, became_empty) = queue_cache.dequeue::<T>(dispatchee);
			if let Some(upward_message) = upward_message {
				used_weight_so_far +=
					T::UmpSink::process_upward_message(dispatchee, upward_message);
			}

			if became_empty {
				// the queue is empty now - this para doesn't need attention anymore.
				cursor.remove();
			} else {
				cursor.advance();
			}
		}

		cursor.flush::<T>();
		queue_cache.flush::<T>();
	}
}

/// To avoid constant fetching, deserializing and serialization the queues are cached.
///
/// After an item dequeued from a queue for the first time, the queue is stored in this struct rather
/// than being serialized and persisted.
///
/// This implementation works best when:
///
/// 1. when the queues are shallow
/// 2. the dispatcher makes more than one cycle
///
/// if the queues are deep and there are many we would load and keep the queues for a long time,
/// thus increasing the peak memory consumption of the wasm runtime. Under such conditions persisting
/// queues might play better since it's unlikely that they are going to be requested once more.
///
/// On the other hand, the situation when deep queues exist and it takes more than one dipsatcher
/// cycle to traverse the queues is already sub-optimal and better be avoided.
///
/// This struct is not supposed to be dropped but rather to be consumed by [`flush`].
struct QueueCache(BTreeMap<ParaId, QueueCacheEntry>);

struct QueueCacheEntry {
	queue: VecDeque<UpwardMessage>,
	count: u32,
	total_size: u32,
}

impl QueueCache {
	fn new() -> Self {
		Self(BTreeMap::new())
	}

	/// Dequeues one item from the upward message queue of the given para.
	///
	/// Returns `(upward_message, became_empty)`, where
	///
	/// - `upward_message` a dequeued message or `None` if the queue _was_ empty.
	/// - `became_empty` is true if the queue _became_ empty.
404
	fn dequeue<T: Config>(&mut self, para: ParaId) -> (Option<UpwardMessage>, bool) {
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
		let cache_entry = self.0.entry(para).or_insert_with(|| {
			let queue = <Module<T> as Store>::RelayDispatchQueues::get(&para);
			let (count, total_size) = <Module<T> as Store>::RelayDispatchQueueSize::get(&para);
			QueueCacheEntry {
				queue,
				count,
				total_size,
			}
		});
		let upward_message = cache_entry.queue.pop_front();
		if let Some(ref msg) = upward_message {
			cache_entry.count -= 1;
			cache_entry.total_size -= msg.len() as u32;
		}

		let became_empty = cache_entry.queue.is_empty();
		(upward_message, became_empty)
	}

	/// Flushes the updated queues into the storage.
425
	fn flush<T: Config>(self) {
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
		// NOTE we use an explicit method here instead of Drop impl because it has unwanted semantics
		// within runtime. It is dangerous to use because of double-panics and flushing on a panic
		// is not necessary as well.
		for (
			para,
			QueueCacheEntry {
				queue,
				count,
				total_size,
			},
		) in self.0
		{
			if queue.is_empty() {
				// remove the entries altogether.
				<Module<T> as Store>::RelayDispatchQueues::remove(&para);
				<Module<T> as Store>::RelayDispatchQueueSize::remove(&para);
			} else {
				<Module<T> as Store>::RelayDispatchQueues::insert(&para, queue);
				<Module<T> as Store>::RelayDispatchQueueSize::insert(&para, (count, total_size));
			}
		}
	}
}

/// A cursor that iterates over all entries in `NeedsDispatch`.
///
/// This cursor will start with the para indicated by `NextDispatchRoundStartWith` storage entry.
/// This cursor is cyclic meaning that after reaching the end it will jump to the beginning. Unlike
/// an iterator, this cursor allows removing items during the iteration.
///
/// Each iteration cycle *must be* concluded with a call to either `advance` or `remove`.
///
/// This struct is not supposed to be dropped but rather to be consumed by [`flush`].
#[derive(Debug)]
struct NeedsDispatchCursor {
	needs_dispatch: Vec<ParaId>,
	cur_idx: usize,
}

impl NeedsDispatchCursor {
466
	fn new<T: Config>() -> Self {
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
		let needs_dispatch: Vec<ParaId> = <Module<T> as Store>::NeedsDispatch::get();
		let start_with = <Module<T> as Store>::NextDispatchRoundStartWith::get();

		let start_with_idx = match start_with {
			Some(para) => match needs_dispatch.binary_search(&para) {
				Ok(found_idx) => found_idx,
				Err(_supposed_idx) => {
					// well that's weird because we maintain an invariant that
					// `NextDispatchRoundStartWith` must point into one of the items in
					// `NeedsDispatch`.
					//
					// let's select 0 as the starting index as a safe bet.
					debug_assert!(false);
					0
				}
			},
			None => 0,
		};

		Self {
			needs_dispatch,
			cur_idx: start_with_idx,
		}
	}

	/// Returns the item the cursor points to.
	fn peek(&self) -> Option<ParaId> {
		self.needs_dispatch.get(self.cur_idx).cloned()
	}

	/// Moves the cursor to the next item.
	fn advance(&mut self) {
		if self.needs_dispatch.is_empty() {
			return;
		}
		self.cur_idx = (self.cur_idx + 1) % self.needs_dispatch.len();
	}

	/// Removes the item under the cursor.
	fn remove(&mut self) {
		if self.needs_dispatch.is_empty() {
			return;
		}
		let _ = self.needs_dispatch.remove(self.cur_idx);

		// we might've removed the last element and that doesn't necessarily mean that `needs_dispatch`
		// became empty. Reposition the cursor in this case to the beginning.
		if self.needs_dispatch.get(self.cur_idx).is_none() {
			self.cur_idx = 0;
		}
	}

	/// Flushes the dispatcher state into the persistent storage.
520
	fn flush<T: Config>(self) {
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
		let next_one = self.peek();
		<Module<T> as Store>::NextDispatchRoundStartWith::set(next_one);
		<Module<T> as Store>::NeedsDispatch::put(self.needs_dispatch);
	}
}

#[cfg(test)]
pub(crate) mod mock_sink {
	//! An implementation of a mock UMP sink that allows attaching a probe for mocking the weights
	//! and checking the sent messages.
	//!
	//! A default behavior of the UMP sink is to ignore an incoming message and return 0 weight.
	//!
	//! A probe can be attached to the mock UMP sink. When attached, the mock sink would consult the
	//! probe to check whether the received message was expected and what weight it should return.
	//!
	//! There are two rules on how to use a probe:
	//!
	//! 1. There can be only one active probe at a time. Creation of another probe while there is
	//!    already an active one leads to a panic. The probe is scoped to a thread where it was created.
	//!
	//! 2. All messages expected by the probe must be received by the time of dropping it. Unreceived
	//!    messages will lead to a panic while dropping a probe.

	use super::{UmpSink, UpwardMessage, ParaId};
	use std::cell::RefCell;
	use std::collections::vec_deque::VecDeque;
	use frame_support::weights::Weight;

	#[derive(Debug)]
	struct UmpExpectation {
		expected_origin: ParaId,
		expected_msg: UpwardMessage,
		mock_weight: Weight,
	}

	std::thread_local! {
		// `Some` here indicates that there is an active probe.
		static HOOK: RefCell<Option<VecDeque<UmpExpectation>>> = RefCell::new(None);
	}

	pub struct MockUmpSink;
	impl UmpSink for MockUmpSink {
		fn process_upward_message(actual_origin: ParaId, actual_msg: Vec<u8>) -> Weight {
			HOOK.with(|opt_hook| match &mut *opt_hook.borrow_mut() {
				Some(hook) => {
					let UmpExpectation {
						expected_origin,
						expected_msg,
						mock_weight,
					} = match hook.pop_front() {
						Some(expectation) => expectation,
						None => {
							panic!(
								"The probe is active but didn't expect the message:\n\n\t{:?}.",
								actual_msg,
							);
						}
					};
					assert_eq!(expected_origin, actual_origin);
					assert_eq!(expected_msg, actual_msg);
					mock_weight
				}
				None => 0,
			})
		}
	}

	pub struct Probe {
		_private: (),
	}

	impl Probe {
		pub fn new() -> Self {
			HOOK.with(|opt_hook| {
				let prev = opt_hook.borrow_mut().replace(VecDeque::default());

				// that can trigger if there were two probes were created during one session which
				// is may be a bit strict, but may save time figuring out what's wrong.
				// if you land here and you do need the two probes in one session consider
				// dropping the the existing probe explicitly.
				assert!(prev.is_none());
			});
			Self { _private: () }
		}

		/// Add an expected message.
		///
		/// The enqueued messages are processed in FIFO order.
		pub fn assert_msg(
			&mut self,
			expected_origin: ParaId,
			expected_msg: UpwardMessage,
			mock_weight: Weight,
		) {
			HOOK.with(|opt_hook| {
				opt_hook
					.borrow_mut()
					.as_mut()
					.unwrap()
					.push_back(UmpExpectation {
						expected_origin,
						expected_msg,
						mock_weight,
					})
			});
		}
	}

	impl Drop for Probe {
		fn drop(&mut self) {
			let _ = HOOK.try_with(|opt_hook| {
				let prev = opt_hook.borrow_mut().take().expect(
					"this probe was created and hasn't been yet destroyed;
					the probe cannot be replaced;
					there is only one probe at a time allowed;
					thus it cannot be `None`;
					qed",
				);

				if !prev.is_empty() {
					// some messages are left unchecked. We should notify the developer about this.
					// however, we do so only if the thread doesn't panic already. Otherwise, the
					// developer would get a SIGILL or SIGABRT without a meaningful error message.
					if !std::thread::panicking() {
						panic!(
							"the probe is dropped and not all expected messages arrived: {:?}",
							prev
						);
					}
				}
			});
			// an `Err` here signals here that the thread local was already destroyed.
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use super::mock_sink::Probe;
662
	use crate::mock::{Configuration, Ump, new_test_ext, GenesisConfig as MockGenesisConfig};
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
	use frame_support::IterableStorageMap;
	use std::collections::HashSet;

	struct GenesisConfigBuilder {
		max_upward_message_size: u32,
		max_upward_message_num_per_candidate: u32,
		max_upward_queue_count: u32,
		max_upward_queue_size: u32,
		preferred_dispatchable_upward_messages_step_weight: Weight,
	}

	impl Default for GenesisConfigBuilder {
		fn default() -> Self {
			Self {
				max_upward_message_size: 16,
				max_upward_message_num_per_candidate: 2,
				max_upward_queue_count: 4,
				max_upward_queue_size: 64,
				preferred_dispatchable_upward_messages_step_weight: 1000,
			}
		}
	}

	impl GenesisConfigBuilder {
		fn build(self) -> crate::mock::GenesisConfig {
			let mut genesis = default_genesis_config();
			let config = &mut genesis.configuration.config;

			config.max_upward_message_size = self.max_upward_message_size;
			config.max_upward_message_num_per_candidate = self.max_upward_message_num_per_candidate;
			config.max_upward_queue_count = self.max_upward_queue_count;
			config.max_upward_queue_size = self.max_upward_queue_size;
			config.preferred_dispatchable_upward_messages_step_weight =
				self.preferred_dispatchable_upward_messages_step_weight;
			genesis
		}
	}

701
702
703
704
705
706
707
708
709
710
711
712
	fn default_genesis_config() -> MockGenesisConfig {
		MockGenesisConfig {
			configuration: crate::configuration::GenesisConfig {
				config: crate::configuration::HostConfiguration {
					max_downward_message_size: 1024,
					..Default::default()
				},
			},
			..Default::default()
		}
	}

713
714
	fn queue_upward_msg(para: ParaId, msg: UpwardMessage) {
		let msgs = vec![msg];
715
716
		assert!(Ump::check_upward_messages(&Configuration::config(), para, &msgs).is_ok());
		let _ = Ump::enact_upward_messages(para, msgs);
717
718
719
720
	}

	fn assert_storage_consistency_exhaustive() {
		// check that empty queues don't clutter the storage.
721
		for (_para, queue) in <Ump as Store>::RelayDispatchQueues::iter() {
722
723
724
725
			assert!(!queue.is_empty());
		}

		// actually count the counts and sizes in queues and compare them to the bookkeeped version.
726
727
		for (para, queue) in <Ump as Store>::RelayDispatchQueues::iter() {
			let (expected_count, expected_size) = <Ump as Store>::RelayDispatchQueueSize::get(para);
728
729
730
731
732
733
734
735
736
737
738
			let (actual_count, actual_size) =
				queue.into_iter().fold((0, 0), |(acc_count, acc_size), x| {
					(acc_count + 1, acc_size + x.len() as u32)
				});

			assert_eq!(expected_count, actual_count);
			assert_eq!(expected_size, actual_size);
		}

		// since we wipe the empty queues the sets of paras in queue contents, queue sizes and
		// need dispatch set should all be equal.
739
		let queue_contents_set = <Ump as Store>::RelayDispatchQueues::iter()
740
741
			.map(|(k, _)| k)
			.collect::<HashSet<ParaId>>();
742
		let queue_sizes_set = <Ump as Store>::RelayDispatchQueueSize::iter()
743
744
			.map(|(k, _)| k)
			.collect::<HashSet<ParaId>>();
745
		let needs_dispatch_set = <Ump as Store>::NeedsDispatch::get()
746
747
748
749
750
751
			.into_iter()
			.collect::<HashSet<ParaId>>();
		assert_eq!(queue_contents_set, queue_sizes_set);
		assert_eq!(queue_contents_set, needs_dispatch_set);

		// `NextDispatchRoundStartWith` should point into a para that is tracked.
752
		if let Some(para) = <Ump as Store>::NextDispatchRoundStartWith::get() {
753
754
755
756
			assert!(queue_contents_set.contains(&para));
		}

		// `NeedsDispatch` is always sorted.
757
758
759
760
761
		assert!(
			<Ump as Store>::NeedsDispatch::get()
				.windows(2)
				.all(|xs| xs[0] <= xs[1])
		);
762
763
764
765
766
767
768
769
	}

	#[test]
	fn dispatch_empty() {
		new_test_ext(default_genesis_config()).execute_with(|| {
			assert_storage_consistency_exhaustive();

			// make sure that the case with empty queues is handled properly
770
			Ump::process_pending_upward_messages();
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786

			assert_storage_consistency_exhaustive();
		});
	}

	#[test]
	fn dispatch_single_message() {
		let a = ParaId::from(228);
		let msg = vec![1, 2, 3];

		new_test_ext(GenesisConfigBuilder::default().build()).execute_with(|| {
			let mut probe = Probe::new();

			probe.assert_msg(a, msg.clone(), 0);
			queue_upward_msg(a, msg);

787
			Ump::process_pending_upward_messages();
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825

			assert_storage_consistency_exhaustive();
		});
	}

	#[test]
	fn dispatch_resume_after_exceeding_dispatch_stage_weight() {
		let a = ParaId::from(128);
		let c = ParaId::from(228);
		let q = ParaId::from(911);

		let a_msg_1 = vec![1, 2, 3];
		let a_msg_2 = vec![3, 2, 1];
		let c_msg_1 = vec![4, 5, 6];
		let c_msg_2 = vec![9, 8, 7];
		let q_msg = b"we are Q".to_vec();

		new_test_ext(
			GenesisConfigBuilder {
				preferred_dispatchable_upward_messages_step_weight: 500,
				..Default::default()
			}
			.build(),
		)
		.execute_with(|| {
			queue_upward_msg(q, q_msg.clone());
			queue_upward_msg(c, c_msg_1.clone());
			queue_upward_msg(a, a_msg_1.clone());
			queue_upward_msg(a, a_msg_2.clone());

			assert_storage_consistency_exhaustive();

			// we expect only two first messages to fit in the first iteration.
			{
				let mut probe = Probe::new();

				probe.assert_msg(a, a_msg_1.clone(), 300);
				probe.assert_msg(c, c_msg_1.clone(), 300);
826
				Ump::process_pending_upward_messages();
827
828
829
830
831
832
833
834
835
836
837
838
839
				assert_storage_consistency_exhaustive();

				drop(probe);
			}

			queue_upward_msg(c, c_msg_2.clone());
			assert_storage_consistency_exhaustive();

			// second iteration should process the second message.
			{
				let mut probe = Probe::new();

				probe.assert_msg(q, q_msg.clone(), 500);
840
				Ump::process_pending_upward_messages();
841
842
843
844
845
846
847
848
849
850
851
				assert_storage_consistency_exhaustive();

				drop(probe);
			}

			// 3rd iteration.
			{
				let mut probe = Probe::new();

				probe.assert_msg(a, a_msg_2.clone(), 100);
				probe.assert_msg(c, c_msg_2.clone(), 100);
852
				Ump::process_pending_upward_messages();
853
854
855
856
857
858
859
860
861
				assert_storage_consistency_exhaustive();

				drop(probe);
			}

			// finally, make sure that the queue is empty.
			{
				let probe = Probe::new();

862
				Ump::process_pending_upward_messages();
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
				assert_storage_consistency_exhaustive();

				drop(probe);
			}
		});
	}

	#[test]
	fn dispatch_correctly_handle_remove_of_latest() {
		let a = ParaId::from(1991);
		let b = ParaId::from(1999);

		let a_msg_1 = vec![1, 2, 3];
		let a_msg_2 = vec![3, 2, 1];
		let b_msg_1 = vec![4, 5, 6];

		new_test_ext(
			GenesisConfigBuilder {
				preferred_dispatchable_upward_messages_step_weight: 900,
				..Default::default()
			}
			.build(),
		)
		.execute_with(|| {
			// We want to test here an edge case, where we remove the queue with the highest
			// para id (i.e. last in the needs_dispatch order).
			//
			// If the last entry was removed we should proceed execution, assuming we still have
			// weight available.

			queue_upward_msg(a, a_msg_1.clone());
			queue_upward_msg(a, a_msg_2.clone());
			queue_upward_msg(b, b_msg_1.clone());

			{
				let mut probe = Probe::new();

				probe.assert_msg(a, a_msg_1.clone(), 300);
				probe.assert_msg(b, b_msg_1.clone(), 300);
				probe.assert_msg(a, a_msg_2.clone(), 300);

904
				Ump::process_pending_upward_messages();
905
906
907
908
909
910

				drop(probe);
			}
		});
	}
}