session: add handler for genesis session

node/cli/src/chain_spec.rs

@@ -166,13 +166,13 @@ fn staging_testnet_config_genesis() -> GenesisConfig {
 			key: endowed_accounts[0].clone(),
 		}),
 		babe: Some(BabeConfig {
-			authorities: initial_authorities.iter().map(|x| (x.3.clone(), 1)).collect(),
+			authorities: vec![],
 		}),
 		im_online: Some(ImOnlineConfig {
-			keys: initial_authorities.iter().map(|x| x.4.clone()).collect(),
+			keys: vec![],
 		}),
 		grandpa: Some(GrandpaConfig {
-			authorities: initial_authorities.iter().map(|x| (x.2.clone(), 1)).collect(),
+			authorities: vec![],
 		}),
 		membership_Instance1: Some(Default::default()),
 	}

srml/babe/src/lib.rs

@@ -292,6 +292,15 @@ impl<T: Trait> Module<T> {
 		this_randomness
 	}
 
+	fn set_authorities<'a, I: 'a>(authorities: I)
+		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
+	{
+		let authorities = authorities.map(|(_account, k)| {
+			(k, 1)
+		}).collect::<Vec<_>>();
+
+		Authorities::put(authorities);
+	}
 }
 
 impl<T: Trait> OnTimestampSet<T::Moment> for Module<T> {
@@ -300,6 +309,13 @@ impl<T: Trait> OnTimestampSet<T::Moment> for Module<T> {
 
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = AuthorityId;
+
+	fn on_genesis_session<'a, I: 'a>(validators: I)
+		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
+	{
+		Self::set_authorities(validators);
+	}
+
 	fn on_new_session<'a, I: 'a>(_changed: bool, validators: I, queued_validators: I)
 		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
 	{
@@ -311,11 +327,7 @@ impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 		EpochIndex::put(epoch_index);
 
 		// Update authorities.
-		let authorities = validators.map(|(_account, k)| {
-			(k, 1)
-		}).collect::<Vec<_>>();
-
-		Authorities::put(authorities);
+		Self::set_authorities(validators);
 
 		// Update epoch start slot.
 		let now = CurrentSlot::get();

srml/grandpa/src/lib.rs

@@ -346,12 +346,19 @@ impl<T: Trait> Module<T> {
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = AuthorityId;
 
+	fn on_genesis_session<'a, I: 'a>(validators: I)
+		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
+	{
+		let authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
+		Authorities::put(authorities);
+	}
+
 	fn on_new_session<'a, I: 'a>(changed: bool, validators: I, _queued_validators: I)
 		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
 	{
 		// instant changes
 		if changed {
-			let next_authorities = validators.map(|(_, k)| (k, 1u64)).collect::<Vec<_>>();
+			let next_authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
 			let last_authorities = <Module<T>>::grandpa_authorities();
 			if next_authorities != last_authorities {
 				if let Some((further_wait, median)) = <Stalled<T>>::take() {

srml/im-online/src/lib.rs

@@ -360,6 +360,12 @@ impl<T: Trait> Module<T> {
 impl<T: Trait> session::OneSessionHandler<T::AccountId> for Module<T> {
 	type Key = AuthorityId;
 
+	fn on_genesis_session<'a, I: 'a>(validators: I)
+		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
+	{
+		Keys::put(validators.map(|x| x.1).collect::<Vec<_>>());
+	}
+
 	fn on_new_session<'a, I: 'a>(_changed: bool, _validators: I, next_validators: I)
 		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
 	{

srml/session/src/lib.rs

@@ -187,6 +187,12 @@ impl<A> OnSessionEnding<A> for () {
 
 /// Handler for when a session keys set changes.
 pub trait SessionHandler<ValidatorId> {
+	/// The given validator set will be used for the genesis session.
+	/// It is guaranteed that the given validator set will also be used
+	/// for the second session, therefore the first call to `on_new_session`
+	/// should provide the same validator set.
+	fn on_genesis_session<Ks: OpaqueKeys>(validators: &[(ValidatorId, Ks)]);
+
 	/// Session set has changed; act appropriately.
 	fn on_new_session<Ks: OpaqueKeys>(
 		changed: bool,
@@ -203,21 +209,35 @@ pub trait OneSessionHandler<ValidatorId> {
 	/// The key type expected.
 	type Key: Decode + Default + AppKey;
 
+	fn on_genesis_session<'a, I: 'a>(validators: I)
+		where I: Iterator<Item=(&'a ValidatorId, Self::Key)>, ValidatorId: 'a;
+
 	fn on_new_session<'a, I: 'a>(changed: bool, validators: I, queued_validators: I)
 		where I: Iterator<Item=(&'a ValidatorId, Self::Key)>, ValidatorId: 'a;
+
 	fn on_disabled(i: usize);
 }
 
 macro_rules! impl_session_handlers {
 	() => (
 		impl<AId> SessionHandler<AId> for () {
+			fn on_genesis_session<Ks: OpaqueKeys>(_: &[(AId, Ks)]) {}
 			fn on_new_session<Ks: OpaqueKeys>(_: bool, _: &[(AId, Ks)], _: &[(AId, Ks)]) {}
 			fn on_disabled(_: usize) {}
 		}
 	);
 
 	( $($t:ident)* ) => {
 		impl<AId, $( $t: OneSessionHandler<AId> ),*> SessionHandler<AId> for ( $( $t , )* ) {
+			fn on_genesis_session<Ks: OpaqueKeys>(validators: &[(AId, Ks)]) {
+				$(
+					let our_keys: Box<dyn Iterator<Item=_>> = Box::new(validators.iter()
+						.map(|k| (&k.0, k.1.get::<$t::Key>(<$t::Key as AppKey>::ID)
+							.unwrap_or_default())));
+
+					$t::on_genesis_session(our_keys);
+				)*
+			}
 			fn on_new_session<Ks: OpaqueKeys>(
 				changed: bool,
 				validators: &[(AId, Ks)],
@@ -348,6 +368,9 @@ decl_storage! {
 					))
 					.collect();
 
+				// Tell everyone about the genesis session keys
+				T::SessionHandler::on_genesis_session::<T::Keys>(&queued_keys);
+
 				<Validators<T>>::put(initial_validators);
 				<QueuedKeys<T>>::put(queued_keys);
 			});