fix: manage nodeByFarmID map when node is updated (#433)

Author: DylanVerstraete

substrate-node/pallets/pallet-tfgrid/src/grid_migration.rs

@@ -112,7 +112,7 @@ pub mod v7 {
 
         #[cfg(feature = "try-runtime")]
         fn post_upgrade() -> Result<(), &'static str> {
-            assert!(PalletVersion::<T>::get() == types::StorageVersion::V7Struct);
+            assert!(PalletVersion::<T>::get() == types::StorageVersion::V8Struct);
 
             info!(
                 "👥  TFGrid pallet migration to {:?} passes POST migrate checks ✅",
@@ -210,6 +210,11 @@ pub fn migrate_nodes<T: Config>() -> frame_support::weights::Weight {
     );
 
     for (farm_id, nodes) in farms_with_nodes.iter() {
+        info!(
+            "inserting nodes: {:?} with farm id: {:?}",
+            nodes.clone(),
+            farm_id
+        );
         NodesByFarmID::<T>::insert(farm_id, nodes);
         writes += 1;
     }
@@ -275,7 +280,7 @@ pub fn migrate_farms<T: Config>() -> frame_support::weights::Weight {
     );
 
     // Update pallet storage version
-    PalletVersion::<T>::set(types::StorageVersion::V7Struct);
+    PalletVersion::<T>::set(types::StorageVersion::V8Struct);
     info!(" <<< Storage version upgraded");
 
     // Return the weight consumed by the migration.

substrate-node/pallets/pallet-tfgrid/src/lib.rs

@@ -35,6 +35,7 @@ pub mod farm;
 pub mod grid_migration;
 pub mod interface;
 pub mod nodes_migration;
+pub mod nodes_migration_v2;
 pub mod pub_config;
 pub mod pub_ip;
 pub mod twin;
@@ -1044,6 +1045,18 @@ pub mod pallet {
 
             let old_node = Nodes::<T>::get(node_id).ok_or(Error::<T>::NodeNotExists)?;
 
+            // If the farm ID changed on the node,
+            // remove the node from the old map from the farm and insert into the correct one
+            if old_node.farm_id != farm_id {
+                let mut old_nodes_by_farm = NodesByFarmID::<T>::get(old_node.farm_id);
+                old_nodes_by_farm.retain(|&id| id != node_id);
+                NodesByFarmID::<T>::insert(farm_id, old_nodes_by_farm);
+
+                let mut nodes_by_farm = NodesByFarmID::<T>::get(farm_id);
+                nodes_by_farm.push(node_id);
+                NodesByFarmID::<T>::insert(farm_id, nodes_by_farm);
+            };
+
             stored_node.farm_id = farm_id;
             stored_node.resources = resources;
             stored_node.location = location;

substrate-node/pallets/pallet-tfgrid/src/nodes_migration.rs

@@ -22,7 +22,7 @@ pub mod v7patch {
         }
 
         fn on_runtime_upgrade() -> Weight {
-            if PalletVersion::<T>::get() == types::StorageVersion::V6Struct {
+            if PalletVersion::<T>::get() == types::StorageVersion::V7Struct {
                 add_farm_nodes_index::<T>()
             } else {
                 info!(" >>> Unused migration");
@@ -32,7 +32,7 @@ pub mod v7patch {
 
         #[cfg(feature = "try-runtime")]
         fn post_upgrade() -> Result<(), &'static str> {
-            assert!(PalletVersion::<T>::get() == types::StorageVersion::V7Struct);
+            assert!(PalletVersion::<T>::get() == types::StorageVersion::V8Struct);
 
             info!(
                 "👥  TFGrid pallet migration to {:?} passes POST migrate checks ✅",
@@ -62,12 +62,17 @@ pub fn add_farm_nodes_index<T: Config>() -> frame_support::weights::Weight {
     }
 
     for (farm_id, nodes) in farms_with_nodes.iter() {
+        info!(
+            "inserting nodes: {:?} with farm id: {:?}",
+            nodes.clone(),
+            farm_id
+        );
         NodesByFarmID::<T>::insert(farm_id, nodes);
         writes += 1;
     }
 
     // Update pallet storage version
-    PalletVersion::<T>::set(types::StorageVersion::V7Struct);
+    PalletVersion::<T>::set(types::StorageVersion::V8Struct);
     info!(" <<< Storage version upgraded");
 
     // Return the weight consumed by the migration.

substrate-node/pallets/pallet-tfgrid/src/nodes_migration_v2.rs

@@ -0,0 +1,82 @@
+use super::Config;
+use super::*;
+use frame_support::{traits::Get, weights::Weight};
+use log::info;
+use sp_std::collections::btree_map::BTreeMap;
+
+pub mod v8patch {
+    use super::*;
+    use crate::Config;
+
+    use frame_support::{pallet_prelude::Weight, traits::OnRuntimeUpgrade};
+    use sp_std::marker::PhantomData;
+    pub struct FixFarmNodeIndexMap<T: Config>(PhantomData<T>);
+
+    impl<T: Config> OnRuntimeUpgrade for FixFarmNodeIndexMap<T> {
+        #[cfg(feature = "try-runtime")]
+        fn pre_upgrade() -> Result<(), &'static str> {
+            assert!(PalletVersion::<T>::get() == types::StorageVersion::V7Struct);
+
+            info!("👥  TFGrid pallet to v4 passes PRE migrate checks ✅",);
+            Ok(())
+        }
+
+        fn on_runtime_upgrade() -> Weight {
+            if PalletVersion::<T>::get() == types::StorageVersion::V7Struct {
+                add_farm_nodes_index::<T>()
+            } else {
+                info!(" >>> Unused migration");
+                return 0;
+            }
+        }
+
+        #[cfg(feature = "try-runtime")]
+        fn post_upgrade() -> Result<(), &'static str> {
+            assert!(PalletVersion::<T>::get() == types::StorageVersion::V8Struct);
+
+            info!(
+                "👥  TFGrid pallet migration to {:?} passes POST migrate checks ✅",
+                Pallet::<T>::pallet_version()
+            );
+
+            Ok(())
+        }
+    }
+}
+
+pub fn add_farm_nodes_index<T: Config>() -> frame_support::weights::Weight {
+    info!(" >>> Migrating nodes storage...");
+
+    NodesByFarmID::<T>::remove_all(None);
+
+    let mut reads = 0;
+    let mut writes = 0;
+
+    let mut farms_with_nodes: BTreeMap<u32, Vec<u32>> = BTreeMap::new();
+    for (_, node) in Nodes::<T>::iter() {
+        // Add index of farm - list (nodes)
+        farms_with_nodes
+            .entry(node.farm_id)
+            .or_insert(vec![])
+            .push(node.id);
+
+        reads += 1;
+    }
+
+    for (farm_id, nodes) in farms_with_nodes.iter() {
+        info!(
+            "inserting nodes: {:?} with farm id: {:?}",
+            nodes.clone(),
+            farm_id
+        );
+        NodesByFarmID::<T>::insert(farm_id, nodes);
+        writes += 1;
+    }
+
+    // Update pallet storage version
+    PalletVersion::<T>::set(types::StorageVersion::V8Struct);
+    info!(" <<< Storage version upgraded");
+
+    // Return the weight consumed by the migration.
+    T::DbWeight::get().reads_writes(reads as Weight + 1, writes as Weight + 1)
+}

substrate-node/pallets/pallet-tfgrid/src/types.rs

@@ -14,6 +14,7 @@ pub enum StorageVersion {
     V5Struct,
     V6Struct,
     V7Struct,
+    V8Struct,
 }
 
 impl Default for StorageVersion {

substrate-node/runtime/src/lib.rs

@@ -759,6 +759,7 @@ pub type Executive = frame_executive::Executive<
         pallet_smart_contract::contract_migration::v5::ContractMigrationV5<Runtime>,
         pallet_tfgrid::grid_migration::v7::GridMigration<Runtime>,
         pallet_tfgrid::nodes_migration::v7patch::NodesMigration<Runtime>,
+        pallet_tfgrid::nodes_migration_v2::v8patch::FixFarmNodeIndexMap<Runtime>,
     ),
 >;