1use super::*;
21use alloc::collections::btree_map::BTreeMap;
22use frame_support::{
23 storage_alias,
24 traits::{ConstU32, OnRuntimeUpgrade},
25};
26
27#[cfg(feature = "try-runtime")]
28use frame_support::ensure;
29#[cfg(feature = "try-runtime")]
30use sp_runtime::TryRuntimeError;
31
32const TARGET: &'static str = "runtime::preimage::migration::v1";
34
35mod v0 {
37 use super::*;
38
39 #[derive(Clone, Eq, PartialEq, Encode, Decode, TypeInfo, MaxEncodedLen, RuntimeDebug)]
40 pub enum OldRequestStatus<AccountId, Balance> {
41 Unrequested(Option<(AccountId, Balance)>),
42 Requested(u32),
43 }
44
45 #[storage_alias]
46 pub type PreimageFor<T: Config> = StorageMap<
47 Pallet<T>,
48 Identity,
49 <T as frame_system::Config>::Hash,
50 BoundedVec<u8, ConstU32<MAX_SIZE>>,
51 >;
52
53 #[storage_alias]
54 pub type StatusFor<T: Config> = StorageMap<
55 Pallet<T>,
56 Identity,
57 <T as frame_system::Config>::Hash,
58 OldRequestStatus<<T as frame_system::Config>::AccountId, BalanceOf<T>>,
59 >;
60
61 #[cfg(feature = "try-runtime")]
63 pub fn image_count<T: Config>() -> Option<u32> {
64 let images = v0::PreimageFor::<T>::iter_values().count() as u32;
65 let status = v0::StatusFor::<T>::iter_values().count() as u32;
66
67 if images == status {
68 Some(images)
69 } else {
70 None
71 }
72 }
73}
74
75pub mod v1 {
76 use super::*;
77
78 pub struct Migration<T>(core::marker::PhantomData<T>);
83
84 impl<T: Config> OnRuntimeUpgrade for Migration<T> {
85 #[cfg(feature = "try-runtime")]
86 fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
87 ensure!(StorageVersion::get::<Pallet<T>>() == 0, "can only upgrade from version 0");
88
89 let images = v0::image_count::<T>().expect("v0 storage corrupted");
90 log::info!(target: TARGET, "Migrating {} images", &images);
91 Ok((images as u32).encode())
92 }
93
94 fn on_runtime_upgrade() -> Weight {
95 let mut weight = T::DbWeight::get().reads(1);
96 if StorageVersion::get::<Pallet<T>>() != 0 {
97 log::warn!(
98 target: TARGET,
99 "skipping MovePreimagesIntoBuckets: executed on wrong storage version.\
100 Expected version 0"
101 );
102 return weight
103 }
104
105 let status = v0::StatusFor::<T>::drain().collect::<Vec<_>>();
106 weight.saturating_accrue(T::DbWeight::get().reads(status.len() as u64));
107
108 let preimages = v0::PreimageFor::<T>::drain().collect::<BTreeMap<_, _>>();
109 weight.saturating_accrue(T::DbWeight::get().reads(preimages.len() as u64));
110
111 for (hash, status) in status.into_iter() {
112 let preimage = if let Some(preimage) = preimages.get(&hash) {
113 preimage
114 } else {
115 log::error!(target: TARGET, "preimage not found for hash {:?}", &hash);
116 continue
117 };
118 let len = preimage.len() as u32;
119 if len > MAX_SIZE {
120 log::error!(
121 target: TARGET,
122 "preimage too large for hash {:?}, len: {}",
123 &hash,
124 len
125 );
126 continue
127 }
128
129 let status = match status {
130 v0::OldRequestStatus::Unrequested(deposit) => match deposit {
131 Some(deposit) => OldRequestStatus::Unrequested { deposit, len },
132 None =>
134 OldRequestStatus::Requested { deposit: None, count: 1, len: Some(len) },
135 },
136 v0::OldRequestStatus::Requested(0) => {
137 log::error!(target: TARGET, "preimage has counter of zero: {:?}", hash);
138 continue
139 },
140 v0::OldRequestStatus::Requested(count) =>
141 OldRequestStatus::Requested { deposit: None, count, len: Some(len) },
142 };
143 log::trace!(target: TARGET, "Moving preimage {:?} with len {}", hash, len);
144
145 #[allow(deprecated)]
146 crate::StatusFor::<T>::insert(hash, status);
147 crate::PreimageFor::<T>::insert(&(hash, len), preimage);
148
149 weight.saturating_accrue(T::DbWeight::get().writes(2));
150 }
151 StorageVersion::new(1).put::<Pallet<T>>();
152
153 weight.saturating_add(T::DbWeight::get().writes(1))
154 }
155
156 #[cfg(feature = "try-runtime")]
157 fn post_upgrade(state: Vec<u8>) -> DispatchResult {
158 let old_images: u32 =
159 Decode::decode(&mut &state[..]).expect("pre_upgrade provides a valid state; qed");
160 let new_images = image_count::<T>().expect("V1 storage corrupted");
161
162 if new_images != old_images {
163 log::error!(
164 target: TARGET,
165 "migrated {} images, expected {}",
166 new_images,
167 old_images
168 );
169 }
170 ensure!(StorageVersion::get::<Pallet<T>>() == 1, "must upgrade");
171 Ok(())
172 }
173 }
174
175 #[cfg(feature = "try-runtime")]
177 pub fn image_count<T: Config>() -> Option<u32> {
178 let images = crate::PreimageFor::<T>::iter_values().count() as u32;
180 #[allow(deprecated)]
181 let status = crate::StatusFor::<T>::iter_values().count() as u32;
182
183 if images == status {
184 Some(images)
185 } else {
186 None
187 }
188 }
189}
190
191#[cfg(test)]
192#[cfg(feature = "try-runtime")]
193mod test {
194 #![allow(deprecated)]
195 use super::*;
196 use crate::mock::{Test as T, *};
197
198 use sp_runtime::bounded_vec;
199
200 #[test]
201 fn migration_works() {
202 new_test_ext().execute_with(|| {
203 assert_eq!(StorageVersion::get::<Pallet<T>>(), 0);
204 let (p, h) = preimage::<T>(128);
208 v0::PreimageFor::<T>::insert(h, p);
209 v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(None));
210 let (p, h) = preimage::<T>(1024);
212 v0::PreimageFor::<T>::insert(h, p);
213 v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Unrequested(Some((1, 1))));
214 let (p, h) = preimage::<T>(8192);
216 v0::PreimageFor::<T>::insert(h, p);
217 v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(0));
218 let (p, h) = preimage::<T>(65536);
220 v0::PreimageFor::<T>::insert(h, p);
221 v0::StatusFor::<T>::insert(h, v0::OldRequestStatus::Requested(10));
222
223 assert_eq!(v0::image_count::<T>(), Some(4));
224 assert_eq!(v1::image_count::<T>(), None, "V1 storage should be corrupted");
225
226 let state = v1::Migration::<T>::pre_upgrade().unwrap();
227 let _w = v1::Migration::<T>::on_runtime_upgrade();
228 v1::Migration::<T>::post_upgrade(state).unwrap();
229
230 assert_eq!(v0::image_count::<T>(), Some(3));
232 assert_eq!(v1::image_count::<T>(), Some(3)); assert_eq!(StorageVersion::get::<Pallet<T>>(), 1);
234
235 let (p, h) = preimage::<T>(128);
237 assert_eq!(crate::PreimageFor::<T>::get(&(h, 128)), Some(p));
238 assert_eq!(
239 crate::StatusFor::<T>::get(h),
240 Some(OldRequestStatus::Requested { deposit: None, count: 1, len: Some(128) })
241 );
242 let (p, h) = preimage::<T>(1024);
244 assert_eq!(crate::PreimageFor::<T>::get(&(h, 1024)), Some(p));
245 assert_eq!(
246 crate::StatusFor::<T>::get(h),
247 Some(OldRequestStatus::Unrequested { deposit: (1, 1), len: 1024 })
248 );
249 let (_, h) = preimage::<T>(8192);
251 assert_eq!(crate::PreimageFor::<T>::get(&(h, 8192)), None);
252 assert_eq!(crate::StatusFor::<T>::get(h), None);
253 let (p, h) = preimage::<T>(65536);
255 assert_eq!(crate::PreimageFor::<T>::get(&(h, 65536)), Some(p));
256 assert_eq!(
257 crate::StatusFor::<T>::get(h),
258 Some(OldRequestStatus::Requested { deposit: None, count: 10, len: Some(65536) })
259 );
260 });
261 }
262
263 fn preimage<T: Config>(
265 len: usize,
266 ) -> (BoundedVec<u8, ConstU32<MAX_SIZE>>, <T as frame_system::Config>::Hash) {
267 let p = bounded_vec![1; len];
268 let h = <T as frame_system::Config>::Hashing::hash_of(&p);
269 (p, h)
270 }
271}