1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
16

            
17
use crate::{
18
    Candidate, CandidateSummaries, CandidateSummary, Config, Delegator,
19
    DelegatorCandidateSummaries, DelegatorCandidateSummary, Pallet, PausePoolsExtrinsics, PoolKind,
20
    Pools, PoolsKey,
21
};
22
use core::marker::PhantomData;
23
use frame_support::{
24
    migrations::{MigrationId, SteppedMigration, SteppedMigrationError},
25
    pallet_prelude::{StorageVersion, Zero},
26
    traits::GetStorageVersion,
27
    weights::WeightMeter,
28
    BoundedVec, Deserialize, Serialize,
29
};
30
use parity_scale_codec::{Decode, Encode};
31
use scale_info::TypeInfo;
32
use sp_core::{ConstU32, Get, MaxEncodedLen, RuntimeDebug};
33
use sp_runtime::Saturating;
34

            
35
const LOG_TARGET: &str = "pallet_pooled_staking::migrations::stepped_generate_summaries";
36
pub const PALLET_MIGRATIONS_ID: &[u8; 21] = b"pallet-pooled-staking";
37

            
38
pub struct MigrationGenerateSummaries<T: Config>(PhantomData<T>);
39
impl<T: Config> SteppedMigration for MigrationGenerateSummaries<T> {
40
    type Cursor = MigrationState;
41
    type Identifier = MigrationId<21>;
42

            
43
2
    fn id() -> Self::Identifier {
44
2
        MigrationId {
45
2
            pallet_id: *PALLET_MIGRATIONS_ID,
46
2
            version_from: 0,
47
2
            version_to: 1,
48
2
        }
49
2
    }
50

            
51
1
    fn step(
52
1
        cursor: Option<Self::Cursor>,
53
1
        meter: &mut WeightMeter,
54
1
    ) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
55
        // Consumes in advance the cost of reading and writing the storage version.
56
1
        meter.consume(T::DbWeight::get().reads_writes(1, 1));
57

            
58
1
        if Pallet::<T>::on_chain_storage_version() != u16::from(Self::id().version_from) {
59
            return Ok(None);
60
1
        }
61

            
62
        // We make a weight meter with 70% of allowed weight to be extra sure the migration will not
63
        // cause issues
64
1
        let mut meter2 = WeightMeter::with_limit(meter.remaining() * 70 / 100);
65

            
66
1
        let new_state = stepped_generate_summaries::<
67
1
            T,
68
1
            DelegatorCandidateSummaries<T>,
69
1
            CandidateSummaries<T>,
70
1
        >(cursor, &mut meter2)?;
71
1
        if new_state.is_none() {
72
1
            // migration is finished, we update the on chain storage version
73
1
            StorageVersion::new(u16::from(Self::id().version_to)).put::<Pallet<T>>();
74
1
        }
75

            
76
1
        meter.consume(meter2.consumed());
77
1
        Ok(new_state)
78
1
    }
79

            
80
    #[cfg(feature = "try-runtime")]
81
    fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
82
        // Can we test it somehow without performing the same process? (which would be useless)
83
        Ok(Default::default())
84
    }
85

            
86
    #[cfg(feature = "try-runtime")]
87
    fn post_upgrade(_state: alloc::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
88
        Ok(())
89
    }
90
}
91

            
92
#[derive(
93
    RuntimeDebug,
94
    PartialEq,
95
    Eq,
96
    Encode,
97
    Decode,
98
    Clone,
99
    TypeInfo,
100
    Serialize,
101
    Deserialize,
102
    MaxEncodedLen,
103
)]
104
pub struct MigrationState {
105
    /// Last key processed by the migration.
106
    /// Migration will resume from there on next step.
107
    last_raw_key: BoundedVec<u8, ConstU32<1024>>,
108
}
109

            
110
/// Generate summaries from `Pools` into provided `D` and `C`, which can be used both for
111
/// a migration (where pooled_staking was used before the introduction of summaries) and for
112
/// try_state (generate the summaries from `Pools` into separate storages, then check they
113
/// match the versions dynamically maintained).
114
1
pub fn stepped_generate_summaries<
115
1
    T: Config,
116
1
    D: frame_support::storage::StorageDoubleMap<
117
1
        Delegator<T>,
118
1
        Candidate<T>,
119
1
        DelegatorCandidateSummary,
120
1
        Query = DelegatorCandidateSummary,
121
1
    >,
122
1
    C: frame_support::storage::StorageMap<Candidate<T>, CandidateSummary, Query = CandidateSummary>,
123
1
>(
124
1
    state: Option<MigrationState>,
125
1
    meter: &mut WeightMeter,
126
1
) -> Result<Option<MigrationState>, SteppedMigrationError> {
127
    // If there is not yet a state, it means we're starting the migration process.
128
    // In this case we pause the pools extrinsics
129
1
    if state.is_none() {
130
1
        log::info!(target: LOG_TARGET, "Starting migration. Pools extrinsics will be paused until the end of the migration.");
131
1
        PausePoolsExtrinsics::<T>::put(true);
132
    } else {
133
        log::info!(target: LOG_TARGET, "Resuming migration.");
134
    }
135

            
136
    // One migration step is to read from the Pools iterator (1 read). If it is
137
    // a XXXShares PoolsKey, it will then read+write
138
    // an entry from both the candidates and delegators summaries (2 reads+writes).
139
    // We need to check before performing each step that we can consume that weight.
140
1
    let step_weight = T::DbWeight::get().reads_writes(3, 2);
141

            
142
    // If the available weight is less than the cost of 1 step, we'll never be able to migrate.
143
1
    if meter.remaining().any_lt(step_weight) {
144
        return Err(SteppedMigrationError::InsufficientWeight {
145
            required: step_weight,
146
        });
147
1
    }
148

            
149
    // Documentation if `StorageDoubleMap::iter_keys` warn sabout inserting/deleting keys
150
    // giving undefined results. While the risk is reduced by pausing the pools extrinsics,
151
    // it may happen while distributing manual rewards if this is the first time a candidate gets
152
    // manual rewards distributed (counter goes from absent to non-zero). However by looking at the
153
    // implementation of those iterators, they iterate the keys by lexicographic order. So even in
154
    // the case this entry is inserted, it will not affect iterating over the non-explored keys, and
155
    // we don't care missing reward counter entries since they are not used by the migration.
156
    //
157
    // We also prefer to use `iter_keys` over `iter` as it might optimize the POV for keys we skip.
158
1
    let mut iterator = match state {
159
        Some(MigrationState { last_raw_key }) => Pools::<T>::iter_keys_from(last_raw_key.to_vec()),
160
1
        None => Pools::<T>::iter_keys(),
161
    };
162

            
163
1
    let mut count = 0;
164

            
165
    loop {
166
5
        if !meter.can_consume(step_weight) {
167
            log::warn!(target: LOG_TARGET, "Migration limit reached. Processed {count} delegations.");
168
            return Ok(Some(MigrationState {
169
                last_raw_key: iterator
170
                    .last_raw_key()
171
                    .to_vec()
172
                    .try_into()
173
                    .expect("size should be larger than key length"),
174
            }));
175
5
        }
176

            
177
5
        meter.consume(T::DbWeight::get().reads_writes(1, 0));
178

            
179
5
        let Some((candidate, key)) = iterator.next() else {
180
1
            break;
181
        };
182

            
183
4
        count += 1;
184
4
        let (pool, delegator) = match key.clone() {
185
1
            PoolsKey::JoiningShares { delegator } => (PoolKind::Joining, delegator),
186
1
            PoolsKey::AutoCompoundingShares { delegator } => (PoolKind::AutoCompounding, delegator),
187
2
            PoolsKey::ManualRewardsShares { delegator } => (PoolKind::ManualRewards, delegator),
188
            PoolsKey::LeavingShares { delegator } => (PoolKind::Leaving, delegator),
189
            _ => continue, // we only care about share values
190
        };
191

            
192
4
        let value = Pools::<T>::get(&candidate, key);
193
        // Are 0/Default values automatically removed from storage?
194
        // In case they aren't we check the amount of shares isn't 0.
195
4
        if value.is_zero() {
196
1
            continue;
197
3
        }
198

            
199
        // We first modify the delegator summary. If the summary is empty it means we have not yet
200
        // encountered that delegator for this candidate, so we'll consider it a new delegator that
201
        // will increase the `delegators` count in the candidate summary.
202
3
        let mut new_delegator = false;
203
3
        meter.consume(T::DbWeight::get().reads_writes(2, 2));
204

            
205
3
        D::mutate(&delegator, &candidate, |summary| {
206
3
            if summary.is_empty() {
207
3
                new_delegator = true;
208
3
            }
209

            
210
3
            summary.set_pool(pool, true);
211
3
        });
212

            
213
3
        C::mutate(&candidate, |summary| {
214
3
            if new_delegator {
215
3
                summary.delegators.saturating_inc();
216
3
            }
217

            
218
3
            summary.pool_delegators_mut(pool).saturating_inc();
219
3
        });
220
    }
221

            
222
1
    log::info!(target: LOG_TARGET, "Migration finished. Processed {count} delegations");
223
1
    PausePoolsExtrinsics::<T>::put(false);
224

            
225
1
    Ok(None)
226
1
}