Clean up session key rotation (#1911)

* Clean up session key rotation

* Fix build

* Bump version

substrate/node/runtime/src/lib.rs

@@ -60,8 +60,8 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("node"),
 	impl_name: create_runtime_str!("substrate-node"),
 	authoring_version: 10,
-	spec_version: 31,
-	impl_version: 33,
+	spec_version: 32,
+	impl_version: 32,
 	apis: RUNTIME_API_VERSIONS,
 };
 
@@ -101,7 +101,7 @@ impl indices::Trait for Runtime {
 
 impl balances::Trait for Runtime {
 	type Balance = Balance;
-	type OnFreeBalanceZero = ((Staking, Contract), Democracy);
+	type OnFreeBalanceZero = (((Staking, Contract), Democracy), Session);
 	type OnNewAccount = Indices;
 	type EnsureAccountLiquid = (Staking, Democracy);
 	type Event = Event;

substrate/srml/session/src/lib.rs

@@ -22,7 +22,7 @@
 use rstd::prelude::*;
 use primitives::traits::{As, Zero, One, Convert};
 use srml_support::{StorageValue, StorageMap, for_each_tuple, decl_module, decl_event, decl_storage};
-use srml_support::dispatch::Result;
+use srml_support::{dispatch::Result, traits::OnFreeBalanceZero};
 use system::ensure_signed;
 use rstd::ops::Mul;
 
@@ -141,13 +141,10 @@ impl<T: Trait> Module<T> {
 
 	/// Set the current set of validators.
 	///
-	/// Called by `staking::new_era()` only. `next_session` should be called after this in order to
+	/// Called by `staking::new_era()` only. `rotate_session` must be called after this in order to
 	/// update the session keys to the next validator set.
 	pub fn set_validators(new: &[T::AccountId]) {
 		<Validators<T>>::put(&new.to_vec());
-		<consensus::Module<T>>::set_authorities(
-			&new.iter().cloned().map(T::ConvertAccountIdToSessionKey::convert).collect::<Vec<_>>()
-		);
 	}
 
 	/// Hook to be called after transaction processing.
@@ -190,11 +187,13 @@ impl<T: Trait> Module<T> {
 		T::OnSessionChange::on_session_change(time_elapsed, apply_rewards);
 
 		// Update any changes in session keys.
-		Self::validators().iter().enumerate().for_each(|(i, v)| {
-			if let Some(n) = <NextKeyFor<T>>::take(v) {
-				<consensus::Module<T>>::set_authority(i as u32, &n);
-			}
-		});
+		for (i, v) in Self::validators().into_iter().enumerate() {
+			<consensus::Module<T>>::set_authority(
+				i as u32,
+				&<NextKeyFor<T>>::get(&v)
+					.unwrap_or_else(|| T::ConvertAccountIdToSessionKey::convert(v))
+			);
+		};
 	}
 
 	/// Get the time that should have elapsed over a session if everything was working perfectly.
@@ -215,6 +214,12 @@ impl<T: Trait> Module<T> {
 	}
 }
 
+impl<T: Trait> OnFreeBalanceZero<T::AccountId> for Module<T> {
+	fn on_free_balance_zero(who: &T::AccountId) {
+		<NextKeyFor<T>>::remove(who);
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;