2323// Ensure we're `no_std` when compiling for Wasm.
2424#![ cfg_attr( not( feature = "std" ) , no_std) ]
2525
26- use sp_std:: prelude:: * ;
26+ use sp_std:: { collections :: btree_set :: BTreeSet , prelude:: * } ;
2727use frame_support:: { decl_module, decl_storage} ;
2828use sp_authority_discovery:: AuthorityId ;
2929
@@ -32,7 +32,7 @@ pub trait Trait: frame_system::Trait + pallet_session::Trait {}
3232
3333decl_storage ! {
3434 trait Store for Module <T : Trait > as AuthorityDiscovery {
35- /// Keys of the current authority set.
35+ /// Keys of the current and next authority set.
3636 Keys get( fn keys) : Vec <AuthorityId >;
3737 }
3838 add_extra_genesis {
@@ -47,7 +47,7 @@ decl_module! {
4747}
4848
4949impl < T : Trait > Module < T > {
50- /// Retrieve authority identifiers of the current authority set.
50+ /// Retrieve authority identifiers of the current and next authority set.
5151 pub fn authorities ( ) -> Vec < AuthorityId > {
5252 Keys :: get ( )
5353 }
@@ -71,17 +71,17 @@ impl<T: Trait> pallet_session::OneSessionHandler<T::AccountId> for Module<T> {
7171 where
7272 I : Iterator < Item = ( & ' a T :: AccountId , Self :: Key ) > ,
7373 {
74- let keys = authorities. map ( |x| x. 1 ) . collect :: < Vec < _ > > ( ) ;
75- Self :: initialize_keys ( & keys) ;
74+ Self :: initialize_keys ( & authorities. map ( |x| x. 1 ) . collect :: < Vec < _ > > ( ) ) ;
7675 }
7776
78- fn on_new_session < ' a , I : ' a > ( changed : bool , validators : I , _queued_validators : I )
77+ fn on_new_session < ' a , I : ' a > ( changed : bool , validators : I , queued_validators : I )
7978 where
8079 I : Iterator < Item = ( & ' a T :: AccountId , Self :: Key ) > ,
8180 {
82- // Remember who the authorities are for the new session.
81+ // Remember who the authorities are for the new and next session.
8382 if changed {
84- Keys :: put ( validators. map ( |x| x. 1 ) . collect :: < Vec < _ > > ( ) ) ;
83+ let keys = validators. chain ( queued_validators) . map ( |x| x. 1 ) . collect :: < BTreeSet < _ > > ( ) ;
84+ Keys :: put ( keys. into_iter ( ) . collect :: < Vec < _ > > ( ) ) ;
8585 }
8686 }
8787
@@ -192,12 +192,13 @@ mod tests {
192192 }
193193
194194 #[ test]
195- fn authorities_returns_current_authority_set ( ) {
196- // The whole authority discovery module ignores account ids, but we still need it for
197- // `pallet_session::OneSessionHandler::on_new_session`, thus its safe to use the same value everywhere.
195+ fn authorities_returns_current_and_next_authority_set ( ) {
196+ // The whole authority discovery module ignores account ids, but we still need them for
197+ // `pallet_session::OneSessionHandler::on_new_session`, thus its safe to use the same value
198+ // everywhere.
198199 let account_id = AuthorityPair :: from_seed_slice ( vec ! [ 10 ; 32 ] . as_ref ( ) ) . unwrap ( ) . public ( ) ;
199200
200- let first_authorities: Vec < AuthorityId > = vec ! [ 0 , 1 ] . into_iter ( )
201+ let mut first_authorities: Vec < AuthorityId > = vec ! [ 0 , 1 ] . into_iter ( )
201202 . map ( |i| AuthorityPair :: from_seed_slice ( vec ! [ i; 32 ] . as_ref ( ) ) . unwrap ( ) . public ( ) )
202203 . map ( AuthorityId :: from)
203204 . collect ( ) ;
@@ -206,12 +207,21 @@ mod tests {
206207 . map ( |i| AuthorityPair :: from_seed_slice ( vec ! [ i; 32 ] . as_ref ( ) ) . unwrap ( ) . public ( ) )
207208 . map ( AuthorityId :: from)
208209 . collect ( ) ;
209-
210210 // Needed for `pallet_session::OneSessionHandler::on_new_session`.
211- let second_authorities_and_account_ids: Vec < ( & AuthorityId , AuthorityId ) > = second_authorities. clone ( )
211+ let second_authorities_and_account_ids = second_authorities. clone ( )
212212 . into_iter ( )
213213 . map ( |id| ( & account_id, id) )
214+ . collect :: < Vec < ( & AuthorityId , AuthorityId ) > > ( ) ;
215+
216+ let mut third_authorities: Vec < AuthorityId > = vec ! [ 4 , 5 ] . into_iter ( )
217+ . map ( |i| AuthorityPair :: from_seed_slice ( vec ! [ i; 32 ] . as_ref ( ) ) . unwrap ( ) . public ( ) )
218+ . map ( AuthorityId :: from)
214219 . collect ( ) ;
220+ // Needed for `pallet_session::OneSessionHandler::on_new_session`.
221+ let third_authorities_and_account_ids = third_authorities. clone ( )
222+ . into_iter ( )
223+ . map ( |id| ( & account_id, id) )
224+ . collect :: < Vec < ( & AuthorityId , AuthorityId ) > > ( ) ;
215225
216226 // Build genesis.
217227 let mut t = frame_system:: GenesisConfig :: default ( )
@@ -233,23 +243,55 @@ mod tests {
233243 AuthorityDiscovery :: on_genesis_session (
234244 first_authorities. iter ( ) . map ( |id| ( id, id. clone ( ) ) )
235245 ) ;
236- assert_eq ! ( first_authorities, AuthorityDiscovery :: authorities( ) ) ;
246+ first_authorities. sort ( ) ;
247+ let mut authorities_returned = AuthorityDiscovery :: authorities ( ) ;
248+ authorities_returned. sort ( ) ;
249+ assert_eq ! ( first_authorities, authorities_returned) ;
237250
238251 // When `changed` set to false, the authority set should not be updated.
239252 AuthorityDiscovery :: on_new_session (
240253 false ,
241254 second_authorities_and_account_ids. clone ( ) . into_iter ( ) ,
242- vec ! [ ] . into_iter ( ) ,
255+ third_authorities_and_account_ids. clone ( ) . into_iter ( ) ,
256+ ) ;
257+ let mut authorities_returned = AuthorityDiscovery :: authorities ( ) ;
258+ authorities_returned. sort ( ) ;
259+ assert_eq ! (
260+ first_authorities,
261+ authorities_returned,
262+ "Expected authority set not to change as `changed` was set to false." ,
243263 ) ;
244- assert_eq ! ( first_authorities, AuthorityDiscovery :: authorities( ) ) ;
245264
246265 // When `changed` set to true, the authority set should be updated.
247266 AuthorityDiscovery :: on_new_session (
248267 true ,
249268 second_authorities_and_account_ids. into_iter ( ) ,
250- vec ! [ ] . into_iter ( ) ,
269+ third_authorities_and_account_ids. clone ( ) . into_iter ( ) ,
270+ ) ;
271+ let mut second_and_third_authorities = second_authorities. iter ( )
272+ . chain ( third_authorities. iter ( ) )
273+ . cloned ( )
274+ . collect :: < Vec < AuthorityId > > ( ) ;
275+ second_and_third_authorities. sort ( ) ;
276+ assert_eq ! (
277+ second_and_third_authorities,
278+ AuthorityDiscovery :: authorities( ) ,
279+ "Expected authority set to contain both the authorities of the new as well as the \
280+ next session."
281+ ) ;
282+
283+ // With overlapping authority sets, `authorities()` should return a deduplicated set.
284+ AuthorityDiscovery :: on_new_session (
285+ true ,
286+ third_authorities_and_account_ids. clone ( ) . into_iter ( ) ,
287+ third_authorities_and_account_ids. clone ( ) . into_iter ( ) ,
288+ ) ;
289+ third_authorities. sort ( ) ;
290+ assert_eq ! (
291+ third_authorities,
292+ AuthorityDiscovery :: authorities( ) ,
293+ "Expected authority set to be deduplicated."
251294 ) ;
252- assert_eq ! ( second_authorities, AuthorityDiscovery :: authorities( ) ) ;
253295 } ) ;
254296 }
255297}
0 commit comments