From 676b295322afebb45b472d202d034a2992f336cb Mon Sep 17 00:00:00 2001
From: Xavier Lau <xavier@inv.cafe>
Date: Sat, 11 Mar 2023 10:46:06 +0800
Subject: [PATCH] Adjustable collator count

---
 pallet/ecdsa-authority/tests/tests.rs |  2 +-
 pallet/staking/src/lib.rs             | 25 +++++++++++++++++++++++++
 pallet/staking/tests/mock.rs          |  6 +++++-
 pallet/staking/tests/tests.rs         | 17 ++++++++++++++++-
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/pallet/ecdsa-authority/tests/tests.rs b/pallet/ecdsa-authority/tests/tests.rs
index ee63e1214..10dc3eadb 100644
--- a/pallet/ecdsa-authority/tests/tests.rs
+++ b/pallet/ecdsa-authority/tests/tests.rs
@@ -430,7 +430,7 @@ fn submit_new_message_root_signature() {
 		// Case 3.
 		let s_3 = sign(&k_3, &message.0);
 		assert_noop!(
-			EcdsaAuthority::submit_new_message_root_signature(RuntimeOrigin::signed(a_3), s_3,),
+			EcdsaAuthority::submit_new_message_root_signature(RuntimeOrigin::signed(a_3), s_3),
 			<Error<Runtime>>::NotAuthority
 		);
 
diff --git a/pallet/staking/src/lib.rs b/pallet/staking/src/lib.rs
index d38d9d861..5d2924c85 100644
--- a/pallet/staking/src/lib.rs
+++ b/pallet/staking/src/lib.rs
@@ -245,6 +245,8 @@ pub mod pallet {
 		NotStaker,
 		/// Target is not a collator.
 		TargetNotCollator,
+		/// Collator count mustn't be zero.
+		ZeroCollatorCount,
 	}
 
 	/// All staking ledgers.
@@ -327,6 +329,10 @@ pub mod pallet {
 	#[pallet::genesis_build]
 	impl<T: Config> GenesisBuild<T> for GenesisConfig<T> {
 		fn build(&self) {
+			if self.collator_count == 0 {
+				panic!("[pallet::staking] collator count mustn't be 0");
+			}
+
 			<SessionStartTime<T>>::put(self.now);
 			<ElapsedTime<T>>::put(self.elapsed_time);
 			<CollatorCount<T>>::put(self.collator_count);
@@ -570,6 +576,25 @@ pub mod pallet {
 
 			Ok(())
 		}
+
+		/// Set collator count.
+		///
+		/// This will apply to the incoming session.
+		///
+		/// Require root origin.
+		#[pallet::call_index(7)]
+		#[pallet::weight(0)]
+		pub fn set_collator_count(origin: OriginFor<T>, count: u32) -> DispatchResult {
+			ensure_root(origin)?;
+
+			if count == 0 {
+				return Err(<Error<T>>::ZeroCollatorCount)?;
+			}
+
+			<CollatorCount<T>>::put(count);
+
+			Ok(())
+		}
 	}
 	impl<T> Pallet<T>
 	where
diff --git a/pallet/staking/tests/mock.rs b/pallet/staking/tests/mock.rs
index 5a3e1325c..8ed031c89 100644
--- a/pallet/staking/tests/mock.rs
+++ b/pallet/staking/tests/mock.rs
@@ -290,7 +290,6 @@ fn initialize_block(number: u64) {
 	<AllPalletsWithSystem as OnInitialize<u64>>::on_initialize(number);
 }
 
-#[derive(Default)]
 pub struct ExtBuilder {
 	collator_count: u32,
 }
@@ -332,3 +331,8 @@ impl ExtBuilder {
 		ext
 	}
 }
+impl Default for ExtBuilder {
+	fn default() -> Self {
+		Self { collator_count: 1 }
+	}
+}
diff --git a/pallet/staking/tests/tests.rs b/pallet/staking/tests/tests.rs
index 209282298..1c1054d14 100644
--- a/pallet/staking/tests/tests.rs
+++ b/pallet/staking/tests/tests.rs
@@ -27,7 +27,7 @@ use darwinia_staking::*;
 use dc_types::{Balance, UNIT};
 // substrate
 use frame_support::{assert_noop, assert_ok, BoundedVec};
-use sp_runtime::{assert_eq_error_rate, Perbill};
+use sp_runtime::{assert_eq_error_rate, DispatchError, Perbill};
 
 #[test]
 fn stake_should_work() {
@@ -403,6 +403,21 @@ fn chill_should_work() {
 	});
 }
 
+#[test]
+fn set_collator_count_should_work() {
+	ExtBuilder::default().build().execute_with(|| {
+		assert_noop!(
+			Staking::set_collator_count(RuntimeOrigin::signed(1), 1),
+			DispatchError::BadOrigin
+		);
+		assert_noop!(
+			Staking::set_collator_count(RuntimeOrigin::root(), 0),
+			<Error<Runtime>>::ZeroCollatorCount
+		);
+		assert_ok!(Staking::set_collator_count(RuntimeOrigin::root(), 1));
+	});
+}
+
 #[test]
 fn power_should_work() {
 	ExtBuilder::default().build().execute_with(|| {