[AHM] Xcm send weights and Batch helper types

integration-tests/ahm/src/bench_rc.rs

@@ -71,3 +71,10 @@ fn test_bench_update_ah_msg_processed_count() {
 		test_update_ah_msg_processed_count::<RelayChain>();
 	});
 }
+
+#[test]
+fn test_bench_send_chunked_xcm_and_track() {
+	new_test_ext().execute_with(|| {
+		test_send_chunked_xcm_and_track::<RelayChain>();
+	});
+}

pallets/rc-migrator/src/accounts.rs

@@ -213,14 +213,7 @@ impl<T: Config> PalletMigration for AccountsMigrator<T> {
 		// we should not send more than we allocated on AH for the migration.
 		let mut ah_weight = WeightMeter::with_limit(T::MaxAhWeight::get());
 		// accounts batch for the current iteration.
-		let mut batch = Vec::new();
-
-		// TODO transport weight. probably we need to leave some buffer since we do not know how
-		// many send batches the one migrate_many will require.
-		let xcm_weight = Weight::from_all(1);
-		if weight_counter.try_consume(xcm_weight).is_err() {
-			return Err(Error::OutOfWeight);
-		}
+		let mut batch = XcmBatchAndMeter::new_from_config::<T>();
 
 		let mut iter = if let Some(ref last_key) = last_key {
 			SystemAccount::<T>::iter_from_key(last_key)
@@ -231,7 +224,9 @@ impl<T: Config> PalletMigration for AccountsMigrator<T> {
 		let mut maybe_last_key = last_key;
 		loop {
 			// account the weight for migrating a single account on Relay Chain.
-			if weight_counter.try_consume(T::RcWeightInfo::withdraw_account()).is_err() {
+			if weight_counter.try_consume(T::RcWeightInfo::withdraw_account()).is_err() ||
+				weight_counter.try_consume(batch.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", batch.len());
 				if batch.is_empty() {
 					return Err(Error::OutOfWeight);
@@ -670,6 +665,7 @@ impl<T: Config> AccountsMigrator<T> {
 	///
 	/// Should be executed once before the migration starts.
 	pub fn obtain_rc_accounts() -> Weight {
+		let mut weight = Weight::zero();
 		let mut reserves = sp_std::collections::btree_map::BTreeMap::new();
 		let mut update_reserves = |id, deposit| {
 			if deposit == 0 {
@@ -679,6 +675,7 @@ impl<T: Config> AccountsMigrator<T> {
 		};
 
 		for (channel_id, info) in hrmp::HrmpChannels::<T>::iter() {
+			weight += T::DbWeight::get().reads(1);
 			// source: https://github.com/paritytech/polkadot-sdk/blob/3dc3a11cd68762c2e5feb0beba0b61f448c4fc92/polkadot/runtime/parachains/src/hrmp.rs#L1475
 			let sender: T::AccountId = channel_id.sender.into_account_truncating();
 			update_reserves(sender, info.sender_deposit);
@@ -689,16 +686,19 @@ impl<T: Config> AccountsMigrator<T> {
 		}
 
 		for (channel_id, info) in hrmp::HrmpOpenChannelRequests::<T>::iter() {
+			weight += T::DbWeight::get().reads(1);
 			// source: https://github.com/paritytech/polkadot-sdk/blob/3dc3a11cd68762c2e5feb0beba0b61f448c4fc92/polkadot/runtime/parachains/src/hrmp.rs#L1475
 			let sender: T::AccountId = channel_id.sender.into_account_truncating();
 			update_reserves(sender, info.sender_deposit);
 		}
 
 		for (_, info) in Paras::<T>::iter() {
+			weight += T::DbWeight::get().reads(1);
 			update_reserves(info.manager, info.deposit);
 		}
 
 		for (id, rc_reserved) in reserves {
+			weight += T::DbWeight::get().reads(4);
 			let account_entry = SystemAccount::<T>::get(&id);
 			let free = <T as Config>::Currency::balance(&id);
 			let total_frozen = account_entry.data.frozen;
@@ -728,6 +728,7 @@ impl<T: Config> AccountsMigrator<T> {
 			}
 
 			if ah_free < ah_ed && rc_reserved >= total_reserved && total_frozen.is_zero() {
+				weight += T::DbWeight::get().writes(1);
 				// when there is no much free balance and the account is used only for reserves
 				// for parachains registering or hrmp channels we will keep the entire account on
 				// the RC.
@@ -738,6 +739,7 @@ impl<T: Config> AccountsMigrator<T> {
 				);
 				RcAccounts::<T>::insert(&id, AccountState::Preserve);
 			} else {
+				weight += T::DbWeight::get().writes(1);
 				log::debug!(
 					target: LOG_TARGET,
 					"Keep part of account: {:?} reserve: {:?} on the RC",
@@ -747,9 +749,7 @@ impl<T: Config> AccountsMigrator<T> {
 				RcAccounts::<T>::insert(&id, AccountState::Part { reserved: rc_reserved });
 			}
 		}
-
-		// TODO: define actual weight
-		Weight::from_all(1)
+		weight
 	}
 
 	/// Try to translate a Parachain sovereign account to the Parachain AH sovereign account.

pallets/rc-migrator/src/asset_rate.rs

@@ -31,12 +31,12 @@ impl<T: Config> PalletMigration for AssetRateMigrator<T> {
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		log::info!(target: LOG_TARGET, "Migrating asset rates");
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 
 		loop {
 			if weight_counter
 				.try_consume(<T as frame_system::Config>::DbWeight::get().reads_writes(1, 1))
-				.is_err()
+				.is_err() || weight_counter.try_consume(messages.consume_weight()).is_err()
 			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if messages.is_empty() {

pallets/rc-migrator/src/benchmarking.rs

@@ -101,6 +101,22 @@ pub mod benchmarks {
 		);
 	}
 
+	#[benchmark]
+	fn send_chunked_xcm_and_track() {
+		let mut batches = XcmBatch::new();
+		batches.push(vec![0u8; (MAX_XCM_SIZE / 2 - 10) as usize]);
+		batches.push(vec![1u8; (MAX_XCM_SIZE / 2 - 10) as usize]);
+		#[block]
+		{
+			let res = Pallet::<T>::send_chunked_xcm_and_track(
+				batches,
+				|batch| types::AhMigratorCall::<T>::TestCall { data: batch },
+				|_| Weight::from_all(1),
+			);
+			assert_eq!(res.unwrap(), 1);
+		}
+	}
+
 	#[benchmark]
 	fn update_ah_msg_processed_count() {
 		let new_processed = 100;
@@ -137,4 +153,9 @@ pub mod benchmarks {
 	pub fn test_update_ah_msg_processed_count<T: Config>() {
 		_update_ah_msg_processed_count::<T>(true /* enable checks */);
 	}
+
+	#[cfg(feature = "std")]
+	pub fn test_send_chunked_xcm_and_track<T: Config>() {
+		_send_chunked_xcm_and_track::<T>(true /* enable checks */);
+	}
 }

pallets/rc-migrator/src/bounties.rs

@@ -62,12 +62,14 @@ impl<T: Config> PalletMigration for BountiesMigrator<T> {
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		let mut last_key = last_key.unwrap_or(BountiesStage::BountyCount);
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 
 		log::info!(target: LOG_TARGET, "Migrating Bounties at stage {:?}", &last_key);
 
 		loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(messages.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if messages.is_empty() {
 					return Err(Error::OutOfWeight);
@@ -150,7 +152,7 @@ impl<T: Config> PalletMigration for BountiesMigrator<T> {
 
 		if !messages.is_empty() {
 			Pallet::<T>::send_chunked_xcm_and_track(
-				messages,
+				messages.into_inner(),
 				|messages| types::AhMigratorCall::<T>::ReceiveBountiesMessages { messages },
 				|len| T::AhWeightInfo::receive_bounties_messages(len),
 			)?;

pallets/rc-migrator/src/claims.rs

@@ -59,10 +59,12 @@ impl<T: Config> PalletMigration for ClaimsMigrator<T> {
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		let mut inner_key = current_key.unwrap_or(ClaimsStage::StorageValues);
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 
 		loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(messages.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if messages.is_empty() {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/conviction_voting.rs

@@ -55,11 +55,13 @@ impl<T: Config> PalletMigration for ConvictionVotingMigrator<T> {
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		let mut last_key = last_key.unwrap_or(ConvictionVotingStage::VotingFor(None));
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 		let mut made_progress = false;
 
 		loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(messages.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if !made_progress {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/crowdloan.rs

@@ -104,12 +104,12 @@ impl<T: Config> PalletMigration for CrowdloanMigrator<T>
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		let mut inner_key = current_key.unwrap_or(CrowdloanStage::Setup);
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 
 		loop {
 			if weight_counter
 				.try_consume(T::DbWeight::get().reads_writes(2, 1))
-				.is_err()
+				.is_err() || weight_counter.try_consume(messages.consume_weight()).is_err()
 			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if messages.is_empty() {

pallets/rc-migrator/src/indices.rs

@@ -51,10 +51,12 @@ impl<T: Config> PalletMigration for IndicesMigrator<T> {
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
 		let mut inner_key = current_key;
-		let mut messages = Vec::new();
+		let mut messages = XcmBatchAndMeter::new_from_config::<T>();
 
 		loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(messages.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", messages.len());
 				if messages.is_empty() {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/lib.rs

@@ -61,7 +61,8 @@ pub mod xcm_config;
 pub use weights::*;
 
 use crate::{
-	accounts::MigratedBalances, types::MigrationFinishedData,
+	accounts::MigratedBalances,
+	types::{MigrationFinishedData, XcmBatch, XcmBatchAndMeter},
 	xcm_config::TrustedTeleportersBeforeAndAfter,
 };
 use accounts::AccountsMigrator;
@@ -621,8 +622,8 @@ pub mod pallet {
 					return weight_counter.consumed();
 				},
 				MigrationStage::AccountsMigrationInit => {
-					// TODO: weights
-					let _ = AccountsMigrator::<T>::obtain_rc_accounts();
+					let weight = AccountsMigrator::<T>::obtain_rc_accounts();
+					weight_counter.consume(weight);
 					RcMigratedBalance::<T>::mutate(|tracker| {
 						// initialize `kept` balance as total issuance, we'll substract from it as
 						// we migrate accounts
@@ -1452,31 +1453,16 @@ pub mod pallet {
 		/// Will modify storage in the error path.
 		/// This is done to avoid exceeding the XCM message size limit.
 		pub fn send_chunked_xcm<E: Encode>(
-			mut items: Vec<E>,
+			items: impl Into<XcmBatch<E>>,
 			create_call: impl Fn(Vec<E>) -> types::AhMigratorCall<T>,
 			weight_at_most: impl Fn(u32) -> Weight,
 		) -> Result<u32, Error<T>> {
+			let mut items = items.into();
 			log::info!(target: LOG_TARGET, "Batching {} items to send via XCM", items.len());
 			defensive_assert!(items.len() > 0, "Sending XCM with empty items");
-			items.reverse();
 			let mut batch_count = 0;
 
-			while !items.is_empty() {
-				let mut remaining_size: u32 = MAX_XCM_SIZE;
-				let mut batch = Vec::new();
-
-				while !items.is_empty() {
-					// Taking from the back as optimization is fine since we reversed
-					let item = items.last().unwrap(); // FAIL-CI no unwrap
-					let msg_size = item.encoded_size() as u32;
-					if msg_size > remaining_size {
-						break;
-					}
-					remaining_size -= msg_size;
-
-					batch.push(items.pop().unwrap()); // FAIL-CI no unwrap
-				}
-
+			while let Some(batch) = items.pop_front() {
 				let batch_len = batch.len() as u32;
 				log::info!(target: LOG_TARGET, "Sending XCM batch of {} items", batch_len);
 				let call = types::AssetHubPalletConfig::<T>::AhmController(create_call(batch));
@@ -1575,7 +1561,7 @@ pub mod pallet {
 		///
 		/// Check the `send_chunked_xcm` function for the documentation.
 		pub fn send_chunked_xcm_and_track<E: Encode>(
-			items: Vec<E>,
+			items: impl Into<XcmBatch<E>>,
 			create_call: impl Fn(Vec<E>) -> types::AhMigratorCall<T>,
 			weight_at_most: impl Fn(u32) -> Weight,
 		) -> Result<u32, Error<T>> {

pallets/rc-migrator/src/multisig.rs

@@ -21,7 +21,6 @@ use frame_support::traits::Currency;
 
 extern crate alloc;
 use crate::{types::*, *};
-use alloc::vec::Vec;
 
 mod aliases {
 	use super::*;
@@ -113,15 +112,17 @@ impl<T: Config> PalletMigration for MultisigMigrator<T> {
 		mut last_key: Option<Self::Key>,
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Error<T>> {
-		let mut batch = Vec::new();
+		let mut batch = XcmBatchAndMeter::new_from_config::<T>();
 		let mut iter = match last_key.clone() {
 			Some((k1, k2)) =>
 				aliases::Multisigs::<T>::iter_from(aliases::Multisigs::<T>::hashed_key_for(k1, k2)),
 			None => aliases::Multisigs::<T>::iter(),
 		};
 
 		loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(batch.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", batch.len());
 				if batch.is_empty() {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/preimage/chunks.rs

@@ -50,11 +50,13 @@ impl<T: Config> PalletMigration for PreimageChunkMigrator<T> {
 		mut next_key: Option<Self::Key>,
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
-		let mut batch = Vec::new();
+		let mut batch = XcmBatchAndMeter::new_from_config::<T>();
 		let mut ah_weight_counter = WeightMeter::new();
 
 		let last_key = loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 2)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 2)).is_err() ||
+				weight_counter.try_consume(batch.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", batch.len());
 				if batch.is_empty() {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/preimage/legacy_request_status.rs

@@ -46,10 +46,12 @@ impl<T: Config> PalletMigration for PreimageLegacyRequestStatusMigrator<T> {
 		mut next_key: Option<Self::Key>,
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
-		let mut batch = Vec::new();
+		let mut batch = XcmBatchAndMeter::new_from_config::<T>();
 
 		let new_next_key = loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(batch.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", batch.len());
 				if batch.is_empty() {
 					return Err(Error::OutOfWeight);

pallets/rc-migrator/src/preimage/request_status.rs

@@ -42,10 +42,12 @@ impl<T: Config> PalletMigration for PreimageRequestStatusMigrator<T> {
 		mut next_key: Option<Self::Key>,
 		weight_counter: &mut WeightMeter,
 	) -> Result<Option<Self::Key>, Self::Error> {
-		let mut batch = Vec::new();
+		let mut batch = XcmBatchAndMeter::new_from_config::<T>();
 
 		let new_next_key = loop {
-			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() {
+			if weight_counter.try_consume(T::DbWeight::get().reads_writes(1, 1)).is_err() ||
+				weight_counter.try_consume(batch.consume_weight()).is_err()
+			{
 				log::info!("RC weight limit reached at batch length {}, stopping", batch.len());
 				if batch.is_empty() {
 					return Err(Error::OutOfWeight);