Era change can be forced on and forced off (#3374)

* Forcing can be on or off.

* Add a testcase.

substrate/srml/staking/src/lib.rs

@@ -555,6 +555,22 @@ pub trait Trait: system::Trait {
 	type SessionInterface: self::SessionInterface<Self::AccountId>;
 }
 
+/// Mode of era-forcing.
+#[derive(Copy, Clone, PartialEq, Eq, Encode, Decode)]
+#[cfg_attr(feature = "std", derive(Debug, Serialize, Deserialize))]
+pub enum Forcing {
+	/// Not forcing anything - just let whatever happen.
+	NotForcing,
+	/// Force a new era, then reset to `NotForcing` as soon as it is done.
+	ForceNew,
+	/// Avoid a new era indefinitely.
+	ForceNone,
+}
+
+impl Default for Forcing {
+	fn default() -> Self { Forcing::NotForcing }
+}
+
 decl_storage! {
 	trait Store for Module<T: Trait> as Staking {
 
@@ -625,7 +641,7 @@ decl_storage! {
 		pub RecentlyOffline get(recently_offline): Vec<(T::AccountId, T::BlockNumber, u32)>;
 
 		/// True if the next session change will be a new era regardless of index.
-		pub ForceNewEra get(forcing_new_era): bool;
+		pub ForceEra get(force_era) config(): Forcing;
 
 		/// A mapping from still-bonded eras to the first session index of that era.
 		BondedEras: Vec<(EraIndex, SessionIndex)>;
@@ -988,18 +1004,27 @@ decl_module! {
 
 		// ----- Root calls.
 
-		/// Force there to be a new era. This also forces a new session immediately after.
-		/// `apply_rewards` should be true for validators to get the session reward.
+		/// Force there to be no new eras indefinitely.
 		///
 		/// # <weight>
-		/// - Independent of the arguments.
-		/// - Triggers the Phragmen election. Expensive but not user-controlled.
-		/// - Depends on state: `O(|edges| * |validators|)`.
+		/// - No arguments.
+		/// # </weight>
+		#[weight = SimpleDispatchInfo::FixedOperational(10_000)]
+		fn force_no_eras(origin) {
+			ensure_root(origin)?;
+			ForceEra::put(Forcing::ForceNone);
+		}
+
+		/// Force there to be a new era at the end of the next session. After this, it will be
+		/// reset to normal (non-forced) behaviour.
+		///
+		/// # <weight>
+		/// - No arguments.
 		/// # </weight>
 		#[weight = SimpleDispatchInfo::FixedOperational(10_000)]
 		fn force_new_era(origin) {
 			ensure_root(origin)?;
-			Self::apply_force_new_era()
+			ForceEra::put(Forcing::ForceNew);
 		}
 
 		/// Set the offline slash grace period.
@@ -1129,16 +1154,17 @@ impl<T: Trait> Module<T> {
 	fn new_session(session_index: SessionIndex)
 		-> Option<(Vec<T::AccountId>, Vec<(T::AccountId, Exposure<T::AccountId, BalanceOf<T>>)>)>
 	{
-		if ForceNewEra::take() || session_index % T::SessionsPerEra::get() == 0 {
-			let validators = T::SessionInterface::validators();
-			let prior = validators.into_iter()
-				.map(|v| { let e = Self::stakers(&v); (v, e) })
-				.collect();
-
-			Self::new_era(session_index).map(move |new| (new, prior))
-		} else {
-			None
+		match ForceEra::get() {
+			Forcing::ForceNew => ForceEra::kill(),
+			Forcing::NotForcing if session_index % T::SessionsPerEra::get() == 0 => (),
+			_ => return None,
 		}
+		let validators = T::SessionInterface::validators();
+		let prior = validators.into_iter()
+			.map(|v| { let e = Self::stakers(&v); (v, e) })
+			.collect();
+
+		Self::new_era(session_index).map(move |new| (new, prior))
 	}
 
 	/// The era has changed - enact new staking set.
@@ -1335,10 +1361,6 @@ impl<T: Trait> Module<T> {
 		}
 	}
 
-	fn apply_force_new_era() {
-		ForceNewEra::put(true);
-	}
-
 	/// Remove all associated data of a stash account from the staking system.
 	///
 	/// This is called :

substrate/srml/staking/src/mock.rs

@@ -310,6 +310,7 @@ impl ExtBuilder {
 			offline_slash: Perbill::from_percent(5),
 			offline_slash_grace: 0,
 			invulnerables: vec![],
+			.. Default::default()
 		}.assimilate_storage(&mut storage);
 
 		let _ = session::GenesisConfig::<Test> {

substrate/srml/staking/src/tests.rs

@@ -108,7 +108,7 @@ fn no_offline_should_work() {
 		assert_eq!(Staking::slash_count(&10), 0);
 		assert_eq!(Balances::free_balance(&10), 1);
 		// New era is not being forced
-		assert!(!Staking::forcing_new_era());
+		assert_eq!(Staking::force_era(), Forcing::NotForcing);
 	});
 }
 
@@ -163,7 +163,7 @@ fn invulnerability_should_work() {
 		assert!(<Validators<Test>>::exists(&11));
 		// New era not being forced
 		// NOTE: new era is always forced once slashing happens -> new validators need to be chosen.
-		assert!(!Staking::forcing_new_era());
+		assert_eq!(Staking::force_era(), Forcing::NotForcing);
 	});
 }
 
@@ -884,6 +884,43 @@ fn session_and_eras_work() {
 	});
 }
 
+#[test]
+fn forcing_new_era_works() {
+	with_externalities(&mut ExtBuilder::default().build(),|| {
+		// normal flow of session.
+		assert_eq!(Staking::current_era(), 0);
+		start_session(0);
+		assert_eq!(Staking::current_era(), 0);
+		start_session(1);
+		assert_eq!(Staking::current_era(), 0);
+		start_session(2);
+		assert_eq!(Staking::current_era(), 1);
+
+		// no era change.
+		ForceEra::put(Forcing::ForceNone);
+		start_session(3);
+		assert_eq!(Staking::current_era(), 1);
+		start_session(4);
+		assert_eq!(Staking::current_era(), 1);
+		start_session(5);
+		assert_eq!(Staking::current_era(), 1);
+		start_session(6);
+		assert_eq!(Staking::current_era(), 1);
+
+		// back to normal
+		ForceEra::put(Forcing::NotForcing);
+		start_session(7);
+		assert_eq!(Staking::current_era(), 1);
+		start_session(8);
+		assert_eq!(Staking::current_era(), 2);
+
+		// forceful change
+		ForceEra::put(Forcing::ForceNew);
+		start_session(9);
+		assert_eq!(Staking::current_era(), 3);
+	});
+}
+
 #[test]
 fn cannot_transfer_staked_balance() {
 	// Tests that a stash account cannot transfer funds