diff --git a/pallets/collator-selection/src/lib.rs b/pallets/collator-selection/src/lib.rs index efcb97ddc..ead4a768a 100644 --- a/pallets/collator-selection/src/lib.rs +++ b/pallets/collator-selection/src/lib.rs @@ -81,6 +81,7 @@ #![cfg_attr(not(feature = "std"), no_std)] pub use pallet::*; +pub mod migrations; #[cfg(test)] mod mock; diff --git a/pallets/collator-selection/src/migrations.rs b/pallets/collator-selection/src/migrations.rs new file mode 100644 index 000000000..ad341f963 --- /dev/null +++ b/pallets/collator-selection/src/migrations.rs @@ -0,0 +1,96 @@ +// This file is part of Astar. + +// Copyright (C) Stake Technologies Pte.Ltd. +// SPDX-License-Identifier: GPL-3.0-or-later + +// Astar is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. + +// Astar is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. + +// You should have received a copy of the GNU General Public License +// along with Astar. If not, see . + +use super::*; +use frame_support::{pallet_prelude::*, traits::OnRuntimeUpgrade}; +use sp_std::{marker::PhantomData, vec::Vec}; + +/// One-time migration that removes outdated LastAuthoredBlock entries. +/// It keeps entries only for accounts that are currently: +/// - active candidates +/// - invulnerables +/// +/// All other accounts are removed. +pub struct LastAuthoredBlockCleanup(PhantomData); + +impl OnRuntimeUpgrade for LastAuthoredBlockCleanup { + fn on_runtime_upgrade() -> Weight { + log::info!("Running CollatorSelection::LastAuthoredBlockCleanup..."); + + // Snapshot active identifiers for faster membership checks + let invulnerables = Invulnerables::::get(); + let candidates: sp_std::collections::btree_set::BTreeSet = + Candidates::::get().into_iter().map(|c| c.who).collect(); + + let mut read = 0u64; + let mut write = 0u64; + let mut stale = Vec::new(); + + // Sanity limit + const MAX_SCAN: u64 = 200; + + for (account, _) in LastAuthoredBlock::::iter() { + if read >= MAX_SCAN { + log::warn!("LastAuthoredBlockCleanup: scan limit {} reached.", MAX_SCAN); + break; + } + read += 1; + + let keep = invulnerables.contains(&account) || candidates.contains(&account); + if !keep { + stale.push(account); + } + } + + for account in stale { + LastAuthoredBlock::::remove(account); + write += 1; + } + + log::info!( + "LastAuthoredBlockCleanup completed: removed {} entries (reads {}, writes {}).", + write, + read, + write + ); + + ::DbWeight::get().reads_writes(read, write) + } + + #[cfg(feature = "try-runtime")] + fn pre_upgrade() -> Result, sp_runtime::TryRuntimeError> { + let old_count = LastAuthoredBlock::::iter().count() as u64; + Ok(old_count.encode()) + } + + #[cfg(feature = "try-runtime")] + fn post_upgrade(data: Vec) -> Result<(), sp_runtime::TryRuntimeError> { + let old_count: u64 = Decode::decode(&mut &data[..]).map_err(|_| { + sp_runtime::TryRuntimeError::Other("Failed to decode pre-upgrade count") + })?; + + let new_count = LastAuthoredBlock::::iter().count() as u64; + + assert!( + new_count < old_count, + "LastAuthoredBlockCleanup: new count > old count (should only decrease)" + ); + + Ok(()) + } +}