Pallet session new API (#4609)

* Initial work

* Fix most things

* fix test

* fix old comment

* migration

* fix

* remove useless stuff

* fix

* less spaghetti implementation

* fix initial session

* fix

bin/node/runtime/src/lib.rs

@@ -236,14 +236,13 @@ parameter_types! {
 }
 
 impl pallet_session::Trait for Runtime {
-	type OnSessionEnding = Staking;
+	type SessionManager = Staking;
 	type SessionHandler = <SessionKeys as OpaqueKeys>::KeyTypeIdProviders;
 	type ShouldEndSession = Babe;
 	type Event = Event;
 	type Keys = SessionKeys;
 	type ValidatorId = <Self as frame_system::Trait>::AccountId;
 	type ValidatorIdOf = pallet_staking::StashOf<Self>;
-	type SelectInitialValidators = Staking;
 	type DisabledValidatorsThreshold = DisabledValidatorsThreshold;
 }
 

frame/authority-discovery/src/lib.rs

@@ -109,26 +109,18 @@ mod tests {
 	pub struct Test;
 	impl Trait for Test {}
 
-	pub struct TestOnSessionEnding;
-	impl pallet_session::OnSessionEnding<AuthorityId> for TestOnSessionEnding {
-		fn on_session_ending(_: SessionIndex, _: SessionIndex) -> Option<Vec<AuthorityId>> {
-			None
-		}
-	}
-
 	parameter_types! {
 		pub const DisabledValidatorsThreshold: Perbill = Perbill::from_percent(33);
 	}
 
 	impl pallet_session::Trait for Test {
-		type OnSessionEnding = TestOnSessionEnding;
+		type SessionManager = ();
 		type Keys = UintAuthorityId;
 		type ShouldEndSession = pallet_session::PeriodicSessions<Period, Offset>;
 		type SessionHandler = TestSessionHandler;
 		type Event = ();
 		type ValidatorId = AuthorityId;
 		type ValidatorIdOf = ConvertInto;
-		type SelectInitialValidators = ();
 		type DisabledValidatorsThreshold = DisabledValidatorsThreshold;
 	}
 

frame/babe/src/mock.rs

@@ -79,9 +79,8 @@ impl pallet_session::Trait for Test {
 	type ValidatorId = <Self as frame_system::Trait>::AccountId;
 	type ShouldEndSession = Babe;
 	type SessionHandler = (Babe,Babe,);
-	type OnSessionEnding = ();
+	type SessionManager = ();
 	type ValidatorIdOf = ();
-	type SelectInitialValidators = ();
 	type Keys = MockSessionKeys;
 	type DisabledValidatorsThreshold = DisabledValidatorsThreshold;
 }

frame/im-online/src/mock.rs

@@ -43,28 +43,25 @@ thread_local! {
 	pub static VALIDATORS: RefCell<Option<Vec<u64>>> = RefCell::new(Some(vec![1, 2, 3]));
 }
 
-pub struct TestOnSessionEnding;
-impl pallet_session::OnSessionEnding<u64> for TestOnSessionEnding {
-	fn on_session_ending(_ending_index: SessionIndex, _will_apply_at: SessionIndex)
-		-> Option<Vec<u64>>
-	{
+pub struct TestSessionManager;
+impl pallet_session::SessionManager<u64> for TestSessionManager {
+	fn new_session(_new_index: SessionIndex) -> Option<Vec<u64>> {
 		VALIDATORS.with(|l| l.borrow_mut().take())
 	}
+	fn end_session(_: SessionIndex) {}
 }
 
-impl pallet_session::historical::OnSessionEnding<u64, u64> for TestOnSessionEnding {
-	fn on_session_ending(_ending_index: SessionIndex, _will_apply_at: SessionIndex)
-		-> Option<(Vec<u64>, Vec<(u64, u64)>)>
-	{
+impl pallet_session::historical::SessionManager<u64, u64> for TestSessionManager {
+	fn new_session(_new_index: SessionIndex) -> Option<Vec<(u64, u64)>> {
 		VALIDATORS.with(|l| l
 			.borrow_mut()
 			.take()
 			.map(|validators| {
-				let full_identification = validators.iter().map(|v| (*v, *v)).collect();
-				(validators, full_identification)
+				validators.iter().map(|v| (*v, *v)).collect()
 			})
 		)
 	}
+	fn end_session(_: SessionIndex) {}
 }
 
 /// An extrinsic type used for tests.
@@ -131,13 +128,12 @@ parameter_types! {
 
 impl pallet_session::Trait for Runtime {
 	type ShouldEndSession = pallet_session::PeriodicSessions<Period, Offset>;
-	type OnSessionEnding = pallet_session::historical::NoteHistoricalRoot<Runtime, TestOnSessionEnding>;
+	type SessionManager = pallet_session::historical::NoteHistoricalRoot<Runtime, TestSessionManager>;
 	type SessionHandler = (ImOnline, );
 	type ValidatorId = u64;
 	type ValidatorIdOf = ConvertInto;
 	type Keys = UintAuthorityId;
 	type Event = ();
-	type SelectInitialValidators = ();
 	type DisabledValidatorsThreshold = DisabledValidatorsThreshold;
 }
 

frame/session/src/historical.rs

@@ -48,25 +48,27 @@ pub trait Trait: super::Trait {
 	/// validator, since they may be outdated by the time this is queried from a
 	/// historical trie.
 	///
-	/// This mapping is expected to remain stable in between calls to
-	/// `Self::OnSessionEnding::on_session_ending` which return new validators.
+	/// It must return the identification for the current session index.
 	type FullIdentificationOf: Convert<Self::ValidatorId, Option<Self::FullIdentification>>;
 }
 
 decl_storage! {
 	trait Store for Module<T: Trait> as Session {
 		/// Mapping from historical session indices to session-data root hash and validator count.
 		HistoricalSessions get(fn historical_root): map SessionIndex => Option<(T::Hash, ValidatorCount)>;
-		/// Queued full identifications for queued sessions whose validators have become obsolete.
-		CachedObsolete get(fn cached_obsolete): map SessionIndex
-			=> Option<Vec<(T::ValidatorId, T::FullIdentification)>>;
 		/// The range of historical sessions we store. [first, last)
 		StoredRange: Option<(SessionIndex, SessionIndex)>;
+		/// Deprecated.
+		CachedObsolete: map SessionIndex => Option<Vec<(T::ValidatorId, T::FullIdentification)>>;
 	}
 }
 
 decl_module! {
-	pub struct Module<T: Trait> for enum Call where origin: T::Origin { }
+	pub struct Module<T: Trait> for enum Call where origin: T::Origin {
+		fn on_initialize(_n: T::BlockNumber) {
+			CachedObsolete::<T>::remove_all();
+		}
+	}
 }
 
 impl<T: Trait> Module<T> {
@@ -97,52 +99,52 @@ impl<T: Trait> Module<T> {
 	}
 }
 
-/// Specialization of the crate-level `OnSessionEnding` which returns the old
-/// set of full identification when changing the validator set.
-pub trait OnSessionEnding<ValidatorId, FullIdentification>: crate::OnSessionEnding<ValidatorId> {
-	/// If there was a validator set change, its returns the set of new validators along with the
-	/// old validators and their full identifications.
-	fn on_session_ending(ending: SessionIndex, will_apply_at: SessionIndex)
-		-> Option<(Vec<ValidatorId>, Vec<(ValidatorId, FullIdentification)>)>;
+/// Specialization of the crate-level `SessionManager` which returns the set of full identification
+/// when creating a new session.
+pub trait SessionManager<ValidatorId, FullIdentification>: crate::SessionManager<ValidatorId> {
+	/// If there was a validator set change, its returns the set of new validators along with their
+	/// full identifications.
+	fn new_session(new_index: SessionIndex) -> Option<Vec<(ValidatorId, FullIdentification)>>;
+	fn end_session(end_index: SessionIndex);
 }
 
-/// An `OnSessionEnding` implementation that wraps an inner `I` and also
+/// An `SessionManager` implementation that wraps an inner `I` and also
 /// sets the historical trie root of the ending session.
 pub struct NoteHistoricalRoot<T, I>(sp_std::marker::PhantomData<(T, I)>);
 
-impl<T: Trait, I> crate::OnSessionEnding<T::ValidatorId> for NoteHistoricalRoot<T, I>
-	where I: OnSessionEnding<T::ValidatorId, T::FullIdentification>
+impl<T: Trait, I> crate::SessionManager<T::ValidatorId> for NoteHistoricalRoot<T, I>
+	where I: SessionManager<T::ValidatorId, T::FullIdentification>
 {
-	fn on_session_ending(ending: SessionIndex, applied_at: SessionIndex) -> Option<Vec<T::ValidatorId>> {
+	fn new_session(new_index: SessionIndex) -> Option<Vec<T::ValidatorId>> {
 		StoredRange::mutate(|range| {
-			range.get_or_insert_with(|| (ending, ending)).1 = ending + 1;
+			range.get_or_insert_with(|| (new_index, new_index)).1 = new_index + 1;
 		});
 
-		// do all of this _before_ calling the other `on_session_ending` impl
-		// so that we have e.g. correct exposures from the _current_.
+		let new_validators_and_id = <I as SessionManager<_, _>>::new_session(new_index);
+		let new_validators = new_validators_and_id.as_ref().map(|new_validators| {
+			new_validators.iter().map(|(v, _id)| v.clone()).collect()
+		});
 
-		let count = <SessionModule<T>>::validators().len() as u32;
-		match ProvingTrie::<T>::generate_for(ending) {
-			Ok(trie) => <HistoricalSessions<T>>::insert(ending, &(trie.root, count)),
-			Err(reason) => {
-				print("Failed to generate historical ancestry-inclusion proof.");
-				print(reason);
+		if let Some(new_validators) = new_validators_and_id {
+			let count = new_validators.len() as u32;
+			match ProvingTrie::<T>::generate_for(new_validators) {
+				Ok(trie) => <HistoricalSessions<T>>::insert(new_index, &(trie.root, count)),
+				Err(reason) => {
+					print("Failed to generate historical ancestry-inclusion proof.");
+					print(reason);
+				}
+			};
+		} else {
+			let previous_index = new_index.saturating_sub(1);
+			if let Some(previous_session) = <HistoricalSessions<T>>::get(previous_index) {
+				<HistoricalSessions<T>>::insert(new_index, previous_session);
 			}
-		};
-
-		// trie has been generated for this session, so it's no longer queued.
-		<CachedObsolete<T>>::remove(&ending);
-
-		let (new_validators, old_exposures) = <I as OnSessionEnding<_, _>>::on_session_ending(ending, applied_at)?;
-
-		// every session from `ending+1 .. applied_at` now has obsolete `FullIdentification`
-		// now that a new validator election has occurred.
-		// we cache these in the trie until those sessions themselves end.
-		for obsolete in (ending + 1) .. applied_at {
-			<CachedObsolete<T>>::insert(obsolete, &old_exposures);
 		}
 
-		Some(new_validators)
+		new_validators
+	}
+	fn end_session(end_index: SessionIndex) {
+		<I as SessionManager<_, _>>::end_session(end_index)
 	}
 }
 
@@ -158,27 +160,22 @@ pub struct ProvingTrie<T: Trait> {
 }
 
 impl<T: Trait> ProvingTrie<T> {
-	fn generate_for(now: SessionIndex) -> Result<Self, &'static str> {
+	fn generate_for<I>(validators: I) -> Result<Self, &'static str>
+		where I: IntoIterator<Item=(T::ValidatorId, T::FullIdentification)>
+	{
 		let mut db = MemoryDB::default();
 		let mut root = Default::default();
 
-		fn build<T: Trait, I>(root: &mut T::Hash, db: &mut MemoryDB<HasherOf<T>>, validators: I)
-			-> Result<(), &'static str>
-			where I: IntoIterator<Item=(T::ValidatorId, Option<T::FullIdentification>)>
 		{
-			let mut trie = TrieDBMut::new(db, root);
+			let mut trie = TrieDBMut::new(&mut db, &mut root);
 			for (i, (validator, full_id)) in validators.into_iter().enumerate() {
 				let i = i as u32;
 				let keys = match <SessionModule<T>>::load_keys(&validator) {
 					None => continue,
 					Some(k) => k,
 				};
 
-				let full_id = full_id.or_else(|| T::FullIdentificationOf::convert(validator.clone()));
-				let full_id = match full_id {
-					None => return Err("no full identification for a current validator"),
-					Some(full) => (validator, full),
-				};
+				let full_id = (validator, full_id);
 
 				// map each key to the owner index.
 				for key_id in T::Keys::key_ids() {
@@ -194,17 +191,6 @@ impl<T: Trait> ProvingTrie<T> {
 				let _ = i.using_encoded(|k| full_id.using_encoded(|v| trie.insert(k, v)))
 					.map_err(|_| "failed to insert into trie")?;
 			}
-
-			Ok(())
-		}
-
-		// if the current session's full identifications are obsolete but cached,
-		// use those.
-		if let Some(obsolete) = <CachedObsolete<T>>::get(&now) {
-			build::<T, _>(&mut root, &mut db, obsolete.into_iter().map(|(v, f)| (v, Some(f))))?
-		} else {
-			let validators = <SessionModule<T>>::validators();
-			build::<T, _>(&mut root, &mut db, validators.into_iter().map(|v| (v, None)))?
 		}
 
 		Ok(ProvingTrie {
@@ -281,7 +267,12 @@ impl<T: Trait, D: AsRef<[u8]>> frame_support::traits::KeyOwnerProofSystem<(KeyTy
 
 	fn prove(key: (KeyTypeId, D)) -> Option<Self::Proof> {
 		let session = <SessionModule<T>>::current_index();
-		let trie = ProvingTrie::<T>::generate_for(session).ok()?;
+		let validators = <SessionModule<T>>::validators().into_iter()
+			.filter_map(|validator| {
+				T::FullIdentificationOf::convert(validator.clone())
+					.map(|full_id| (validator, full_id))
+			});
+		let trie = ProvingTrie::<T>::generate_for(validators).ok()?;
 
 		let (id, data) = key;
 
@@ -348,13 +339,9 @@ mod tests {
 			set_next_validators(vec![1, 2, 4]);
 			force_new_session();
 
-			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_none());
-
 			System::set_block_number(2);
 			Session::on_initialize(2);
 
-			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_some());
-
 			assert!(Historical::historical_root(proof.session).is_some());
 			assert!(Session::current_index() > proof.session);
 
@@ -366,15 +353,13 @@ mod tests {
 			force_new_session();
 			System::set_block_number(3);
 			Session::on_initialize(3);
-
-			assert!(Historical::cached_obsolete(&(proof.session + 1)).is_none());
 		});
 	}
 
 	#[test]
 	fn prune_up_to_works() {
 		new_test_ext().execute_with(|| {
-			for i in 1..101u64 {
+			for i in 1..99u64 {
 				set_next_validators(vec![i]);
 				force_new_session();
 
@@ -385,7 +370,7 @@ mod tests {
 
 			assert_eq!(StoredRange::get(), Some((0, 100)));
 
-			for i in 1..100 {
+			for i in 0..100 {
 				assert!(Historical::historical_root(i).is_some())
 			}
 
@@ -405,7 +390,7 @@ mod tests {
 			Historical::prune_up_to(100);
 			assert_eq!(StoredRange::get(), None);
 
-			for i in 101..201u64 {
+			for i in 99..199u64 {
 				set_next_validators(vec![i]);
 				force_new_session();
 
@@ -416,14 +401,14 @@ mod tests {
 
 			assert_eq!(StoredRange::get(), Some((100, 200)));
 
-			for i in 101..200 {
+			for i in 100..200 {
 				assert!(Historical::historical_root(i).is_some())
 			}
 
 			Historical::prune_up_to(9999);
 			assert_eq!(StoredRange::get(), None);
 
-			for i in 101..200 {
+			for i in 100..200 {
 				assert!(Historical::historical_root(i).is_none())
 			}
 		});