[AHM] Preimage migration

integration-tests/ahm/src/tests.rs

@@ -33,7 +33,7 @@
 
 use cumulus_primitives_core::{AggregateMessageOrigin, ParaId};
 use frame_support::{pallet_prelude::*, traits::*, weights::WeightMeter};
-use pallet_rc_migrator::{MigrationStage, RcMigrationStage};
+use pallet_rc_migrator::{types::PalletMigrationChecks, MigrationStage, RcMigrationStage};
 use polkadot_primitives::InboundDownwardMessage;
 use remote_externalities::RemoteExternalities;
 
@@ -47,8 +47,10 @@ async fn account_migration_works() {
 	let para_id = ParaId::from(1000);
 
 	// Simulate relay blocks and grab the DMP messages
-	let dmp_messages = rc.execute_with(|| {
+	let (dmp_messages, pre_check_payload) = rc.execute_with(|| {
 		let mut dmps = Vec::new();
+		let pre_check_payload =
+			pallet_rc_migrator::preimage::PreimageChunkMigrator::<Polkadot>::pre_check();
 
 		// Loop until no more DMPs are added and we had at least 1
 		loop {
@@ -59,10 +61,10 @@ async fn account_migration_works() {
 			dmps.extend(new_dmps);
 
 			if RcMigrationStage::<Polkadot>::get() ==
-				pallet_rc_migrator::MigrationStage::ProxyMigrationDone
+				pallet_rc_migrator::MigrationStage::AllMigrationsDone
 			{
 				log::info!("Migration done");
-				break dmps;
+				break (dmps, pre_check_payload);
 			}
 		}
 	});
@@ -88,6 +90,10 @@ async fn account_migration_works() {
 			log::debug!("AH DMP messages left: {}", fp.storage.count);
 			next_block_ah();
 		}
+
+		pallet_rc_migrator::preimage::PreimageChunkMigrator::<Polkadot>::post_check(
+			pre_check_payload,
+		);
 		// NOTE that the DMP queue is probably not empty because the snapshot that we use contains
 		// some overweight ones.
 	});

pallets/ah-migrator/src/lib.rs

@@ -33,6 +33,7 @@
 
 pub mod account;
 pub mod multisig;
+pub mod preimage;
 pub mod proxy;
 pub mod types;
 
@@ -48,8 +49,9 @@ use frame_support::{
 };
 use frame_system::pallet_prelude::*;
 use pallet_balances::{AccountData, Reasons as LockReasons};
-use pallet_rc_migrator::{accounts::Account as RcAccount, multisig::*, proxy::*};
+use pallet_rc_migrator::{accounts::Account as RcAccount, multisig::*, preimage::*, proxy::*};
 use sp_application_crypto::Ss58Codec;
+use sp_core::H256;
 use sp_runtime::{
 	traits::{Convert, TryConvert},
 	AccountId32,
@@ -71,6 +73,7 @@ pub mod pallet {
 		+ pallet_balances::Config<Balance = u128>
 		+ pallet_multisig::Config
 		+ pallet_proxy::Config
+		+ pallet_preimage::Config<Hash = H256>
 	{
 		/// The overarching event type.
 		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
@@ -148,6 +151,30 @@ pub mod pallet {
 			/// How many proxy announcements failed to integrate.
 			count_bad: u32,
 		},
+		/// Received a batch of `RcPreimageChunk` that are going to be integrated.
+		PreimageChunkBatchReceived {
+			/// How many preimage chunks are in the batch.
+			count: u32,
+		},
+		/// We processed a batch of `RcPreimageChunk` that we received.
+		PreimageChunkBatchProcessed {
+			/// How many preimage chunks were successfully integrated.
+			count_good: u32,
+			/// How many preimage chunks failed to integrate.
+			count_bad: u32,
+		},
+		/// We received a batch of `RcPreimageRequestStatus` that we are going to integrate.
+		PreimageRequestStatusBatchReceived {
+			/// How many preimage request status are in the batch.
+			count: u32,
+		},
+		/// We processed a batch of `RcPreimageRequestStatus` that we received.
+		PreimageRequestStatusBatchProcessed {
+			/// How many preimage request status were successfully integrated.
+			count_good: u32,
+			/// How many preimage request status failed to integrate.
+			count_bad: u32,
+		},
 	}
 
 	#[pallet::pallet]
@@ -207,8 +234,29 @@ pub mod pallet {
 			announcements: Vec<RcProxyAnnouncementOf<T>>,
 		) -> DispatchResult {
 			ensure_root(origin)?;
+
 			Self::do_receive_proxy_announcements(announcements).map_err(Into::into)
 		}
+
+		#[pallet::call_index(4)]
+		pub fn receive_preimage_chunks(
+			origin: OriginFor<T>,
+			chunks: Vec<RcPreimageChunk>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_receive_preimage_chunks(chunks).map_err(Into::into)
+		}
+
+		#[pallet::call_index(5)]
+		pub fn receive_preimage_request_status(
+			origin: OriginFor<T>,
+			request_status: Vec<RcPreimageRequestStatusOf<T>>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_receive_preimage_request_statuses(request_status).map_err(Into::into)
+		}
 	}
 
 	#[pallet::hooks]

pallets/ah-migrator/src/preimage.rs

@@ -0,0 +1,169 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{types::*, *};
+use frame_support::traits::{Consideration, Footprint};
+use pallet_rc_migrator::preimage::{chunks::*, *};
+
+impl<T: Config> Pallet<T> {
+	pub fn do_receive_preimage_chunks(chunks: Vec<RcPreimageChunk>) -> Result<(), Error<T>> {
+		Self::deposit_event(Event::PreimageChunkBatchReceived { count: chunks.len() as u32 });
+		let (mut count_good, mut count_bad) = (0, 0);
+		log::info!(target: LOG_TARGET, "Integrating {} preimage chunks", chunks.len());
+
+		for chunk in chunks {
+			match Self::do_receive_preimage_chunk(chunk) {
+				Ok(()) => count_good += 1,
+				Err(e) => {
+					count_bad += 1;
+					log::error!(target: LOG_TARGET, "Error while integrating preimage chunk: {:?}", e);
+				},
+			}
+		}
+		Self::deposit_event(Event::PreimageChunkBatchProcessed { count_good, count_bad });
+
+		Ok(())
+	}
+
+	pub fn do_receive_preimage_chunk(chunk: RcPreimageChunk) -> Result<(), Error<T>> {
+		log::debug!(target: LOG_TARGET, "Integrating preimage chunk {} offset {}/{}", chunk.preimage_hash, chunk.chunk_byte_offset + chunk.chunk_bytes.len() as u32, chunk.preimage_len);
+		let key = (chunk.preimage_hash, chunk.preimage_len);
+
+		// First check that we did not miss a chunk
+		let preimage = match alias::PreimageFor::<T>::get(&key) {
+			Some(mut preimage) => {
+				if preimage.len() != chunk.chunk_byte_offset as usize {
+					defensive!("Preimage chunk missing");
+					return Err(Error::<T>::TODO);
+				}
+
+				match preimage.try_mutate(|p| {
+					p.extend(chunk.chunk_bytes.clone());
+				}) {
+					Some(preimage) => {
+						alias::PreimageFor::<T>::insert(&key, &preimage);
+						preimage
+					},
+					None => {
+						defensive!("Preimage too big");
+						return Err(Error::<T>::TODO);
+					},
+				}
+			},
+			None => {
+				if chunk.chunk_byte_offset != 0 {
+					defensive!("Preimage chunk missing");
+					return Err(Error::<T>::TODO);
+				}
+
+				let preimage: BoundedVec<u8, ConstU32<{ CHUNK_SIZE }>> = chunk.chunk_bytes;
+				debug_assert!(CHUNK_SIZE <= pallet_rc_migrator::preimage::alias::MAX_SIZE);
+				let bounded_preimage: BoundedVec<
+					u8,
+					ConstU32<{ pallet_rc_migrator::preimage::alias::MAX_SIZE }>,
+				> = preimage.into_inner().try_into().expect("Asserted");
+				alias::PreimageFor::<T>::insert(key, &bounded_preimage);
+				bounded_preimage
+			},
+		};
+
+		if preimage.len() == chunk.preimage_len as usize + chunk.chunk_byte_offset as usize {
+			log::debug!(target: LOG_TARGET, "Preimage complete: {}", chunk.preimage_hash);
+		}
+
+		Ok(())
+	}
+
+	pub fn do_receive_preimage_request_statuses(
+		request_status: Vec<RcPreimageRequestStatusOf<T>>,
+	) -> Result<(), Error<T>> {
+		Self::deposit_event(Event::PreimageRequestStatusBatchReceived {
+			count: request_status.len() as u32,
+		});
+		log::info!(target: LOG_TARGET, "Integrating {} preimage request status", request_status.len());
+		let (mut count_good, mut count_bad) = (0, 0);
+
+		for request_status in request_status {
+			match Self::do_receive_preimage_request_status(request_status) {
+				Ok(()) => count_good += 1,
+				Err(e) => {
+					count_bad += 1;
+					log::error!(target: LOG_TARGET, "Error while integrating preimage request status: {:?}", e);
+				},
+			}
+		}
+
+		Self::deposit_event(Event::PreimageRequestStatusBatchProcessed { count_good, count_bad });
+		Ok(())
+	}
+
+	pub fn do_receive_preimage_request_status(
+		mut request_status: RcPreimageRequestStatusOf<T>,
+	) -> Result<(), Error<T>> {
+		if alias::RequestStatusFor::<T>::contains_key(&request_status.hash) {
+			log::warn!(target: LOG_TARGET, "Request status already migrated: {:?}", request_status.hash);
+			return Ok(());
+		}
+
+		let new_ticket = match request_status.request_status {
+			alias::RequestStatus::Unrequested { ticket: (ref who, ref ticket), len } => {
+				let fp = Footprint::from_parts(1, len as usize);
+				ticket.clone().update(&who, fp).ok()
+			},
+			alias::RequestStatus::Requested {
+				maybe_ticket: Some((ref who, ref ticket)),
+				maybe_len: Some(len),
+				..
+			} => {
+				let fp = Footprint::from_parts(1, len as usize);
+				ticket.clone().update(&who, fp).ok()
+			},
+			alias::RequestStatus::Requested { maybe_ticket: Some(_), maybe_len: None, .. } => {
+				defensive!("Ticket cannot be re-evaluated");
+				// I think this is unreachable, but not exactly sure. Either way, nothing that we
+				// could do about it.
+				None
+			},
+			_ => None,
+		};
+
+		let new_request_status = match (new_ticket, request_status.request_status.clone()) {
+			(
+				Some(new_ticket),
+				alias::RequestStatus::Unrequested { ticket: (who, ref mut ticket), len },
+			) => alias::RequestStatus::Unrequested { ticket: (who, new_ticket), len },
+			(
+				Some(new_ticket),
+				alias::RequestStatus::Requested {
+					maybe_ticket: Some((who, ref mut ticket)),
+					maybe_len: Some(len),
+					count: count,
+				},
+			) => alias::RequestStatus::Requested {
+				maybe_ticket: Some((who, new_ticket)),
+				maybe_len: Some(len),
+				count,
+			},
+			_ => request_status.request_status,
+		};
+
+		alias::RequestStatusFor::<T>::insert(&request_status.hash, &new_request_status);
+		log::debug!(target: LOG_TARGET, "Integrating preimage request status: {:?}", new_request_status);
+
+		Ok(())
+	}
+}

pallets/rc-migrator/Cargo.toml

@@ -23,6 +23,7 @@ pallet-balances = { workspace = true }
 pallet-staking = { workspace = true }
 pallet-proxy = { workspace = true }
 pallet-multisig = { workspace = true }
+pallet-preimage = { workspace = true }
 polkadot-runtime-common = { workspace = true }
 runtime-parachains = { workspace = true }
 polkadot-parachain-primitives = { workspace = true }
@@ -51,6 +52,7 @@ std = [
 	"sp-runtime/std",
 	"sp-std/std",
 	"xcm/std",
+	"pallet-preimage/std"
 ]
 runtime-benchmarks = [
 	"frame-benchmarking/runtime-benchmarks",
@@ -64,6 +66,7 @@ runtime-benchmarks = [
 	"polkadot-runtime-common/runtime-benchmarks",
 	"runtime-parachains/runtime-benchmarks",
 	"sp-runtime/runtime-benchmarks",
+	"pallet-preimage/runtime-benchmarks"
 ]
 try-runtime = [
 	"frame-support/try-runtime",
@@ -75,4 +78,5 @@ try-runtime = [
 	"polkadot-runtime-common/try-runtime",
 	"runtime-parachains/try-runtime",
 	"sp-runtime/try-runtime",
+	"pallet-preimage/try-runtime"
 ]