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

substrate/frame/session/src/historical.rs

@@ -48,8 +48,7 @@ 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>>;
 }
 
@@ -57,16 +56,19 @@ 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,15 +160,14 @@ 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) {
@@ -174,11 +175,7 @@ impl<T: Trait> ProvingTrie<T> {
 					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())
 			}
 		});

substrate/frame/session/src/lib.rs

@@ -162,29 +162,30 @@ impl<
 	}
 }
 
-/// An event handler for when the session is ending.
-/// TODO [slashing] consider renaming to OnSessionStarting
-pub trait OnSessionEnding<ValidatorId> {
-	/// Handle the fact that the session is ending, and optionally provide the new validator set.
+/// A trait for managing creation of new validator set.
+pub trait SessionManager<ValidatorId> {
+	/// Plan a new session, and optionally provide the new validator set.
 	///
 	/// Even if the validator-set is the same as before, if any underlying economic
 	/// conditions have changed (i.e. stake-weights), the new validator set must be returned.
 	/// This is necessary for consensus engines making use of the session module to
 	/// issue a validator-set change so misbehavior can be provably associated with the new
 	/// economic conditions as opposed to the old.
+	/// The returned validator set, if any, will not be applied until `new_index`.
+	/// `new_index` is strictly greater than from previous call.
 	///
-	/// `ending_index` is the index of the currently ending session.
-	/// The returned validator set, if any, will not be applied until `will_apply_at`.
-	/// `will_apply_at` is guaranteed to be at least `ending_index + 1`, since session indices don't
-	/// repeat, but it could be some time after in case we are staging authority set changes.
-	fn on_session_ending(
-		ending_index: SessionIndex,
-		will_apply_at: SessionIndex
-	) -> Option<Vec<ValidatorId>>;
+	/// The first session start at index 0.
+	fn new_session(new_index: SessionIndex) -> Option<Vec<ValidatorId>>;
+	/// End the session.
+	///
+	/// Because the session pallet can queue validator set the ending session can be lower than the
+	/// last new session index.
+	fn end_session(end_index: SessionIndex);
 }
 
-impl<A> OnSessionEnding<A> for () {
-	fn on_session_ending(_: SessionIndex, _: SessionIndex) -> Option<Vec<A>> { None }
+impl<A> SessionManager<A> for () {
+	fn new_session(_: SessionIndex) -> Option<Vec<A>> { None }
+	fn end_session(_: SessionIndex) {}
 }
 
 /// Handler for session lifecycle events.
@@ -214,7 +215,7 @@ pub trait SessionHandler<ValidatorId> {
 
 	/// A notification for end of the session.
 	///
-	/// Note it is triggered before any `OnSessionEnding` handlers,
+	/// Note it is triggered before any `SessionManager::end_session` handlers,
 	/// so we can still affect the validator set.
 	fn on_before_session_ending() {}
 
@@ -248,7 +249,7 @@ pub trait OneSessionHandler<ValidatorId>: BoundToRuntimeAppPublic {
 
 	/// A notification for end of the session.
 	///
-	/// Note it is triggered before any `OnSessionEnding` handlers,
+	/// Note it is triggered before any `SessionManager::end_session` handlers,
 	/// so we can still affect the validator set.
 	fn on_before_session_ending() {}
 
@@ -318,21 +319,6 @@ impl<AId> SessionHandler<AId> for TestSessionHandler {
 	fn on_disabled(_: usize) {}
 }
 
-/// Handler for selecting the genesis validator set.
-pub trait SelectInitialValidators<ValidatorId> {
-	/// Returns the initial validator set. If `None` is returned
-	/// all accounts that have session keys set in the genesis block
-	/// will be validators.
-	fn select_initial_validators() -> Option<Vec<ValidatorId>>;
-}
-
-/// Implementation of `SelectInitialValidators` that does nothing.
-impl<V> SelectInitialValidators<V> for () {
-	fn select_initial_validators() -> Option<Vec<V>> {
-		None
-	}
-}
-
 impl<T: Trait> ValidatorRegistration<T::ValidatorId> for Module<T> {
 	fn is_registered(id: &T::ValidatorId) -> bool {
 		Self::load_keys(id).is_some()
@@ -352,8 +338,8 @@ pub trait Trait: frame_system::Trait {
 	/// Indicator for when to end the session.
 	type ShouldEndSession: ShouldEndSession<Self::BlockNumber>;
 
-	/// Handler for when a session is about to end.
-	type OnSessionEnding: OnSessionEnding<Self::ValidatorId>;
+	/// Handler for managing new session.
+	type SessionManager: SessionManager<Self::ValidatorId>;
 
 	/// Handler when a session has changed.
 	type SessionHandler: SessionHandler<Self::ValidatorId>;
@@ -366,9 +352,6 @@ pub trait Trait: frame_system::Trait {
 	/// After the threshold is reached `disabled` method starts to return true,
 	/// which in combination with `pallet_staking` forces a new era.
 	type DisabledValidatorsThreshold: Get<Perbill>;
-
-	/// Select initial validators.
-	type SelectInitialValidators: SelectInitialValidators<Self::ValidatorId>;
 }
 
 const DEDUP_KEY_PREFIX: &[u8] = b":session:keys";
@@ -435,12 +418,19 @@ decl_storage! {
 					.expect("genesis config must not contain duplicates; qed");
 			}
 
-			let initial_validators = T::SelectInitialValidators::select_initial_validators()
-				.unwrap_or_else(|| config.keys.iter().map(|(ref v, _)| v.clone()).collect());
+			let initial_validators_0 = T::SessionManager::new_session(0)
+				.unwrap_or_else(|| {
+					frame_support::print("No initial validator provided by `SessionManager`, use \
+						session config keys to generate initial validator set.");
+					config.keys.iter().map(|(ref v, _)| v.clone()).collect()
+				});
+			assert!(!initial_validators_0.is_empty(), "Empty validator set for session 0 in genesis block!");
 
-			assert!(!initial_validators.is_empty(), "Empty validator set in genesis block!");
+			let initial_validators_1 = T::SessionManager::new_session(1)
+				.unwrap_or_else(|| initial_validators_0.clone());
+			assert!(!initial_validators_1.is_empty(), "Empty validator set for session 1 in genesis block!");
 
-			let queued_keys: Vec<_> = initial_validators
+			let queued_keys: Vec<_> = initial_validators_1
 				.iter()
 				.cloned()
 				.map(|v| (
@@ -452,7 +442,7 @@ decl_storage! {
 			// Tell everyone about the genesis session keys
 			T::SessionHandler::on_genesis_session::<T::Keys>(&queued_keys);
 
-			<Validators<T>>::put(initial_validators);
+			<Validators<T>>::put(initial_validators_0);
 			<QueuedKeys<T>>::put(queued_keys);
 		});
 	}
@@ -545,10 +535,14 @@ impl<T: Trait> Module<T> {
 			DisabledValidators::take();
 		}
 
-		let applied_at = session_index + 2;
+		T::SessionManager::end_session(session_index);
+
+		// Increment session index.
+		let session_index = session_index + 1;
+		CurrentIndex::put(session_index);
 
 		// Get next validator set.
-		let maybe_next_validators = T::OnSessionEnding::on_session_ending(session_index, applied_at);
+		let maybe_next_validators = T::SessionManager::new_session(session_index + 1);
 		let (next_validators, next_identities_changed)
 			= if let Some(validators) = maybe_next_validators
 		{
@@ -560,10 +554,6 @@ impl<T: Trait> Module<T> {
 			(<Validators<T>>::get(), false)
 		};
 
-		// Increment session index.
-		let session_index = session_index + 1;
-		CurrentIndex::put(session_index);
-
 		// Queue next session keys.
 		let (queued_amalgamated, next_changed) = {
 			// until we are certain there has been a change, iterate the prior

substrate/frame/session/src/mock.rs

@@ -88,9 +88,10 @@ impl SessionHandler<u64> for TestSessionHandler {
 	}
 }
 
-pub struct TestOnSessionEnding;
-impl OnSessionEnding<u64> for TestOnSessionEnding {
-	fn on_session_ending(_: SessionIndex, _: SessionIndex) -> Option<Vec<u64>> {
+pub struct TestSessionManager;
+impl SessionManager<u64> for TestSessionManager {
+	fn end_session(_: SessionIndex) {}
+	fn new_session(_: SessionIndex) -> Option<Vec<u64>> {
 		if !TEST_SESSION_CHANGED.with(|l| *l.borrow()) {
 			VALIDATORS.with(|v| {
 				let mut v = v.borrow_mut();
@@ -108,14 +109,13 @@ impl OnSessionEnding<u64> for TestOnSessionEnding {
 }
 
 #[cfg(feature = "historical")]
-impl crate::historical::OnSessionEnding<u64, u64> for TestOnSessionEnding {
-	fn on_session_ending(ending_index: SessionIndex, will_apply_at: SessionIndex)
-		-> Option<(Vec<u64>, Vec<(u64, u64)>)>
+impl crate::historical::SessionManager<u64, u64> for TestSessionManager {
+	fn end_session(_: SessionIndex) {}
+	fn new_session(new_index: SessionIndex)
+		-> Option<Vec<(u64, u64)>>
 	{
-		let pair_with_ids = |vals: &[u64]| vals.iter().map(|&v| (v, v)).collect::<Vec<_>>();
-		<Self as OnSessionEnding<_>>::on_session_ending(ending_index, will_apply_at)
-			.map(|vals| (pair_with_ids(&vals), vals))
-			.map(|(ids, vals)| (vals, ids))
+		<Self as SessionManager<_>>::new_session(new_index)
+			.map(|vals| vals.into_iter().map(|val| (val, val)).collect())
 	}
 }
 
@@ -190,15 +190,14 @@ parameter_types! {
 impl Trait for Test {
 	type ShouldEndSession = TestShouldEndSession;
 	#[cfg(feature = "historical")]
-	type OnSessionEnding = crate::historical::NoteHistoricalRoot<Test, TestOnSessionEnding>;
+	type SessionManager = crate::historical::NoteHistoricalRoot<Test, TestSessionManager>;
 	#[cfg(not(feature = "historical"))]
-	type OnSessionEnding = TestOnSessionEnding;
+	type SessionManager = TestSessionManager;
 	type SessionHandler = TestSessionHandler;
 	type ValidatorId = u64;
 	type ValidatorIdOf = ConvertInto;
 	type Keys = MockSessionKeys;
 	type Event = ();
-	type SelectInitialValidators = ();
 	type DisabledValidatorsThreshold = DisabledValidatorsThreshold;
 }