Init agents and channels with migration

parachain/pallets/system/src/benchmarking.rs

@@ -157,5 +157,9 @@ mod benchmarks {
 		Ok(())
 	}
 
-	impl_benchmark_test_suite!(SnowbridgeControl, crate::mock::new_test_ext(), crate::mock::Test);
+	impl_benchmark_test_suite!(
+		SnowbridgeControl,
+		crate::mock::new_test_ext(true),
+		crate::mock::Test
+	);
 }

parachain/pallets/system/src/lib.rs

@@ -47,6 +47,7 @@ mod tests;
 
 #[cfg(feature = "runtime-benchmarks")]
 mod benchmarking;
+pub mod migration;
 
 pub mod api;
 pub mod weights;
@@ -256,34 +257,7 @@ pub mod pallet {
 	#[pallet::genesis_build]
 	impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
 		fn build(&self) {
-			let bridge_hub_agent_id =
-				agent_id_of::<T>(&MultiLocation::here()).expect("infallible; qed");
-			// Agent for BridgeHub
-			Agents::<T>::insert(bridge_hub_agent_id, ());
-
-			// Primary governance channel
-			Channels::<T>::insert(
-				PRIMARY_GOVERNANCE_CHANNEL,
-				Channel { agent_id: bridge_hub_agent_id, para_id: self.para_id },
-			);
-
-			// Secondary governance channel
-			Channels::<T>::insert(
-				SECONDARY_GOVERNANCE_CHANNEL,
-				Channel { agent_id: bridge_hub_agent_id, para_id: self.para_id },
-			);
-
-			// Asset Hub
-			let asset_hub_location: MultiLocation =
-				ParentThen(X1(Parachain(self.asset_hub_para_id.into()))).into();
-			let asset_hub_agent_id =
-				agent_id_of::<T>(&asset_hub_location).expect("infallible; qed");
-			let asset_hub_channel_id: ChannelId = self.asset_hub_para_id.into();
-			Agents::<T>::insert(asset_hub_agent_id, ());
-			Channels::<T>::insert(
-				asset_hub_channel_id,
-				Channel { agent_id: asset_hub_agent_id, para_id: self.asset_hub_para_id },
-			);
+			Pallet::<T>::initialize(self.para_id, self.asset_hub_para_id).expect("infallible; qed");
 		}
 	}
 
@@ -643,6 +617,49 @@ pub mod pallet {
 			});
 			Ok(())
 		}
+
+		/// Checks if the pallet has been initialized.
+		pub(crate) fn is_initialized() -> bool {
+			let primary_exists = Channels::<T>::contains_key(PRIMARY_GOVERNANCE_CHANNEL);
+			let secondary_exists = Channels::<T>::contains_key(SECONDARY_GOVERNANCE_CHANNEL);
+			primary_exists && secondary_exists
+		}
+
+		/// Initializes agents and channels.
+		pub(crate) fn initialize(
+			para_id: ParaId,
+			asset_hub_para_id: ParaId,
+		) -> Result<(), DispatchError> {
+			// Asset Hub
+			let asset_hub_location: MultiLocation =
+				ParentThen(X1(Parachain(asset_hub_para_id.into()))).into();
+			let asset_hub_agent_id = agent_id_of::<T>(&asset_hub_location)?;
+			let asset_hub_channel_id: ChannelId = asset_hub_para_id.into();
+			Agents::<T>::insert(asset_hub_agent_id, ());
+			Channels::<T>::insert(
+				asset_hub_channel_id,
+				Channel { agent_id: asset_hub_agent_id, para_id: asset_hub_para_id },
+			);
+
+			// Governance channels
+			let bridge_hub_agent_id = agent_id_of::<T>(&MultiLocation::here())?;
+			// Agent for BridgeHub
+			Agents::<T>::insert(bridge_hub_agent_id, ());
+
+			// Primary governance channel
+			Channels::<T>::insert(
+				PRIMARY_GOVERNANCE_CHANNEL,
+				Channel { agent_id: bridge_hub_agent_id, para_id },
+			);
+
+			// Secondary governance channel
+			Channels::<T>::insert(
+				SECONDARY_GOVERNANCE_CHANNEL,
+				Channel { agent_id: bridge_hub_agent_id, para_id },
+			);
+
+			Ok(())
+		}
 	}
 
 	impl<T: Config> StaticLookup for Pallet<T> {

parachain/pallets/system/src/migration.rs

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
+//! Governance API for controlling the Ethereum side of the bridge
+use super::*;
+use frame_support::traits::OnRuntimeUpgrade;
+use log;
+
+#[cfg(feature = "try-runtime")]
+use sp_runtime::TryRuntimeError;
+
+pub mod v0 {
+	use frame_support::{pallet_prelude::*, weights::Weight};
+
+	use super::*;
+
+	const LOG_TARGET: &str = "ethereum_system::migration";
+
+	pub struct InitializeOnUpgrade<T, BridgeHubParaId, AssetHubParaId>(
+		sp_std::marker::PhantomData<(T, BridgeHubParaId, AssetHubParaId)>,
+	);
+	impl<T, BridgeHubParaId, AssetHubParaId> OnRuntimeUpgrade
+		for InitializeOnUpgrade<T, BridgeHubParaId, AssetHubParaId>
+	where
+		T: Config,
+		BridgeHubParaId: Get<u32>,
+		AssetHubParaId: Get<u32>,
+	{
+		fn on_runtime_upgrade() -> Weight {
+			if !Pallet::<T>::is_initialized() {
+				Pallet::<T>::initialize(
+					BridgeHubParaId::get().into(),
+					AssetHubParaId::get().into(),
+				)
+				.expect("infallible; qed");
+				log::info!(
+					target: LOG_TARGET,
+					"Ethereum system initialized."
+				);
+				T::DbWeight::get().reads_writes(2, 5)
+			} else {
+				log::info!(
+					target: LOG_TARGET,
+					"Ethereum system already initialized. Skipping."
+				);
+				T::DbWeight::get().reads(2)
+			}
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, TryRuntimeError> {
+			if !Pallet::<T>::is_initialized() {
+				log::info!(
+					target: LOG_TARGET,
+					"Agents and channels not initialized. Initialization will run."
+				);
+			} else {
+				log::info!(
+					target: LOG_TARGET,
+					"Agents and channels are initialized. Initialization will not run."
+				);
+			}
+			Ok(vec![])
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_: Vec<u8>) -> Result<(), TryRuntimeError> {
+			frame_support::ensure!(
+				Pallet::<T>::is_initialized(),
+				"Agents and channels were not initialized."
+			);
+			Ok(())
+		}
+	}
+}

parachain/pallets/system/src/mock.rs

@@ -232,16 +232,18 @@ impl crate::Config for Test {
 }
 
 // Build genesis storage according to the mock runtime.
-pub fn new_test_ext() -> sp_io::TestExternalities {
+pub fn new_test_ext(genesis_build: bool) -> sp_io::TestExternalities {
 	let mut storage = frame_system::GenesisConfig::<Test>::default().build_storage().unwrap();
 
-	crate::GenesisConfig::<Test> {
-		para_id: OwnParaId::get(),
-		asset_hub_para_id: AssetHubParaId::get(),
-		_config: Default::default(),
+	if genesis_build {
+		crate::GenesisConfig::<Test> {
+			para_id: OwnParaId::get(),
+			asset_hub_para_id: AssetHubParaId::get(),
+			_config: Default::default(),
+		}
+		.assimilate_storage(&mut storage)
+		.unwrap();
 	}
-	.assimilate_storage(&mut storage)
-	.unwrap();
 
 	let mut ext: sp_io::TestExternalities = storage.into();
 	let initial_amount = InitialFunding::get();