grandpa: cleanup stale entries in set id session mapping (#13237)

* grandpa: cleanup stale entries in set id session mapping

* Update frame/grandpa/src/migrations.rs

Co-authored-by: Bastian Köcher <[email protected]>

* grandpa: remove unused import

* grandpa: migration off-by-one

* Update frame/grandpa/src/lib.rs

Co-authored-by: Anton <[email protected]>

* Update frame/grandpa/src/lib.rs

Co-authored-by: Anton <[email protected]>

* grandpa: MaxSetIdSessionEntries as u64

* node-template: fix MaxSetIdSessionEntries type

---------

Co-authored-by: Bastian Köcher <[email protected]>
Co-authored-by: Anton <[email protected]>

bin/node-template/runtime/src/lib.rs

@@ -230,6 +230,7 @@ impl pallet_grandpa::Config for Runtime {
 
 	type WeightInfo = ();
 	type MaxAuthorities = ConstU32<32>;
+	type MaxSetIdSessionEntries = ConstU64<0>;
 }
 
 impl pallet_timestamp::Config for Runtime {

bin/node/runtime/src/lib.rs

@@ -1310,6 +1310,10 @@ impl pallet_authority_discovery::Config for Runtime {
 	type MaxAuthorities = MaxAuthorities;
 }
 
+parameter_types! {
+	pub const MaxSetIdSessionEntries: u32 = BondingDuration::get() * SessionsPerEra::get();
+}
+
 impl pallet_grandpa::Config for Runtime {
 	type RuntimeEvent = RuntimeEvent;
 
@@ -1331,6 +1335,7 @@ impl pallet_grandpa::Config for Runtime {
 
 	type WeightInfo = ();
 	type MaxAuthorities = MaxAuthorities;
+	type MaxSetIdSessionEntries = MaxSetIdSessionEntries;
 }
 
 parameter_types! {

frame/grandpa/src/lib.rs

@@ -121,6 +121,15 @@ pub mod pallet {
 		/// Max Authorities in use
 		#[pallet::constant]
 		type MaxAuthorities: Get<u32>;
+
+		/// The maximum number of entries to keep in the set id to session index mapping.
+		///
+		/// Since the `SetIdSession` map is only used for validating equivocations this
+		/// value should relate to the bonding duration of whatever staking system is
+		/// being used (if any). If equivocation handling is not enabled then this value
+		/// can be zero.
+		#[pallet::constant]
+		type MaxSetIdSessionEntries: Get<u64>;
 	}
 
 	#[pallet::hooks]
@@ -323,6 +332,12 @@ pub mod pallet {
 	/// A mapping from grandpa set ID to the index of the *most recent* session for which its
 	/// members were responsible.
 	///
+	/// This is only used for validating equivocation proofs. An equivocation proof must
+	/// contains a key-ownership proof for a given session, therefore we need a way to tie
+	/// together sessions and GRANDPA set ids, i.e. we need to validate that a validator
+	/// was the owner of a given key on a given session, and what the active set ID was
+	/// during that session.
+	///
 	/// TWOX-NOTE: `SetId` is not under user control.
 	#[pallet::storage]
 	#[pallet::getter(fn session_for_set)]
@@ -643,10 +658,17 @@ where
 			};
 
 			if res.is_ok() {
-				CurrentSetId::<T>::mutate(|s| {
+				let current_set_id = CurrentSetId::<T>::mutate(|s| {
 					*s += 1;
 					*s
-				})
+				});
+
+				let max_set_id_session_entries = T::MaxSetIdSessionEntries::get().max(1);
+				if current_set_id >= max_set_id_session_entries {
+					SetIdSession::<T>::remove(current_set_id - max_set_id_session_entries);
+				}
+
+				current_set_id
 			} else {
 				// either the session module signalled that the validators have changed
 				// or the set was stalled. but since we didn't successfully schedule
@@ -659,8 +681,8 @@ where
 			Self::current_set_id()
 		};
 
-		// if we didn't issue a change, we update the mapping to note that the current
-		// set corresponds to the latest equivalent session (i.e. now).
+		// update the mapping to note that the current set corresponds to the
+		// latest equivalent session (i.e. now).
 		let session_index = <pallet_session::Pallet<T>>::current_index();
 		SetIdSession::<T>::insert(current_set_id, &session_index);
 	}

frame/grandpa/src/migrations.rs

@@ -15,5 +15,51 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use frame_support::{
+	traits::{Get, OnRuntimeUpgrade},
+	weights::Weight,
+};
+
+use crate::{Config, CurrentSetId, SetIdSession, LOG_TARGET};
+
 /// Version 4.
 pub mod v4;
+
+/// This migration will clean up all stale set id -> session entries from the
+/// `SetIdSession` storage map, only the latest `max_set_id_session_entries`
+/// will be kept.
+///
+/// This migration should be added with a runtime upgrade that introduces the
+/// `MaxSetIdSessionEntries` constant to the pallet (although it could also be
+/// done later on).
+pub struct CleanupSetIdSessionMap<T>(sp_std::marker::PhantomData<T>);
+impl<T: Config> OnRuntimeUpgrade for CleanupSetIdSessionMap<T> {
+	fn on_runtime_upgrade() -> Weight {
+		// NOTE: since this migration will loop over all stale entries in the
+		// map we need to set some cutoff value, otherwise the migration might
+		// take too long to run. for scenarios where there are that many entries
+		// to cleanup a multiblock migration will be needed instead.
+		if CurrentSetId::<T>::get() > 25_000 {
+			log::warn!(
+				target: LOG_TARGET,
+				"CleanupSetIdSessionMap migration was aborted since there are too many entries to cleanup."
+			);
+
+			return T::DbWeight::get().reads(1)
+		}
+
+		cleanup_set_id_sesion_map::<T>()
+	}
+}
+
+fn cleanup_set_id_sesion_map<T: Config>() -> Weight {
+	let until_set_id = CurrentSetId::<T>::get().saturating_sub(T::MaxSetIdSessionEntries::get());
+
+	for set_id in 0..=until_set_id {
+		SetIdSession::<T>::remove(set_id);
+	}
+
+	T::DbWeight::get()
+		.reads(1)
+		.saturating_add(T::DbWeight::get().writes(until_set_id + 1))
+}

frame/grandpa/src/mock.rs

@@ -219,6 +219,7 @@ impl pallet_offences::Config for Test {
 parameter_types! {
 	pub const ReportLongevity: u64 =
 		BondingDuration::get() as u64 * SessionsPerEra::get() as u64 * Period::get();
+	pub const MaxSetIdSessionEntries: u32 = BondingDuration::get() * SessionsPerEra::get();
 }
 
 impl Config for Test {
@@ -239,6 +240,7 @@ impl Config for Test {
 
 	type WeightInfo = ();
 	type MaxAuthorities = ConstU32<100>;
+	type MaxSetIdSessionEntries = MaxSetIdSessionEntries;
 }
 
 pub fn grandpa_log(log: ConsensusLog<u64>) -> DigestItem {

frame/grandpa/src/tests.rs

@@ -781,6 +781,33 @@ fn on_new_session_doesnt_start_new_set_if_schedule_change_failed() {
 	});
 }
 
+#[test]
+fn cleans_up_old_set_id_session_mappings() {
+	new_test_ext(vec![(1, 1), (2, 1), (3, 1)]).execute_with(|| {
+		let max_set_id_session_entries = MaxSetIdSessionEntries::get();
+
+		start_era(max_set_id_session_entries);
+
+		// we should have a session id mapping for all the set ids from
+		// `max_set_id_session_entries` eras we have observed
+		for i in 1..=max_set_id_session_entries {
+			assert!(Grandpa::session_for_set(i as u64).is_some());
+		}
+
+		start_era(max_set_id_session_entries * 2);
+
+		// we should keep tracking the new mappings for new eras
+		for i in max_set_id_session_entries + 1..=max_set_id_session_entries * 2 {
+			assert!(Grandpa::session_for_set(i as u64).is_some());
+		}
+
+		// but the old ones should have been pruned by now
+		for i in 1..=max_set_id_session_entries {
+			assert!(Grandpa::session_for_set(i as u64).is_none());
+		}
+	});
+}
+
 #[test]
 fn always_schedules_a_change_on_new_session_when_stalled() {
 	new_test_ext(vec![(1, 1), (2, 1), (3, 1)]).execute_with(|| {