Fix session rotation properly and add test

substrate/runtime/session/src/lib.rs

@@ -79,6 +79,7 @@ decl_module! {
 		fn force_new_session(normal_rotation: bool) -> Result = 1;
 	}
 }
+
 decl_storage! {
 	trait Store for Module<T: Trait>;
 
@@ -93,6 +94,8 @@ decl_storage! {
 	// Percent by which the session must necessarily finish late before we early-exit the session.
 	pub BrokenPercentLate get(broken_percent_late): b"ses:broken_percent_late" => required T::Moment;
 
+	// New session is being forced.
+	pub ForcingNewSession get(forcing_new_session): b"ses:forcing_new_session" => bool;
 	// Block at which the session length last changed.
 	LastLengthChange: b"ses:llc" => T::BlockNumber;
 	// The next key for a given validator.
@@ -127,8 +130,8 @@ impl<T: Trait> Module<T> {
 	}
 
 	/// Forces a new session.
-	fn force_new_session(normal_rotation: bool) -> Result {
-		Self::rotate_session(normal_rotation);
+	pub fn force_new_session(normal_rotation: bool) -> Result {
+		<ForcingNewSession<T>>::put(normal_rotation);
 		Ok(())
 	}
 
@@ -153,13 +156,16 @@ impl<T: Trait> Module<T> {
 		let block_number = <system::Module<T>>::block_number();
 		let is_final_block = ((block_number - Self::last_length_change()) % Self::length()).is_zero();
 		let broken_validation = Self::broken_validation();
-		if is_final_block || broken_validation {
-			Self::rotate_session(!broken_validation);
+		if let Some(normal_rotation) = Self::forcing_new_session() {
+			Self::rotate_session(normal_rotation, is_final_block);
+			<ForcingNewSession<T>>::kill();
+		} else if is_final_block || broken_validation {
+			Self::rotate_session(!broken_validation, is_final_block);
 		}
 	}
 
 	/// Move onto next session: register the new authority set.
-	pub fn rotate_session(normal_rotation: bool) {
+	pub fn rotate_session(normal_rotation: bool, is_final_block: bool) {
 		let now = <timestamp::Module<T>>::get();
 		let time_elapsed = now.clone() - Self::current_start();
 
@@ -174,7 +180,7 @@ impl<T: Trait> Module<T> {
 		} else {
 			false
 		};
-		if len_changed || !normal_rotation {
+		if len_changed || !is_final_block {
 			let block_number = <system::Module<T>>::block_number();
 			<LastLengthChange<T>>::put(block_number);
 		}
@@ -359,6 +365,40 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn should_work_with_early_exit() {
+		with_externalities(&mut new_test_ext(), || {
+			System::set_block_number(1);
+			assert_ok!(Session::set_length(10));
+			assert_eq!(Session::blocks_remaining(), 1);
+			Session::check_rotate_session();
+
+			System::set_block_number(2);
+			assert_eq!(Session::blocks_remaining(), 0);
+			Session::check_rotate_session();
+			assert_eq!(Session::length(), 10);
+
+			System::set_block_number(7);
+			assert_eq!(Session::current_index(), 1);
+			assert_eq!(Session::blocks_remaining(), 5);
+			assert_ok!(Session::force_new_session(false));
+			Session::check_rotate_session();
+
+			System::set_block_number(8);
+			assert_eq!(Session::current_index(), 2);
+			assert_eq!(Session::blocks_remaining(), 9);
+			Session::check_rotate_session();
+
+			System::set_block_number(17);
+			assert_eq!(Session::current_index(), 2);
+			assert_eq!(Session::blocks_remaining(), 0);
+			Session::check_rotate_session();
+
+			System::set_block_number(18);
+			assert_eq!(Session::current_index(), 3);
+		});
+	}
+
 	#[test]
 	fn session_length_change_should_work() {
 		with_externalities(&mut new_test_ext(), || {

substrate/runtime/staking/src/lib.rs

@@ -147,7 +147,7 @@ decl_module! {
 		fn set_sessions_per_era(new: T::BlockNumber) -> Result = 0;
 		fn set_bonding_duration(new: T::BlockNumber) -> Result = 1;
 		fn set_validator_count(new: u32) -> Result = 2;
-		fn force_new_era() -> Result = 3;
+		fn force_new_era(should_slash: bool) -> Result = 3;
 	}
 }
 
@@ -204,6 +204,9 @@ decl_storage! {
 	// The enumeration sets.
 	pub EnumSet get(enum_set): b"sta:enum_set" => default map [ T::AccountIndex => Vec<T::AccountId> ];
 
+	// We are forcing a new era.
+	pub ForcingNewEra get(forcing_new_era): b"sta:forcing_new_era" => ();
+
 	// The "free" balance of a given account.
 	//
 	// This is the only balance that matters in terms of most operations on tokens. It is
@@ -414,10 +417,11 @@ impl<T: Trait> Module<T> {
 		Ok(())
 	}
 
-	/// Force there to be a new era. This also forces a new session immediately after.
-	fn force_new_era() -> Result {
-		<session::Module<T>>::rotate_session(false);
-		Ok(())
+	/// Force there to be a new era. This also forces a new session immediately after by
+	/// setting `normal_rotation` to be false. Validators will get slashed.
+	fn force_new_era(should_slash: bool) -> Result {
+		<ForcingNewEra<T>>::put(());
+		<session::Module<T>>::force_new_session(!should_slash)
 	}
 
 	// PUBLIC MUTABLES (DANGEROUS)
@@ -590,7 +594,10 @@ impl<T: Trait> Module<T> {
 				}
 			}
 		}
-		if ((session_index - Self::last_era_length_change()) % Self::sessions_per_era()).is_zero() || !normal_rotation {
+		if <ForcingNewEra<T>>::take().is_some()
+			|| ((session_index - Self::last_era_length_change()) % Self::sessions_per_era()).is_zero()
+			|| !normal_rotation
+		{
 			Self::new_era();
 		}
 	}