[AHM] refactor events and enable scheduled migration start

integration-tests/ahm/Cargo.toml

@@ -49,3 +49,4 @@ hex = { workspace = true }
 pallet-whitelist = { workspace = true, default-features = true }
 pallet-bounties = { workspace = true, default-features = true }
 cumulus-pallet-parachain-system = { workspace = true }
+pallet-xcm = { workspace = true, default-features = true }

integration-tests/ahm/src/mock.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
-use asset_hub_polkadot_runtime::{Block as AssetHubBlock, Runtime as AssetHub};
+use asset_hub_polkadot_runtime::{AhMigrator, Block as AssetHubBlock, Runtime as AssetHub};
 use codec::Decode;
 use cumulus_primitives_core::{
 	AggregateMessageOrigin as ParachainMessageOrigin, InboundDownwardMessage, ParaId,
@@ -25,7 +25,7 @@ use pallet_rc_migrator::{
 	MigrationStageOf as RcMigrationStageOf, RcMigrationStage as RcMigrationStageStorage,
 };
 use polkadot_primitives::UpwardMessage;
-use polkadot_runtime::{Block as PolkadotBlock, Runtime as Polkadot};
+use polkadot_runtime::{Block as PolkadotBlock, RcMigrator, Runtime as Polkadot};
 use remote_externalities::{Builder, Mode, OfflineConfig, RemoteExternalities};
 use runtime_parachains::{
 	dmp::DownwardMessageQueues,
@@ -75,10 +75,9 @@ pub fn next_block_rc() {
 	frame_system::Pallet::<Polkadot>::set_block_number(now);
 	frame_system::Pallet::<Polkadot>::reset_events();
 	let weight = <polkadot_runtime::MessageQueue as OnInitialize<_>>::on_initialize(now);
-	let weight = <polkadot_runtime::RcMigrator as OnInitialize<_>>::on_initialize(now)
-		.saturating_add(weight);
+	let weight = <RcMigrator as OnInitialize<_>>::on_initialize(now).saturating_add(weight);
 	<polkadot_runtime::MessageQueue as OnFinalize<_>>::on_finalize(now);
-	<polkadot_runtime::RcMigrator as OnFinalize<_>>::on_finalize(now);
+	<RcMigrator as OnFinalize<_>>::on_finalize(now);
 
 	let limit = <Polkadot as frame_system::Config>::BlockWeights::get().max_block;
 	assert!(
@@ -96,10 +95,9 @@ pub fn next_block_ah() {
 	frame_system::Pallet::<AssetHub>::set_block_number(now);
 	frame_system::Pallet::<Polkadot>::reset_events();
 	let weight = <asset_hub_polkadot_runtime::MessageQueue as OnInitialize<_>>::on_initialize(now);
-	let weight = <asset_hub_polkadot_runtime::AhMigrator as OnInitialize<_>>::on_initialize(now)
-		.saturating_add(weight);
+	let weight = <AhMigrator as OnInitialize<_>>::on_initialize(now).saturating_add(weight);
 	<asset_hub_polkadot_runtime::MessageQueue as OnFinalize<_>>::on_finalize(now);
-	<asset_hub_polkadot_runtime::AhMigrator as OnFinalize<_>>::on_finalize(now);
+	<AhMigrator as OnFinalize<_>>::on_finalize(now);
 
 	let limit = <AssetHub as frame_system::Config>::BlockWeights::get().max_block;
 	assert!(
@@ -150,19 +148,20 @@ pub fn enqueue_ump(msgs: Vec<UpwardMessage>) {
 // Sets the initial migration stage on the Relay Chain.
 //
 // If the `START_STAGE` environment variable is set, it will be used to set the initial migration
-// stage. Otherwise, the current migration stage will be returned.
+// stage. Otherwise, the `AccountsMigrationInit` stage will be set bypassing the `Scheduled` stage.
+// The `Scheduled` stage is tested separately by the `scheduled_migration_works` test.
 pub fn set_initial_migration_stage(
 	relay_chain: &mut RemoteExternalities<PolkadotBlock>,
 ) -> RcMigrationStageOf<Polkadot> {
 	let stage = relay_chain.execute_with(|| {
-		if let Ok(stage) = std::env::var("START_STAGE") {
+		let stage = if let Ok(stage) = std::env::var("START_STAGE") {
 			log::info!("Setting start stage: {:?}", &stage);
-			let stage = RcMigrationStage::from_str(&stage).expect("Invalid start stage");
-			RcMigrationStageStorage::<Polkadot>::put(stage.clone());
-			stage
+			RcMigrationStage::from_str(&stage).expect("Invalid start stage")
 		} else {
-			RcMigrationStageStorage::<Polkadot>::get()
-		}
+			RcMigrationStage::AccountsMigrationInit
+		};
+		RcMigrationStageStorage::<Polkadot>::put(stage.clone());
+		stage
 	});
 	relay_chain.commit_all().unwrap();
 	stage

integration-tests/ahm/src/tests.rs

@@ -35,20 +35,25 @@ use super::{mock::*, proxy_test::ProxiesStillWork};
 use asset_hub_polkadot_runtime::Runtime as AssetHub;
 use cumulus_pallet_parachain_system::PendingUpwardMessages;
 use cumulus_primitives_core::{BlockT, Junction, Location, ParaId};
+use frame_support::traits::schedule::DispatchTime;
 use frame_system::pallet_prelude::BlockNumberFor;
-use pallet_ah_migrator::types::AhMigrationCheck;
+use pallet_ah_migrator::{
+	types::AhMigrationCheck, AhMigrationStage as AhMigrationStageStorage,
+	MigrationStage as AhMigrationStage,
+};
 use pallet_rc_migrator::{
 	types::RcMigrationCheck, MigrationStage as RcMigrationStage,
 	RcMigrationStage as RcMigrationStageStorage,
 };
-use polkadot_runtime::{Block as PolkadotBlock, Runtime as Polkadot};
+use polkadot_runtime::{Block as PolkadotBlock, RcMigrator, Runtime as Polkadot};
 use polkadot_runtime_common::{paras_registrar, slots as pallet_slots};
 use remote_externalities::RemoteExternalities;
 use runtime_parachains::dmp::DownwardMessageQueues;
 use sp_core::crypto::Ss58Codec;
 use sp_runtime::AccountId32;
 use std::{collections::BTreeMap, str::FromStr};
-use xcm_emulator::{ConvertLocation, WeightMeter};
+use xcm::latest::*;
+use xcm_emulator::{assert_ok, ConvertLocation, WeightMeter};
 
 type RcChecks = (
 	pallet_rc_migrator::preimage::PreimageChunkMigrator<Polkadot>,
@@ -350,6 +355,118 @@ async fn migration_works() {
 	println!("Migration done in {} RC blocks", rc_block_count);
 }
 
+#[tokio::test(flavor = "current_thread")]
+async fn scheduled_migration_works() {
+	let Some((mut rc, mut ah)) = load_externalities().await else { return };
+
+	// Check that the migration is pending on the RC.
+	rc.execute_with(|| {
+		log::info!("Asserting the initial state on RC");
+		next_block_rc();
+
+		assert_eq!(RcMigrationStageStorage::<Polkadot>::get(), RcMigrationStage::Pending);
+
+		// clear the DMP queue.
+		let _ = DownwardMessageQueues::<Polkadot>::take(AH_PARA_ID);
+	});
+	rc.commit_all().unwrap();
+
+	// Check that the migration is pending on the AH.
+	ah.execute_with(|| {
+		log::info!("Asserting the initial state on AH");
+		next_block_ah();
+
+		assert_eq!(AhMigrationStageStorage::<AssetHub>::get(), AhMigrationStage::Pending);
+
+		// clear the UMP queue.
+		let _ = PendingUpwardMessages::<AssetHub>::take();
+	});
+	ah.commit_all().unwrap();
+
+	// Schedule the migration on RC.
+	let dmp_messages = rc.execute_with(|| {
+		log::info!("Scheduling the migration on RC");
+		next_block_rc();
+
+		let now = frame_system::Pallet::<Polkadot>::block_number();
+		let scheduled_at = now + 2;
+
+		// Fellowship Origin
+		let origin = pallet_xcm::Origin::Xcm(Location::new(
+			0,
+			[
+				Junction::Parachain(1001),
+				Junction::Plurality { id: BodyId::Technical, part: BodyPart::Voice },
+			],
+		));
+		assert_ok!(RcMigrator::schedule_migration(origin.into(), DispatchTime::At(scheduled_at)));
+		assert_eq!(
+			RcMigrationStageStorage::<Polkadot>::get(),
+			RcMigrationStage::Scheduled { block_number: scheduled_at }
+		);
+
+		next_block_rc();
+		// migrating not yet started
+		assert_eq!(
+			RcMigrationStageStorage::<Polkadot>::get(),
+			RcMigrationStage::Scheduled { block_number: scheduled_at }
+		);
+		assert_eq!(DownwardMessageQueues::<Polkadot>::take(AH_PARA_ID).len(), 0);
+
+		next_block_rc();
+
+		// migration started
+		assert_eq!(RcMigrationStageStorage::<Polkadot>::get(), RcMigrationStage::Initializing);
+		let dmp_messages = DownwardMessageQueues::<Polkadot>::take(AH_PARA_ID);
+		assert!(dmp_messages.len() > 0);
+
+		dmp_messages
+	});
+
+	// enqueue DMP messages from RC to AH.
+	ah.execute_with(|| {
+		enqueue_dmp(dmp_messages);
+	});
+	ah.commit_all().unwrap();
+
+	// Asset Hub receives the message from the Relay Chain to start the migration and the
+	// acknowledges it by sending the message back to the Relay Chain.
+	let ump_messages = ah.execute_with(|| {
+		log::info!("Acknowledging the start of the migration on AH");
+		assert_eq!(AhMigrationStageStorage::<AssetHub>::get(), AhMigrationStage::Pending);
+
+		next_block_ah();
+
+		assert_eq!(
+			AhMigrationStageStorage::<AssetHub>::get(),
+			AhMigrationStage::DataMigrationOngoing
+		);
+
+		PendingUpwardMessages::<AssetHub>::take()
+	});
+	ah.commit_all().unwrap();
+
+	// enqueue UMP messages from AH to RC.
+	rc.execute_with(|| {
+		enqueue_ump(ump_messages);
+	});
+	rc.commit_all().unwrap();
+
+	// Relay Chain receives the acknowledgement from the Asset Hub and starts sending the data.
+	rc.execute_with(|| {
+		log::info!("Receiving the acknowledgement from AH on RC");
+		assert_eq!(RcMigrationStageStorage::<Polkadot>::get(), RcMigrationStage::Initializing);
+
+		next_block_rc();
+
+		assert_eq!(
+			RcMigrationStageStorage::<Polkadot>::get(),
+			RcMigrationStage::AccountsMigrationOngoing { last_key: None }
+		);
+	});
+	rc.commit_all().unwrap();
+}
+
 #[tokio::test]
 async fn some_account_migration_works() {
 	use frame_system::Account as SystemAccount;

pallets/ah-migrator/src/account.rs

@@ -25,7 +25,10 @@ impl<T: Config> Pallet<T> {
 	) -> Result<(), Error<T>> {
 		log::info!(target: LOG_TARGET, "Integrating {} accounts", accounts.len());
 
-		Self::deposit_event(Event::<T>::AccountBatchReceived { count: accounts.len() as u32 });
+		Self::deposit_event(Event::<T>::BatchReceived {
+			pallet: PalletEventName::Balances,
+			count: accounts.len() as u32,
+		});
 		let (mut count_good, mut count_bad) = (0, 0);
 
 		for account in accounts {
@@ -48,7 +51,11 @@ impl<T: Config> Pallet<T> {
 			}
 		}
 
-		Self::deposit_event(Event::<T>::AccountBatchProcessed { count_good, count_bad });
+		Self::deposit_event(Event::<T>::BatchProcessed {
+			pallet: PalletEventName::Balances,
+			count_good,
+			count_bad,
+		});
 		Ok(())
 	}
 

pallets/ah-migrator/src/asset_rate.rs

@@ -24,14 +24,18 @@ impl<T: Config> Pallet<T> {
 		log::info!(target: LOG_TARGET, "Processing {} asset rates", rates.len());
 
 		let count = rates.len() as u32;
-		Self::deposit_event(Event::AssetRatesReceived { count });
+		Self::deposit_event(Event::BatchReceived { pallet: PalletEventName::AssetRates, count });
 
 		for rate in rates {
 			Self::do_receive_asset_rate(rate)?;
 		}
 
 		log::info!(target: LOG_TARGET, "Processed {} asset rates", count);
-		Self::deposit_event(Event::AssetRatesProcessed { count_good: count, count_bad: 0 });
+		Self::deposit_event(Event::BatchProcessed {
+			pallet: PalletEventName::AssetRates,
+			count_good: count,
+			count_bad: 0,
+		});
 
 		Ok(())
 	}

pallets/ah-migrator/src/claims.rs

@@ -21,7 +21,10 @@ use pallet_rc_migrator::claims::{alias, RcClaimsMessage, RcClaimsMessageOf};
 impl<T: Config> Pallet<T> {
 	pub fn do_receive_claims(messages: Vec<RcClaimsMessageOf<T>>) -> Result<(), Error<T>> {
 		log::info!(target: LOG_TARGET, "Integrating {} claims", messages.len());
-		Self::deposit_event(Event::ClaimsBatchReceived { count: messages.len() as u32 });
+		Self::deposit_event(Event::BatchReceived {
+			pallet: PalletEventName::Claims,
+			count: messages.len() as u32,
+		});
 		let (mut count_good, mut count_bad) = (0, 0);
 
 		for message in messages {
@@ -33,7 +36,11 @@ impl<T: Config> Pallet<T> {
 				},
 			}
 		}
-		Self::deposit_event(Event::ClaimsBatchProcessed { count_good, count_bad });
+		Self::deposit_event(Event::BatchProcessed {
+			pallet: PalletEventName::Claims,
+			count_good,
+			count_bad,
+		});
 
 		Ok(())
 	}

pallets/ah-migrator/src/conviction_voting.rs

@@ -27,13 +27,20 @@ impl<T: Config> Pallet<T> {
 	) -> Result<(), Error<T>> {
 		log::info!(target: LOG_TARGET, "Processing {} conviction voting messages", messages.len());
 		let count = messages.len() as u32;
-		Self::deposit_event(Event::ConvictionVotingMessagesReceived { count });
+		Self::deposit_event(Event::BatchReceived {
+			pallet: PalletEventName::ConvictionVoting,
+			count,
+		});
 
 		for message in messages {
 			Self::do_receive_conviction_voting_message(message);
 		}
 
-		Self::deposit_event(Event::ConvictionVotingMessagesProcessed { count_good: count });
+		Self::deposit_event(Event::BatchProcessed {
+			pallet: PalletEventName::ConvictionVoting,
+			count_good: count,
+			count_bad: 0,
+		});
 
 		Ok(())
 	}