pallet_collator_selection/
migrations.rs

1// This file is part of Astar.
2
3// Copyright (C) Stake Technologies Pte.Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later
5
6// Astar is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// Astar is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with Astar. If not, see <http://www.gnu.org/licenses/>.
18
19use super::*;
20use frame_support::{pallet_prelude::*, traits::OnRuntimeUpgrade};
21use sp_std::{marker::PhantomData, vec::Vec};
22
23/// One-time migration that removes outdated LastAuthoredBlock entries.
24/// It keeps entries only for accounts that are currently:
25///   - active candidates
26///   - invulnerables
27///
28/// All other accounts are removed.
29pub struct LastAuthoredBlockCleanup<T: Config>(PhantomData<T>);
30
31impl<T: Config> OnRuntimeUpgrade for LastAuthoredBlockCleanup<T> {
32    fn on_runtime_upgrade() -> Weight {
33        log::info!("Running CollatorSelection::LastAuthoredBlockCleanup...");
34
35        // Snapshot active identifiers for faster membership checks
36        let invulnerables = Invulnerables::<T>::get();
37        let candidates: sp_std::collections::btree_set::BTreeSet<T::AccountId> =
38            Candidates::<T>::get().into_iter().map(|c| c.who).collect();
39
40        let mut read = 0u64;
41        let mut write = 0u64;
42        let mut stale = Vec::new();
43
44        // Sanity limit
45        const MAX_SCAN: u64 = 200;
46
47        for (account, _) in LastAuthoredBlock::<T>::iter() {
48            if read >= MAX_SCAN {
49                log::warn!("LastAuthoredBlockCleanup: scan limit {MAX_SCAN} reached.");
50                break;
51            }
52            read += 1;
53
54            let keep = invulnerables.contains(&account) || candidates.contains(&account);
55            if !keep {
56                stale.push(account);
57            }
58        }
59
60        for account in stale {
61            LastAuthoredBlock::<T>::remove(account);
62            write += 1;
63        }
64
65        log::info!(
66            "LastAuthoredBlockCleanup completed: removed {write:?} entries (reads {read:?}, writes {write:?})."
67        );
68
69        <T as frame_system::Config>::DbWeight::get().reads_writes(read, write)
70    }
71
72    #[cfg(feature = "try-runtime")]
73    fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
74        let old_count = LastAuthoredBlock::<T>::iter().count() as u64;
75        Ok(old_count.encode())
76    }
77
78    #[cfg(feature = "try-runtime")]
79    fn post_upgrade(data: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
80        let old_count: u64 = Decode::decode(&mut &data[..]).map_err(|_| {
81            sp_runtime::TryRuntimeError::Other("Failed to decode pre-upgrade count")
82        })?;
83
84        let new_count = LastAuthoredBlock::<T>::iter().count() as u64;
85
86        assert!(
87            new_count < old_count,
88            "LastAuthoredBlockCleanup: new count > old count (should only decrease)"
89        );
90
91        Ok(())
92    }
93}
pallet_collator_selection/migrations.rs

pallet_collator_selection/
migrations.rs
