Fix Election when ForceNone V1 (paritytech#6166)

* Clean

* Better doc

* Better better doc

* Again better doc

* Fix indemt

* Update frame/staking/src/lib.rs

* Update frame/staking/src/lib.rs

* Better test

Co-authored-by: Gavin Wood <gavin@parity.io>

frame/staking/src/lib.rs

@@ -1288,7 +1288,7 @@ decl_module! {
 				// either current session final based on the plan, or we're forcing.
 				(Self::is_current_session_final() || Self::will_era_be_forced())
 			{
-				if let Some(next_session_change) = T::NextNewSession::estimate_next_new_session(now){
+				if let Some(next_session_change) = T::NextNewSession::estimate_next_new_session(now) {
 					if let Some(remaining) = next_session_change.checked_sub(&now) {
 						if remaining <= T::ElectionLookahead::get() && !remaining.is_zero() {
 							// create snapshot.
@@ -2569,16 +2569,19 @@ impl<T: Trait> Module<T> {
 				Forcing::ForceAlways => (),
 				Forcing::NotForcing if era_length >= T::SessionsPerEra::get() => (),
 				_ => {
-					// not forcing, not a new era either. If final, set the flag.
-					if era_length + 1 >= T::SessionsPerEra::get() {
+					// Either `ForceNone`, or `NotForcing && era_length < T::SessionsPerEra::get()`.
+					if era_length + 1 == T::SessionsPerEra::get() {
 						IsCurrentSessionFinal::put(true);
+					} else if era_length >= T::SessionsPerEra::get() {
+						// Should only happen when we are ready to trigger an era but we have ForceNone,
+						// otherwise previous arm would short circuit.
+						Self::close_election_window();
 					}
 					return None
 				},
 			}
 
 			// new era.
-			IsCurrentSessionFinal::put(false);
 			Self::new_era(session_index)
 		} else {
 			// Set initial era
@@ -2912,6 +2915,17 @@ impl<T: Trait> Module<T> {
 		maybe_new_validators
 	}
 
+
+	/// Remove all the storage items associated with the election.
+	fn close_election_window() {
+		// Close window.
+		<EraElectionStatus<T>>::put(ElectionStatus::Closed);
+		// Kill snapshots.
+		Self::kill_stakers_snapshot();
+		// Don't track final session.
+		IsCurrentSessionFinal::put(false);
+	}
+
 	/// Select the new validator set at the end of the era.
 	///
 	/// Runs [`try_do_phragmen`] and updates the following storage items:
@@ -2933,11 +2947,8 @@ impl<T: Trait> Module<T> {
 			exposures,
 			compute,
 		}) = Self::try_do_phragmen() {
-			// We have chosen the new validator set. Submission is no longer allowed.
-			<EraElectionStatus<T>>::put(ElectionStatus::Closed);
-
-			// kill the snapshots.
-			Self::kill_stakers_snapshot();
+			// Totally close the election round and data.
+			Self::close_election_window();
 
 			// Populate Stakers and write slot stake.
 			let mut total_stake: BalanceOf<T> = Zero::zero();

frame/staking/src/tests.rs

@@ -2859,7 +2859,7 @@ mod offchain_phragmen {
 	}
 
 	#[test]
-	fn offchain_election_flag_is_triggered() {
+	fn offchain_window_is_triggered() {
 		ExtBuilder::default()
 			.session_per_era(5)
 			.session_length(10)
@@ -2919,28 +2919,104 @@ mod offchain_phragmen {
 	}
 
 	#[test]
-	fn offchain_election_flag_is_triggered_when_forcing() {
+	fn offchain_window_is_triggered_when_forcing() {
 		ExtBuilder::default()
 			.session_per_era(5)
 			.session_length(10)
 			.election_lookahead(3)
 			.build()
 			.execute_with(|| {
-				run_to_block(7);
-				assert_session_era!(0, 0);
-
 				run_to_block(12);
 				ForceEra::put(Forcing::ForceNew);
 				run_to_block(13);
 				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
 
 				run_to_block(17); // instead of 47
 				assert_eq!(Staking::era_election_status(), ElectionStatus::Open(17));
+
+				run_to_block(20);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+			})
+	}
+
+	#[test]
+	fn offchain_window_is_triggered_when_force_always() {
+		ExtBuilder::default()
+			.session_per_era(5)
+			.session_length(10)
+			.election_lookahead(3)
+			.build()
+			.execute_with(|| {
+
+				ForceEra::put(Forcing::ForceAlways);
+				run_to_block(16);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+
+				run_to_block(17); // instead of 37
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Open(17));
+
+				run_to_block(20);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+
+				run_to_block(26);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+
+				run_to_block(27); // next one again
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Open(27));
+			})
+	}
+
+	#[test]
+	fn offchain_window_closes_when_forcenone() {
+		ExtBuilder::default()
+			.session_per_era(5)
+			.session_length(10)
+			.election_lookahead(3)
+			.build()
+			.execute_with(|| {
+				ForceEra::put(Forcing::ForceNone);
+
+				run_to_block(36);
+				assert_session_era!(3, 0);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+
+				// opens
+				run_to_block(37);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Open(37));
+				assert!(Staking::is_current_session_final());
+				assert!(Staking::snapshot_validators().is_some());
+
+				// closes normally
+				run_to_block(40);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+				assert!(!Staking::is_current_session_final());
+				assert!(Staking::snapshot_validators().is_none());
+				assert_session_era!(4, 0);
+
+				run_to_block(47);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+				assert_session_era!(4, 0);
+
+				run_to_block(57);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+				assert_session_era!(5, 0);
+
+				run_to_block(67);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+
+				// Will not open again as scheduled
+				run_to_block(87);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+				assert_session_era!(8, 0);
+
+				run_to_block(90);
+				assert_eq!(Staking::era_election_status(), ElectionStatus::Closed);
+				assert_session_era!(9, 0);
 			})
 	}
 
 	#[test]
-	fn election_on_chain_fallback_works() {
+	fn offchain_window_on_chain_fallback_works() {
 		ExtBuilder::default().build_and_execute(|| {
 			start_session(1);
 			start_session(2);