Forcing-aware offchain Phragmén. (#5580)

* Make it force-aware

* Fix merge issues

Author: kianenigma

frame/staking/src/lib.rs

@@ -1010,7 +1010,8 @@ decl_storage! {
 		/// solutions to be submitted.
 		pub EraElectionStatus get(fn era_election_status): ElectionStatus<T::BlockNumber>;
 
-		/// True if the current planned session is final.
+		/// True if the current **planned** session is final. Note that this does not take era
+		/// forcing into account.
 		pub IsCurrentSessionFinal get(fn is_current_session_final): bool = false;
 
 		/// True if network has been upgraded to this version.
@@ -1166,7 +1167,8 @@ decl_module! {
 			if
 				// if we don't have any ongoing offchain compute.
 				Self::era_election_status().is_closed() &&
-				Self::is_current_session_final()
+				// 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(remaining) = next_session_change.checked_sub(&now) {
@@ -2895,6 +2897,13 @@ impl<T: Trait> Module<T> {
 		}
 	}
 
+	fn will_era_be_forced() -> bool {
+		match ForceEra::get() {
+			Forcing::ForceAlways | Forcing::ForceNew => true,
+			Forcing::ForceNone | Forcing::NotForcing => false,
+		}
+	}
+
 	#[cfg(feature = "runtime-benchmarks")]
 	pub fn add_era_stakers(current_era: EraIndex, controller: T::AccountId, exposure: Exposure<T::AccountId, BalanceOf<T>>) {
 		<ErasStakers<T>>::insert(&current_era, &controller, &exposure);

frame/staking/src/tests.rs

@@ -2872,6 +2872,27 @@ mod offchain_phragmen {
 			})
 	}
 
+	#[test]
+	fn offchain_election_flag_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));
+			})
+	}
+
 	#[test]
 	fn election_on_chain_fallback_works() {
 		ExtBuilder::default().build().execute_with(|| {