From a3bca4bb65fdbfef99b52b06181779c0f681d3ea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=C3=B3nal=20Murray?= <donal.murray@parity.io>
Date: Mon, 21 Oct 2024 19:19:22 +0200
Subject: [PATCH] [Coretime chain] Add high assignment count mitigation to
 testnets (#6022)

Fixed in Polkadot and Kusama in
https://github.com/polkadot-fellows/runtimes/pull/434 and this PR just
adds to testnets.

We can handle a maximum of 28 assignments inside one XCM, while it's
possible to have 80 (if a region is interlaced 79 times).
This can be chunked on the coretime chain side but currently the
scheduler on the relay makes assumptions that means we can't send more
than one chunk for a given core.

This just truncates the additional assignments until we can extend the
relay to support this. This means that the first 27 assignments are
taken, the final 28th is used to pad with idle to complete the mask (the
relay also assumes that every schedule is complete). Any other
assignments are dropped.

---------

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
---
 .../src/tests/coretime_interface.rs           |  2 +-
 .../src/tests/coretime_interface.rs           |  2 +-
 .../coretime/coretime-rococo/src/coretime.rs  | 30 ++++++++++++++++
 .../coretime/coretime-westend/src/coretime.rs | 34 +++++++++++++++++--
 prdoc/pr_6022.prdoc                           | 14 ++++++++
 5 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 prdoc/pr_6022.prdoc

diff --git a/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-rococo/src/tests/coretime_interface.rs b/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-rococo/src/tests/coretime_interface.rs
index 584bce8f1df..9915b1753ef 100644
--- a/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-rococo/src/tests/coretime_interface.rs
+++ b/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-rococo/src/tests/coretime_interface.rs
@@ -46,7 +46,7 @@ fn transact_hardcoded_weights_are_sane() {
 
 		// Create and populate schedule with the worst case assignment on this core.
 		let mut schedule = Vec::new();
-		for i in 0..27 {
+		for i in 0..80 {
 			schedule.push(ScheduleItem {
 				mask: CoreMask::void().set(i),
 				assignment: CoreAssignment::Task(2000 + i),
diff --git a/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-westend/src/tests/coretime_interface.rs b/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-westend/src/tests/coretime_interface.rs
index f61bc4285a0..00530f80b95 100644
--- a/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-westend/src/tests/coretime_interface.rs
+++ b/cumulus/parachains/integration-tests/emulated/tests/coretime/coretime-westend/src/tests/coretime_interface.rs
@@ -46,7 +46,7 @@ fn transact_hardcoded_weights_are_sane() {
 
 		// Create and populate schedule with the worst case assignment on this core.
 		let mut schedule = Vec::new();
-		for i in 0..27 {
+		for i in 0..80 {
 			schedule.push(ScheduleItem {
 				mask: CoreMask::void().set(i),
 				assignment: CoreAssignment::Task(2000 + i),
diff --git a/cumulus/parachains/runtimes/coretime/coretime-rococo/src/coretime.rs b/cumulus/parachains/runtimes/coretime/coretime-rococo/src/coretime.rs
index 76ee06a87e8..3910a747e9b 100644
--- a/cumulus/parachains/runtimes/coretime/coretime-rococo/src/coretime.rs
+++ b/cumulus/parachains/runtimes/coretime/coretime-rococo/src/coretime.rs
@@ -218,6 +218,36 @@ impl CoretimeInterface for CoretimeAllocator {
 		end_hint: Option<RCBlockNumberOf<Self>>,
 	) {
 		use crate::coretime::CoretimeProviderCalls::AssignCore;
+
+		// The relay chain currently only allows `assign_core` to be called with a complete mask
+		// and only ever with increasing `begin`. The assignments must be truncated to avoid
+		// dropping that core's assignment completely.
+
+		// This shadowing of `assignment` is temporary and can be removed when the relay can accept
+		// multiple messages to assign a single core.
+		let assignment = if assignment.len() > 28 {
+			let mut total_parts = 0u16;
+			// Account for missing parts with a new `Idle` assignment at the start as
+			// `assign_core` on the relay assumes this is sorted. We'll add the rest of the
+			// assignments and sum the parts in one pass, so this is just initialized to 0.
+			let mut assignment_truncated = vec![(CoreAssignment::Idle, 0)];
+			// Truncate to first 27 non-idle assignments.
+			assignment_truncated.extend(
+				assignment
+					.into_iter()
+					.filter(|(a, _)| *a != CoreAssignment::Idle)
+					.take(27)
+					.inspect(|(_, parts)| total_parts += *parts)
+					.collect::<Vec<_>>(),
+			);
+
+			// Set the parts of the `Idle` assignment we injected at the start of the vec above.
+			assignment_truncated[0].1 = 57_600u16.saturating_sub(total_parts);
+			assignment_truncated
+		} else {
+			assignment
+		};
+
 		let assign_core_call =
 			RelayRuntimePallets::Coretime(AssignCore(core, begin, assignment, end_hint));
 
diff --git a/cumulus/parachains/runtimes/coretime/coretime-westend/src/coretime.rs b/cumulus/parachains/runtimes/coretime/coretime-westend/src/coretime.rs
index 865ff68d4c6..86769cb2da1 100644
--- a/cumulus/parachains/runtimes/coretime/coretime-westend/src/coretime.rs
+++ b/cumulus/parachains/runtimes/coretime/coretime-westend/src/coretime.rs
@@ -224,8 +224,6 @@ impl CoretimeInterface for CoretimeAllocator {
 		end_hint: Option<RCBlockNumberOf<Self>>,
 	) {
 		use crate::coretime::CoretimeProviderCalls::AssignCore;
-		let assign_core_call =
-			RelayRuntimePallets::Coretime(AssignCore(core, begin, assignment, end_hint));
 
 		// Weight for `assign_core` from westend benchmarks:
 		// `ref_time` = 10177115 + (1 * 25000000) + (2 * 100000000) + (57600 * 13932) = 937660315
@@ -233,6 +231,38 @@ impl CoretimeInterface for CoretimeAllocator {
 		// Add 5% to each component and round to 2 significant figures.
 		let call_weight = Weight::from_parts(980_000_000, 3800);
 
+		// The relay chain currently only allows `assign_core` to be called with a complete mask
+		// and only ever with increasing `begin`. The assignments must be truncated to avoid
+		// dropping that core's assignment completely.
+
+		// This shadowing of `assignment` is temporary and can be removed when the relay can accept
+		// multiple messages to assign a single core.
+		let assignment = if assignment.len() > 28 {
+			let mut total_parts = 0u16;
+			// Account for missing parts with a new `Idle` assignment at the start as
+			// `assign_core` on the relay assumes this is sorted. We'll add the rest of the
+			// assignments and sum the parts in one pass, so this is just initialized to 0.
+			let mut assignment_truncated = vec![(CoreAssignment::Idle, 0)];
+			// Truncate to first 27 non-idle assignments.
+			assignment_truncated.extend(
+				assignment
+					.into_iter()
+					.filter(|(a, _)| *a != CoreAssignment::Idle)
+					.take(27)
+					.inspect(|(_, parts)| total_parts += *parts)
+					.collect::<Vec<_>>(),
+			);
+
+			// Set the parts of the `Idle` assignment we injected at the start of the vec above.
+			assignment_truncated[0].1 = 57_600u16.saturating_sub(total_parts);
+			assignment_truncated
+		} else {
+			assignment
+		};
+
+		let assign_core_call =
+			RelayRuntimePallets::Coretime(AssignCore(core, begin, assignment, end_hint));
+
 		let message = Xcm(vec![
 			Instruction::UnpaidExecution {
 				weight_limit: WeightLimit::Unlimited,
diff --git a/prdoc/pr_6022.prdoc b/prdoc/pr_6022.prdoc
new file mode 100644
index 00000000000..804d46af661
--- /dev/null
+++ b/prdoc/pr_6022.prdoc
@@ -0,0 +1,14 @@
+title: '[Coretime chain] Add high assignment count mitigation to testnets'
+doc:
+- audience: Runtime User
+  description: |
+    We can handle a maximum of 28 assignments inside one XCM, while it's possible to have 80 (if a
+    region is interlaced 79 times). This can be chunked on the coretime chain side but currently the
+    relay does not support this. This PR truncates the additional assignments on Rococo and Westend
+    to mitigate this until the relay is fixed. The first 27 assignments are taken, the final 28th is
+    used to pad with idle to complete the mask. Any other assignments are dropped.
+crates:
+- name: coretime-rococo-runtime
+  bump: patch
+- name: coretime-westend-runtime
+  bump: patch
-- 
GitLab