From 48f537a954b7e2c11ef9409f408afffe53303798 Mon Sep 17 00:00:00 2001
From: Xavier Lau <x@acg.box>
Date: Wed, 9 Oct 2024 10:26:56 +0800
Subject: [PATCH] Improve migration

---
 node/src/command.rs               |  2 +-
 pallet/staking/src/lib.rs         |  2 +-
 runtime/crab/src/migration.rs     | 44 +------------------------------
 runtime/darwinia/src/migration.rs | 31 +++++++++++++---------
 runtime/koi/src/migration.rs      |  2 +-
 5 files changed, 22 insertions(+), 59 deletions(-)

diff --git a/node/src/command.rs b/node/src/command.rs
index fc7b4f00a..96a019e18 100644
--- a/node/src/command.rs
+++ b/node/src/command.rs
@@ -234,7 +234,7 @@ impl CliConfiguration<Self> for RelayChainCli {
 /// Parse command line arguments into service configuration.
 pub fn run() -> Result<()> {
 	#[cfg(feature = "runtime-benchmarks")]
-	/// Creates partial components for the runtimes that are supported by the benchmarks.
+	/// Creates partial components for the runtime that are supported by the benchmarks.
 	macro_rules! construct_benchmark_partials {
 		($config:expr, $cli:ident, |$partials:ident| $code:expr) => {{
 			#[cfg(feature = "crab-runtime")]
diff --git a/pallet/staking/src/lib.rs b/pallet/staking/src/lib.rs
index 9bc28850b..be6b32a8d 100644
--- a/pallet/staking/src/lib.rs
+++ b/pallet/staking/src/lib.rs
@@ -134,7 +134,7 @@ pub mod pallet {
 		/// Override the [`frame_system::Config::RuntimeEvent`].
 		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
 
-		/// Weight information for extrinsics in this pallet.
+		/// Weight information for extrinsic in this pallet.
 		type WeightInfo: WeightInfo;
 
 		/// Unix time interface.
diff --git a/runtime/crab/src/migration.rs b/runtime/crab/src/migration.rs
index 75e85b19b..198778bb6 100644
--- a/runtime/crab/src/migration.rs
+++ b/runtime/crab/src/migration.rs
@@ -45,47 +45,5 @@ impl frame_support::traits::OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
 }
 
 fn migrate() -> frame_support::weights::Weight {
-	// dawinia
-	use darwinia_staking::CacheState;
-	if let Some(s) = migration::get_storage_value::<(CacheState, CacheState, CacheState)>(
-		b"DarwinaStaking",
-		b"ExposureCacheStates",
-		&[],
-	) {
-		migration::put_storage_value(b"DarwinaStaking", b"CacheStates", &[], s);
-	}
-
-	if let Ok(dao) =
-		array_bytes::hex_n_into::<_, AccountId, 20>("0x08837De0Ae21C270383D9F2de4DB03c7b1314632")
-	{
-		let _ = <pallet_assets::Pallet<Runtime>>::transfer_ownership(
-			RuntimeOrigin::signed(ROOT),
-			codec::Compact(AssetIds::CKton as AssetId),
-			dao,
-		);
-
-		if let Ok(deposit) = array_bytes::hex_n_into::<_, AccountId, 20>(
-			"0x46275d29113f065c2aac262f34C7a3d8a8B7377D",
-		) {
-			let _ = <pallet_assets::Pallet<Runtime>>::set_team(
-				RuntimeOrigin::signed(dao),
-				codec::Compact(AssetIds::CKton as AssetId),
-				deposit,
-				deposit,
-				dao,
-			);
-
-			<darwinia_deposit::DepositContract<Runtime>>::put(deposit);
-		}
-	}
-	if let Ok(who) =
-		array_bytes::hex_n_into::<_, AccountId, 20>("0xa4fFAC7A5Da311D724eD47393848f694Baee7930")
-	{
-		<darwinia_staking::RingStakingContract<Runtime>>::put(who);
-	}
-
-	<darwinia_staking::MigrationStartPoint<Runtime>>::put(darwinia_staking::now::<Runtime>());
-
-	// frame_support::weights::Weight::zero()
-	<Runtime as frame_system::Config>::DbWeight::get().reads_writes(7, 7)
+	<Runtime as frame_system::Config>::DbWeight::get().reads_writes(0, 0)
 }
diff --git a/runtime/darwinia/src/migration.rs b/runtime/darwinia/src/migration.rs
index cd8108886..73d54c66a 100644
--- a/runtime/darwinia/src/migration.rs
+++ b/runtime/darwinia/src/migration.rs
@@ -45,22 +45,14 @@ impl frame_support::traits::OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
 }
 
 fn migrate() -> frame_support::weights::Weight {
-	let _ = migration::clear_storage_prefix(
-		b"BridgeKusamaGrandpa",
-		b"ImportedHeaders",
-		&[],
-		Some(100),
-		None,
-	);
-
-	// dawinia
+	// darwinia
 	use darwinia_staking::CacheState;
 	if let Some(s) = migration::get_storage_value::<(CacheState, CacheState, CacheState)>(
-		b"DarwinaStaking",
+		b"DarwiniaStaking",
 		b"ExposureCacheStates",
 		&[],
 	) {
-		migration::put_storage_value(b"DarwinaStaking", b"CacheStates", &[], s);
+		migration::put_storage_value(b"DarwiniaStaking", b"CacheStates", &[], s);
 	}
 
 	if let Ok(dao) =
@@ -85,15 +77,28 @@ fn migrate() -> frame_support::weights::Weight {
 
 			<darwinia_deposit::DepositContract<Runtime>>::put(deposit);
 		}
+
+		log::info!("successfully transfer ownership of KTON to KTON DAO");
 	}
 	if let Ok(who) =
 		array_bytes::hex_n_into::<_, AccountId, 20>("0xa4fFAC7A5Da311D724eD47393848f694Baee7930")
 	{
 		<darwinia_staking::RingStakingContract<Runtime>>::put(who);
+
+		log::info!("successfully set RING staking contract");
 	}
 
 	<darwinia_staking::MigrationStartPoint<Runtime>>::put(darwinia_staking::now::<Runtime>());
 
-	// frame_support::weights::Weight::zero()
-	<Runtime as frame_system::Config>::DbWeight::get().reads_writes(7, 107)
+	if let Some(k) = migration::take_storage_value::<AccountId>(
+		b"DarwiniaStaking",
+		b"KtonRewardDistributionContract",
+		&[],
+	) {
+		<darwinia_staking::KtonStakingContract<Runtime>>::put(k);
+
+		log::info!("successfully set KTON staking contract");
+	}
+
+	<Runtime as frame_system::Config>::DbWeight::get().reads_writes(7, 9)
 }
diff --git a/runtime/koi/src/migration.rs b/runtime/koi/src/migration.rs
index d6045cb0e..198778bb6 100644
--- a/runtime/koi/src/migration.rs
+++ b/runtime/koi/src/migration.rs
@@ -45,5 +45,5 @@ impl frame_support::traits::OnRuntimeUpgrade for CustomOnRuntimeUpgrade {
 }
 
 fn migrate() -> frame_support::weights::Weight {
-	frame_support::weights::Weight::zero()
+	<Runtime as frame_system::Config>::DbWeight::get().reads_writes(0, 0)
 }